[ENH] Add feature to the rust log service to proxy the go log service.

Following the rollout ADR, this will encode table to go
classic->transitioning->transitioned.

|            | **inactive**  | **active**   |
|------------|---------------|--------------|
| **open**   | classic       | BUG          |
| **sealed** | transitioning | transitioned |

Author: rescrv

rust/log-service/src/lib.rs

@@ -10,14 +10,16 @@ use bytes::Bytes;
 use chroma_cache::CacheConfig;
 use chroma_config::Configurable;
 use chroma_error::ChromaError;
+use chroma_log::{config::GrpcLogConfig, grpc_log::GrpcLog};
 use chroma_storage::config::StorageConfig;
 use chroma_storage::Storage;
 use chroma_types::chroma_proto::{
-    log_service_server::LogService, CollectionInfo, GetAllCollectionInfoToCompactRequest,
-    GetAllCollectionInfoToCompactResponse, InspectDirtyLogRequest, InspectDirtyLogResponse,
-    LogRecord, OperationRecord, PullLogsRequest, PullLogsResponse, PurgeDirtyForCollectionRequest,
-    PurgeDirtyForCollectionResponse, PushLogsRequest, PushLogsResponse, ScoutLogsRequest,
-    ScoutLogsResponse, UpdateCollectionLogOffsetRequest, UpdateCollectionLogOffsetResponse,
+    log_service_client::LogServiceClient, log_service_server::LogService, CollectionInfo,
+    GetAllCollectionInfoToCompactRequest, 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;
@@ -201,7 +203,7 @@ async fn get_log_from_handle<'a>(
         });
     }
     tracing::info!("Opening log at {}", prefix);
-    let opened = LogWriter::open_or_initialize(
+    let opened = LogWriter::open(
         options.clone(),
         Arc::clone(storage),
         prefix,
@@ -630,11 +632,70 @@ pub struct LogServer {
     storage: Arc<Storage>,
     open_logs: Arc<StateHashTable<LogKey, LogStub>>,
     dirty_log: Arc<LogWriter>,
+    #[allow(clippy::type_complexity)]
+    proxy: Option<LogServiceClient<chroma_tracing::GrpcTraceService<tonic::transport::Channel>>>,
     compacting: tokio::sync::Mutex<()>,
     cache: Option<Box<dyn chroma_cache::PersistentCache<String, CachedParquetFragment>>>,
     metrics: Metrics,
 }
 
+impl LogServer {
+    fn should_initialize_log(&self, collection: CollectionUuid) -> bool {
+        todo!();
+    }
+
+    async fn forward_push_logs(
+        &self,
+        request: Request<PushLogsRequest>,
+    ) -> Result<Response<PushLogsResponse>, Status> {
+        todo!();
+    }
+
+    async fn forward_scout_logs(
+        &self,
+        request: Request<ScoutLogsRequest>,
+    ) -> Result<Response<ScoutLogsResponse>, Status> {
+        if let Some(proxy) = self.proxy.as_ref() {
+            proxy.clone().scout_logs(request).await
+        } else {
+            Err(Status::failed_precondition("proxy not initialized"))
+        }
+    }
+
+    async fn forward_pull_logs(
+        &self,
+        request: Request<PullLogsRequest>,
+    ) -> Result<Response<PullLogsResponse>, Status> {
+        if let Some(proxy) = self.proxy.as_ref() {
+            proxy.clone().pull_logs(request).await
+        } else {
+            Err(Status::failed_precondition("proxy not initialized"))
+        }
+    }
+
+    async fn forward_update_collection_log_offset(
+        &self,
+        request: Request<UpdateCollectionLogOffsetRequest>,
+    ) -> Result<Response<UpdateCollectionLogOffsetResponse>, Status> {
+        if let Some(proxy) = self.proxy.as_ref() {
+            proxy.clone().update_collection_log_offset(request).await
+        } else {
+            Err(Status::failed_precondition("proxy not initialized"))
+        }
+    }
+
+    async fn forward_fork_logs(
+        &self,
+        request: Request<ForkLogsRequest>,
+    ) -> Result<Response<ForkLogsResponse>, Status> {
+        if let Some(proxy) = self.proxy.as_ref() {
+            proxy.clone().fork_logs(request).await
+        } else {
+            Err(Status::failed_precondition("proxy not initialized"))
+        }
+    }
+}
+
 #[async_trait::async_trait]
 impl LogService for LogServer {
     async fn push_logs(
@@ -662,16 +723,36 @@ impl LogService for LogServer {
                 collection_id,
                 dirty_log: Arc::clone(&self.dirty_log),
             };
-            let log = get_log_from_handle(
+            let log = match get_log_from_handle(
                 &handle,
                 &self.config.writer,
                 &self.storage,
                 &prefix,
                 mark_dirty,
             )
             .await
-            // TODO(rescrv): better error handling.
-            .map_err(|err| Status::unknown(err.to_string()))?;
+            {
+                Ok(log) => log,
+                Err(wal3::Error::UninitializedLog) => {
+                    if self.should_initialize_log(collection_id) {
+                        if let Err(err) = LogWriter::initialize(
+                            &self.config.writer,
+                            &self.storage,
+                            &prefix,
+                            "push_logs initializer",
+                        )
+                        .await
+                        {
+                            return Err(Status::unknown(err.to_string()));
+                        }
+                        return Box::pin(self.push_logs(Request::new(push_logs))).await;
+                    }
+                    return self.forward_push_logs(Request::new(push_logs)).await;
+                }
+                Err(err) => {
+                    return Err(Status::unknown(err.to_string()));
+                }
+            };
             let mut messages = Vec::with_capacity(push_logs.records.len());
             for record in push_logs.records {
                 let mut buf = vec![];
@@ -718,6 +799,17 @@ impl LogService for LogServer {
             );
             let limit_position = match log_reader.maximum_log_position().await {
                 Ok(limit_position) => limit_position,
+                Err(wal3::Error::UninitializedLog) => {
+                    // NOTE(rescrv):  In this case, what we have is a guarantee that we know the
+                    // log is ours.  This only comes from the by-tenant or by-collection alt-log
+                    // shunts.  The log is uninitialized, so we know it's not because
+                    // has_manifest().
+                    if self.should_initialize_log(collection_id) {
+                        LogPosition::from_offset(1)
+                    } else {
+                        return self.forward_scout_logs(Request::new(scout_logs)).await;
+                    }
+                }
                 Err(err) => {
                     if err.code() == chroma_error::ErrorCodes::FailedPrecondition {
                         LogPosition::from_offset(1)
@@ -768,14 +860,18 @@ impl LogService for LogServer {
                 .await
             {
                 Ok(fragments) => fragments,
-                Err(err) => {
-                    if let wal3::Error::UninitializedLog = err {
+                Err(wal3::Error::UninitializedLog) => {
+                    // NOTE(rescrv):  Same as with ScoutLogs.
+                    if self.should_initialize_log(collection_id) {
                         tracing::info!("Uninitialized log for collection {}", collection_id);
                         return Ok(Response::new(PullLogsResponse { records: vec![] }));
                     } else {
-                        return Err(Status::new(err.code().into(), err.to_string()));
+                        return self.forward_pull_logs(Request::new(pull_logs)).await;
                     }
                 }
+                Err(err) => {
+                    return Err(Status::new(err.code().into(), err.to_string()));
+                }
             };
             let futures = fragments
                 .iter()
@@ -854,6 +950,37 @@ impl LogService for LogServer {
                 Arc::clone(&storage),
                 source_prefix.clone(),
             );
+            if let Err(err) = log_reader.maximum_log_position().await {
+                match err {
+                    wal3::Error::UninitializedLog => {
+                        // NOTE(rescrv):  Same as with ScoutLogs.
+                        if self.should_initialize_log(source_collection_id) {
+                            LogWriter::initialize(
+                                &self.config.writer,
+                                &storage,
+                                &source_prefix,
+                                "fork logs initializer",
+                            )
+                            .await
+                            .map_err(|err| {
+                                Status::new(
+                                    err.code().into(),
+                                    format!("Failed to initialize log for fork: {err:?}"),
+                                )
+                            })?;
+                            return Box::pin(self.fork_logs(Request::new(request))).await;
+                        } else {
+                            return self.forward_fork_logs(Request::new(request)).await;
+                        }
+                    }
+                    _ => {
+                        return Err(Status::new(
+                            err.code().into(),
+                            format!("Failed to load log: {}", err),
+                        ));
+                    }
+                }
+            }
             let cursors = CursorStore::new(
                 CursorStoreOptions::default(),
                 Arc::clone(&storage),
@@ -862,7 +989,6 @@ impl LogService for LogServer {
             );
             let cursor_name = &COMPACTION;
             let witness = cursors.load(cursor_name).await.map_err(|err| {
-                tracing::info!("FINDME");
                 Status::new(err.code().into(), format!("Failed to load cursor: {}", err))
             })?;
             // This is the existing compaction_offset, which is the last record that was compacted.
@@ -879,7 +1005,6 @@ impl LogService for LogServer {
             )
             .await
             .map_err(|err| {
-                tracing::info!("FINDME");
                 Status::new(err.code().into(), format!("Failed to copy log: {}", err))
             })?;
             let log_reader = LogReader::new(
@@ -889,11 +1014,9 @@ impl LogService for LogServer {
             );
             // This is the next record to insert, so we'll have to adjust downwards.
             let max_offset = log_reader.maximum_log_position().await.map_err(|err| {
-                tracing::info!("FINDME");
                 Status::new(err.code().into(), format!("Failed to copy log: {}", err))
             })?;
             if max_offset < offset {
-                tracing::info!("FINDME");
                 return Err(Status::new(
                     chroma_error::ErrorCodes::Internal.into(),
                     format!("max_offset={:?} < offset={:?}", max_offset, offset),
@@ -1074,6 +1197,26 @@ impl LogService for LogServer {
             request.log_offset
         );
         let storage_prefix = storage_prefix_for_log(collection_id);
+        let log_reader = LogReader::new(
+            self.config.reader.clone(),
+            Arc::clone(&self.storage),
+            storage_prefix.clone(),
+        );
+
+        let res = log_reader.maximum_log_position().await;
+        if let Err(wal3::Error::UninitializedLog) = res {
+            if self.should_initialize_log(collection_id) {
+                return Err(Status::failed_precondition(
+                    "uninitialized log has its cursor updated",
+                ));
+            } else {
+                return self
+                    .forward_update_collection_log_offset(Request::new(request))
+                    .await;
+            }
+        }
+        res.map_err(|err| Status::unknown(err.to_string()))?;
+
         let cursor_name = &COMPACTION;
         let cursor_store = CursorStore::new(
             CursorStoreOptions::default(),
@@ -1416,6 +1559,8 @@ pub struct LogServerConfig {
     pub reinsert_threshold: u64,
     #[serde(default = "LogServerConfig::default_timeout_us")]
     pub timeout_us: u64,
+    #[serde(default)]
+    pub proxy_to: Option<GrpcLogConfig>,
 }
 
 impl LogServerConfig {
@@ -1447,6 +1592,7 @@ impl Default for LogServerConfig {
             record_count_threshold: Self::default_record_count_threshold(),
             reinsert_threshold: Self::default_reinsert_threshold(),
             timeout_us: Self::default_timeout_us(),
+            proxy_to: None,
         }
     }
 }
@@ -1484,13 +1630,24 @@ impl Configurable<LogServerConfig> for LogServer {
         .await
         .map_err(|err| -> Box<dyn ChromaError> { Box::new(err) as _ })?;
         let dirty_log = Arc::new(dirty_log);
+        let proxy = if let Some(proxy_to) = config.proxy_to.as_ref() {
+            match GrpcLog::primary_client_from_config(proxy_to).await {
+                Ok(log) => Some(log),
+                Err(err) => {
+                    return Err(err);
+                }
+            }
+        } else {
+            None
+        };
         let compacting = tokio::sync::Mutex::new(());
         let metrics = Metrics::new(opentelemetry::global::meter("chroma"));
         Ok(Self {
             config: config.clone(),
             open_logs: Arc::new(StateHashTable::default()),
             storage,
             dirty_log,
+            proxy,
             compacting,
             cache,
             metrics,

rust/log/src/grpc_log.rs

@@ -172,6 +172,7 @@ impl Configurable<GrpcLogConfig> for GrpcLog {
         my_config: &GrpcLogConfig,
         _registry: &Registry,
     ) -> Result<Self, Box<dyn ChromaError>> {
+        // NOTE(rescrv):  This code is duplicated with primary_client_from_config below.  A transient hack.
         let host = &my_config.host;
         let port = &my_config.port;
         let max_encoding_message_size = my_config.max_encoding_message_size;
@@ -209,6 +210,40 @@ impl Configurable<GrpcLogConfig> for GrpcLog {
 }
 
 impl GrpcLog {
+    // NOTE(rescrv) This is a transient hack, so the code duplication is not worth eliminating.
+    pub async fn primary_client_from_config(
+        my_config: &GrpcLogConfig,
+    ) -> Result<
+        LogServiceClient<chroma_tracing::GrpcTraceService<tonic::transport::Channel>>,
+        Box<dyn ChromaError>,
+    > {
+        let host = &my_config.host;
+        let port = &my_config.port;
+        let max_encoding_message_size = my_config.max_encoding_message_size;
+        let max_decoding_message_size = my_config.max_decoding_message_size;
+        let connection_string = format!("http://{}:{}", host, port);
+        let client_for_conn_str =
+            |connection_string: String| -> Result<LogServiceClient<_>, Box<dyn ChromaError>> {
+                tracing::info!("Connecting to log service at {}", connection_string);
+                let endpoint_res = match Endpoint::from_shared(connection_string) {
+                    Ok(endpoint) => endpoint,
+                    Err(e) => return Err(Box::new(GrpcLogError::FailedToConnect(e))),
+                };
+                let endpoint_res = endpoint_res
+                    .connect_timeout(Duration::from_millis(my_config.connect_timeout_ms))
+                    .timeout(Duration::from_millis(my_config.request_timeout_ms));
+                let channel = endpoint_res.connect_lazy();
+                let channel = ServiceBuilder::new()
+                    .layer(chroma_tracing::GrpcTraceLayer)
+                    .service(channel);
+                let client = LogServiceClient::new(channel)
+                    .max_encoding_message_size(max_encoding_message_size)
+                    .max_decoding_message_size(max_decoding_message_size);
+                Ok(client)
+            };
+        client_for_conn_str(connection_string)
+    }
+
     fn client_for(
         &mut self,
         tenant: &str,