WASI OTel

Cargo.toml

@@ -84,6 +84,7 @@ openssl = { version = "0.10" }
 anyhow = { workspace = true, features = ["backtrace"] }
 conformance = { path = "tests/conformance-tests" }
 conformance-tests = { workspace = true }
+fake-opentelemetry-collector = "0.26"
 hex = "0.4"
 http-body-util = { workspace = true }
 hyper = { workspace = true }
@@ -144,6 +145,10 @@ hyper-util = { version = "0.1", features = ["tokio"] }
 indexmap = "2"
 itertools = "0.14"
 lazy_static = "1.5"
+opentelemetry = "0.28"
+# The default `reqwest-blocking-client` causes a runtime panic
+opentelemetry-otlp = { version = "0.28", default-features = false, features = ["http-proto", "reqwest-client", "logs"]}
+opentelemetry_sdk = {version = "0.28", features = ["experimental_metrics_periodicreader_with_async_runtime", "experimental_trace_batch_span_processor_with_async_runtime", "experimental_logs_batch_log_processor_with_async_runtime"]}
 path-absolutize = "3"
 pin-project-lite = "0.2.16"
 quote = "1"
@@ -172,6 +177,7 @@ toml_edit = "0.22"
 tower-service = "0.3.3"
 tracing = { version = "0.1.41", features = ["log"] }
 url = "2.5.7"
+tracing-opentelemetry = "0.29"
 walkdir = "2"
 wasm-encoder = "0.239.0"
 wasm-metadata = "0.239.0"

crates/factor-key-value/Cargo.toml

@@ -8,6 +8,7 @@ edition = { workspace = true }
 anyhow = { workspace = true }
 serde = { workspace = true }
 spin-core = { path = "../core" }
+spin-factor-otel = { path = "../factor-otel" }
 spin-factors = { path = "../factors" }
 spin-locked-app = { path = "../locked-app" }
 spin-resource-table = { path = "../table" }

crates/factor-key-value/src/host.rs

@@ -1,6 +1,7 @@
 use super::{Cas, SwapError};
 use anyhow::{Context, Result};
 use spin_core::{async_trait, wasmtime::component::Resource};
+use spin_factor_otel::OtelFactorState;
 use spin_resource_table::Table;
 use spin_telemetry::traces::{self, Blame};
 use spin_world::v2::key_value;
@@ -49,23 +50,26 @@ pub struct KeyValueDispatch {
     manager: Arc<dyn StoreManager>,
     stores: Table<Arc<dyn Store>>,
     compare_and_swaps: Table<Arc<dyn Cas>>,
+    otel_context: Option<OtelFactorState>,
 }
 
 impl KeyValueDispatch {
     pub fn new(allowed_stores: HashSet<String>, manager: Arc<dyn StoreManager>) -> Self {
-        Self::new_with_capacity(allowed_stores, manager, DEFAULT_STORE_TABLE_CAPACITY)
+        Self::new_with_capacity(allowed_stores, manager, DEFAULT_STORE_TABLE_CAPACITY, None)
     }
 
     pub fn new_with_capacity(
         allowed_stores: HashSet<String>,
         manager: Arc<dyn StoreManager>,
         capacity: u32,
+        otel_context: Option<OtelFactorState>,
     ) -> Self {
         Self {
             allowed_stores,
             manager,
             stores: Table::new(capacity),
             compare_and_swaps: Table::new(capacity),
+            otel_context,
         }
     }
 
@@ -113,6 +117,9 @@ impl key_value::Host for KeyValueDispatch {}
 impl key_value::HostStore for KeyValueDispatch {
     #[instrument(name = "spin_key_value.open", skip(self), err, fields(otel.kind = "client", kv.backend=self.manager.summary(&name).unwrap_or("unknown".to_string())))]
     async fn open(&mut self, name: String) -> Result<Result<Resource<key_value::Store>, Error>> {
+        if let Some(otel_context) = self.otel_context.as_ref() {
+            otel_context.reparent_tracing_span()
+        }
         Ok(async {
             if self.allowed_stores.contains(&name) {
                 let store = self.manager.get(&name).await?;
@@ -135,6 +142,9 @@ impl key_value::HostStore for KeyValueDispatch {
         store: Resource<key_value::Store>,
         key: String,
     ) -> Result<Result<Option<Vec<u8>>, Error>> {
+        if let Some(otel_context) = self.otel_context.as_ref() {
+            otel_context.reparent_tracing_span()
+        }
         let store = self.get_store(store)?;
         Ok(store.get(&key).await.map_err(track_error_on_span))
     }
@@ -146,6 +156,9 @@ impl key_value::HostStore for KeyValueDispatch {
         key: String,
         value: Vec<u8>,
     ) -> Result<Result<(), Error>> {
+        if let Some(otel_context) = self.otel_context.as_ref() {
+            otel_context.reparent_tracing_span()
+        }
         let store = self.get_store(store)?;
         Ok(store.set(&key, &value).await.map_err(track_error_on_span))
     }
@@ -156,6 +169,9 @@ impl key_value::HostStore for KeyValueDispatch {
         store: Resource<key_value::Store>,
         key: String,
     ) -> Result<Result<(), Error>> {
+        if let Some(otel_context) = self.otel_context.as_ref() {
+            otel_context.reparent_tracing_span()
+        }
         let store = self.get_store(store)?;
         Ok(store.delete(&key).await.map_err(track_error_on_span))
     }
@@ -166,6 +182,9 @@ impl key_value::HostStore for KeyValueDispatch {
         store: Resource<key_value::Store>,
         key: String,
     ) -> Result<Result<bool, Error>> {
+        if let Some(otel_context) = self.otel_context.as_ref() {
+            otel_context.reparent_tracing_span()
+        }
         let store = self.get_store(store)?;
         Ok(store.exists(&key).await.map_err(track_error_on_span))
     }
@@ -175,6 +194,9 @@ impl key_value::HostStore for KeyValueDispatch {
         &mut self,
         store: Resource<key_value::Store>,
     ) -> Result<Result<Vec<String>, Error>> {
+        if let Some(otel_context) = self.otel_context.as_ref() {
+            otel_context.reparent_tracing_span()
+        }
         let store = self.get_store(store)?;
         Ok(store.get_keys().await.map_err(track_error_on_span))
     }

crates/factor-key-value/src/lib.rs

@@ -8,6 +8,7 @@ use std::{
 };
 
 use anyhow::ensure;
+use spin_factor_otel::OtelFactorState;
 use spin_factors::{
     ConfigureAppContext, Factor, FactorData, FactorInstanceBuilder, InitContext, PrepareContext,
     RuntimeFactors,
@@ -87,17 +88,19 @@ impl Factor for KeyValueFactor {
 
     fn prepare<T: RuntimeFactors>(
         &self,
-        ctx: PrepareContext<T, Self>,
+        mut ctx: PrepareContext<T, Self>,
     ) -> anyhow::Result<InstanceBuilder> {
         let app_state = ctx.app_state();
         let allowed_stores = app_state
             .component_allowed_stores
             .get(ctx.app_component().id())
             .expect("component should be in component_stores")
             .clone();
+        let otel_context = OtelFactorState::from_prepare_context(&mut ctx)?;
         Ok(InstanceBuilder {
             store_manager: app_state.store_manager.clone(),
             allowed_stores,
+            otel_context,
         })
     }
 }
@@ -177,6 +180,7 @@ pub struct InstanceBuilder {
     store_manager: Arc<AppStoreManager>,
     /// The allowed stores for this component instance.
     allowed_stores: HashSet<String>,
+    otel_context: OtelFactorState,
 }
 
 impl FactorInstanceBuilder for InstanceBuilder {
@@ -186,11 +190,13 @@ impl FactorInstanceBuilder for InstanceBuilder {
         let Self {
             store_manager,
             allowed_stores,
+            otel_context,
         } = self;
         Ok(KeyValueDispatch::new_with_capacity(
             allowed_stores,
             store_manager,
             u32::MAX,
+            Some(otel_context),
         ))
     }
 }

crates/factor-llm/Cargo.toml

@@ -17,6 +17,7 @@ llm-cublas = ["llm", "spin-llm-local/cublas"]
 anyhow = { workspace = true }
 async-trait = { workspace = true }
 serde = { workspace = true }
+spin-factor-otel = { path = "../factor-otel" }
 spin-factors = { path = "../factors" }
 spin-llm-local = { path = "../llm-local", optional = true }
 spin-llm-remote-http = { path = "../llm-remote-http" }

crates/factor-llm/src/host.rs

@@ -13,6 +13,8 @@ impl v2::Host for InstanceState {
         prompt: String,
         params: Option<v2::InferencingParams>,
     ) -> Result<v2::InferencingResult, v2::Error> {
+        self.otel_context.reparent_tracing_span();
+
         if !self.allowed_models.contains(&model) {
             return Err(access_denied_error(&model));
         }
@@ -40,6 +42,8 @@ impl v2::Host for InstanceState {
         model: v1::EmbeddingModel,
         data: Vec<String>,
     ) -> Result<v2::EmbeddingsResult, v2::Error> {
+        self.otel_context.reparent_tracing_span();
+
         if !self.allowed_models.contains(&model) {
             return Err(access_denied_error(&model));
         }

crates/factor-llm/src/lib.rs

@@ -5,6 +5,7 @@ use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use spin_factor_otel::OtelFactorState;
 use spin_factors::{
     ConfigureAppContext, Factor, FactorData, PrepareContext, RuntimeFactors, SelfInstanceBuilder,
 };
@@ -73,7 +74,7 @@ impl Factor for LlmFactor {
 
     fn prepare<T: RuntimeFactors>(
         &self,
-        ctx: PrepareContext<T, Self>,
+        mut ctx: PrepareContext<T, Self>,
     ) -> anyhow::Result<Self::InstanceBuilder> {
         let allowed_models = ctx
             .app_state()
@@ -82,10 +83,12 @@ impl Factor for LlmFactor {
             .cloned()
             .unwrap_or_default();
         let engine = ctx.app_state().engine.clone();
+        let otel_context = OtelFactorState::from_prepare_context(&mut ctx)?;
 
         Ok(InstanceState {
             engine,
             allowed_models,
+            otel_context,
         })
     }
 }
@@ -100,6 +103,7 @@ pub struct AppState {
 pub struct InstanceState {
     engine: Arc<Mutex<dyn LlmEngine>>,
     pub allowed_models: Arc<HashSet<String>>,
+    otel_context: OtelFactorState,
 }
 
 /// The runtime configuration for the LLM factor.

crates/factor-otel/Cargo.toml

@@ -0,0 +1,25 @@
+[package]
+name = "spin-factor-otel"
+version = { workspace = true }
+authors = { workspace = true }
+edition = { workspace = true }
+
+[dependencies]
+anyhow = { workspace = true }
+indexmap = "2.2.6"
+opentelemetry = { workspace = true }
+opentelemetry_sdk = { workspace = true }
+opentelemetry-otlp = { workspace = true }
+spin-core = { path = "../core" }
+spin-factors = { path = "../factors" }
+spin-resource-table = { path = "../table" }
+spin-telemetry = { path = "../telemetry" }
+wasi-otel = { path = "../wasi-otel" }
+tracing = { workspace = true }
+tracing-opentelemetry = { workspace = true }
+
+[dev-dependencies]
+toml = "0.5"
+
+[lints]
+workspace = true

crates/factor-otel/src/host.rs

@@ -0,0 +1,88 @@
+use crate::InstanceState;
+use anyhow::anyhow;
+use anyhow::Result;
+use opentelemetry::trace::TraceContextExt;
+use opentelemetry_sdk::error::OTelSdkError;
+use opentelemetry_sdk::metrics::exporter::PushMetricExporter;
+use opentelemetry_sdk::trace::SpanProcessor;
+use tracing_opentelemetry::OpenTelemetrySpanExt;
+// TODO: This feels weird. I'm wondering if it can be fixed by wrangling the bindgen macro in `crates/wasi_otel/lib.rs`.
+use wasi_otel::wasi;
+
+impl wasi::otel::tracing::Host for InstanceState {
+    async fn on_start(&mut self, context: wasi::otel::tracing::SpanContext) -> Result<()> {
+        let mut state = self.state.write().unwrap();
+
+        // Before we do anything make sure we track the original host span ID for reparenting
+        if state.original_host_span_id.is_none() {
+            state.original_host_span_id = Some(
+                tracing::Span::current()
+                    .context()
+                    .span()
+                    .span_context()
+                    .span_id(),
+            );
+        }
+
+        // Track the guest spans context in our ordered map
+        let span_context: opentelemetry::trace::SpanContext = context.into();
+        state
+            .guest_span_contexts
+            .insert(span_context.span_id(), span_context);
+
+        Ok(())
+    }
+
+    async fn on_end(&mut self, span_data: wasi::otel::tracing::SpanData) -> Result<()> {
+        let mut state = self.state.write().unwrap();
+
+        let span_context: opentelemetry::trace::SpanContext = span_data.span_context.clone().into();
+        let span_id: opentelemetry::trace::SpanId = span_context.span_id();
+
+        if state.guest_span_contexts.shift_remove(&span_id).is_none() {
+            Err(anyhow!("Trying to end a span that was not started"))?;
+        }
+
+        self.span_processor.on_end(span_data.into());
+
+        Ok(())
+    }
+
+    async fn outer_span_context(&mut self) -> Result<wasi::otel::tracing::SpanContext> {
+        Ok(tracing::Span::current()
+            .context()
+            .span()
+            .span_context()
+            .clone()
+            .into())
+    }
+}
+
+impl wasi::otel::metrics::Host for InstanceState {
+    async fn export(
+        &mut self,
+        metrics: wasi::otel::metrics::ResourceMetrics,
+    ) -> spin_core::wasmtime::Result<std::result::Result<(), wasi::otel::metrics::Error>> {
+        let mut rm: opentelemetry_sdk::metrics::data::ResourceMetrics = metrics.into();
+        match self.metric_exporter.export(&mut rm).await {
+            Ok(_) => Ok(Ok(())),
+            Err(e) => match e {
+                OTelSdkError::AlreadyShutdown => {
+                    let msg = "Shutdown has already been invoked";
+                    tracing::error!(msg);
+                    Ok(Err(msg.to_string()))
+                }
+                OTelSdkError::InternalFailure(e) => {
+                    let detailed_msg = format!("Internal failure: {}", e);
+                    tracing::error!(detailed_msg);
+                    Ok(Err("Internal failure.".to_string()))
+                }
+                OTelSdkError::Timeout(d) => {
+                    let detailed_msg = format!("Operation timed out after {} seconds", d.as_secs());
+                    tracing::error!(detailed_msg);
+                    Ok(Err("Operation timed out.".to_string()))
+                }
+            },
+        }
+    }
+}