Add correct segment name to span

Author: codetheweb

rust/worker/src/execution/operators/apply_log_to_segment_writer.rs

@@ -174,7 +174,8 @@ impl<Writer: SegmentWriter + Send + Sync + Clone>
             .segment_writer
             .apply_materialized_log_chunk(res.clone())
             .instrument(tracing::trace_span!(
-                "Apply materialized logs to record segment"
+                "Apply materialized logs",
+                segment = input.segment_writer.get_name()
             ))
             .await
         {

rust/worker/src/segment/distributed_hnsw_segment.rs

@@ -178,6 +178,10 @@ impl DistributedHNSWSegmentWriter {
 }
 
 impl SegmentWriter for DistributedHNSWSegmentWriter {
+    fn get_name(&self) -> &'static str {
+        "DistributedHNSWSegmentWriter"
+    }
+
     async fn apply_materialized_log_chunk(
         &self,
         records: chroma_types::Chunk<super::MaterializedLogRecord<'_>>,
@@ -241,6 +245,10 @@ impl SegmentWriter for DistributedHNSWSegmentWriter {
 }
 
 impl SegmentWriter for Box<DistributedHNSWSegmentWriter> {
+    fn get_name(&self) -> &'static str {
+        self.as_ref().get_name()
+    }
+
     async fn apply_materialized_log_chunk(
         &self,
         records: chroma_types::Chunk<super::MaterializedLogRecord<'_>>,

rust/worker/src/segment/metadata_segment.rs

@@ -531,6 +531,10 @@ impl<'me> MetadataSegmentWriter<'me> {
 }
 
 impl SegmentWriter for MetadataSegmentWriter<'_> {
+    fn get_name(&self) -> &'static str {
+        "MetadataSegmentWriter"
+    }
+
     async fn apply_materialized_log_chunk(
         &self,
         records: Chunk<MaterializedLogRecord<'_>>,

rust/worker/src/segment/record_segment.rs

@@ -333,6 +333,10 @@ impl ChromaError for ApplyMaterializedLogError {
 }
 
 impl SegmentWriter for RecordSegmentWriter {
+    fn get_name(&self) -> &'static str {
+        "RecordSegmentWriter"
+    }
+
     async fn apply_materialized_log_chunk<'referred_data>(
         &self,
         records: Chunk<MaterializedLogRecord<'referred_data>>,

rust/worker/src/segment/types.rs

@@ -789,6 +789,7 @@ impl<'me> LogMaterializer<'me> {
 // This needs to be public for testing
 #[allow(async_fn_in_trait)]
 pub trait SegmentWriter {
+    fn get_name(&self) -> &'static str;
     fn apply_materialized_log_chunk(
         &self,
         records: Chunk<MaterializedLogRecord>,