mirror of
https://github.com/seaweedfs/seaweedfs.git
synced 2026-11-12 23:09:33 +01:00
vector: reuse the Lance session, version leases and opened tables across reads (#12573)
* vector: share one Lance session across reads Each request opened its table in a session of its own, so every query loaded and decoded the index again. Get, list and query now open in one session the server keeps, sized by -index-cache-mb and -metadata-cache-mb; writes keep a session of their own. Lance keys some of its cached entries, such as a version's index list, by location and version alone, and an index dropped and created again at its location numbers its versions from 1 again: a shared session answered a query on the new table with the old table's index. CreateIndex now marks each table with an id in its schema metadata, and the session starts over when a location turns up holding another table. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * vector: share a version lease among the reads of that version Every get, list and query tagged the version it read and deleted the tag afterwards: two filer writes and a manifest read per request. Reads of one version now share one tag, which goes 10 s after the last read took it; a tag is reused for at most 5 minutes, so a busy version never reaches the reapers' TTLs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * vector: keep opened tables for the next read A read built the table's store and read its latest manifest each time. Reads now keep the opened table by location and encryption, up to 1024, and move it to the latest version, which lists the versions and takes the manifest from the session. A handle opened in an earlier session is not reused. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * vector: key shared leases by manifest, and take one per version at once A table dropped and created again within the linger reused, for a version number the old table had, a lease whose tag the drop had deleted. Leases are now keyed by the manifest's ETag as well. Reads that arrive before a version has a lease wait for the first to take it, rather than each writing a tag. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * vector: drop kept tables when the shared session starts over A kept table holds the session it was opened in, so a replaced session and its caches lived on in the handles of other tables until they were evicted. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * vector: prune shared lease slots only when they double Every lease lookup scanned all slots under the map's lock, a cost that grew with the versions read in the last linger. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * vector: shrink the lease slot map after a prune After a burst the map kept its peak allocation, which every later prune scanned. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * vector: keep a lease slot a read has taken through a prune A read that took a slot and had not yet locked it could see the slot pruned, and then pin through it while the next read pinned through a new slot, writing a second tag. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
1 parent
b924ffabf1
commit
13334a59b9
8 files changed
+537
-23
No files matched your search
Generated
+1
@@ -7256,6 +7256,7 @@ dependencies = [
|
||||
"tonic-prost-build",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"uuid",
|
||||
"weed-lance-worker",
|
||||
"weed-vector-core",
|
||||
]
|
||||
|
||||
@@ -170,7 +170,7 @@ async fn open_dataset(
|
||||
storage_options: Default::default(),
|
||||
managed_versioning: false,
|
||||
};
|
||||
return dataset::open_at(&description, storage.wrapper(id).await?).await;
|
||||
return dataset::open_at(&description, storage.wrapper(id).await?, None).await;
|
||||
}
|
||||
let table = dataset::open(client, &parse_id(id), storage).await?;
|
||||
Ok(table.dataset)
|
||||
|
||||
@@ -2,8 +2,8 @@
|
||||
//! and version tags act as the lease maintenance must respect.
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::{LazyLock, Mutex};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
use std::sync::{Arc, LazyLock, Mutex, Weak};
|
||||
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use anyhow::Result;
|
||||
use lance::Dataset;
|
||||
@@ -21,6 +21,11 @@ const FOREIGN_TTL: Duration = Duration::from_secs(24 * 3600);
|
||||
/// Expired-lease scans run at most once per interval per dataset, not per
|
||||
/// request.
|
||||
const REAP_INTERVAL: Duration = Duration::from_secs(60);
|
||||
/// How long a shared lease stays after the last read that took it.
|
||||
const SHARED_LINGER: Duration = Duration::from_secs(10);
|
||||
/// A shared lease is handed out for this long after it was taken, then
|
||||
/// replaced, so a busy version's tag never ages into a reaper's TTL.
|
||||
const SHARED_FOR: Duration = Duration::from_secs(300);
|
||||
/// A process-unique owner id embedded in lease tags so a reaper only applies
|
||||
/// the short TTL to its own leases.
|
||||
static INSTANCE: LazyLock<String> = LazyLock::new(|| Uuid::new_v4().simple().to_string());
|
||||
@@ -138,6 +143,104 @@ fn untrack(tag: &str) {
|
||||
}
|
||||
}
|
||||
|
||||
/// Leases that reads of one version share: the first read tags the version,
|
||||
/// reads that follow reuse the tag, and the tag goes once no read has taken it
|
||||
/// for [`SHARED_LINGER`]. Spares a tag write and delete per read.
|
||||
pub struct SharedLeases {
|
||||
/// By location, version and manifest ETag: a table dropped and created
|
||||
/// again at its location reuses version numbers, not manifests.
|
||||
held: Mutex<Held>,
|
||||
linger: Duration,
|
||||
share_for: Duration,
|
||||
}
|
||||
|
||||
type LeaseKey = (String, u64, Option<String>);
|
||||
|
||||
#[derive(Default)]
|
||||
struct Held {
|
||||
slots: HashMap<LeaseKey, Slot>,
|
||||
/// Slots are pruned once they reach this many, so a lookup does not
|
||||
/// scan them all.
|
||||
prune_at: usize,
|
||||
}
|
||||
/// Reads of a version that find no lease wait here for the first to take it.
|
||||
type Slot = Arc<tokio::sync::Mutex<Weak<SharedLease>>>;
|
||||
|
||||
impl Default for SharedLeases {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
held: Mutex::default(),
|
||||
linger: SHARED_LINGER,
|
||||
share_for: SHARED_FOR,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A lease held for as long as any read, or the linger after the last one,
|
||||
/// keeps it.
|
||||
pub struct SharedLease {
|
||||
_lease: Lease,
|
||||
taken: Instant,
|
||||
}
|
||||
|
||||
impl SharedLeases {
|
||||
/// [`try_pin`], reusing a lease another read of this version holds.
|
||||
pub async fn try_pin(&self, dataset: &Dataset) -> Result<Option<Arc<SharedLease>>> {
|
||||
let key = (
|
||||
dataset.uri().to_string(),
|
||||
dataset.version().version,
|
||||
dataset.manifest_location().e_tag.clone(),
|
||||
);
|
||||
let slot = match self.held.lock() {
|
||||
Ok(mut held) => {
|
||||
if held.slots.len() >= held.prune_at {
|
||||
// A slot another read has taken but not yet locked stays.
|
||||
held.slots.retain(|_, slot| {
|
||||
Arc::strong_count(slot) > 1
|
||||
|| slot
|
||||
.try_lock()
|
||||
.map_or(true, |lease| lease.strong_count() > 0)
|
||||
});
|
||||
held.prune_at = (2 * held.slots.len()).max(64);
|
||||
// A burst's allocation would otherwise make every later
|
||||
// prune scan it.
|
||||
let capacity = held.prune_at;
|
||||
held.slots.shrink_to(capacity);
|
||||
}
|
||||
held.slots.entry(key).or_default().clone()
|
||||
}
|
||||
Err(_) => Slot::default(),
|
||||
};
|
||||
let mut slot = slot.lock().await;
|
||||
let lease = match slot
|
||||
.upgrade()
|
||||
.filter(|lease| lease.taken.elapsed() < self.share_for)
|
||||
{
|
||||
Some(lease) => lease,
|
||||
None => {
|
||||
let Some(lease) = try_pin(dataset).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let lease = Arc::new(SharedLease {
|
||||
_lease: lease,
|
||||
taken: Instant::now(),
|
||||
});
|
||||
*slot = Arc::downgrade(&lease);
|
||||
lease
|
||||
}
|
||||
};
|
||||
drop(slot);
|
||||
if let Ok(handle) = tokio::runtime::Handle::try_current() {
|
||||
let (lingering, linger) = (lease.clone(), self.linger);
|
||||
handle.spawn(async move {
|
||||
tokio::time::sleep(linger).await;
|
||||
drop(lingering);
|
||||
});
|
||||
}
|
||||
Ok(Some(lease))
|
||||
}
|
||||
}
|
||||
|
||||
/// Drop lease tags whose owner never released them, so an abandoned lease
|
||||
/// cannot keep a version out of cleanup forever.
|
||||
async fn reap_expired(dataset: &Dataset, now: u128) {
|
||||
@@ -324,6 +427,131 @@ mod tests {
|
||||
b.release().await.unwrap();
|
||||
}
|
||||
|
||||
/// Reads of one version share a tag, another version gets its own, and
|
||||
/// the tags go once the linger after the last read has passed.
|
||||
#[tokio::test]
|
||||
async fn reads_of_a_version_share_one_lease() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let mut dataset =
|
||||
testing::write_dataset(&dir, &testing::vectors(4, 4), &["a", "b", "c", "d"]).await;
|
||||
let leases = SharedLeases {
|
||||
linger: Duration::from_millis(50),
|
||||
..Default::default()
|
||||
};
|
||||
let a = leases.try_pin(&dataset).await.unwrap().unwrap();
|
||||
let b = leases.try_pin(&dataset).await.unwrap().unwrap();
|
||||
assert!(Arc::ptr_eq(&a, &b));
|
||||
assert_eq!(dataset.tags().list().await.unwrap().len(), 1);
|
||||
|
||||
testing::append_dataset(&mut dataset, &testing::vectors(1, 4), &["e"]).await;
|
||||
let c = leases.try_pin(&dataset).await.unwrap().unwrap();
|
||||
assert!(!Arc::ptr_eq(&a, &c));
|
||||
assert_eq!(dataset.tags().list().await.unwrap().len(), 2);
|
||||
|
||||
drop((a, b, c));
|
||||
tokio::time::sleep(leases.linger).await;
|
||||
let mut left = usize::MAX;
|
||||
for _ in 0..1000 {
|
||||
// A listing that meets a tag as it is deleted fails; try again.
|
||||
if let Ok(tags) = dataset.tags().list().await {
|
||||
left = tags.len();
|
||||
}
|
||||
if left == 0 {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
assert_eq!(left, 0);
|
||||
}
|
||||
|
||||
/// Reads that arrive together, before any lease exists, share the one the
|
||||
/// first takes.
|
||||
#[tokio::test]
|
||||
async fn concurrent_reads_take_one_lease() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let dataset =
|
||||
testing::write_dataset(&dir, &testing::vectors(4, 4), &["a", "b", "c", "d"]).await;
|
||||
let leases = SharedLeases::default();
|
||||
let pins: Vec<_> = futures::future::join_all((0..8).map(|_| leases.try_pin(&dataset)))
|
||||
.await
|
||||
.into_iter()
|
||||
.map(|lease| lease.unwrap().unwrap())
|
||||
.collect();
|
||||
assert!(pins.iter().all(|lease| Arc::ptr_eq(lease, &pins[0])));
|
||||
assert_eq!(dataset.tags().list().await.unwrap().len(), 1);
|
||||
}
|
||||
|
||||
/// A table dropped and written again at its location gets its own lease
|
||||
/// for a version number the old table had.
|
||||
#[tokio::test]
|
||||
async fn a_recreated_table_takes_its_own_lease() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let rows = testing::vectors(4, 4);
|
||||
let old = testing::write_dataset(&dir, &rows, &["a", "b", "c", "d"]).await;
|
||||
let leases = SharedLeases::default();
|
||||
let first = leases.try_pin(&old).await.unwrap().unwrap();
|
||||
std::fs::remove_dir_all(dir.path()).unwrap();
|
||||
let new = testing::write_dataset(&dir, &rows, &["e", "f", "g", "h"]).await;
|
||||
assert_eq!(new.version().version, old.version().version);
|
||||
let second = leases.try_pin(&new).await.unwrap().unwrap();
|
||||
assert!(!Arc::ptr_eq(&first, &second));
|
||||
assert_eq!(new.tags().list().await.unwrap().len(), 1);
|
||||
}
|
||||
|
||||
/// Slots of released leases are pruned as new versions come in.
|
||||
#[tokio::test]
|
||||
async fn released_lease_slots_are_pruned() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let mut dataset = testing::write_dataset(&dir, &testing::vectors(1, 4), &["a"]).await;
|
||||
let leases = SharedLeases {
|
||||
linger: Duration::ZERO,
|
||||
..Default::default()
|
||||
};
|
||||
for _ in 0..200 {
|
||||
drop(leases.try_pin(&dataset).await.unwrap());
|
||||
tokio::task::yield_now().await;
|
||||
testing::append_dataset(&mut dataset, &testing::vectors(1, 4), &["b"]).await;
|
||||
}
|
||||
assert!(leases.held.lock().unwrap().slots.len() <= 128);
|
||||
}
|
||||
|
||||
/// A slot a read has taken but not yet locked survives a prune, so that
|
||||
/// read and the next share the lease.
|
||||
#[tokio::test]
|
||||
async fn a_taken_slot_survives_a_prune() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let dataset = testing::write_dataset(&dir, &testing::vectors(1, 4), &["a"]).await;
|
||||
let leases = SharedLeases::default();
|
||||
let key: LeaseKey = ("elsewhere".into(), 1, None);
|
||||
let taken = {
|
||||
let mut held = leases.held.lock().unwrap();
|
||||
held.prune_at = 0;
|
||||
held.slots.entry(key.clone()).or_default().clone()
|
||||
};
|
||||
leases.try_pin(&dataset).await.unwrap().unwrap();
|
||||
assert!(Arc::ptr_eq(
|
||||
&leases.held.lock().unwrap().slots[&key],
|
||||
&taken
|
||||
));
|
||||
}
|
||||
|
||||
/// A shared lease is not handed out past SHARED_FOR, so a busy version
|
||||
/// moves to a fresh tag.
|
||||
#[tokio::test]
|
||||
async fn a_shared_lease_is_replaced_when_old() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let dataset =
|
||||
testing::write_dataset(&dir, &testing::vectors(4, 4), &["a", "b", "c", "d"]).await;
|
||||
let leases = SharedLeases {
|
||||
share_for: Duration::from_millis(50),
|
||||
..Default::default()
|
||||
};
|
||||
let a = leases.try_pin(&dataset).await.unwrap().unwrap();
|
||||
tokio::time::sleep(leases.share_for).await;
|
||||
let b = leases.try_pin(&dataset).await.unwrap().unwrap();
|
||||
assert!(!Arc::ptr_eq(&a, &b));
|
||||
}
|
||||
|
||||
/// A stale lease tag left behind is reaped by the next pin; a stale-named
|
||||
/// tag this process still holds is not.
|
||||
#[tokio::test]
|
||||
|
||||
@@ -36,6 +36,7 @@ bytes = "1"
|
||||
clap = { version = "4", features = ["derive", "env"] }
|
||||
futures = "0.3"
|
||||
moka = { version = "0.12", features = ["sync"] }
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
object_store = { version = "0.13", features = ["aws"] }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
|
||||
@@ -41,6 +41,13 @@ struct Args {
|
||||
/// Memory for caching reads of write-once table files, in MiB; 0 disables.
|
||||
#[arg(long, env = "WEED_VECTOR_READ_CACHE_MB", default_value_t = 512)]
|
||||
read_cache_mb: u64,
|
||||
/// Memory for indexes loaded by queries, kept for the next query, in MiB.
|
||||
#[arg(long, env = "WEED_VECTOR_INDEX_CACHE_MB", default_value_t = 1024)]
|
||||
index_cache_mb: usize,
|
||||
/// Memory for table metadata (manifests, file footers) kept across
|
||||
/// queries, in MiB.
|
||||
#[arg(long, env = "WEED_VECTOR_METADATA_CACHE_MB", default_value_t = 256)]
|
||||
metadata_cache_mb: usize,
|
||||
/// Partitions each query probes in a table's vector index. Default: a
|
||||
/// tenth of the index's partitions, at least 8; more under a filter.
|
||||
#[arg(long, env = "WEED_VECTOR_NPROBES")]
|
||||
@@ -163,6 +170,7 @@ async fn main() -> Result<()> {
|
||||
.then(|| ReadCache::new(args.read_cache_mb << 20).with_metrics(metrics.clone()));
|
||||
let mut service = VectorServer::new()
|
||||
.with_nprobes(args.nprobes.map(std::num::NonZeroUsize::get))
|
||||
.with_lance_caches(args.index_cache_mb << 20, args.metadata_cache_mb << 20)
|
||||
.with_storage(reads);
|
||||
if let Some(cache) = &cache {
|
||||
service = service.with_read_cache(cache.clone());
|
||||
|
||||
@@ -353,7 +353,7 @@ mod tests {
|
||||
use super::*;
|
||||
use crate::pb::vector as pb;
|
||||
use crate::pb::vector::vector_service_server::VectorService;
|
||||
use crate::service::VectorServer;
|
||||
use crate::service::{VectorServer, schema_for, table_id};
|
||||
use object_store::ObjectStoreExt;
|
||||
use object_store::memory::InMemory;
|
||||
use tonic::Request;
|
||||
@@ -616,6 +616,117 @@ mod tests {
|
||||
assert_eq!(query(&server, &table, vec![1.0, 0.0]).await, ["z"]);
|
||||
}
|
||||
|
||||
/// A query loads the index once; the next one finds it in the session
|
||||
/// every read shares.
|
||||
#[tokio::test]
|
||||
async fn queries_reuse_the_loaded_index() {
|
||||
let bucket = format!("session-{}", uuid_like());
|
||||
let location = format!("shared-memory://{bucket}/vector/x");
|
||||
let server = VectorServer::new().with_schemes(&["shared-memory"]);
|
||||
let table = pb::Table {
|
||||
id: vec![bucket, "vector".into(), "x".into()],
|
||||
location: location.clone(),
|
||||
..Default::default()
|
||||
};
|
||||
create(&server, &table, 2).await;
|
||||
put(
|
||||
&server,
|
||||
&table,
|
||||
(0..300)
|
||||
.map(|i| vector(&format!("k{i}"), vec![i as f32, 1.0]))
|
||||
.collect(),
|
||||
)
|
||||
.await;
|
||||
let mut dataset = lance::Dataset::open(&location).await.unwrap();
|
||||
weed_lance_worker::vector_index::build(&mut dataset, "vector", "idx", 2, false)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(query(&server, &table, vec![5.0, 1.0]).await[0], "k5");
|
||||
let session = server.reads.lock().unwrap().session.clone();
|
||||
let loaded = session.index_cache_stats().await;
|
||||
assert!(session.approx_num_items() > 0);
|
||||
assert_eq!(query(&server, &table, vec![7.0, 1.0]).await[0], "k7");
|
||||
let again = session.index_cache_stats().await;
|
||||
assert!(server.handles.contains_key(&(location, false)));
|
||||
assert!(again.hits > loaded.hits, "{loaded:?} then {again:?}");
|
||||
assert_eq!(again.misses, loaded.misses, "{loaded:?} then {again:?}");
|
||||
}
|
||||
|
||||
/// An index dropped and created again at the same location reaches the
|
||||
/// version numbers of the old one; the shared session must not answer
|
||||
/// for them with the old index.
|
||||
#[tokio::test]
|
||||
async fn a_recreated_table_does_not_see_the_old_index() {
|
||||
recreate_after(false).await;
|
||||
recreate_after(true).await;
|
||||
}
|
||||
|
||||
/// `unmarked`: the first table is written without a table id, as indexes
|
||||
/// created before ids were are.
|
||||
async fn recreate_after(unmarked: bool) {
|
||||
let bucket = format!("recreate-{}", uuid_like());
|
||||
let location = format!("shared-memory://{bucket}/vector/x");
|
||||
let server = VectorServer::new().with_schemes(&["shared-memory"]);
|
||||
let table = pb::Table {
|
||||
id: vec![bucket, "vector".into(), "x".into()],
|
||||
location: location.clone(),
|
||||
..Default::default()
|
||||
};
|
||||
let rows = |offset: usize| {
|
||||
(0..300)
|
||||
.map(|i| vector(&format!("k{}", i + offset), vec![(i + offset) as f32, 1.0]))
|
||||
.collect::<Vec<_>>()
|
||||
};
|
||||
if unmarked {
|
||||
let schema = schema_for(2, pb::DistanceMetric::L2, &[]);
|
||||
let empty = arrow_array::RecordBatch::new_empty(schema.clone());
|
||||
let reader = arrow_array::RecordBatchIterator::new(vec![Ok(empty)], schema);
|
||||
lance::Dataset::write(reader, &location, None)
|
||||
.await
|
||||
.unwrap();
|
||||
} else {
|
||||
create(&server, &table, 2).await;
|
||||
}
|
||||
put(&server, &table, rows(0)).await;
|
||||
let mut dataset = lance::Dataset::open(&location).await.unwrap();
|
||||
assert_eq!(table_id(&dataset).is_empty(), unmarked);
|
||||
weed_lance_worker::vector_index::build(&mut dataset, "vector", "idx", 2, false)
|
||||
.await
|
||||
.unwrap();
|
||||
let indexed = dataset.version().version;
|
||||
assert_eq!(query(&server, &table, vec![5.0, 1.0]).await[0], "k5");
|
||||
let other = pb::Table {
|
||||
id: vec![table.id[0].clone(), "vector".into(), "y".into()],
|
||||
location: location.replace("/vector/x", "/vector/y"),
|
||||
..Default::default()
|
||||
};
|
||||
create(&server, &other, 2).await;
|
||||
put(&server, &other, rows(0)).await;
|
||||
assert_eq!(query(&server, &other, vec![5.0, 1.0]).await[0], "k5");
|
||||
|
||||
let (store, path) = lance::io::ObjectStore::from_uri(&location).await.unwrap();
|
||||
store.remove_dir_all(path).await.unwrap();
|
||||
create(&server, &table, 2).await;
|
||||
let mut offset = 1000;
|
||||
while lance::Dataset::open(&location)
|
||||
.await
|
||||
.unwrap()
|
||||
.version()
|
||||
.version
|
||||
< indexed
|
||||
{
|
||||
put(&server, &table, rows(offset)).await;
|
||||
offset += 1000;
|
||||
}
|
||||
assert_eq!(query(&server, &table, vec![1005.0, 1.0]).await[0], "k1005");
|
||||
server.handles.run_pending_tasks();
|
||||
let session = server.reads.lock().unwrap().session.clone();
|
||||
for (_, (_, dataset)) in server.handles.iter() {
|
||||
assert!(Arc::ptr_eq(&dataset.session(), &session));
|
||||
}
|
||||
}
|
||||
|
||||
fn uuid_like() -> String {
|
||||
format!(
|
||||
"{:x}",
|
||||
|
||||
@@ -9,16 +9,18 @@ use lance::dataset::ProjectionRequest;
|
||||
use lance::dataset::WriteParams;
|
||||
use lance::dataset::{MergeInsertBuilder, WhenMatched, WhenNotMatched};
|
||||
use lance::io::WrappingObjectStore;
|
||||
use lance::session::Session;
|
||||
use lance::table::format::Fragment;
|
||||
use prost_types::{ListValue, Struct, Value, value::Kind};
|
||||
use std::collections::{BTreeMap, HashMap};
|
||||
use std::sync::Arc;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tonic::{Request, Response, Status};
|
||||
use weed_lance_worker::catalog::TableDescription;
|
||||
use weed_lance_worker::dataset;
|
||||
use weed_lance_worker::direct_read::{DirectReads, Wrappers};
|
||||
use weed_lance_worker::jobs::vector_index::is_vector_table;
|
||||
use weed_vector_core::index;
|
||||
use weed_vector_core::maintenance::SharedLeases;
|
||||
use weed_vector_core::query::{self, Query};
|
||||
use weed_vector_core::table::VectorTable;
|
||||
|
||||
@@ -55,7 +57,9 @@ const KEY_COLUMN: &str = "key";
|
||||
const METADATA_COLUMN: &str = "metadata";
|
||||
const VECTOR_COLUMN: &str = "vector";
|
||||
const ROW_ID_COLUMN: &str = "_rowid";
|
||||
fn schema_for(
|
||||
const DEFAULT_INDEX_CACHE_BYTES: usize = 1 << 30;
|
||||
const DEFAULT_METADATA_CACHE_BYTES: usize = 256 << 20;
|
||||
pub(crate) fn schema_for(
|
||||
dimension: usize,
|
||||
metric: pb::DistanceMetric,
|
||||
non_filterable: &[String],
|
||||
@@ -444,9 +448,85 @@ fn resolved(table: &pb::Table, schemes: &[&str]) -> Option<TableDescription> {
|
||||
})
|
||||
}
|
||||
|
||||
/// Marks the table an index was created as, in its schema metadata, which
|
||||
/// every later version keeps.
|
||||
const TABLE_ID_KEY: &str = "weed_vector_table_id";
|
||||
|
||||
/// The table a dataset was created as; empty for one created before tables
|
||||
/// were marked.
|
||||
pub(crate) fn table_id(dataset: &Dataset) -> &str {
|
||||
dataset
|
||||
.schema()
|
||||
.metadata
|
||||
.get(TABLE_ID_KEY)
|
||||
.map_or("", String::as_str)
|
||||
}
|
||||
|
||||
/// The session reads share, and the table each location held when this
|
||||
/// session first read it. Lance keys some of its caches by location and
|
||||
/// version alone, and an index dropped and created again at its location
|
||||
/// numbers its versions from 1 again, so a session that sees a location hold
|
||||
/// another table starts over.
|
||||
pub(crate) struct Reads {
|
||||
pub(crate) session: Arc<Session>,
|
||||
generation: u64,
|
||||
tables: HashMap<String, String>,
|
||||
}
|
||||
|
||||
/// Opened tables kept for the next read of them.
|
||||
const MAX_HANDLES: u64 = 1024;
|
||||
|
||||
fn handle_key(table: &pb::Table) -> (String, bool) {
|
||||
(table.location.clone(), table.sse_s3)
|
||||
}
|
||||
|
||||
/// Locations a session remembers before it starts over.
|
||||
const MAX_READ_TABLES: usize = 100_000;
|
||||
|
||||
impl Reads {
|
||||
fn new(index_bytes: usize, metadata_bytes: usize) -> Self {
|
||||
Self {
|
||||
session: Arc::new(Session::new(
|
||||
index_bytes,
|
||||
metadata_bytes,
|
||||
Default::default(),
|
||||
)),
|
||||
generation: 0,
|
||||
tables: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether a dataset of `table_id` at `location`, opened in this session,
|
||||
/// may be read. Otherwise the session was replaced and the caller opens
|
||||
/// the table again.
|
||||
fn admit(&mut self, location: &str, table_id: &str, (index, metadata): (usize, usize)) -> bool {
|
||||
match self.tables.get(location) {
|
||||
Some(seen) if seen == table_id => return true,
|
||||
None if self.tables.len() < MAX_READ_TABLES => {
|
||||
self.tables
|
||||
.insert(location.to_string(), table_id.to_string());
|
||||
return true;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
let generation = self.generation + 1;
|
||||
*self = Self::new(index, metadata);
|
||||
self.generation = generation;
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
pub struct VectorServer {
|
||||
storage: Option<DirectReads>,
|
||||
read_cache: Option<ReadCache>,
|
||||
/// Lance's index and metadata caches, shared by every request that only
|
||||
/// reads. Writes open in a session of their own.
|
||||
pub(crate) reads: Mutex<Reads>,
|
||||
/// A table opened by an earlier read, by location and encryption, with
|
||||
/// the generation of the session it was opened in.
|
||||
pub(crate) handles: moka::sync::Cache<(String, bool), (u64, Dataset)>,
|
||||
lance_cache_bytes: (usize, usize),
|
||||
leases: SharedLeases,
|
||||
nprobes: Option<usize>,
|
||||
/// URI schemes a table's location may have.
|
||||
schemes: &'static [&'static str],
|
||||
@@ -463,6 +543,13 @@ impl VectorServer {
|
||||
Self {
|
||||
storage: None,
|
||||
read_cache: None,
|
||||
reads: Mutex::new(Reads::new(
|
||||
DEFAULT_INDEX_CACHE_BYTES,
|
||||
DEFAULT_METADATA_CACHE_BYTES,
|
||||
)),
|
||||
lance_cache_bytes: (DEFAULT_INDEX_CACHE_BYTES, DEFAULT_METADATA_CACHE_BYTES),
|
||||
handles: moka::sync::Cache::new(MAX_HANDLES),
|
||||
leases: SharedLeases::default(),
|
||||
nprobes: None,
|
||||
schemes: &["s3"],
|
||||
}
|
||||
@@ -488,6 +575,13 @@ impl VectorServer {
|
||||
self
|
||||
}
|
||||
|
||||
/// Bytes of loaded indexes and of table metadata kept across reads.
|
||||
pub fn with_lance_caches(mut self, index_bytes: usize, metadata_bytes: usize) -> Self {
|
||||
self.reads = Mutex::new(Reads::new(index_bytes, metadata_bytes));
|
||||
self.lance_cache_bytes = (index_bytes, metadata_bytes);
|
||||
self
|
||||
}
|
||||
|
||||
/// Serve repeated reads of write-once table files from `cache`.
|
||||
pub fn with_read_cache(mut self, cache: ReadCache) -> Self {
|
||||
self.read_cache = Some(cache);
|
||||
@@ -529,9 +623,72 @@ impl VectorServer {
|
||||
}
|
||||
|
||||
async fn open(&self, table: &pb::Table) -> Result<Dataset, Status> {
|
||||
self.open_in(table, None).await
|
||||
}
|
||||
|
||||
/// [`Self::open`] in the shared session, for a request that only reads.
|
||||
async fn open_to_read(&self, table: &pb::Table) -> Result<Dataset, Status> {
|
||||
loop {
|
||||
let (session, generation) = {
|
||||
let reads = self.reads.lock().unwrap();
|
||||
(reads.session.clone(), reads.generation)
|
||||
};
|
||||
let dataset = match self.handle(table, generation).await? {
|
||||
Some(dataset) => dataset,
|
||||
None => {
|
||||
let dataset = self.open_in(table, Some(session)).await?;
|
||||
// Under the lock, so no handle of a replaced session
|
||||
// lands after the replacement cleared them.
|
||||
let reads = self.reads.lock().unwrap();
|
||||
if table.version == 0 && reads.generation == generation {
|
||||
self.handles
|
||||
.insert(handle_key(table), (generation, dataset.clone()));
|
||||
}
|
||||
dataset
|
||||
}
|
||||
};
|
||||
let mut reads = self.reads.lock().unwrap();
|
||||
if reads.generation != generation {
|
||||
continue;
|
||||
}
|
||||
if reads.admit(&table.location, table_id(&dataset), self.lance_cache_bytes) {
|
||||
return Ok(dataset);
|
||||
}
|
||||
// The handles hold the replaced session and its caches.
|
||||
self.handles.invalidate_all();
|
||||
}
|
||||
}
|
||||
|
||||
/// The table from the handle an earlier read of it left in this session,
|
||||
/// moved to the version asked for: the latest costs a listing of its
|
||||
/// versions, not a new store and manifest read.
|
||||
async fn handle(&self, table: &pb::Table, generation: u64) -> Result<Option<Dataset>, Status> {
|
||||
self.describe(table)?;
|
||||
let key = handle_key(table);
|
||||
let Some((at, mut dataset)) = self.handles.get(&key) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if at != generation {
|
||||
return Ok(None);
|
||||
}
|
||||
let moved = match table.version {
|
||||
0 => dataset.checkout_latest().await.map(|()| dataset),
|
||||
version => dataset.checkout_version(version).await,
|
||||
};
|
||||
moved.map(Some).map_err(|err| {
|
||||
self.handles.invalidate(&key);
|
||||
open_status(err.into())
|
||||
})
|
||||
}
|
||||
|
||||
async fn open_in(
|
||||
&self,
|
||||
table: &pb::Table,
|
||||
session: Option<Arc<Session>>,
|
||||
) -> Result<Dataset, Status> {
|
||||
let description = self.describe(table)?;
|
||||
let wrapper = self.wrapper(table);
|
||||
let mut dataset = dataset::open_at(&description, wrapper)
|
||||
let mut dataset = dataset::open_at(&description, wrapper, session)
|
||||
.await
|
||||
.map_err(open_status)?;
|
||||
if table.version != 0 {
|
||||
@@ -564,6 +721,13 @@ impl VectorService for VectorServer {
|
||||
// The gateway has already created the table through the catalog.
|
||||
let description = self.describe(&table)?;
|
||||
let schema = schema_for(dimension, metric, &req.non_filterable_metadata_keys);
|
||||
let schema = Arc::new(Schema::new_with_metadata(
|
||||
schema.fields().clone(),
|
||||
HashMap::from([(
|
||||
TABLE_ID_KEY.to_string(),
|
||||
uuid::Uuid::new_v4().simple().to_string(),
|
||||
)]),
|
||||
));
|
||||
let empty = RecordBatch::new_empty(schema.clone());
|
||||
let reader = arrow_array::RecordBatchIterator::new(vec![Ok(empty)].into_iter(), schema);
|
||||
let dataset = Dataset::write(
|
||||
@@ -680,12 +844,10 @@ impl VectorService for VectorServer {
|
||||
let table = req
|
||||
.table
|
||||
.ok_or_else(|| Status::invalid_argument("missing table"))?;
|
||||
let dataset = self.open(&table).await?;
|
||||
let dataset = self.open_to_read(&table).await?;
|
||||
// Lease the resolved version so maintenance cannot reclaim its files
|
||||
// while this query runs against it.
|
||||
let _lease = weed_vector_core::maintenance::try_pin(&dataset)
|
||||
.await
|
||||
.map_err(internal)?;
|
||||
let _lease = self.leases.try_pin(&dataset).await.map_err(internal)?;
|
||||
let resolved = VectorTable::resolve(&dataset, None, Some(KEY_COLUMN)).map_err(invalid)?;
|
||||
let columns = projection(req.return_data, req.return_metadata);
|
||||
let batches = query::get(&dataset, &resolved, &req.keys, Some(&columns))
|
||||
@@ -730,7 +892,7 @@ impl VectorService for VectorServer {
|
||||
"segment_index must be below segment_count, which is at most {MAX_SEGMENTS}"
|
||||
)));
|
||||
}
|
||||
let dataset = self.open(&table).await?;
|
||||
let dataset = self.open_to_read(&table).await?;
|
||||
let offset: u64 = if req.next_token.is_empty() {
|
||||
0
|
||||
} else {
|
||||
@@ -750,9 +912,7 @@ impl VectorService for VectorServer {
|
||||
// Data and metadata come from a second read of this version, so
|
||||
// lease it as get_vectors does.
|
||||
let _lease = if columns.len() > 1 {
|
||||
weed_vector_core::maintenance::try_pin(&dataset)
|
||||
.await
|
||||
.map_err(internal)?
|
||||
self.leases.try_pin(&dataset).await.map_err(internal)?
|
||||
} else {
|
||||
None
|
||||
};
|
||||
@@ -797,12 +957,10 @@ impl VectorService for VectorServer {
|
||||
let table = req
|
||||
.table
|
||||
.ok_or_else(|| Status::invalid_argument("missing table"))?;
|
||||
let dataset = self.open(&table).await?;
|
||||
let dataset = self.open_to_read(&table).await?;
|
||||
// Lease the resolved version so maintenance cannot reclaim its files
|
||||
// while this query runs against it.
|
||||
let _lease = weed_vector_core::maintenance::try_pin(&dataset)
|
||||
.await
|
||||
.map_err(internal)?;
|
||||
let _lease = self.leases.try_pin(&dataset).await.map_err(internal)?;
|
||||
let resolved = VectorTable::resolve(&dataset, None, Some(KEY_COLUMN)).map_err(invalid)?;
|
||||
if req.query_vector.len() != resolved.dimension {
|
||||
return Err(Status::invalid_argument(format!(
|
||||
|
||||
@@ -15,6 +15,7 @@ use lance::io::{
|
||||
ObjectStore, ObjectStoreParams, ObjectStoreRegistry, StorageOptionsAccessor,
|
||||
WrappingObjectStore,
|
||||
};
|
||||
use lance::session::Session;
|
||||
|
||||
use crate::catalog::{NamespaceClient, TableDescription};
|
||||
use crate::direct_read::DirectReads;
|
||||
@@ -74,7 +75,7 @@ pub fn unsigned_options() -> HashMap<String, String> {
|
||||
pub async fn open(client: &NamespaceClient, id: &[String], storage: &Storage) -> Result<OpenTable> {
|
||||
let description = client.describe_table(id).await?;
|
||||
let wrapper = storage.wrapper(&description.location).await?;
|
||||
let dataset = open_at(&description, wrapper).await?;
|
||||
let dataset = open_at(&description, wrapper, None).await?;
|
||||
Ok(OpenTable {
|
||||
id: id.to_vec(),
|
||||
location: description.location,
|
||||
@@ -82,13 +83,19 @@ pub async fn open(client: &NamespaceClient, id: &[String], storage: &Storage) ->
|
||||
})
|
||||
}
|
||||
|
||||
/// Open a table the caller has already resolved, with the store wrapped.
|
||||
/// Open a table the caller has already resolved, with the store wrapped, in
|
||||
/// `session` when given, otherwise in a session of its own.
|
||||
pub async fn open_at(
|
||||
description: &TableDescription,
|
||||
wrapper: Option<Arc<dyn WrappingObjectStore>>,
|
||||
session: Option<Arc<Session>>,
|
||||
) -> Result<Dataset> {
|
||||
DatasetBuilder::from_uri(&description.location)
|
||||
.with_store_params(store_params(wrapper))
|
||||
let mut builder =
|
||||
DatasetBuilder::from_uri(&description.location).with_store_params(store_params(wrapper));
|
||||
if let Some(session) = session {
|
||||
builder = builder.with_session(session);
|
||||
}
|
||||
builder
|
||||
.load()
|
||||
.await
|
||||
.with_context(|| format!("open lance dataset at {}", description.location))
|
||||
|
||||
Reference in new issue
Block a user