[ENH] A tool to purge the cache. (#5085)

## Description of changes

In anticipation of turning on the disk cache for wal3, we have a tool
that allows us to purge a single entry from the cache.

## Test plan

It seems to work from tilt.

## Documentation Changes

N/A

Author: rescrv

go/pkg/log/server/server.go

@@ -187,6 +187,10 @@ func (s *logServer) GarbageCollectPhase2(ctx context.Context, req *logservicepb.
 	return
 }
 
+func (s *logServer) PurgeFromCache(ctx context.Context, req *logservicepb.PurgeFromCacheRequest) (res *logservicepb.PurgeFromCacheResponse, err error) {
+	return
+}
+
 func NewLogServer(lr *repository.LogRepository) logservicepb.LogServiceServer {
 	return &logServer{
 		lr: lr,

idl/chromadb/proto/logservice.proto

@@ -154,6 +154,22 @@ message GarbageCollectPhase2Request {
 message GarbageCollectPhase2Response {
 }
 
+message PurgeFromCacheRequest {
+  oneof entry_to_evict {
+    string cursor_for_collection_id = 1;
+    string manifest_for_collection_id = 2;
+    FragmentToEvict fragment = 3;
+  };
+}
+
+message FragmentToEvict {
+  string collection_id = 1;
+  string fragment_path = 2;
+};
+
+message PurgeFromCacheResponse {
+}
+
 service LogService {
   rpc PushLogs(PushLogsRequest) returns (PushLogsResponse) {}
   rpc ScoutLogs(ScoutLogsRequest) returns (ScoutLogsResponse) {}
@@ -162,17 +178,19 @@ service LogService {
   rpc GetAllCollectionInfoToCompact(GetAllCollectionInfoToCompactRequest) returns (GetAllCollectionInfoToCompactResponse) {}
   rpc UpdateCollectionLogOffset(UpdateCollectionLogOffsetRequest) returns (UpdateCollectionLogOffsetResponse) {}
   rpc PurgeDirtyForCollection(PurgeDirtyForCollectionRequest) returns (PurgeDirtyForCollectionResponse) {}
-  // RPC endpoints to expose for operator debuggability.
   // This endpoint must route to the rust log service.
   rpc InspectDirtyLog(InspectDirtyLogRequest) returns (InspectDirtyLogResponse) {}
   // This endpoint must route to the go log service.
   rpc SealLog(SealLogRequest) returns (SealLogResponse) {}
   // This endpoint must route to the rust log service.
   rpc MigrateLog(MigrateLogRequest) returns (MigrateLogResponse) {}
+  // RPC endpoints to expose for operator debuggability.
   // This endpoint can be supported by any log service.
   rpc InspectLogState(InspectLogStateRequest) returns (InspectLogStateResponse) {}
   // This endpoint should route to the rust log service.
   rpc ScrubLog(ScrubLogRequest) returns (ScrubLogResponse) {}
   // This endpoint should route to the rust log service.
   rpc GarbageCollectPhase2(GarbageCollectPhase2Request) returns (GarbageCollectPhase2Response) {}
+  // This endpoint will purge from cache the specified items.
+  rpc PurgeFromCache(PurgeFromCacheRequest) returns (PurgeFromCacheResponse) {}
 }

rust/log-service/src/bin/chroma-log-service-purge-cache-entry.rs

@@ -0,0 +1,63 @@
+use tonic::transport::Channel;
+
+use chroma_types::chroma_proto::log_service_client::LogServiceClient;
+use chroma_types::chroma_proto::{
+    purge_from_cache_request::EntryToEvict, FragmentToEvict, PurgeFromCacheRequest,
+};
+
+#[tokio::main]
+async fn main() {
+    let args = std::env::args().skip(1).collect::<Vec<_>>();
+    if args.len() != 3 && args.len() != 4 {
+        eprintln!(
+            "USAGE: chroma-log-service-purge-cache-entry HOST TYPE COLLECTION_UUID [FRAGMENT_PATH]"
+        );
+        std::process::exit(13);
+    }
+    let logservice = Channel::from_shared(args[0].clone())
+        .expect("could not create channel")
+        .connect()
+        .await
+        .expect("could not connect to log service");
+    let req = match args[1].as_str() {
+        "cursor" => {
+            if args.len() != 3 {
+                eprintln!("purge cache entry takes no fragment path");
+                std::process::exit(13);
+            }
+            PurgeFromCacheRequest {
+                entry_to_evict: Some(EntryToEvict::CursorForCollectionId(args[2].clone())),
+            }
+        }
+        "manifest" => {
+            if args.len() != 3 {
+                eprintln!("purge cache entry takes no fragment path");
+                std::process::exit(13);
+            }
+            PurgeFromCacheRequest {
+                entry_to_evict: Some(EntryToEvict::ManifestForCollectionId(args[2].clone())),
+            }
+        }
+        "fragment" => {
+            if args.len() != 4 {
+                eprintln!("purge cache entry takes a fragment path");
+                std::process::exit(13);
+            }
+            PurgeFromCacheRequest {
+                entry_to_evict: Some(EntryToEvict::Fragment(FragmentToEvict {
+                    collection_id: args[2].clone(),
+                    fragment_path: args[3].clone(),
+                })),
+            }
+        }
+        _ => {
+            eprintln!("unknown type: {}", args[1]);
+            std::process::exit(13);
+        }
+    };
+    let mut client = LogServiceClient::new(logservice);
+    let _state = client
+        .purge_from_cache(req)
+        .await
+        .expect("could not purge from cache");
+}

rust/log-service/src/lib.rs

@@ -18,14 +18,15 @@ use chroma_storage::Storage;
 use chroma_tracing::util::wrap_span_with_parent_context;
 use chroma_types::chroma_proto::{
     garbage_collect_phase2_request::LogToCollect, log_service_client::LogServiceClient,
-    log_service_server::LogService, scrub_log_request::LogToScrub, CollectionInfo,
-    GarbageCollectPhase2Request, GarbageCollectPhase2Response,
-    GetAllCollectionInfoToCompactRequest, GetAllCollectionInfoToCompactResponse,
-    InspectDirtyLogRequest, InspectDirtyLogResponse, InspectLogStateRequest,
-    InspectLogStateResponse, LogRecord, MigrateLogRequest, MigrateLogResponse, OperationRecord,
-    PullLogsRequest, PullLogsResponse, PurgeDirtyForCollectionRequest,
-    PurgeDirtyForCollectionResponse, PushLogsRequest, PushLogsResponse, ScoutLogsRequest,
-    ScoutLogsResponse, ScrubLogRequest, ScrubLogResponse, SealLogRequest, SealLogResponse,
+    log_service_server::LogService, purge_from_cache_request::EntryToEvict,
+    scrub_log_request::LogToScrub, CollectionInfo, GarbageCollectPhase2Request,
+    GarbageCollectPhase2Response, GetAllCollectionInfoToCompactRequest,
+    GetAllCollectionInfoToCompactResponse, InspectDirtyLogRequest, InspectDirtyLogResponse,
+    InspectLogStateRequest, InspectLogStateResponse, LogRecord, MigrateLogRequest,
+    MigrateLogResponse, OperationRecord, PullLogsRequest, PullLogsResponse,
+    PurgeDirtyForCollectionRequest, PurgeDirtyForCollectionResponse, PurgeFromCacheRequest,
+    PurgeFromCacheResponse, PushLogsRequest, PushLogsResponse, ScoutLogsRequest, ScoutLogsResponse,
+    ScrubLogRequest, ScrubLogResponse, SealLogRequest, SealLogResponse,
     UpdateCollectionLogOffsetRequest, UpdateCollectionLogOffsetResponse,
 };
 use chroma_types::chroma_proto::{ForkLogsRequest, ForkLogsResponse};
@@ -293,6 +294,10 @@ fn cache_key_for_cursor(collection_id: CollectionUuid, name: &CursorName) -> Str
     format!("{collection_id}::cursor::{}", name.path())
 }
 
+fn cache_key_for_fragment(collection_id: CollectionUuid, fragment_path: &str) -> String {
+    format!("{collection_id}::{}", fragment_path)
+}
+
 ////////////////////////////////////////// CachedFragment //////////////////////////////////////////
 
 #[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
@@ -1473,7 +1478,7 @@ impl LogServer {
                 .map(|fragment| async {
                     let prefix = collection_id.storage_prefix_for_log();
                     if let Some(cache) = self.cache.as_ref() {
-                        let cache_key = format!("{collection_id}::{}", fragment.path);
+                        let cache_key = cache_key_for_fragment(collection_id, &fragment.path);
                         if let Ok(Some(answer)) = cache.get(&cache_key).await {
                             return Ok(Arc::new(answer.bytes));
                         }
@@ -2007,6 +2012,48 @@ impl LogServer {
         .instrument(span)
         .await
     }
+
+    async fn purge_from_cache(
+        &self,
+        request: Request<PurgeFromCacheRequest>,
+    ) -> Result<Response<PurgeFromCacheResponse>, Status> {
+        let span = wrap_span_with_parent_context(
+            tracing::trace_span!("PurgeFromCache",),
+            request.metadata(),
+        );
+        let purge = request.into_inner();
+        async move {
+            let key = match purge.entry_to_evict {
+                Some(EntryToEvict::CursorForCollectionId(x)) => {
+                    let collection_id = Uuid::parse_str(&x)
+                        .map(CollectionUuid)
+                        .map_err(|_| Status::invalid_argument("Failed to parse collection id"))?;
+                    Some(cache_key_for_cursor(collection_id, &COMPACTION))
+                }
+                Some(EntryToEvict::ManifestForCollectionId(x)) => {
+                    let collection_id = Uuid::parse_str(&x)
+                        .map(CollectionUuid)
+                        .map_err(|_| Status::invalid_argument("Failed to parse collection id"))?;
+                    Some(cache_key_for_manifest(collection_id))
+                }
+                Some(EntryToEvict::Fragment(f)) => {
+                    let collection_id = Uuid::parse_str(&f.collection_id)
+                        .map(CollectionUuid)
+                        .map_err(|_| Status::invalid_argument("Failed to parse collection id"))?;
+                    Some(cache_key_for_fragment(collection_id, &f.fragment_path))
+                }
+                None => None,
+            };
+            if let Some(key) = key {
+                if let Some(cache) = self.cache.as_ref() {
+                    cache.remove(&key).await;
+                }
+            }
+            Ok(Response::new(PurgeFromCacheResponse {}))
+        }
+        .instrument(span)
+        .await
+    }
 }
 
 struct LogServerWrapper {
@@ -2107,6 +2154,13 @@ impl LogService for LogServerWrapper {
     ) -> Result<Response<GarbageCollectPhase2Response>, Status> {
         self.log_server.garbage_collect_phase2(request).await
     }
+
+    async fn purge_from_cache(
+        &self,
+        request: Request<PurgeFromCacheRequest>,
+    ) -> Result<Response<PurgeFromCacheResponse>, Status> {
+        self.log_server.purge_from_cache(request).await
+    }
 }
 
 fn parquet_to_records(parquet: Arc<Vec<u8>>) -> Result<Vec<(LogPosition, Vec<u8>)>, Status> {