[ENH] A route and tool to inspect the dirty log. (#4461)

## Description of changes

To get ground truth on the dirty log, add a tool that can print it as it
gets interpreted in the 'get all to compact' call.

## Test plan

- [ ] Tests pass locally with `pytest` for python, `yarn test` for js,
`cargo test` for rust

## Documentation Changes

N/A

Author: rescrv

go/pkg/log/server/server.go

@@ -152,6 +152,12 @@ func (s *logServer) PurgeDirtyForCollection(ctx context.Context, req *logservice
 	return
 }
 
+func (s *logServer) InspectDirtyLog(ctx context.Context, req *logservicepb.InspectDirtyLogRequest) (res *logservicepb.InspectDirtyLogResponse, err error) {
+	// no-op for now
+	return
+}
+
+
 func NewLogServer(lr *repository.LogRepository) logservicepb.LogServiceServer {
 	return &logServer{
 		lr: lr,

idl/chromadb/proto/logservice.proto

@@ -87,6 +87,14 @@ message PurgeDirtyForCollectionResponse {
   // Empty
 }
 
+message InspectDirtyLogRequest {
+  // Empty
+}
+
+message InspectDirtyLogResponse {
+  repeated string markers = 1;
+}
+
 service LogService {
   rpc PushLogs(PushLogsRequest) returns (PushLogsResponse) {}
   rpc ScoutLogs(ScoutLogsRequest) returns (ScoutLogsResponse) {}
@@ -95,4 +103,5 @@ service LogService {
   rpc GetAllCollectionInfoToCompact(GetAllCollectionInfoToCompactRequest) returns (GetAllCollectionInfoToCompactResponse) {}
   rpc UpdateCollectionLogOffset(UpdateCollectionLogOffsetRequest) returns (UpdateCollectionLogOffsetResponse) {}
   rpc PurgeDirtyForCollection(PurgeDirtyForCollectionRequest) returns (PurgeDirtyForCollectionResponse) {}
+  rpc InspectDirtyLog(InspectDirtyLogRequest) returns (InspectDirtyLogResponse) {}
 }

rust/log-service/src/bin/chroma-inspect-dirty-log.rs

@@ -0,0 +1,27 @@
+use tonic::transport::Channel;
+
+use chroma_types::chroma_proto::log_service_client::LogServiceClient;
+use chroma_types::chroma_proto::InspectDirtyLogRequest;
+
+#[tokio::main]
+async fn main() {
+    let args = std::env::args().skip(1).collect::<Vec<_>>();
+    if args.len() != 1 {
+        eprintln!("USAGE: chroma-inspect-dirty-log [HOST]");
+        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 mut client = LogServiceClient::new(logservice);
+    let dirty = client
+        .inspect_dirty_log(InspectDirtyLogRequest {})
+        .await
+        .expect("could not inspect dirty log");
+    let dirty = dirty.into_inner();
+    for line in dirty.markers {
+        println!("{line}");
+    }
+}

rust/log-service/src/lib.rs

@@ -13,10 +13,10 @@ use chroma_storage::config::StorageConfig;
 use chroma_storage::Storage;
 use chroma_types::chroma_proto::{
     log_service_server::LogService, CollectionInfo, GetAllCollectionInfoToCompactRequest,
-    GetAllCollectionInfoToCompactResponse, LogRecord, OperationRecord, PullLogsRequest,
-    PullLogsResponse, PurgeDirtyForCollectionRequest, PurgeDirtyForCollectionResponse,
-    PushLogsRequest, PushLogsResponse, ScoutLogsRequest, ScoutLogsResponse,
-    UpdateCollectionLogOffsetRequest, UpdateCollectionLogOffsetResponse,
+    GetAllCollectionInfoToCompactResponse, InspectDirtyLogRequest, InspectDirtyLogResponse,
+    LogRecord, OperationRecord, PullLogsRequest, PullLogsResponse, PurgeDirtyForCollectionRequest,
+    PurgeDirtyForCollectionResponse, PushLogsRequest, PushLogsResponse, ScoutLogsRequest,
+    ScoutLogsResponse, UpdateCollectionLogOffsetRequest, UpdateCollectionLogOffsetResponse,
 };
 use chroma_types::chroma_proto::{ForkLogsRequest, ForkLogsResponse};
 use chroma_types::CollectionUuid;
@@ -1080,6 +1080,65 @@ impl LogService for LogServer {
             .map_err(|err| Status::new(err.code().into(), err.to_string()))?;
         Ok(Response::new(PurgeDirtyForCollectionResponse {}))
     }
+
+    #[tracing::instrument(skip(self, _request))]
+    async fn inspect_dirty_log(
+        &self,
+        _request: Request<InspectDirtyLogRequest>,
+    ) -> Result<Response<InspectDirtyLogResponse>, Status> {
+        let Some(reader) = self.dirty_log.reader(LogReaderOptions::default()) else {
+            return Err(Status::unavailable("Failed to get dirty log reader"));
+        };
+        let Some(cursors) = self.dirty_log.cursors(CursorStoreOptions::default()) else {
+            return Err(Status::unavailable("Failed to get dirty log cursors"));
+        };
+        let witness = match cursors.load(&STABLE_PREFIX).await {
+            Ok(witness) => witness,
+            Err(err) => {
+                return Err(Status::new(err.code().into(), err.to_string()));
+            }
+        };
+        let default = Cursor::default();
+        let cursor = witness.as_ref().map(|w| w.cursor()).unwrap_or(&default);
+        tracing::info!("cursoring from {cursor:?}");
+        let dirty_fragments = reader
+            .scan(
+                cursor.position,
+                Limits {
+                    max_files: Some(1_000_000),
+                    max_bytes: Some(1_000_000_000),
+                },
+            )
+            .await
+            .map_err(|err| Status::new(err.code().into(), err.to_string()))?;
+        let dirty_futures = dirty_fragments
+            .iter()
+            .map(|fragment| reader.read_parquet(fragment))
+            .collect::<Vec<_>>();
+        let dirty_raw = futures::future::try_join_all(dirty_futures)
+            .await
+            .map_err(|err| {
+                Status::new(
+                    err.code().into(),
+                    format!("Failed to fetch dirty parquet: {}", err),
+                )
+            })?;
+        let mut markers = vec![];
+        for (_, records, _) in dirty_raw {
+            let records = records
+                .into_iter()
+                .map(|x| String::from_utf8(x.1))
+                .collect::<Result<Vec<_>, _>>()
+                .map_err(|err| {
+                    Status::new(
+                        chroma_error::ErrorCodes::DataLoss.into(),
+                        format!("Failed to extract records: {}", err),
+                    )
+                })?;
+            markers.extend(records);
+        }
+        Ok(Response::new(InspectDirtyLogResponse { markers }))
+    }
 }
 
 fn parquet_to_records(parquet: Arc<Vec<u8>>) -> Result<Vec<(LogPosition, Vec<u8>)>, Status> {