[ENH] add panics to tracing

rust/worker/Dockerfile

@@ -34,9 +34,9 @@ RUN if [ "$RELEASE_MODE" = "1" ]; then mv target/release/query_service .; else m
 
 FROM debian:bookworm-slim as query_service
 
+RUN apt-get update && apt-get install -y libssl-dev ca-certificates
 COPY --from=query_service_builder /chroma/query_service .
 COPY --from=query_service_builder /chroma/rust/worker/chroma_config.yaml .
-RUN apt-get update && apt-get install -y libssl-dev ca-certificates
 
 ENTRYPOINT [ "./query_service" ]
 
@@ -54,8 +54,8 @@ RUN if [ "$RELEASE_MODE" = "1" ]; then mv target/release/compaction_service .; e
 
 FROM debian:bookworm-slim as compaction_service
 
+RUN apt-get update && apt-get install -y libssl-dev ca-certificates
 COPY --from=compaction_service_builder /chroma/compaction_service .
 COPY --from=compaction_service_builder /chroma/rust/worker/chroma_config.yaml .
-RUN apt-get update && apt-get install -y libssl-dev ca-certificates
 
 ENTRYPOINT [ "./compaction_service" ]

rust/worker/src/server.rs

@@ -532,18 +532,30 @@ impl chroma_proto::metadata_reader_server::MetadataReader for WorkerServer {
 impl chroma_proto::debug_server::Debug for WorkerServer {
     async fn get_info(
         &self,
-        _: Request<()>,
+        request: 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))
+        // Note: We cannot write a middleware that instruments every service rpc
+        // with a span because of https://github.com/hyperium/tonic/pull/1202.
+        let request_span = trace_span!("Get info");
+
+        wrap_span_with_parent_context(request_span, request.metadata()).in_scope(|| {
+            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");
+    async fn trigger_panic(&self, request: Request<()>) -> Result<Response<()>, Status> {
+        // Note: We cannot write a middleware that instruments every service rpc
+        // with a span because of https://github.com/hyperium/tonic/pull/1202.
+        let request_span = trace_span!("Trigger panic");
+
+        wrap_span_with_parent_context(request_span, request.metadata()).in_scope(|| {
+            panic!("Intentional panic triggered");
+        })
     }
 }
 

rust/worker/src/tracing/opentelemetry_config.rs

@@ -49,4 +49,28 @@ pub(crate) fn init_otel_tracing(service_name: &String, otel_endpoint: &String) {
     tracing::subscriber::set_global_default(subscriber)
         .expect("Set global default subscriber failed");
     println!("Set global subscriber for {}", service_name);
+
+    // Add panics to tracing
+    let prev_hook = std::panic::take_hook();
+    std::panic::set_hook(Box::new(move |panic_info| {
+        let payload = panic_info.payload();
+
+        #[allow(clippy::manual_map)]
+        let payload = if let Some(s) = payload.downcast_ref::<&str>() {
+            Some(&**s)
+        } else if let Some(s) = payload.downcast_ref::<String>() {
+            Some(s.as_str())
+        } else {
+            None
+        };
+
+        tracing::error!(
+            panic.payload = payload,
+            panic.location = panic_info.location().map(|l| l.to_string()),
+            panic.backtrace = tracing::field::display(std::backtrace::Backtrace::capture()),
+            "A panic occurred"
+        );
+
+        prev_hook(panic_info);
+    }));
 }