[ENH]: Tracing for rust sysdb (#6333)

## Description of changes

This change adds tracing spans to various parts of the rust sysdb. All methods in spanner.rs are instrumented. Some queries are instrumented with info spans while most are instrumented with debug spans.

- Improvements & Bug fixes
  - ...
- New functionality
  - ...

## Test plan

Manual testing on Jaeger.

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

## Migration plan

_Are there any migrations, or any forwards/backwards compatibility changes needed in order to make sure this change deploys reliably?_

## Observability plan

_What is the plan to instrument and monitor this change?_

## 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)?_

Author: tanujnay112

rust/rust-sysdb/src/server.rs

@@ -53,6 +53,7 @@ use tokio::{
     signal::unix::{signal, SignalKind},
 };
 use tonic::{transport::Server, Request, Response, Status};
+use tracing::instrument;
 
 pub struct SysdbService {
     port: u16,
@@ -746,7 +747,7 @@ impl SysDb for SysdbService {
         .retry(backoff)
         .when(|e: &SysDbError| {
             if matches!(e, SysDbError::CollectionEntryIsStale) {
-                tracing::info!(
+                tracing::warn!(
                     "Collection entry is stale, retrying flush collection compaction for collection_id: {}",
                     collection_id
                 );
@@ -841,6 +842,7 @@ impl SysDb for SysdbService {
 
 impl SysdbService {
     /// Create a new version file in object storage
+    #[instrument(skip(self, storage, segments), level = "info", fields(collection_id = %collection.collection_id))]
     async fn create_new_version_file(
         &self,
         storage: &Storage,

rust/rust-sysdb/src/spanner.rs

@@ -27,6 +27,8 @@ use google_cloud_spanner::session::SessionConfig;
 use google_cloud_spanner::statement::Statement;
 use google_cloud_spanner::transaction_rw::ReadWriteTransaction;
 use thiserror::Error;
+use tracing::instrument;
+use tracing::Instrument;
 use uuid::Uuid;
 
 use crate::config::{SpannerBackendConfig, SpannerConfig};
@@ -141,6 +143,7 @@ impl SpannerBackend {
     /// Uses commit timestamps for created_at and updated_at.
     /// Sets last_compaction_time to Unix epoch (0) by default.
     /// If the tenant already exists, does nothing (insert on conflict do nothing).
+    #[instrument(skip(self), level = "info")]
     pub async fn create_tenant(
         &self,
         req: CreateTenantRequest,
@@ -156,7 +159,7 @@ impl SpannerBackend {
                     );
                     check_stmt.add_param("id", &tenant_id);
 
-                    let mut iter = tx.query(check_stmt).await?;
+                    let mut iter = tx.query(check_stmt).instrument(tracing::debug_span!("create_tenant query")).await?;
 
                     // If tenant doesn't exist, insert it otherwise ignore for idempotency
                     // Set last_compaction_time to Unix epoch (0) by default
@@ -184,6 +187,7 @@ impl SpannerBackend {
     /// Get tenants by ids.
     ///
     /// Returns `SysDbError::NotFound` if any tenant does not exist or is marked as deleted.
+    #[instrument(skip(self), level = "info")]
     pub async fn get_tenants(
         &self,
         req: GetTenantsRequest,
@@ -195,7 +199,10 @@ impl SpannerBackend {
 
         let mut tx = self.client.single().await?;
 
-        let mut iter = tx.query(stmt).await?;
+        let mut iter = tx
+            .query(stmt)
+            .instrument(tracing::debug_span!("get_tenants query"))
+            .await?;
         let mut tenants = Vec::new();
 
         // Get all rows
@@ -214,6 +221,7 @@ impl SpannerBackend {
         }
     }
 
+    #[instrument(skip(self, tx), level = "info")]
     async fn update_tenant(
         &self,
         tx: &mut ReadWriteTransaction,
@@ -240,6 +248,7 @@ impl SpannerBackend {
     ///
     /// This mimics the Go RegisterFilePaths function which updates the file_paths
     /// column for each segment with the new file paths from compaction.
+    #[instrument(skip(self, tx, req), fields(collection_id = %req.collection_id), level = "info")]
     async fn update_segments(
         &self,
         tx: &mut ReadWriteTransaction,
@@ -276,7 +285,13 @@ impl SpannerBackend {
             update_stmt.add_param("region", &self.local_region().as_str());
             update_stmt.add_param("file_paths", &file_paths_json);
 
-            let rows_affected = tx.update(update_stmt).await?;
+            let rows_affected = tx
+                .update(update_stmt)
+                .instrument(tracing::info_span!(
+                    "update_segment_file_paths",
+                    segment_id = %flush_segment_compaction.segment_id
+                ))
+                .await?;
             if rows_affected == 0 {
                 return Err(SysDbError::NotFound(format!(
                     "segment '{}' not found",
@@ -302,6 +317,7 @@ impl SpannerBackend {
     /// Validates that the database name is not empty and that the tenant exists.
     /// Uses commit timestamps for created_at and updated_at.
     /// All checks and the insert are done atomically in a single transaction.
+    #[instrument(skip(self), level = "info")]
     pub async fn create_database(
         &self,
         req: CreateDatabaseRequest,
@@ -320,7 +336,10 @@ impl SpannerBackend {
                     );
                     tenant_check_stmt.add_param("id", &tenant_id);
 
-                    let mut tenant_iter = tx.query(tenant_check_stmt).await?;
+                    let mut tenant_iter = tx
+                        .query(tenant_check_stmt)
+                        .instrument(tracing::debug_span!("check_tenant_exists"))
+                        .await?;
                     if tenant_iter.next().await?.is_none() {
                         return Err(SysDbError::NotFound(format!(
                             "tenant '{}' not found",
@@ -335,7 +354,7 @@ impl SpannerBackend {
                     name_check_stmt.add_param("name", &db_name);
                     name_check_stmt.add_param("tenant_id", &tenant_id);
 
-                    let mut name_iter = tx.query(name_check_stmt).await?;
+                    let mut name_iter = tx.query(name_check_stmt).instrument(tracing::debug_span!("check_database_name_exists")).await?;
                     if name_iter.next().await?.is_some() {
                         return Err(SysDbError::AlreadyExists(format!(
                             "database with name '{}' already exists for tenant '{}'",
@@ -349,7 +368,10 @@ impl SpannerBackend {
                     );
                     check_stmt.add_param("id", &db_id);
 
-                    let mut iter = tx.query(check_stmt).await?;
+                    let mut iter = tx
+                        .query(check_stmt)
+                        .instrument(tracing::debug_span!("check_database_exists"))
+                        .await?;
                     if iter.next().await?.is_some() {
                         return Err(SysDbError::AlreadyExists(format!(
                             "database with id '{}' already exists",
@@ -366,7 +388,9 @@ impl SpannerBackend {
                     insert_stmt.add_param("tenant_id", &tenant_id);
                     insert_stmt.add_param("is_deleted", &false);
 
-                    tx.update(insert_stmt).await?;
+                    tx.update(insert_stmt)
+                        .instrument(tracing::info_span!("insert_database"))
+                        .await?;
                     tracing::info!("Created database: {} for tenant: {}", db_name, tenant_id);
 
                     Ok(())
@@ -383,6 +407,7 @@ impl SpannerBackend {
     /// Get a database by name and tenant.
     ///
     /// Returns `SysDbError::NotFound` if the database does not exist or is marked as deleted.
+    #[instrument(skip(self), level = "info")]
     pub async fn get_database(
         &self,
         req: GetDatabaseRequest,
@@ -396,7 +421,10 @@ impl SpannerBackend {
 
         let mut tx = self.client.single().await?;
 
-        let mut iter = tx.query(stmt).await?;
+        let mut iter = tx
+            .query(stmt)
+            .instrument(tracing::info_span!("get_database query"))
+            .await?;
 
         // Get the first row if it exists
         if let Some(row) = iter.next().await? {
@@ -411,6 +439,7 @@ impl SpannerBackend {
         }
     }
 
+    #[instrument(skip(self), level = "info")]
     pub async fn list_databases(&self, tenant: &str) -> Result<Vec<Database>, SysDbError> {
         let query = "SELECT id, name, tenant_id FROM databases \
                      WHERE tenant_id = @tenant_id AND is_deleted = FALSE \
@@ -420,7 +449,10 @@ impl SpannerBackend {
         stmt.add_param("tenant_id", &tenant);
 
         let mut tx = self.client.single().await?;
-        let mut result_set = tx.query(stmt).await?;
+        let mut result_set = tx
+            .query(stmt)
+            .instrument(tracing::info_span!("list_databases query"))
+            .await?;
 
         let mut databases = Vec::new();
         while let Some(row) = result_set.next().await? {
@@ -439,6 +471,7 @@ impl SpannerBackend {
     // Collection Operations
     // ============================================================
 
+    #[instrument(skip(self), level = "info")]
     pub async fn create_collection(
         &self,
         req: CreateCollectionRequest,
@@ -512,7 +545,7 @@ impl SpannerBackend {
                     db_stmt.add_param("name", &database_name);
                     db_stmt.add_param("tenant_id", &tenant_id_str);
 
-                    let mut db_iter = tx.query(db_stmt).await?;
+                    let mut db_iter = tx.query(db_stmt).instrument(tracing::debug_span!("get_database_by_name")).await?;
                     let db_row = db_iter.next().await?;
                     let database_id = match db_row {
                         Some(row) => {
@@ -535,7 +568,7 @@ impl SpannerBackend {
                     );
                     id_check_stmt.add_param("collection_id", &collection_id);
 
-                    let mut id_iter = tx.query(id_check_stmt).await?;
+                    let mut id_iter = tx.query(id_check_stmt).instrument(tracing::debug_span!("check_collection_exists_by_id")).await?;
                     if id_iter.next().await?.is_some() {
                         if get_or_create {
                             // Return the existing collection
@@ -558,7 +591,7 @@ impl SpannerBackend {
                     check_stmt.add_param("database_name", &database_name);
                     check_stmt.add_param("name", &collection_name);
 
-                    let mut check_iter = tx.query(check_stmt).await?;
+                    let mut check_iter = tx.query(check_stmt).instrument(tracing::debug_span!("check_collection_exists_by_name")).await?;
                     if let Some(existing_row) = check_iter.next().await? {
                         // Collection with same name exists
                         if get_or_create {
@@ -775,6 +808,7 @@ impl SpannerBackend {
     /// - `limit` and `offset`: Pagination
     ///
     /// Returns a list of matching collections.
+    #[instrument(skip(self), level = "info")]
     pub async fn get_collections(
         &self,
         req: GetCollectionsRequest,
@@ -898,7 +932,10 @@ impl SpannerBackend {
         }
 
         let mut tx = self.client.single().await?;
-        let mut result_set = tx.query(stmt).await?;
+        let mut result_set = tx
+            .query(stmt)
+            .instrument(tracing::info_span!("get_collections query"))
+            .await?;
 
         // Collect all rows, grouped by collection_id, preserving query order (created_at ASC)
         let mut collection_order: Vec<String> = Vec::new();
@@ -1023,7 +1060,10 @@ impl SpannerBackend {
         fetch_stmt.add_param("collection_id", &collection_id);
         fetch_stmt.add_param("region", &region);
 
-        let mut fetch_iter = tx.query(fetch_stmt).await?;
+        let mut fetch_iter = tx
+            .query(fetch_stmt)
+            .instrument(tracing::info_span!("fetch_collection_in_tx query"))
+            .await?;
 
         // Collect all rows and convert to Collection using TryFrom<Vec<Row>>
         let mut rows = Vec::new();
@@ -1047,6 +1087,7 @@ impl SpannerBackend {
     /// - Segments: deduplicated using HashMap
     ///
     /// This approach fetches all data in a single network round trip.
+    #[instrument(skip(self), level = "info")]
     pub async fn get_collection_with_segments(
         &self,
         req: GetCollectionWithSegmentsRequest,
@@ -1098,7 +1139,10 @@ impl SpannerBackend {
         stmt.add_param("region", &region);
 
         let mut tx = self.client.single().await?;
-        let mut result_set = tx.query(stmt).await?;
+        let mut result_set = tx
+            .query(stmt)
+            .instrument(tracing::debug_span!("get_collection_with_segments query"))
+            .await?;
 
         // Collect all rows
         let mut rows: Vec<Row> = Vec::new();
@@ -1161,6 +1205,7 @@ impl SpannerBackend {
     /// - `new_configuration`: New configuration (selective update of hnsw, spann, or embedding function)
     ///
     /// Returns the updated collection.
+    #[instrument(skip(self))]
     pub async fn update_collection(
         &self,
         req: UpdateCollectionRequest,
@@ -1197,7 +1242,7 @@ impl SpannerBackend {
                     );
                     check_stmt.add_param("collection_id", &collection_id);
 
-                    let mut check_iter = tx.query(check_stmt).await?;
+                    let mut check_iter = tx.query(check_stmt).instrument(tracing::debug_span!("check_collection_metadata")).await?;
                     let (tenant_id, database_name): (String, String) = match check_iter.next().await? {
                         Some(row) => {
                             let tenant: String = row.column_by_name("tenant_id").map_err(SysDbError::FailedToReadColumn)?;
@@ -1222,7 +1267,7 @@ impl SpannerBackend {
                         name_check_stmt.add_param("name", new_name);
                         name_check_stmt.add_param("collection_id", &collection_id);
 
-                        let mut name_iter = tx.query(name_check_stmt).await?;
+                        let mut name_iter = tx.query(name_check_stmt).instrument(tracing::debug_span!("check_collection_name_exists")).await?;
                         if name_iter.next().await?.is_some() {
                             return Err(SysDbError::AlreadyExists(format!(
                                 "collection with name '{}' already exists in database '{}'",
@@ -1263,7 +1308,7 @@ impl SpannerBackend {
                             "SELECT name, dimension FROM collections WHERE collection_id = @collection_id",
                         );
                         current_stmt.add_param("collection_id", &collection_id);
-                        let mut current_iter = tx.query(current_stmt).await?;
+                        let mut current_iter = tx.query(current_stmt).instrument(tracing::debug_span!("get_current_collection_info")).await?;
                         let (current_name, current_dimension): (String, Option<i64>) = match current_iter.next().await? {
                             Some(row) => {
                                 let n: String = row.column_by_name("name").map_err(SysDbError::FailedToReadColumn)?;
@@ -1290,7 +1335,7 @@ impl SpannerBackend {
                             "SELECT key FROM collection_metadata WHERE collection_id = @collection_id",
                         );
                         keys_stmt.add_param("collection_id", &collection_id);
-                        let mut keys_iter = tx.query(keys_stmt).await?;
+                        let mut keys_iter = tx.query(keys_stmt).instrument(tracing::debug_span!("get_collection_metadata_keys")).await?;
                         while let Some(row) = keys_iter.next().await? {
                             let key: String = row.column_by_name("key").map_err(SysDbError::FailedToReadColumn)?;
                             mutations.push(delete(
@@ -1364,7 +1409,7 @@ impl SpannerBackend {
                         );
                         schema_stmt.add_param("collection_id", &collection_id);
 
-                        let mut schema_iter = tx.query(schema_stmt).await?;
+                        let mut schema_iter = tx.query(schema_stmt).instrument(tracing::debug_span!("get_collection_schemas")).await?;
                         let mut region_schemas: Vec<(String, String)> = Vec::new();
 
                         while let Some(row) = schema_iter.next().await? {
@@ -1453,6 +1498,7 @@ impl SpannerBackend {
 
     /// Flush collection compaction results to the database.
     /// This mimics the logic in go/pkg/sysdb/coordinator/table_catalog.go::FlushCollectionCompactionForVersionedCollection
+    #[instrument(skip(self, req))]
     pub async fn flush_collection_compaction(
         &self,
         req: FlushCompactionRequest,
@@ -1605,7 +1651,15 @@ impl SpannerBackend {
                     update_stmt.add_param("schema_str", &schema_str);
                     update_stmt.add_param("region", &region.to_string());
                     update_stmt.add_param("old_version_file_name", &old_version_file_name);
-                    let rows_affected = tx.update(update_stmt).await?;
+                    let rows_affected = tx
+                        .update(update_stmt)
+                        .instrument(tracing::info_span!(
+                            "flush_compaction update query",
+                            collection_id = %collection_id,
+                            version_file_path = %version_file_path,
+                            region = %region,
+                        ))
+                        .await?;
 
                     if rows_affected == 0 {
                         // CAS operation failed - collection was updated by another transaction
@@ -1659,13 +1713,17 @@ impl SpannerBackend {
 
     /// Reset the database state by dropping all tables and re-running migrations.
     /// This provides a completely clean slate for testing.
+    #[instrument(skip(self), level = "info")]
     pub async fn reset(&self) -> Result<(), SysDbError> {
         // Step 1: Get all indexes first
         let get_indexes_stmt = Statement::new(
             "SELECT table_name, index_name FROM INFORMATION_SCHEMA.INDEXES WHERE table_schema = ''",
         );
         let mut tx = self.client.single().await?;
-        let mut indexes_result = tx.query(get_indexes_stmt).await?;
+        let mut indexes_result = tx
+            .query(get_indexes_stmt)
+            .instrument(tracing::debug_span!("get_spanner_indexes"))
+            .await?;
         let mut table_names = HashSet::new();
 
         let mut indexes: Vec<(String, String)> = Vec::new();