[ENH] Add a scout-logs function to find the max log position.

go/pkg/log/repository/log.go

@@ -169,6 +169,18 @@ func (r *LogRepository) GetLastCompactedOffsetForCollection(ctx context.Context,
 	return
 }
 
+func (r *LogRepository) GetBoundsForCollection(ctx context.Context, collectionId string) (start, limit int64, err error) {
+	bounds, err := r.queries.GetBoundsForCollection(ctx, collectionId)
+	if err != nil {
+		trace_log.Error("Error in getting minimum and maximum offset for collection", zap.Error(err), zap.String("collectionId", collectionId))
+		return
+	}
+	start = bounds.RecordCompactionOffsetPosition
+	limit = bounds.RecordEnumerationOffsetPosition + 1
+	err = nil
+	return
+}
+
 func (r *LogRepository) GarbageCollection(ctx context.Context) error {
 	collectionToCompact, err := r.queries.GetAllCollections(ctx)
 	if err != nil {

go/pkg/log/server/server.go

@@ -44,6 +44,23 @@ func (s *logServer) PushLogs(ctx context.Context, req *logservicepb.PushLogsRequ
 	return
 }
 
+func (s *logServer) ScoutLogs(ctx context.Context, req *logservicepb.ScoutLogsRequest) (res *logservicepb.ScoutLogsResponse, err error) {
+	var collectionID types.UniqueID
+	collectionID, err = types.ToUniqueID(&req.CollectionId)
+	if err != nil {
+		return
+	}
+	var limit int64
+	_, limit, err = s.lr.GetBoundsForCollection(ctx, collectionID.String())
+	if err != nil {
+		return
+	}
+	res = &logservicepb.ScoutLogsResponse {
+		LimitOffset: int64(limit),
+	}
+	return
+}
+
 func (s *logServer) PullLogs(ctx context.Context, req *logservicepb.PullLogsRequest) (res *logservicepb.PullLogsResponse, err error) {
 	var collectionID types.UniqueID
 	collectionID, err = types.ToUniqueID(&req.CollectionId)

go/pkg/log/store/queries/queries.sql

@@ -45,6 +45,11 @@ SELECT CAST(COALESCE(MIN(r.offset), 0) as bigint) AS min_offset, CAST(COALESCE(M
 FROM record_log r
 WHERE r.collection_id = $1;
 
+-- name: GetBoundsForCollection :one
+SELECT record_compaction_offset_position, record_enumeration_offset_position
+FROM collection
+WHERE id = $1;
+
 -- name: DeleteCollection :exec
 DELETE FROM collection c where c.id = ANY(@collection_ids::text[]);
 

idl/chromadb/proto/logservice.proto

@@ -14,6 +14,14 @@ message PushLogsResponse {
   int32 record_count = 1;
 }
 
+message ScoutLogsRequest {
+  string collection_id = 1;
+}
+
+message ScoutLogsResponse {
+  int64 limit_offset = 1;
+}
+
 message PullLogsRequest {
   string collection_id = 1;
   int64 start_from_offset = 2;
@@ -68,6 +76,7 @@ message PurgeDirtyForCollectionResponse {
 
 service LogService {
   rpc PushLogs(PushLogsRequest) returns (PushLogsResponse) {}
+  rpc ScoutLogs(ScoutLogsRequest) returns (ScoutLogsResponse) {}
   rpc PullLogs(PullLogsRequest) returns (PullLogsResponse) {}
   rpc GetAllCollectionInfoToCompact(GetAllCollectionInfoToCompactRequest) returns (GetAllCollectionInfoToCompactResponse) {}
   rpc UpdateCollectionLogOffset(UpdateCollectionLogOffsetRequest) returns (UpdateCollectionLogOffsetResponse) {}

rust/log-service/src/lib.rs

@@ -15,8 +15,8 @@ use chroma_types::chroma_proto::{
     log_service_server::LogService, CollectionInfo, GetAllCollectionInfoToCompactRequest,
     GetAllCollectionInfoToCompactResponse, LogRecord, OperationRecord, PullLogsRequest,
     PullLogsResponse, PurgeDirtyForCollectionRequest, PurgeDirtyForCollectionResponse,
-    PushLogsRequest, PushLogsResponse, UpdateCollectionLogOffsetRequest,
-    UpdateCollectionLogOffsetResponse,
+    PushLogsRequest, PushLogsResponse, ScoutLogsRequest, ScoutLogsResponse,
+    UpdateCollectionLogOffsetRequest, UpdateCollectionLogOffsetResponse,
 };
 use chroma_types::CollectionUuid;
 use figment::providers::{Env, Format, Yaml};
@@ -643,6 +643,30 @@ impl LogService for LogServer {
         Ok(Response::new(PushLogsResponse { record_count }))
     }
 
+    #[tracing::instrument(skip(self, request), err(Display))]
+    async fn scout_logs(
+        &self,
+        request: Request<ScoutLogsRequest>,
+    ) -> Result<Response<ScoutLogsResponse>, Status> {
+        let scout_logs = request.into_inner();
+        let collection_id = Uuid::parse_str(&scout_logs.collection_id)
+            .map(CollectionUuid)
+            .map_err(|_| Status::invalid_argument("Failed to parse collection id"))?;
+        tracing::info!("Scouting logs for collection {}", collection_id,);
+        let prefix = storage_prefix_for_log(collection_id);
+        let log_reader = LogReader::new(
+            self.config.reader.clone(),
+            Arc::clone(&self.storage),
+            prefix,
+        );
+        let limit_position = log_reader
+            .maximum_log_position()
+            .await
+            .map_err(|err| Status::new(err.code().into(), err.to_string()))?;
+        let limit_offset = limit_position.offset() as i64;
+        Ok(Response::new(ScoutLogsResponse { limit_offset }))
+    }
+
     #[tracing::instrument(skip(self, request), err(Display))]
     async fn pull_logs(
         &self,

rust/log/src/grpc_log.rs

@@ -18,6 +18,8 @@ use uuid::Uuid;
 pub enum GrpcPullLogsError {
     #[error("Failed to fetch")]
     FailedToPullLogs(#[from] tonic::Status),
+    #[error("Failed to scout logs: {0}")]
+    FailedToScoutLogs(tonic::Status),
     #[error("Failed to convert proto embedding record into EmbeddingRecord")]
     ConversionError(#[from] RecordConversionError),
 }
@@ -26,6 +28,7 @@ impl ChromaError for GrpcPullLogsError {
     fn code(&self) -> ErrorCodes {
         match self {
             GrpcPullLogsError::FailedToPullLogs(err) => err.code().into(),
+            GrpcPullLogsError::FailedToScoutLogs(err) => err.code().into(),
             GrpcPullLogsError::ConversionError(_) => ErrorCodes::Internal,
         }
     }
@@ -194,6 +197,29 @@ impl GrpcLog {
         &mut self.client
     }
 
+    pub(super) async fn scout_logs(
+        &mut self,
+        collection_id: CollectionUuid,
+        starting_offset: u64,
+    ) -> Result<u64, Box<dyn ChromaError>> {
+        let request = self
+            .client_for(collection_id)
+            .scout_logs(chroma_proto::ScoutLogsRequest {
+                collection_id: collection_id.0.to_string(),
+            });
+        let response = request.await;
+        let response = match response {
+            Ok(response) => response,
+            Err(err) => {
+                tracing::error!("Failed to scout logs: {}", err);
+                return Err(Box::new(GrpcPullLogsError::FailedToScoutLogs(err)));
+            }
+        };
+        let scout = response.into_inner();
+        tracing::info!("scout logs: {} -> {}", starting_offset, scout.limit_offset);
+        Ok(scout.limit_offset as u64)
+    }
+
     pub(super) async fn read(
         &mut self,
         collection_id: CollectionUuid,

rust/log/src/in_memory_log.rs

@@ -1,8 +1,11 @@
-use crate::types::CollectionInfo;
-use chroma_types::{CollectionUuid, LogRecord};
 use std::collections::HashMap;
 use std::fmt::Debug;
 
+use chroma_error::ChromaError;
+use chroma_types::{CollectionUuid, LogRecord};
+
+use crate::types::CollectionInfo;
+
 // This is used for testing only, it represents a log record that is stored in memory
 // internal to a mock log implementation
 #[derive(Clone)]
@@ -123,6 +126,25 @@ impl InMemoryLog {
     ) {
         self.offsets.insert(collection_id, new_offset);
     }
+
+    pub(super) async fn scout_logs(
+        &mut self,
+        collection_id: CollectionUuid,
+        starting_offset: u64,
+    ) -> Result<u64, Box<dyn ChromaError>> {
+        let answer = self
+            .collection_to_log
+            .get(&collection_id)
+            .iter()
+            .flat_map(|x| x.iter().map(|rec| rec.log_offset + 1).max())
+            .max()
+            .unwrap_or(starting_offset as i64) as u64;
+        if answer >= starting_offset {
+            Ok(answer)
+        } else {
+            Ok(starting_offset)
+        }
+    }
 }
 
 impl Default for InMemoryLog {

rust/log/src/log.rs

@@ -50,6 +50,28 @@ impl Log {
         }
     }
 
+    #[tracing::instrument(skip(self))]
+    pub async fn scout_logs(
+        &mut self,
+        collection_id: CollectionUuid,
+        starting_offset: u64,
+    ) -> Result<u64, Box<dyn ChromaError>> {
+        match self {
+            Log::Sqlite(log) => log
+                .scout_logs(collection_id, starting_offset as i64)
+                .await
+                .map_err(|e| Box::new(e) as Box<dyn ChromaError>),
+            Log::Grpc(log) => log
+                .scout_logs(collection_id, starting_offset)
+                .await
+                .map_err(|e| Box::new(e) as Box<dyn ChromaError>),
+            Log::InMemory(log) => log
+                .scout_logs(collection_id, starting_offset)
+                .await
+                .map_err(|e| Box::new(e) as Box<dyn ChromaError>),
+        }
+    }
+
     #[tracing::instrument(skip(self, records))]
     pub async fn push_logs(
         &mut self,

rust/log/src/sqlite_log.rs

@@ -27,6 +27,8 @@ pub enum SqlitePullLogsError {
     InvalidEmbedding(bytemuck::PodCastError),
     #[error("Failed to parse metadata: {0}")]
     InvalidMetadata(#[from] serde_json::Error),
+    #[error("Method {0} is not implemented")]
+    NotImplemented(String),
 }
 
 impl ChromaError for SqlitePullLogsError {
@@ -36,6 +38,7 @@ impl ChromaError for SqlitePullLogsError {
             SqlitePullLogsError::InvalidEncoding(_) => ErrorCodes::InvalidArgument,
             SqlitePullLogsError::InvalidEmbedding(_) => ErrorCodes::InvalidArgument,
             SqlitePullLogsError::InvalidMetadata(_) => ErrorCodes::InvalidArgument,
+            SqlitePullLogsError::NotImplemented(_) => ErrorCodes::Internal,
         }
     }
 }
@@ -162,6 +165,16 @@ impl SqliteLog {
             .map_err(|_| SqlitePushLogsError::CompactorHandleSetError)
     }
 
+    pub(super) async fn scout_logs(
+        &mut self,
+        _collection_id: CollectionUuid,
+        _starting_offset: i64,
+    ) -> Result<u64, SqlitePullLogsError> {
+        Err(SqlitePullLogsError::NotImplemented(
+            "scout_logs".to_string(),
+        ))
+    }
+
     pub(super) async fn read(
         &mut self,
         collection_id: CollectionUuid,