[ENH] Scout logs issues a HEAD for cached manifests.

Author: rescrv

rust/log-service/src/lib.rs

@@ -1128,28 +1128,64 @@ impl LogServer {
         let collection_id = Uuid::parse_str(&scout_logs.collection_id)
             .map(CollectionUuid)
             .map_err(|_| Status::invalid_argument("Failed to parse collection id"))?;
-
         let prefix = collection_id.storage_prefix_for_log();
         let log_reader = LogReader::new(
             self.config.reader.clone(),
             Arc::clone(&self.storage),
             prefix,
         );
-        let (start_position, limit_position) = match log_reader.manifest().await {
-            Ok(Some(manifest)) => (manifest.oldest_timestamp(), manifest.next_write_timestamp()),
-            Ok(None) => (LogPosition::from_offset(1), LogPosition::from_offset(1)),
-            Err(wal3::Error::UninitializedLog) => {
-                return Err(Status::not_found(format!(
-                    "collection {collection_id} not found"
-                )));
+        let cache_key = cache_key_for_manifest_and_etag(collection_id);
+        let mut cached_manifest_and_e_tag = None;
+        if let Some(cache) = self.cache.as_ref() {
+            if let Some(cache_bytes) = cache.get(&cache_key).await.ok().flatten() {
+                let met = serde_json::from_slice::<ManifestAndETag>(&cache_bytes.bytes).ok();
+                cached_manifest_and_e_tag = met;
             }
-            Err(err) => {
-                return Err(Status::new(
-                    err.code().into(),
-                    format!("could not scout logs: {err:?}"),
-                ));
+        }
+        // NOTE(rescrv):  We verify and if verification fails, we take the cached manifest to fall
+        // back to the uncached path.
+        if let Some(cached) = cached_manifest_and_e_tag.as_ref() {
+            if !log_reader.verify(cached).await.unwrap_or_default() {
+                cached_manifest_and_e_tag.take();
             }
-        };
+        }
+        let (start_position, limit_position) =
+            if let Some(manifest_and_e_tag) = cached_manifest_and_e_tag {
+                (
+                    manifest_and_e_tag.manifest.oldest_timestamp(),
+                    manifest_and_e_tag.manifest.next_write_timestamp(),
+                )
+            } else {
+                let (start_position, limit_position) = match log_reader.manifest_and_e_tag().await {
+                    Ok(Some(manifest_and_e_tag)) => {
+                        if let Some(cache) = self.cache.as_ref() {
+                            let json = serde_json::to_string(&manifest_and_e_tag)
+                                .map_err(|err| Status::unknown(err.to_string()))?;
+                            let cached_bytes = CachedBytes {
+                                bytes: Vec::from(json),
+                            };
+                            cache.insert(cache_key, cached_bytes).await;
+                        }
+                        (
+                            manifest_and_e_tag.manifest.oldest_timestamp(),
+                            manifest_and_e_tag.manifest.next_write_timestamp(),
+                        )
+                    }
+                    Ok(None) => (LogPosition::from_offset(1), LogPosition::from_offset(1)),
+                    Err(wal3::Error::UninitializedLog) => {
+                        return Err(Status::not_found(format!(
+                            "collection {collection_id} not found"
+                        )));
+                    }
+                    Err(err) => {
+                        return Err(Status::new(
+                            err.code().into(),
+                            format!("could not scout logs: {err:?}"),
+                        ));
+                    }
+                };
+                (start_position, limit_position)
+            };
         let start_offset = start_position.offset() as i64;
         let limit_offset = limit_position.offset() as i64;
         Ok(Response::new(ScoutLogsResponse {

rust/storage/src/admissioncontrolleds3.rs

@@ -458,6 +458,23 @@ impl AdmissionControlledS3Storage {
             .await
     }
 
+    pub async fn confirm_same(&self, key: &str, e_tag: &ETag) -> Result<bool, StorageError> {
+        self.metrics.nac_outstanding_read_requests.record(
+            self.metrics
+                .outstanding_read_requests
+                .load(Ordering::Relaxed) as u64,
+            &self.metrics.hostname_attribute,
+        );
+        self.metrics
+            .outstanding_read_requests
+            .fetch_add(1, Ordering::Relaxed);
+        let res = self.storage.confirm_same(key, e_tag).await;
+        self.metrics
+            .outstanding_read_requests
+            .fetch_sub(1, Ordering::Relaxed);
+        res
+    }
+
     async fn execute_fetch<FetchReturn, FetchFn, FetchFut>(
         fetch_fn: FetchFn,
         input: Result<(Arc<Vec<u8>>, Option<ETag>), StorageError>,

rust/storage/src/lib.rs

@@ -284,6 +284,19 @@ impl Storage {
         }
     }
 
+    // NOTE(rescrv):  Returns Ok(true) if the file is definitely the same.  Returns Ok(false) if
+    // the file cannot be confirmed to be the same but it exists.  Returns Err on error.  It is up
+    // to the user to know how they are confirming the same and to react to Ok(false) even if the
+    // file is definitely the same file on storage.
+    pub async fn confirm_same(&self, key: &str, e_tag: &ETag) -> Result<bool, StorageError> {
+        match self {
+            Storage::ObjectStore(object_store) => object_store.confirm_same(key, e_tag).await,
+            Storage::S3(s3) => s3.confirm_same(key, e_tag).await,
+            Storage::Local(local) => local.confirm_same(key, e_tag).await,
+            Storage::AdmissionControlledS3(as3) => as3.confirm_same(key, e_tag).await,
+        }
+    }
+
     pub async fn put_file(
         &self,
         key: &str,

rust/storage/src/local.rs

@@ -67,6 +67,10 @@ impl LocalStorage {
         Ok((bytes, Some(etag)))
     }
 
+    pub async fn confirm_same(&self, _: &str, _: &ETag) -> Result<bool, StorageError> {
+        Err(StorageError::NotImplemented)
+    }
+
     pub async fn put_bytes(
         &self,
         key: &str,

rust/storage/src/object_store.rs

@@ -153,6 +153,10 @@ impl ObjectStore {
         Err(StorageError::NotImplemented)
     }
 
+    pub async fn confirm_same(&self, _: &str, _: &ETag) -> Result<bool, StorageError> {
+        Err(StorageError::NotImplemented)
+    }
+
     pub async fn get_parallel(&self, key: &str) -> Result<Arc<Vec<u8>>, StorageError> {
         let meta = self.object_store.head(&Path::from(key)).await?;
         let file_size = meta.size;

rust/storage/src/s3.rs

@@ -211,6 +211,31 @@ impl S3Storage {
         }
     }
 
+    #[allow(clippy::type_complexity)]
+    pub async fn confirm_same(&self, key: &str, e_tag: &ETag) -> Result<bool, StorageError> {
+        let res = self
+            .client
+            .head_object()
+            .bucket(self.bucket.clone())
+            .key(key)
+            .send()
+            .await;
+        match res {
+            Ok(res) => Ok(res.e_tag() == Some(&e_tag.0)),
+            Err(e) => match e {
+                SdkError::ServiceError(err) => {
+                    let inner = err.into_err();
+                    Err(StorageError::Generic {
+                        source: Arc::new(inner),
+                    })
+                }
+                _ => Err(StorageError::Generic {
+                    source: Arc::new(e),
+                }),
+            },
+        }
+    }
+
     #[allow(clippy::type_complexity)]
     async fn get_stream_and_e_tag(
         &self,

rust/wal3/src/manifest.rs

@@ -723,6 +723,17 @@ impl Manifest {
         Ok(())
     }
 
+    /// Validate the e_tag against the manifest on object storage.
+    pub async fn head(
+        _: &ThrottleOptions,
+        storage: &Storage,
+        prefix: &str,
+        e_tag: &ETag,
+    ) -> Result<bool, Error> {
+        let path = manifest_path(prefix);
+        Ok(storage.confirm_same(&path, e_tag).await.map_err(Arc::new)?)
+    }
+
     /// Load the latest manifest from object storage.
     pub async fn load(
         options: &ThrottleOptions,

rust/wal3/src/reader.rs

@@ -13,8 +13,8 @@ use chroma_storage::{
 };
 
 use crate::{
-    parse_fragment_path, Error, Fragment, LogPosition, LogReaderOptions, Manifest, ScrubError,
-    ScrubSuccess, Snapshot, SnapshotCache,
+    parse_fragment_path, Error, Fragment, LogPosition, LogReaderOptions, Manifest, ManifestAndETag,
+    ScrubError, ScrubSuccess, Snapshot, SnapshotCache,
 };
 
 fn ranges_overlap(lhs: (LogPosition, LogPosition), rhs: (LogPosition, LogPosition)) -> bool {
@@ -74,6 +74,16 @@ impl LogReader {
         self.cache = Some(cache);
     }
 
+    pub async fn verify(&self, manifest_and_etag: &ManifestAndETag) -> Result<bool, Error> {
+        Manifest::head(
+            &self.options.throttle,
+            &self.storage,
+            &self.prefix,
+            &manifest_and_etag.e_tag,
+        )
+        .await
+    }
+
     pub async fn manifest(&self) -> Result<Option<Manifest>, Error> {
         Ok(
             Manifest::load(&self.options.throttle, &self.storage, &self.prefix)
@@ -82,6 +92,14 @@ impl LogReader {
         )
     }
 
+    pub async fn manifest_and_e_tag(&self) -> Result<Option<ManifestAndETag>, Error> {
+        match Manifest::load(&self.options.throttle, &self.storage, &self.prefix).await {
+            Ok(Some((manifest, e_tag))) => Ok(Some(ManifestAndETag { manifest, e_tag })),
+            Ok(None) => Ok(None),
+            Err(err) => Err(err),
+        }
+    }
+
     pub async fn oldest_timestamp(&self) -> Result<LogPosition, Error> {
         let Some((manifest, _)) =
             Manifest::load(&self.options.throttle, &self.storage, &self.prefix).await?