[ENH] handle panics in query service

idl/chromadb/proto/debug.proto

@@ -0,0 +1,14 @@
+syntax = "proto3";
+
+package chroma;
+
+import "google/protobuf/empty.proto";
+
+message GetInfoResponse {
+  string version = 1;
+}
+
+service Debug {
+  rpc GetInfo(google.protobuf.Empty) returns (GetInfoResponse) {}
+  rpc TriggerPanic(google.protobuf.Empty) returns (google.protobuf.Empty) {}
+}

rust/worker/Cargo.toml

@@ -48,9 +48,14 @@ tracing = "0.1"
 tracing-bunyan-formatter = "0.3.3"
 tracing-opentelemetry = "0.19.0"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
-opentelemetry = { version = "0.19.0", default-features = false, features = ["trace", "rt-tokio"] }
+opentelemetry = { version = "0.19.0", default-features = false, features = [
+  "trace",
+  "rt-tokio",
+] }
 opentelemetry-otlp = "0.12.0"
 shuttle = "0.7.1"
+tower = "0.4.13"
+hyper = "0.14"
 
 
 [dev-dependencies]
@@ -59,6 +64,7 @@ proptest-state-machine = "0.1.0"
 "rand" = "0.8.5"
 rayon = "1.8.0"
 criterion = "0.3"
+random-port = "0.1.1"
 
 [build-dependencies]
 tonic-build = "0.10"

rust/worker/build.rs

@@ -1,15 +1,17 @@
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // Compile the protobuf files in the chromadb proto directory.
+    let mut proto_paths = vec![
+        "../../idl/chromadb/proto/chroma.proto",
+        "../../idl/chromadb/proto/coordinator.proto",
+        "../../idl/chromadb/proto/logservice.proto",
+    ];
+
+    #[cfg(debug_assertions)]
+    proto_paths.push("../../idl/chromadb/proto/debug.proto");
+
     tonic_build::configure()
         .emit_rerun_if_changed(true)
-        .compile(
-            &[
-                "../../idl/chromadb/proto/chroma.proto",
-                "../../idl/chromadb/proto/coordinator.proto",
-                "../../idl/chromadb/proto/logservice.proto",
-            ],
-            &["../../idl/"],
-        )?;
+        .compile(&proto_paths, &["../../idl/"])?;
 
     // Compile the hnswlib bindings.
     cc::Build::new()

rust/worker/src/server.rs

@@ -94,18 +94,23 @@ impl WorkerServer {
             ))
             .add_service(
                 chroma_proto::metadata_reader_server::MetadataReaderServer::new(worker.clone()),
-            )
-            .serve_with_shutdown(addr, async {
-                let mut sigterm = match signal(SignalKind::terminate()) {
-                    Ok(sigterm) => sigterm,
-                    Err(e) => {
-                        tracing::error!("Failed to create signal handler: {:?}", e);
-                        return;
-                    }
-                };
-                sigterm.recv().await;
-                tracing::info!("Received SIGTERM, shutting down");
-            });
+            );
+
+        #[cfg(debug_assertions)]
+        let server =
+            server.add_service(chroma_proto::debug_server::DebugServer::new(worker.clone()));
+
+        let server = server.serve_with_shutdown(addr, async {
+            let mut sigterm = match signal(SignalKind::terminate()) {
+                Ok(sigterm) => sigterm,
+                Err(e) => {
+                    tracing::error!("Failed to create signal handler: {:?}", e);
+                    return;
+                }
+            };
+            sigterm.recv().await;
+            tracing::info!("Received SIGTERM, shutting down");
+        });
 
         server.await?;
         Ok(())
@@ -522,3 +527,81 @@ impl chroma_proto::metadata_reader_server::MetadataReader for WorkerServer {
             .await
     }
 }
+
+#[tonic::async_trait]
+impl chroma_proto::debug_server::Debug for WorkerServer {
+    async fn get_info(
+        &self,
+        _: Request<()>,
+    ) -> Result<Response<chroma_proto::GetInfoResponse>, Status> {
+        let response = chroma_proto::GetInfoResponse {
+            version: option_env!("CARGO_PKG_VERSION")
+                .unwrap_or("unknown")
+                .to_string(),
+        };
+        Ok(Response::new(response))
+    }
+
+    async fn trigger_panic(&self, _: Request<()>) -> Result<Response<()>, Status> {
+        panic!("Intentional panic triggered");
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::execution::dispatcher;
+    use crate::log::log::InMemoryLog;
+    use crate::storage::local::LocalStorage;
+    use crate::storage::Storage;
+    use crate::sysdb::test_sysdb::TestSysDb;
+    use crate::system;
+
+    use super::*;
+    use chroma_proto::debug_client::DebugClient;
+    use tempfile::tempdir;
+
+    #[tokio::test]
+    async fn gracefully_handles_panics() {
+        let sysdb = TestSysDb::new();
+        let log = InMemoryLog::new();
+        let tmp_dir = tempdir().unwrap();
+        let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
+
+        let port = random_port::PortPicker::new().pick().unwrap();
+        let mut server = WorkerServer {
+            dispatcher: None,
+            system: None,
+            sysdb: Box::new(SysDb::Test(sysdb)),
+            log: Box::new(Log::InMemory(log)),
+            hnsw_index_provider: HnswIndexProvider::new(
+                storage.clone(),
+                tmp_dir.path().to_path_buf(),
+            ),
+            blockfile_provider: BlockfileProvider::new_arrow(storage),
+            port,
+        };
+
+        let system: system::System = system::System::new();
+        let dispatcher = dispatcher::Dispatcher::new(4, 10, 10);
+        let dispatcher_handle = system.start_component(dispatcher);
+
+        server.set_system(system.clone());
+        server.set_dispatcher(dispatcher_handle.receiver());
+
+        tokio::spawn(async move {
+            let _ = crate::server::WorkerServer::run(server).await;
+        });
+
+        let mut client = DebugClient::connect(format!("http://localhost:{}", port))
+            .await
+            .unwrap();
+
+        // Test response when handler panics
+        let err_response = client.trigger_panic(Request::new(())).await.unwrap_err();
+        assert_eq!(err_response.code(), tonic::Code::Cancelled);
+
+        // The server should still work, even after a panic was thrown
+        let response = client.get_info(Request::new(())).await;
+        assert!(response.is_ok());
+    }
+}

rust/worker/src/sysdb/test_sysdb.rs

@@ -91,15 +91,8 @@ impl TestSysDb {
         if id.is_some() && id.unwrap() != segment.id {
             return false;
         }
-        if r#type.is_some() {
-            match r#type.unwrap().as_str() {
-                "hnsw" => {
-                    if segment.r#type != SegmentType::HnswDistributed {
-                        return false;
-                    }
-                }
-                _ => return false,
-            }
+        if let Some(r#type) = r#type {
+            return segment.r#type == SegmentType::try_from(r#type.as_str()).unwrap();
         }
         if scope.is_some() && scope.unwrap() != segment.scope {
             return false;

rust/worker/src/types/segment.rs

@@ -28,6 +28,20 @@ impl From<SegmentType> for String {
     }
 }
 
+impl TryFrom<&str> for SegmentType {
+    type Error = SegmentConversionError;
+
+    fn try_from(segment_type: &str) -> Result<Self, Self::Error> {
+        match segment_type {
+            "urn:chroma:segment/vector/hnsw-distributed" => Ok(SegmentType::HnswDistributed),
+            "urn:chroma:segment/record/blockfile" => Ok(SegmentType::BlockfileRecord),
+            "urn:chroma:segment/metadata/sqlite" => Ok(SegmentType::Sqlite),
+            "urn:chroma:segment/metadata/blockfile" => Ok(SegmentType::BlockfileMetadata),
+            _ => Err(SegmentConversionError::InvalidSegmentType),
+        }
+    }
+}
+
 #[derive(Clone, Debug, PartialEq)]
 pub(crate) struct Segment {
     pub(crate) id: Uuid,