diff --git a/seaweed-vector/Cargo.lock b/seaweed-vector/Cargo.lock index e1a3225934..da09df9035 100644 --- a/seaweed-vector/Cargo.lock +++ b/seaweed-vector/Cargo.lock @@ -6486,11 +6486,15 @@ dependencies = [ "arrow-array", "arrow-cast", "arrow-schema", + "async-trait", + "bytes", "clap", "futures", "lance", "lance-arrow", "lance-linalg", + "moka", + "object_store", "prost", "prost-types", "protoc-bin-vendored", diff --git a/seaweed-vector/crates/server/Cargo.toml b/seaweed-vector/crates/server/Cargo.toml index 5f938f66a4..8d65d590d1 100644 --- a/seaweed-vector/crates/server/Cargo.toml +++ b/seaweed-vector/crates/server/Cargo.toml @@ -27,8 +27,12 @@ reqwest = { version = "0.12", default-features = false } tonic = { version = "0.14", features = ["tls-aws-lc"] } tonic-prost = "0.14" anyhow.workspace = true +async-trait = "0.1" +bytes = "1" clap = { version = "4", features = ["derive", "env"] } futures = "0.3" +moka = { version = "0.12", features = ["sync"] } +object_store = "0.13" serde = { version = "1", features = ["derive"] } serde_json = "1" tokio.workspace = true diff --git a/seaweed-vector/crates/server/src/lib.rs b/seaweed-vector/crates/server/src/lib.rs index 322ec27dcc..84db0b59cd 100644 --- a/seaweed-vector/crates/server/src/lib.rs +++ b/seaweed-vector/crates/server/src/lib.rs @@ -1,5 +1,6 @@ //! The VectorService gRPC server the S3 gateway talks to. +pub mod read_cache; pub mod service; pub mod remote_pb { diff --git a/seaweed-vector/crates/server/src/main.rs b/seaweed-vector/crates/server/src/main.rs index 453aa8ed5f..10d46f372d 100644 --- a/seaweed-vector/crates/server/src/main.rs +++ b/seaweed-vector/crates/server/src/main.rs @@ -4,6 +4,7 @@ use tonic::transport::Server; use weed_lance_worker::catalog::NamespaceClient; use weed_lance_worker::dataset::FallbackOptions; use weed_vector_server::pb::vector::vector_service_server::VectorServiceServer; +use weed_vector_server::read_cache::ReadCache; use weed_vector_server::service::VectorServer; /// The VectorService the S3 gateway fans vector calls out to. @@ -27,6 +28,9 @@ struct Args { /// gateway discovers it via ListClusterNodes(client_type=vector). #[arg(long, alias = "master", env = "WEED_VECTOR_MASTERS")] masters: Option, + /// 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, } fn parse_option(s: &str) -> Result<(String, String), String> { @@ -108,7 +112,10 @@ async fn main() -> Result<()> { } let fallback: FallbackOptions = args.storage_options.into_iter().collect(); let client = NamespaceClient::new(args.namespace.clone()).with_fallback(&fallback); - let service = VectorServer::new(client, fallback); + let mut service = VectorServer::new(client, fallback); + if args.read_cache_mb > 0 { + service = service.with_read_cache(ReadCache::new(args.read_cache_mb << 20)); + } tracing::info!(namespace = %args.namespace, listen = %args.listen, "weed-vector"); Server::builder() // Match the gateway's 1 GiB cap — put batches over the tonic 4 MiB diff --git a/seaweed-vector/crates/server/src/read_cache.rs b/seaweed-vector/crates/server/src/read_cache.rs new file mode 100644 index 0000000000..748bd4b66d --- /dev/null +++ b/seaweed-vector/crates/server/src/read_cache.rs @@ -0,0 +1,565 @@ +//! A bounded in-memory cache of reads from the files Lance never rewrites. +//! +//! Data, index and deletion files are written once under fresh random names, +//! so a cached read cannot go stale, even when an index is dropped and +//! recreated at the same location. Manifests, tags and every other path are +//! read from the store each time, which keeps read-your-writes intact. + +use std::fmt; +use std::ops::Range; +use std::sync::Arc; + +use bytes::Bytes; +use futures::StreamExt; +use futures::stream::BoxStream; +use lance::io::WrappingObjectStore; +use object_store::path::Path; +use object_store::{ + Attributes, CopyOptions, GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, + ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, + Result, +}; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct Key { + store: Arc, + path: Path, + /// `None` for a whole-object read. + range: Option>, +} + +#[derive(Debug, Clone)] +struct Entry { + bytes: Bytes, + /// Set for a whole-object read, which answers with its metadata. + object: Option>, +} + +/// Wraps each object store Lance opens; clones share one cache. +#[derive(Clone)] +pub struct ReadCache(moka::sync::Cache); + +impl ReadCache { + pub fn new(capacity_bytes: u64) -> Self { + Self( + moka::sync::Cache::builder() + .max_capacity(capacity_bytes) + .weigher(|key: &Key, entry: &Entry| { + u32::try_from(key.path.as_ref().len() + entry.bytes.len()).unwrap_or(u32::MAX) + }) + .build(), + ) + } +} + +impl fmt::Debug for ReadCache { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ReadCache") + .field("entries", &self.0.entry_count()) + .field("bytes", &self.0.weighted_size()) + .finish() + } +} + +impl WrappingObjectStore for ReadCache { + fn wrap(&self, store_prefix: &str, original: Arc) -> Arc { + Arc::new(CachedStore { + inner: original, + store: store_prefix.into(), + cache: self.0.clone(), + }) + } +} + +/// Lance names data files `data/.lance`, deletion files +/// `_deletions/--.` and index files +/// `_indices//`; nothing else under a table is write-once. +fn immutable(path: &Path) -> bool { + let parts: Vec<_> = path.parts().collect(); + let Some((name, dirs)) = parts.split_last() else { + return false; + }; + match dirs { + [.., dir] if dir.as_ref() == "data" => name.as_ref().ends_with(".lance"), + [.., dir] if dir.as_ref() == "_deletions" => true, + [.., dir, id] if dir.as_ref() == "_indices" => !id.as_ref().starts_with('_'), + _ => false, + } +} + +/// Lance reads files up to its block size whole; larger whole reads stream +/// through uncached. +const MAX_WHOLE_OBJECT: u64 = 4 << 20; + +/// A read with no conditions, so its answer depends only on the object. +fn plain(options: &GetOptions) -> bool { + options.if_match.is_none() + && options.if_none_match.is_none() + && options.if_modified_since.is_none() + && options.if_unmodified_since.is_none() + && options.range.is_none() + && options.version.is_none() + && !options.head +} + +/// A fetched range can be a slice of a larger response buffer; the cache +/// keeps its own copy so it holds only the bytes it counts. +fn owned(bytes: Bytes) -> Bytes { + Bytes::copy_from_slice(&bytes) +} + +fn whole(bytes: Bytes, object: &(ObjectMeta, Attributes)) -> GetResult { + GetResult { + range: 0..bytes.len() as u64, + payload: GetResultPayload::Stream(futures::stream::once(async { Ok(bytes) }).boxed()), + meta: object.0.clone(), + attributes: object.1.clone(), + } +} + +#[derive(Debug)] +struct CachedStore { + inner: Arc, + store: Arc, + cache: moka::sync::Cache, +} + +impl CachedStore { + fn key(&self, location: &Path, range: Option>) -> Key { + Key { + store: self.store.clone(), + path: location.clone(), + range, + } + } +} + +impl fmt::Display for CachedStore { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "ReadCache({})", self.inner) + } +} + +#[async_trait::async_trait] +#[deny(clippy::missing_trait_methods)] +impl ObjectStore for CachedStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> Result { + self.inner.put_opts(location, payload, opts).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> Result> { + self.inner.put_multipart_opts(location, opts).await + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> Result { + if !immutable(location) || !plain(&options) { + return self.inner.get_opts(location, options).await; + } + let key = self.key(location, None); + if let Some(Entry { + bytes, + object: Some(object), + }) = self.cache.get(&key) + { + return Ok(whole(bytes, &object)); + } + let result = self.inner.get_opts(location, options).await?; + if result.meta.size > MAX_WHOLE_OBJECT { + return Ok(result); + } + let object = Arc::new((result.meta.clone(), result.attributes.clone())); + let bytes = owned(result.bytes().await?); + self.cache.insert( + key, + Entry { + bytes: bytes.clone(), + object: Some(object.clone()), + }, + ); + Ok(whole(bytes, &object)) + } + + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> Result> { + if !immutable(location) { + return self.inner.get_ranges(location, ranges).await; + } + let mut out: Vec> = ranges + .iter() + .map(|range| { + self.cache + .get(&self.key(location, Some(range.clone()))) + .map(|entry| entry.bytes) + }) + .collect(); + let missing: Vec> = ranges + .iter() + .zip(&out) + .filter(|(_, hit)| hit.is_none()) + .map(|(range, _)| range.clone()) + .collect(); + if !missing.is_empty() { + let mut fetched = self.inner.get_ranges(location, &missing).await?.into_iter(); + for (range, slot) in ranges.iter().zip(out.iter_mut()) { + if slot.is_some() { + continue; + } + let Some(bytes) = fetched.next() else { + break; + }; + let bytes = owned(bytes); + self.cache.insert( + self.key(location, Some(range.clone())), + Entry { + bytes: bytes.clone(), + object: None, + }, + ); + *slot = Some(bytes); + } + } + out.into_iter() + .map(|bytes| { + bytes.ok_or_else(|| object_store::Error::Generic { + store: "ReadCache", + source: "get_ranges returned fewer ranges than requested".into(), + }) + }) + .collect() + } + + fn delete_stream( + &self, + locations: BoxStream<'static, Result>, + ) -> BoxStream<'static, Result> { + self.inner.delete_stream(locations) + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result> { + self.inner.list(prefix) + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, Result> { + self.inner.list_with_offset(prefix, offset) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result { + self.inner.list_with_delimiter(prefix).await + } + + async fn copy_opts(&self, from: &Path, to: &Path, options: CopyOptions) -> Result<()> { + self.inner.copy_opts(from, to, options).await + } + + async fn rename_opts(&self, from: &Path, to: &Path, options: RenameOptions) -> Result<()> { + self.inner.rename_opts(from, to, options).await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::pb::vector as pb; + use crate::pb::vector::vector_service_server::VectorService; + use crate::service::VectorServer; + use object_store::ObjectStoreExt; + use object_store::memory::InMemory; + use tonic::Request; + use weed_lance_worker::catalog::NamespaceClient; + use weed_lance_worker::dataset::FallbackOptions; + + fn cached(inner: Arc) -> Arc { + ReadCache::new(1 << 20).wrap("memory", inner) + } + + #[test] + fn only_write_once_files_are_immutable() { + for path in [ + "vector/x/data/0101abcd.lance", + "vector/data/data/0101abcd.lance", + "vector/x/_deletions/3-7-12345.arrow", + "vector/x/_indices/1b6f/index.idx", + ] { + assert!(immutable(&Path::from(path)), "{path}"); + } + for path in [ + "vector/data/_versions/1.manifest", + "vector/data/_latest.manifest", + "vector/x/_versions/18446744073709551614.manifest", + "vector/x/_refs/tags/weed-vector-query-1.json", + "vector/x/_transactions/0-abc.txn", + "vector/x/data/notes.txt", + "vector/_indices/_versions/1.manifest", + "data/x.lance/_versions/1.manifest", + ] { + assert!(!immutable(&Path::from(path)), "{path}"); + } + } + + async fn read(store: &Arc, path: &Path, range: Range) -> Bytes { + let mut got = store + .get_ranges(path, std::slice::from_ref(&range)) + .await + .unwrap(); + got.remove(0) + } + + #[tokio::test] + async fn write_once_reads_come_from_the_cache() { + let inner = Arc::new(InMemory::new()); + let store = cached(inner.clone()); + let data = Path::from("vector/x/data/a.lance"); + let manifest = Path::from("vector/x/_versions/1.manifest"); + for path in [&data, &manifest] { + inner + .put(path, Bytes::from_static(b"first").into()) + .await + .unwrap(); + assert_eq!(read(&store, path, 0..3).await, &b"fir"[..]); + assert_eq!( + store.get(path).await.unwrap().bytes().await.unwrap(), + &b"first"[..] + ); + inner + .put(path, Bytes::from_static(b"second").into()) + .await + .unwrap(); + } + // The overwrite stands in for a store answering differently: the data + // file keeps its cached bytes, the manifest is read again. + assert_eq!( + store.get_ranges(&data, &[0..3, 1..4]).await.unwrap(), + [&b"fir"[..], &b"eco"[..]] + ); + let got = store.get(&data).await.unwrap(); + assert_eq!(got.meta.size, 5); + assert_eq!(got.bytes().await.unwrap(), &b"first"[..]); + assert_eq!(read(&store, &manifest, 0..3).await, &b"sec"[..]); + assert_eq!( + store.get(&manifest).await.unwrap().bytes().await.unwrap(), + &b"second"[..] + ); + // A conditional or ranged get is the store's to answer. + let ranged = GetOptions { + range: Some((0..2).into()), + ..Default::default() + }; + let got = store.get_opts(&data, ranged).await.unwrap(); + assert_eq!(got.bytes().await.unwrap(), &b"se"[..]); + } + + #[tokio::test] + async fn the_cache_stays_within_its_capacity() { + let cache = ReadCache::new(64 << 10); + let inner = Arc::new(InMemory::new()); + let store = cache.wrap("memory", inner.clone()); + for i in 0..64 { + let path = Path::from(format!("vector/x/data/{i}.lance")); + inner + .put(&path, Bytes::from(vec![0u8; 8 << 10]).into()) + .await + .unwrap(); + read(&store, &path, 0..(8 << 10)).await; + } + cache.0.run_pending_tasks(); + assert!(cache.0.weighted_size() <= 64 << 10, "{cache:?}"); + } + + #[tokio::test] + async fn a_cached_range_does_not_pin_its_response_buffer() { + let cache = ReadCache::new(1 << 20); + let inner = Arc::new(InMemory::new()); + let store = cache.wrap("memory", inner.clone()); + let path = Path::from("vector/x/data/big.lance"); + inner + .put(&path, Bytes::from(vec![7u8; 256 << 10]).into()) + .await + .unwrap(); + // InMemory answers a range with a slice of the whole object. + let whole = inner.get(&path).await.unwrap().bytes().await.unwrap(); + let backing = whole.as_ptr_range(); + read(&store, &path, 10..20).await; + let key = Key { + store: "memory".into(), + path, + range: Some(10..20), + }; + let cached = cache.0.get(&key).unwrap().bytes; + assert_eq!(cached.len(), 10); + assert!(!backing.contains(&cached.as_ptr())); + } + + /// A catalog that answers every request with `location`. + async fn catalog_at(location: String) -> String { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let body = serde_json::json!({ "location": location }).to_string(); + tokio::spawn(async move { + loop { + let (mut conn, _) = listener.accept().await.unwrap(); + let body = body.clone(); + tokio::spawn(async move { + let mut request = [0u8; 8192]; + let _ = conn.read(&mut request).await; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + let _ = conn.write_all(response.as_bytes()).await; + let _ = conn.shutdown().await; + }); + } + }); + url + } + + fn vector(key: &str, data: Vec) -> pb::Vector { + pb::Vector { + key: key.to_string(), + data, + metadata: None, + } + } + + async fn get(server: &VectorServer, table: &pb::Table, keys: &[&str]) -> Vec { + let mut vectors = server + .get_vectors(Request::new(pb::GetVectorsRequest { + table: Some(table.clone()), + keys: keys.iter().map(|k| k.to_string()).collect(), + return_data: true, + return_metadata: false, + })) + .await + .unwrap() + .into_inner() + .vectors; + vectors.sort_by(|a, b| a.key.cmp(&b.key)); + vectors + } + + async fn query(server: &VectorServer, table: &pb::Table, v: Vec) -> Vec { + server + .query_vectors(Request::new(pb::QueryVectorsRequest { + table: Some(table.clone()), + query_vector: v, + top_k: 10, + ..Default::default() + })) + .await + .unwrap() + .into_inner() + .vectors + .into_iter() + .map(|r| r.key) + .collect() + } + + async fn create(server: &VectorServer, table: &pb::Table, dimension: u32) { + server + .create_index(Request::new(pb::CreateIndexRequest { + table: Some(table.clone()), + dimension, + ..Default::default() + })) + .await + .unwrap(); + } + + async fn put(server: &VectorServer, table: &pb::Table, vectors: Vec) { + server + .put_vectors(Request::new(pb::PutVectorsRequest { + table: Some(table.clone()), + vectors, + })) + .await + .unwrap(); + } + + /// Reads through the cache see every committed write, and an index + /// dropped and recreated at the same location answers with its new rows. + #[tokio::test] + async fn cached_reads_follow_writes_and_a_recreated_index() { + let location = format!("shared-memory://cache-{}/vector/x", uuid_like(),); + let server = VectorServer::new( + NamespaceClient::new(catalog_at(location.clone()).await), + FallbackOptions::new(), + ) + .with_read_cache(ReadCache::new(16 << 20)); + let table = pb::Table { + id: vec!["bkt".into(), "vector".into(), "x".into()], + version: 0, + }; + + create(&server, &table, 4).await; + put( + &server, + &table, + vec![ + vector("a", vec![1.0, 0.0, 0.0, 0.0]), + vector("b", vec![0.0, 1.0, 0.0, 0.0]), + ], + ) + .await; + assert_eq!( + query(&server, &table, vec![1.0, 0.0, 0.0, 0.0]).await, + ["a", "b"] + ); + assert_eq!( + get(&server, &table, &["a"]).await[0].data, + [1.0, 0.0, 0.0, 0.0] + ); + + put( + &server, + &table, + vec![ + vector("a", vec![0.0, 0.0, 2.0, 0.0]), + vector("c", vec![1.0, 0.1, 0.0, 0.0]), + ], + ) + .await; + assert_eq!( + query(&server, &table, vec![1.0, 0.0, 0.0, 0.0]).await, + ["c", "b", "a"] + ); + assert_eq!( + get(&server, &table, &["a"]).await[0].data, + [0.0, 0.0, 2.0, 0.0] + ); + + let (store, path) = lance::io::ObjectStore::from_uri(&location).await.unwrap(); + store.remove_dir_all(path).await.unwrap(); + create(&server, &table, 2).await; + put(&server, &table, vec![vector("z", vec![0.5, 0.5])]).await; + let vectors = get(&server, &table, &["a", "b", "c", "z"]).await; + assert_eq!(vectors.len(), 1); + assert_eq!( + (vectors[0].key.as_str(), vectors[0].data.as_slice()), + ("z", &[0.5, 0.5][..]) + ); + assert_eq!(query(&server, &table, vec![1.0, 0.0]).await, ["z"]); + } + + fn uuid_like() -> String { + format!( + "{:x}", + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + ) + } +} diff --git a/seaweed-vector/crates/server/src/service.rs b/seaweed-vector/crates/server/src/service.rs index fb8ab22264..cb9f51b1c4 100644 --- a/seaweed-vector/crates/server/src/service.rs +++ b/seaweed-vector/crates/server/src/service.rs @@ -10,6 +10,7 @@ use lance::dataset::WriteParams; use lance::dataset::{MergeInsertBuilder, WhenMatched, WhenNotMatched}; use lance::io::ObjectStoreParams; use lance::io::StorageOptionsAccessor; +use lance::io::WrappingObjectStore; use lance::table::format::Fragment; use lance_linalg::distance::DistanceType; use prost_types::{ListValue, Struct, Value, value::Kind}; @@ -22,6 +23,7 @@ use weed_vector_core::query::{self, Query}; use weed_vector_core::table::VectorTable; use crate::pb::vector::{self as pb, vector_service_server::VectorService}; +use crate::read_cache::ReadCache; fn internal(err: impl Into) -> Status { Status::internal(format!("{:#}", err.into())) @@ -371,15 +373,32 @@ fn rows_as_vectors( pub struct VectorServer { client: NamespaceClient, fallback: FallbackOptions, + read_cache: Option>, } impl VectorServer { pub fn new(client: NamespaceClient, fallback: FallbackOptions) -> Self { - Self { client, fallback } + Self { + client, + fallback, + read_cache: None, + } + } + + /// Serve repeated reads of write-once table files from `cache`. + pub fn with_read_cache(mut self, cache: ReadCache) -> Self { + self.read_cache = Some(Arc::new(cache)); + self } async fn open(&self, table: &pb::Table) -> Result { - let opened = dataset::open(&self.client, &table.id, &self.fallback).await?; + let opened = dataset::open_with( + &self.client, + &table.id, + &self.fallback, + self.read_cache.clone(), + ) + .await?; let mut dataset = opened.dataset; if table.version != 0 { dataset = dataset.checkout_version(table.version).await?; diff --git a/seaweed-worker/crates/lance/src/dataset.rs b/seaweed-worker/crates/lance/src/dataset.rs index 19eb1b82e4..d28bcf2d4c 100644 --- a/seaweed-worker/crates/lance/src/dataset.rs +++ b/seaweed-worker/crates/lance/src/dataset.rs @@ -11,7 +11,10 @@ use std::sync::Arc; use anyhow::{Context, Result}; use lance::dataset::Dataset; use lance::dataset::builder::DatasetBuilder; -use lance::io::{ObjectStore, ObjectStoreParams, ObjectStoreRegistry, StorageOptionsAccessor}; +use lance::io::{ + ObjectStore, ObjectStoreParams, ObjectStoreRegistry, StorageOptionsAccessor, + WrappingObjectStore, +}; use crate::catalog::{NamespaceClient, TableDescription}; @@ -42,9 +45,19 @@ pub async fn open( client: &NamespaceClient, id: &[String], fallback: &FallbackOptions, +) -> Result { + open_with(client, id, fallback, None).await +} + +/// [`open`] with the dataset's object store wrapped, e.g. by a read cache. +pub async fn open_with( + client: &NamespaceClient, + id: &[String], + fallback: &FallbackOptions, + wrapper: Option>, ) -> Result { let description = client.describe_table(id).await?; - let dataset = open_at(&description, fallback).await?; + let dataset = open_at(&description, fallback, wrapper).await?; Ok(OpenTable { id: id.to_vec(), location: description.location, @@ -56,11 +69,21 @@ pub async fn open( /// through. What it vends always wins: the fallback exists because a deployment /// without STS vends no credentials at all, and then the worker has no other way /// to reach the data. -async fn open_at(description: &TableDescription, fallback: &FallbackOptions) -> Result { +async fn open_at( + description: &TableDescription, + fallback: &FallbackOptions, + wrapper: Option>, +) -> Result { let mut options: HashMap = fallback.clone(); options.extend(description.storage_options.clone()); DatasetBuilder::from_uri(&description.location) - .with_storage_options(options) + .with_store_params(ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + options, + ))), + object_store_wrapper: wrapper, + ..Default::default() + }) .load() .await .with_context(|| format!("open lance dataset at {}", description.location))