[ENH] Add otel support for query and compaction service

rust/worker/Cargo.toml

@@ -45,7 +45,11 @@ arrow = "50.0.0"
 roaring = "0.10.3"
 tantivy = "0.21.1"
 tracing = "0.1"
-tracing-subscriber = "0.3"
+tracing-bunyan-formatter = "0.3.3"
+tracing-opentelemetry = "0.19.0"
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+opentelemetry = { version = "0.19.0", default-features = false, features = ["trace", "rt-tokio"] }
+opentelemetry-otlp = "0.12.0"
 
 [dev-dependencies]
 proptest = "1.4.0"

rust/worker/chroma_config.yaml

@@ -4,6 +4,8 @@
 # for now we nest it in the worker directory
 
 query_service:
+    service_name: "query-service"
+    otel_endpoint: "http://jaeger:4317"
     my_ip: "10.244.0.9"
     my_port: 50051
     assignment_policy:
@@ -32,6 +34,8 @@ query_service:
         worker_queue_size: 100
 
 compaction_service:
+    service_name: "compaction-service"
+    otel_endpoint: "http://jaeger:4317"
     my_ip: "10.244.0.9"
     my_port: 50051
     assignment_policy:

rust/worker/src/bin/query_service.rs

@@ -2,8 +2,5 @@ use worker::query_service_entrypoint;
 
 #[tokio::main]
 async fn main() {
-    tracing_subscriber::fmt()
-        .with_max_level(tracing::Level::INFO)
-        .init();
     query_service_entrypoint().await;
 }

rust/worker/src/config.rs

@@ -93,6 +93,8 @@ impl RootConfig {
 /// Each submodule that needs to be configured from the config object should implement the Configurable trait and
 /// have its own field in this struct for its Config struct.
 pub(crate) struct QueryServiceConfig {
+    pub(crate) service_name: String,
+    pub(crate) otel_endpoint: String,
     pub(crate) my_ip: String,
     pub(crate) my_port: u16,
     pub(crate) assignment_policy: crate::assignment::config::AssignmentPolicyConfig,
@@ -115,6 +117,8 @@ pub(crate) struct QueryServiceConfig {
 /// Each submodule that needs to be configured from the config object should implement the Configurable trait and
 /// have its own field in this struct for its Config struct.
 pub(crate) struct CompactionServiceConfig {
+    pub(crate) service_name: String,
+    pub(crate) otel_endpoint: String,
     pub(crate) my_ip: String,
     pub(crate) my_port: u16,
     pub(crate) assignment_policy: crate::assignment::config::AssignmentPolicyConfig,
@@ -150,6 +154,8 @@ mod tests {
                 "chroma_config.yaml",
                 r#"
                 query_service:
+                    service_name: "query-service"
+                    otel_endpoint: "http://jaeger:4317"
                     my_ip: "192.0.0.1"
                     my_port: 50051
                     assignment_policy:
@@ -178,6 +184,8 @@ mod tests {
                         worker_queue_size: 100
 
                 compaction_service:
+                    service_name: "compaction-service"
+                    otel_endpoint: "http://jaeger:4317"
                     my_ip: "192.0.0.1"
                     my_port: 50051
                     assignment_policy:
@@ -227,6 +235,8 @@ mod tests {
                 "random_path.yaml",
                 r#"
                 query_service:
+                    service_name: "query-service"
+                    otel_endpoint: "http://jaeger:4317"
                     my_ip: "192.0.0.1"
                     my_port: 50051
                     assignment_policy:
@@ -255,6 +265,8 @@ mod tests {
                         worker_queue_size: 100
 
                 compaction_service:
+                    service_name: "compaction-service"
+                    otel_endpoint: "http://jaeger:4317"
                     my_ip: "192.0.0.1"
                     my_port: 50051
                     assignment_policy:
@@ -322,6 +334,8 @@ mod tests {
                 "chroma_config.yaml",
                 r#"
                 query_service:
+                    service_name: "query-service"
+                    otel_endpoint: "http://jaeger:4317"
                     my_ip: "192.0.0.1"
                     my_port: 50051
                     assignment_policy:
@@ -350,6 +364,8 @@ mod tests {
                         worker_queue_size: 100
 
                 compaction_service:
+                    service_name: "compaction-service"
+                    otel_endpoint: "http://jaeger:4317"
                     my_ip: "192.0.0.1"
                     my_port: 50051
                     assignment_policy:
@@ -401,6 +417,8 @@ mod tests {
                 "chroma_config.yaml",
                 r#"
                 query_service:
+                    service_name: "query-service"
+                    otel_endpoint: "http://jaeger:4317"
                     assignment_policy:
                         RendezvousHashing:
                             hasher: Murmur3
@@ -427,6 +445,8 @@ mod tests {
                         worker_queue_size: 100
 
                 compaction_service:
+                    service_name: "compaction-service"
+                    otel_endpoint: "http://jaeger:4317"
                     assignment_policy:
                         RendezvousHashing:
                             hasher: Murmur3

rust/worker/src/lib.rs

@@ -13,6 +13,7 @@ mod server;
 mod storage;
 mod sysdb;
 mod system;
+mod tracing;
 mod types;
 
 use config::Configurable;
@@ -35,6 +36,12 @@ pub async fn query_service_entrypoint() {
     };
 
     let config = config.query_service;
+
+    crate::tracing::opentelemetry_config::init_otel_tracing(
+        &config.service_name,
+        &config.otel_endpoint,
+    );
+
     let system: system::System = system::System::new();
     let dispatcher =
         match execution::dispatcher::Dispatcher::try_from_config(&config.dispatcher).await {
@@ -94,6 +101,12 @@ pub async fn compaction_service_entrypoint() {
     };
 
     let config = config.compaction_service;
+
+    crate::tracing::opentelemetry_config::init_otel_tracing(
+        &config.service_name,
+        &config.otel_endpoint,
+    );
+
     let system: system::System = system::System::new();
 
     let mut memberlist = match memberlist::CustomResourceMemberlistProvider::try_from_config(

rust/worker/src/tracing/mod.rs

@@ -0,0 +1 @@
+pub(crate) mod opentelemetry_config;

rust/worker/src/tracing/opentelemetry_config.rs

@@ -0,0 +1,51 @@
+use opentelemetry::global;
+use opentelemetry::sdk::propagation::TraceContextPropagator;
+use opentelemetry::sdk::trace;
+use opentelemetry_otlp::WithExportConfig;
+use tracing_bunyan_formatter::BunyanFormattingLayer;
+use tracing_subscriber::{layer::SubscriberExt, EnvFilter, Layer};
+
+pub(crate) fn init_otel_tracing(service_name: &String, otel_endpoint: &String) {
+    println!(
+        "Registering jaeger subscriber for {} at endpoint {}",
+        service_name, otel_endpoint
+    );
+    let resource = opentelemetry::sdk::Resource::new(vec![opentelemetry::KeyValue::new(
+        "service.name",
+        service_name.clone(),
+    )]);
+    // Prepare trace config.
+    let trace_config = trace::config()
+        .with_sampler(opentelemetry::sdk::trace::Sampler::AlwaysOn)
+        .with_resource(resource);
+    // Prepare exporter.
+    let exporter = opentelemetry_otlp::new_exporter()
+        .tonic()
+        .with_endpoint(otel_endpoint);
+    let otlp_tracer = opentelemetry_otlp::new_pipeline()
+        .tracing()
+        .with_exporter(exporter)
+        .with_trace_config(trace_config)
+        .install_batch(opentelemetry::runtime::Tokio)
+        .expect("Error - Failed to create tracer.");
+    // Layer for adding our configured tracer.
+    // Export everything at this layer. The backend i.e. honeycomb or jaeger will filter at its end.
+    let exporter_layer = tracing_opentelemetry::layer()
+        .with_tracer(otlp_tracer)
+        .with_filter(tracing_subscriber::filter::LevelFilter::TRACE);
+    // Layer for printing spans to stdout. Only print INFO logs by default.
+    let stdout_layer =
+        BunyanFormattingLayer::new(service_name.clone().to_string(), std::io::stdout)
+            .with_filter(tracing_subscriber::filter::LevelFilter::INFO);
+    // global filter layer. Don't filter anything at global layer.
+    let global_layer = EnvFilter::new("TRACE");
+    // Create subscriber.
+    let subscriber = tracing_subscriber::registry()
+        .with(global_layer)
+        .with(stdout_layer)
+        .with(exporter_layer);
+    global::set_text_map_propagator(TraceContextPropagator::new());
+    tracing::subscriber::set_global_default(subscriber)
+        .expect("Set global default subscriber failed");
+    println!("Set global subscriber for {}", service_name);
+}