[ENH]: add readiness probe for garbage collector (#4662)

## Description of changes

Adds a readiness probe for GC.

## Test plan

_How are these changes tested?_

Verified that the service was marked ready using kubectl.

## Documentation Changes

_Are all docstrings for user-facing APIs updated if required? Do we need
to make documentation changes in the [docs
section](https://github.com/chroma-core/chroma/tree/main/docs/docs.trychroma.com)?_
n/a

Author: codetheweb

Cargo.lock

@@ -0,0 +1,6 @@
+syntax = "proto3";
+
+package chroma;
+
+// This is a placeholder service. The garbage collector currently only exposes a health service.
+service GarbageCollector {}

idl/chromadb/proto/garbage_collector.proto

@@ -16,7 +16,7 @@ apiVersion: v2
 name: distributed-chroma
 description: A helm chart for distributed Chroma
 type: application
-version: 0.1.43
+version: 0.1.44
 appVersion: "0.4.24"
 keywords:
   - chroma

k8s/distributed-chroma/Chart.yaml

@@ -40,6 +40,9 @@ spec:
       - name: garbage-collector
         image: "{{ .Values.garbageCollector.image.repository }}:{{ .Values.garbageCollector.image.tag }}"
         imagePullPolicy: IfNotPresent
+        readinessProbe:
+          grpc:
+            port: 50055
         {{ if .Values.garbageCollector.resources }}
         resources:
           limits:

k8s/distributed-chroma/templates/garbage-collector.yaml

@@ -27,6 +27,7 @@ tokio = { workspace = true }
 uuid = { workspace = true }
 futures = { workspace = true }
 tonic = { workspace = true }
+tonic-health = { workspace = true }
 tempfile = { workspace = true }
 tracing = { workspace = true }
 thiserror = { workspace = true }

rust/garbage_collector/Cargo.toml

@@ -38,6 +38,8 @@ pub(super) struct GarbageCollectorConfig {
     pub(super) assignment_policy: chroma_config::assignment::config::AssignmentPolicyConfig,
     pub(super) memberlist_provider: chroma_memberlist::config::MemberlistProviderConfig,
     pub my_member_id: String,
+    #[serde(default = "GarbageCollectorConfig::default_port")]
+    pub port: u16,
 }
 
 impl GarbageCollectorConfig {
@@ -61,6 +63,10 @@ impl GarbageCollectorConfig {
             Err(e) => panic!("Error loading config: {}", e),
         }
     }
+
+    fn default_port() -> u16 {
+        50055
+    }
 }
 
 #[cfg(test)]

rust/garbage_collector/src/config.rs

@@ -633,6 +633,7 @@ mod tests {
             assignment_policy: chroma_config::assignment::config::AssignmentPolicyConfig::default(),
             my_member_id: "test-gc".to_string(),
             memberlist_provider: chroma_memberlist::config::MemberlistProviderConfig::default(),
+            port: 50055,
         };
         let registry = Registry::new();
 
@@ -753,6 +754,7 @@ mod tests {
             assignment_policy: chroma_config::assignment::config::AssignmentPolicyConfig::default(),
             my_member_id: "test-gc".to_string(),
             memberlist_provider: chroma_memberlist::config::MemberlistProviderConfig::default(),
+            port: 50055,
         };
         let registry = Registry::new();
 

rust/garbage_collector/src/garbage_collector_component.rs

@@ -6,9 +6,11 @@ use chroma_tracing::{
     init_global_filter_layer, init_otel_layer, init_panic_tracing_hook, init_stdout_layer,
     init_tracing,
 };
+use chroma_types::chroma_proto::garbage_collector_server::GarbageCollectorServer;
 use config::GarbageCollectorConfig;
 use garbage_collector_component::GarbageCollector;
 use tokio::signal::unix::{signal, SignalKind};
+use tonic::transport::Server;
 use tracing::{debug, error, info};
 
 mod config;
@@ -24,7 +26,20 @@ pub mod types;
 
 const CONFIG_PATH_ENV_VAR: &str = "CONFIG_PATH";
 
+// This is a placeholder service so that we can expose a health service
+struct GarbageCollectorService {}
+
+impl chroma_types::chroma_proto::garbage_collector_server::GarbageCollector
+    for GarbageCollectorService
+{
+}
+
 pub async fn garbage_collector_service_entrypoint() -> Result<(), Box<dyn std::error::Error>> {
+    let (mut health_reporter, health_service) = tonic_health::server::health_reporter();
+    health_reporter
+        .set_not_serving::<GarbageCollectorServer<GarbageCollectorService>>()
+        .await;
+
     debug!("Loading configuration from environment");
     // Parse configuration. Configuration includes sysdb connection details, and
     // gc run details amongst others.
@@ -49,6 +64,27 @@ pub async fn garbage_collector_service_entrypoint() -> Result<(), Box<dyn std::e
 
     info!("Loaded configuration successfully: {:#?}", config);
 
+    let addr = format!("[::]:{}", config.port)
+        .parse()
+        .expect("Invalid address format");
+    let server_join_handle = tokio::spawn(async move {
+        let server = Server::builder().add_service(health_service);
+        server
+            .serve_with_shutdown(addr, async {
+                match signal(SignalKind::terminate()) {
+                    Ok(mut sigterm) => {
+                        sigterm.recv().await;
+                        tracing::info!("Received SIGTERM, shutting down gRPC server");
+                    }
+                    Err(err) => {
+                        tracing::error!("Failed to create SIGTERM handler: {err}")
+                    }
+                }
+            })
+            .await
+            .expect("Failed to start gRPC server");
+    });
+
     let registry = chroma_config::registry::Registry::new();
     let system = System::new();
 
@@ -79,6 +115,10 @@ pub async fn garbage_collector_service_entrypoint() -> Result<(), Box<dyn std::e
     memberlist.subscribe(garbage_collector_handle.receiver());
     let mut memberlist_handle = system.start_component(memberlist);
 
+    health_reporter
+        .set_serving::<GarbageCollectorServer<GarbageCollectorService>>()
+        .await;
+
     // Keep the service running and handle shutdown signals
     let mut sigterm = signal(SignalKind::terminate())?;
     let mut sigint = signal(SignalKind::interrupt())?;
@@ -112,6 +152,7 @@ pub async fn garbage_collector_service_entrypoint() -> Result<(), Box<dyn std::e
         .expect("Dispatcher should be stoppable");
     system.stop().await;
     system.join().await;
+    let _ = server_join_handle.await;
 
     info!("Shutting down garbage collector service");
     Ok(())

rust/garbage_collector/src/lib.rs

@@ -6,6 +6,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         "idl/chromadb/proto/coordinator.proto",
         "idl/chromadb/proto/logservice.proto",
         "idl/chromadb/proto/query_executor.proto",
+        "idl/chromadb/proto/garbage_collector.proto",
     ];
 
     // Can't use #[cfg(test)] here because a build for tests is technically a regular debug build, meaning that #[cfg(test)] is useless in build.rs.

rust/types/build.rs

@@ -48,6 +48,7 @@ pub use tenant::*;
 pub use types::*;
 pub use where_parsing::*;
 
+#[allow(clippy::all)]
 pub mod chroma_proto {
     tonic::include_proto!("chroma");
 }