[ENH] Gracefully shutdown GC system (#4246)

## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
   - We would like to wait until existing GC jobs to finish on shutdown. Thus we will call `system.stop()` to gracefully shutdown all components when SIGINT/SIGTERM is received. Notice that the system component will always be shutdown once the existing message has been handled, so we don't need to worry about partial progress in the GC job.
 - New functionality
   - N/A

## Test plan
*How are these changes tested?*

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

## Documentation Changes
*Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs repository](https://github.com/chroma-core/docs)?*

Author: Sicheng-Pan

Diff for: rust/garbage_collector/src/garbage_collector_component.rs

@@ -178,6 +178,11 @@ impl Component for GarbageCollector {
             || Some(span!(parent: None, tracing::Level::INFO, "Scheduled garbage collection")),
         );
     }
+
+    fn on_stop_timeout(&self) -> Duration {
+        // NOTE: Increased timeout for remaining jobs to finish
+        Duration::from_secs(60)
+    }
 }
 
 #[derive(Debug)]

Diff for: rust/garbage_collector/src/lib.rs

@@ -7,7 +7,6 @@ use chroma_tracing::{
 use config::GarbageCollectorConfig;
 use garbage_collector_component::GarbageCollector;
 use tokio::signal::unix::{signal, SignalKind};
-use tokio::time::{sleep, Duration};
 use tracing::{debug, error, info};
 
 mod config;
@@ -56,7 +55,7 @@ pub async fn garbage_collector_service_entrypoint() -> Result<(), Box<dyn std::e
         })?;
 
     let system = System::new();
-    let dispatcher_handle = system.start_component(dispatcher);
+    let mut dispatcher_handle = system.start_component(dispatcher);
 
     // Start a background task to periodically check for garbage.
     // Garbage collector is a component that gets notified every
@@ -68,36 +67,40 @@ pub async fn garbage_collector_service_entrypoint() -> Result<(), Box<dyn std::e
             e
         })?;
 
-    garbage_collector_component.set_dispatcher(dispatcher_handle);
+    garbage_collector_component.set_dispatcher(dispatcher_handle.clone());
     garbage_collector_component.set_system(system.clone());
 
-    let _ = system.start_component(garbage_collector_component);
+    let mut garbage_collector_handle = system.start_component(garbage_collector_component);
 
     // Keep the service running and handle shutdown signals
     let mut sigterm = signal(SignalKind::terminate())?;
     let mut sigint = signal(SignalKind::interrupt())?;
 
     info!("Service running, waiting for signals");
-    loop {
-        tokio::select! {
-            _ = sigterm.recv() => {
-                info!("Received SIGTERM signal");
-                break;
-            }
-            _ = sigint.recv() => {
-                info!("Received SIGINT signal");
-                break;
-            }
-            _ = sleep(Duration::from_secs(1)) => {
-                // Keep the service running
-                continue;
-            }
+    tokio::select! {
+        _ = sigterm.recv() => {
+            info!("Received SIGTERM signal");
+        }
+        _ = sigint.recv() => {
+            info!("Received SIGINT signal");
         }
     }
-
-    // Give some time for any in-progress garbage collection to complete
     info!("Starting graceful shutdown, waiting for in-progress tasks");
-    sleep(Duration::from_secs(5)).await;
+    // NOTE: We should first stop the garbage collector. The garbage collector will finish the remaining jobs before shutdown.
+    // We cannot directly shutdown the dispatcher and system because that will fail remaining jobs.
+    garbage_collector_handle.stop();
+    garbage_collector_handle
+        .join()
+        .await
+        .expect("Garbage collector should be stoppable");
+    dispatcher_handle.stop();
+    dispatcher_handle
+        .join()
+        .await
+        .expect("Dispatcher should be stoppable");
+    system.stop().await;
+    system.join().await;
+
     info!("Shutting down garbage collector service");
     Ok(())
 }