Pipe observe factor through all the other factors

Signed-off-by: Caleb Schoepp <[email protected]>

Author: calebschoepp

Cargo.lock

@@ -9,6 +9,7 @@ anyhow = "1.0"
 lru = "0.9.0"
 serde = { version = "1.0", features = ["rc"] }
 spin-core = { path = "../core" }
+spin-factor-observe = { path = "../factor-observe" }
 spin-factors = { path = "../factors" }
 spin-locked-app = { path = "../locked-app" }
 spin-world = { path = "../world" }

crates/factor-key-value/Cargo.toml

@@ -1,6 +1,7 @@
 use crate::util::EmptyStoreManager;
 use anyhow::{Context, Result};
 use spin_core::{async_trait, wasmtime::component::Resource};
+use spin_factor_observe::ObserveContext;
 use spin_locked_app::MetadataKey;
 use spin_world::v2::key_value;
 use std::{collections::HashSet, sync::Arc};
@@ -40,6 +41,7 @@ pub struct KeyValueDispatch {
     allowed_stores: HashSet<String>,
     manager: Arc<dyn StoreManager>,
     stores: Table<Arc<dyn Store>>,
+    observe_context: Option<ObserveContext>,
 }
 
 impl KeyValueDispatch {
@@ -52,12 +54,19 @@ impl KeyValueDispatch {
             allowed_stores: HashSet::new(),
             manager: Arc::new(EmptyStoreManager),
             stores: Table::new(capacity),
+            observe_context: None,
         }
     }
 
-    pub fn init(&mut self, allowed_stores: HashSet<String>, manager: Arc<dyn StoreManager>) {
+    pub fn init(
+        &mut self,
+        allowed_stores: HashSet<String>,
+        manager: Arc<dyn StoreManager>,
+        observe_context: Option<ObserveContext>,
+    ) {
         self.allowed_stores = allowed_stores;
         self.manager = manager;
+        self.observe_context = observe_context;
     }
 
     pub fn get_store(&self, store: Resource<key_value::Store>) -> anyhow::Result<&Arc<dyn Store>> {
@@ -82,6 +91,9 @@ impl key_value::Host for KeyValueDispatch {}
 impl key_value::HostStore for KeyValueDispatch {
     #[instrument(name = "spin_key_value.open", skip(self), err(level = Level::INFO), 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(observe_context) = self.observe_context.as_ref() {
+            observe_context.reparent_tracing_span()
+        }
         Ok(async {
             if self.allowed_stores.contains(&name) {
                 let store = self
@@ -102,6 +114,9 @@ impl key_value::HostStore for KeyValueDispatch {
         store: Resource<key_value::Store>,
         key: String,
     ) -> Result<Result<Option<Vec<u8>>, Error>> {
+        if let Some(observe_context) = self.observe_context.as_ref() {
+            observe_context.reparent_tracing_span()
+        }
         let store = self.get_store(store)?;
         Ok(store.get(&key).await)
     }
@@ -113,6 +128,9 @@ impl key_value::HostStore for KeyValueDispatch {
         key: String,
         value: Vec<u8>,
     ) -> Result<Result<(), Error>> {
+        if let Some(observe_context) = self.observe_context.as_ref() {
+            observe_context.reparent_tracing_span()
+        }
         let store = self.get_store(store)?;
         Ok(store.set(&key, &value).await)
     }
@@ -123,6 +141,9 @@ impl key_value::HostStore for KeyValueDispatch {
         store: Resource<key_value::Store>,
         key: String,
     ) -> Result<Result<(), Error>> {
+        if let Some(observe_context) = self.observe_context.as_ref() {
+            observe_context.reparent_tracing_span()
+        }
         let store = self.get_store(store)?;
         Ok(store.delete(&key).await)
     }
@@ -133,6 +154,9 @@ impl key_value::HostStore for KeyValueDispatch {
         store: Resource<key_value::Store>,
         key: String,
     ) -> Result<Result<bool, Error>> {
+        if let Some(observe_context) = self.observe_context.as_ref() {
+            observe_context.reparent_tracing_span()
+        }
         let store = self.get_store(store)?;
         Ok(store.exists(&key).await)
     }
@@ -142,6 +166,9 @@ impl key_value::HostStore for KeyValueDispatch {
         &mut self,
         store: Resource<key_value::Store>,
     ) -> Result<Result<Vec<String>, Error>> {
+        if let Some(observe_context) = self.observe_context.as_ref() {
+            observe_context.reparent_tracing_span()
+        }
         let store = self.get_store(store)?;
         Ok(store.get_keys().await)
     }

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

@@ -9,6 +9,7 @@ use std::{
 
 use anyhow::ensure;
 use host::KEY_VALUE_STORES_KEY;
+use spin_factor_observe::{ObserveContext, ObserveFactor};
 use spin_factors::{
     ConfigureAppContext, Factor, FactorInstanceBuilder, InitContext, PrepareContext, RuntimeFactors,
 };
@@ -88,17 +89,23 @@ 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 observe_context = match ctx.instance_builder::<ObserveFactor>() {
+            Ok(factor) => Some(factor.get_observe_context()),
+            Err(spin_factors::Error::NoSuchFactor(_)) => None,
+            Err(e) => return Err(e.into()),
+        };
         Ok(InstanceBuilder {
             store_manager: app_state.store_manager.clone(),
             allowed_stores,
+            observe_context,
         })
     }
 }
@@ -147,6 +154,7 @@ pub struct InstanceBuilder {
     store_manager: Arc<AppStoreManager>,
     /// The allowed stores for this component instance.
     allowed_stores: HashSet<String>,
+    observe_context: Option<ObserveContext>,
 }
 
 impl FactorInstanceBuilder for InstanceBuilder {
@@ -156,9 +164,10 @@ impl FactorInstanceBuilder for InstanceBuilder {
         let Self {
             store_manager,
             allowed_stores,
+            observe_context,
         } = self;
         let mut dispatch = KeyValueDispatch::new_with_capacity(u32::MAX);
-        dispatch.init(allowed_stores, store_manager);
+        dispatch.init(allowed_stores, store_manager, observe_context);
         Ok(dispatch)
     }
 }

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

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

crates/factor-llm/Cargo.toml

@@ -15,6 +15,10 @@ impl v2::Host for InstanceState {
         prompt: String,
         params: Option<v2::InferencingParams>,
     ) -> Result<v2::InferencingResult, v2::Error> {
+        if let Some(observe_context) = self.observe_context.as_ref() {
+            observe_context.reparent_tracing_span()
+        }
+
         if !self.allowed_models.contains(&model) {
             return Err(access_denied_error(&model));
         }
@@ -42,6 +46,10 @@ impl v2::Host for InstanceState {
         model: v1::EmbeddingModel,
         data: Vec<String>,
     ) -> Result<v2::EmbeddingsResult, v2::Error> {
+        if let Some(observe_context) = self.observe_context.as_ref() {
+            observe_context.reparent_tracing_span()
+        }
+
         if !self.allowed_models.contains(&model) {
             return Err(access_denied_error(&model));
         }

crates/factor-llm/src/host.rs

@@ -5,6 +5,7 @@ use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use spin_factor_observe::{ObserveContext, ObserveFactor};
 use spin_factors::{
     ConfigureAppContext, Factor, PrepareContext, RuntimeFactors, SelfInstanceBuilder,
 };
@@ -76,7 +77,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()
@@ -85,10 +86,16 @@ impl Factor for LlmFactor {
             .cloned()
             .unwrap_or_default();
         let engine = ctx.app_state().engine.clone();
+        let observe_context = match ctx.instance_builder::<ObserveFactor>() {
+            Ok(factor) => Some(factor.get_observe_context()),
+            Err(spin_factors::Error::NoSuchFactor(_)) => None,
+            Err(e) => return Err(e.into()),
+        };
 
         Ok(InstanceState {
             engine,
             allowed_models,
+            observe_context,
         })
     }
 }
@@ -103,6 +110,7 @@ pub struct AppState {
 pub struct InstanceState {
     engine: Arc<Mutex<dyn LlmEngine>>,
     pub allowed_models: Arc<HashSet<String>>,
+    observe_context: Option<ObserveContext>,
 }
 
 /// The runtime configuration for the LLM factor.

crates/factor-llm/src/lib.rs

@@ -76,9 +76,11 @@ impl Factor for OutboundHttpFactor {
         let outbound_networking = ctx.instance_builder::<OutboundNetworkingFactor>()?;
         let allowed_hosts = outbound_networking.allowed_hosts();
         let component_tls_configs = outbound_networking.component_tls_configs().clone();
-        let observe_context = ctx
-            .instance_builder::<ObserveFactor>()?
-            .get_observe_context();
+        let observe_context = match ctx.instance_builder::<ObserveFactor>() {
+            Ok(factor) => Some(factor.get_observe_context()),
+            Err(spin_factors::Error::NoSuchFactor(_)) => None,
+            Err(e) => return Err(e.into()),
+        };
         Ok(InstanceState {
             wasi_http_ctx: WasiHttpCtx::new(),
             allowed_hosts,
@@ -101,7 +103,7 @@ pub struct InstanceState {
     request_interceptor: Option<Arc<dyn OutboundHttpInterceptor>>,
     // Connection-pooling client for 'fermyon:spin/http' interface
     spin_http_client: Option<reqwest::Client>,
-    observe_context: ObserveContext,
+    observe_context: Option<ObserveContext>,
 }
 
 impl InstanceState {

crates/factor-outbound-http/src/lib.rs

@@ -16,7 +16,9 @@ impl spin_http::Host for crate::InstanceState {
         fields(otel.kind = "client", url.full = Empty, http.request.method = Empty,
         http.response.status_code = Empty, otel.name = Empty, server.address = Empty, server.port = Empty))]
     async fn send_request(&mut self, req: Request) -> Result<Response, HttpError> {
-        self.observe_context.reparent_tracing_span();
+        if let Some(observe_context) = self.observe_context.as_ref() {
+            observe_context.reparent_tracing_span()
+        }
 
         let span = Span::current();
         record_request_fields(&span, &req);

crates/factor-outbound-http/src/spin.rs

@@ -87,7 +87,9 @@ impl<'a> WasiHttpView for WasiHttpImplInner<'a> {
         request: Request<wasmtime_wasi_http::body::HyperOutgoingBody>,
         config: wasmtime_wasi_http::types::OutgoingRequestConfig,
     ) -> wasmtime_wasi_http::HttpResult<wasmtime_wasi_http::types::HostFutureIncomingResponse> {
-        self.state.observe_context.reparent_tracing_span();
+        if let Some(observe_context) = self.state.observe_context.as_ref() {
+            observe_context.reparent_tracing_span()
+        }
         Ok(HostFutureIncomingResponse::Pending(
             wasmtime_wasi::runtime::spawn(
                 send_request_impl(