Skip to content

Commit ce91f58

Browse files
committed
feat(wasi-observe): A WASI Observe host component
Signed-off-by: Caleb Schoepp <caleb.schoepp@fermyon.com>
1 parent 24b8f0a commit ce91f58

File tree

52 files changed

+2098
-134
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

52 files changed

+2098
-134
lines changed

Cargo.lock

Lines changed: 161 additions & 35 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ openssl = { version = "0.10" }
7979
anyhow = { workspace = true, features = ["backtrace"] }
8080
conformance = { path = "tests/conformance-tests" }
8181
conformance-tests = { workspace = true }
82+
fake-opentelemetry-collector = "0.21.1"
8283
hex = "0.4"
8384
http-body-util = { workspace = true }
8485
hyper = { workspace = true }

crates/factor-key-value/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ edition = { workspace = true }
88
anyhow = { workspace = true }
99
serde = { workspace = true }
1010
spin-core = { path = "../core" }
11+
spin-factor-observe = { path = "../factor-observe" }
1112
spin-factors = { path = "../factors" }
1213
spin-locked-app = { path = "../locked-app" }
1314
spin-resource-table = { path = "../table" }

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

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use super::{Cas, SwapError};
22
use anyhow::{Context, Result};
33
use spin_core::{async_trait, wasmtime::component::Resource};
4+
use spin_factor_observe::ObserveContext;
45
use spin_resource_table::Table;
56
use spin_world::v2::key_value;
67
use spin_world::wasi::keyvalue as wasi_keyvalue;
@@ -48,23 +49,26 @@ pub struct KeyValueDispatch {
4849
manager: Arc<dyn StoreManager>,
4950
stores: Table<Arc<dyn Store>>,
5051
compare_and_swaps: Table<Arc<dyn Cas>>,
52+
observe_context: Option<ObserveContext>,
5153
}
5254

5355
impl KeyValueDispatch {
5456
pub fn new(allowed_stores: HashSet<String>, manager: Arc<dyn StoreManager>) -> Self {
55-
Self::new_with_capacity(allowed_stores, manager, DEFAULT_STORE_TABLE_CAPACITY)
57+
Self::new_with_capacity(allowed_stores, manager, DEFAULT_STORE_TABLE_CAPACITY, None)
5658
}
5759

5860
pub fn new_with_capacity(
5961
allowed_stores: HashSet<String>,
6062
manager: Arc<dyn StoreManager>,
6163
capacity: u32,
64+
observe_context: Option<ObserveContext>,
6265
) -> Self {
6366
Self {
6467
allowed_stores,
6568
manager,
6669
stores: Table::new(capacity),
6770
compare_and_swaps: Table::new(capacity),
71+
observe_context,
6872
}
6973
}
7074

@@ -110,6 +114,9 @@ impl key_value::Host for KeyValueDispatch {}
110114
impl key_value::HostStore for KeyValueDispatch {
111115
#[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())))]
112116
async fn open(&mut self, name: String) -> Result<Result<Resource<key_value::Store>, Error>> {
117+
if let Some(observe_context) = self.observe_context.as_ref() {
118+
observe_context.reparent_tracing_span()
119+
}
113120
Ok(async {
114121
if self.allowed_stores.contains(&name) {
115122
let store = self.manager.get(&name).await?;
@@ -132,6 +139,9 @@ impl key_value::HostStore for KeyValueDispatch {
132139
store: Resource<key_value::Store>,
133140
key: String,
134141
) -> Result<Result<Option<Vec<u8>>, Error>> {
142+
if let Some(observe_context) = self.observe_context.as_ref() {
143+
observe_context.reparent_tracing_span()
144+
}
135145
let store = self.get_store(store)?;
136146
Ok(store.get(&key).await)
137147
}
@@ -143,6 +153,9 @@ impl key_value::HostStore for KeyValueDispatch {
143153
key: String,
144154
value: Vec<u8>,
145155
) -> Result<Result<(), Error>> {
156+
if let Some(observe_context) = self.observe_context.as_ref() {
157+
observe_context.reparent_tracing_span()
158+
}
146159
let store = self.get_store(store)?;
147160
Ok(store.set(&key, &value).await)
148161
}
@@ -153,6 +166,9 @@ impl key_value::HostStore for KeyValueDispatch {
153166
store: Resource<key_value::Store>,
154167
key: String,
155168
) -> Result<Result<(), Error>> {
169+
if let Some(observe_context) = self.observe_context.as_ref() {
170+
observe_context.reparent_tracing_span()
171+
}
156172
let store = self.get_store(store)?;
157173
Ok(store.delete(&key).await)
158174
}
@@ -163,6 +179,9 @@ impl key_value::HostStore for KeyValueDispatch {
163179
store: Resource<key_value::Store>,
164180
key: String,
165181
) -> Result<Result<bool, Error>> {
182+
if let Some(observe_context) = self.observe_context.as_ref() {
183+
observe_context.reparent_tracing_span()
184+
}
166185
let store = self.get_store(store)?;
167186
Ok(store.exists(&key).await)
168187
}
@@ -172,6 +191,9 @@ impl key_value::HostStore for KeyValueDispatch {
172191
&mut self,
173192
store: Resource<key_value::Store>,
174193
) -> Result<Result<Vec<String>, Error>> {
194+
if let Some(observe_context) = self.observe_context.as_ref() {
195+
observe_context.reparent_tracing_span()
196+
}
175197
let store = self.get_store(store)?;
176198
Ok(store.get_keys().await)
177199
}

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use std::{
88
};
99

1010
use anyhow::ensure;
11+
use spin_factor_observe::ObserveContext;
1112
use spin_factors::{
1213
ConfigureAppContext, Factor, FactorInstanceBuilder, InitContext, PrepareContext, RuntimeFactors,
1314
};
@@ -84,17 +85,19 @@ impl Factor for KeyValueFactor {
8485

8586
fn prepare<T: RuntimeFactors>(
8687
&self,
87-
ctx: PrepareContext<T, Self>,
88+
mut ctx: PrepareContext<T, Self>,
8889
) -> anyhow::Result<InstanceBuilder> {
8990
let app_state = ctx.app_state();
9091
let allowed_stores = app_state
9192
.component_allowed_stores
9293
.get(ctx.app_component().id())
9394
.expect("component should be in component_stores")
9495
.clone();
96+
let observe_context = ObserveContext::from_prepare_context(&mut ctx)?;
9597
Ok(InstanceBuilder {
9698
store_manager: app_state.store_manager.clone(),
9799
allowed_stores,
100+
observe_context,
98101
})
99102
}
100103
}
@@ -174,6 +177,7 @@ pub struct InstanceBuilder {
174177
store_manager: Arc<AppStoreManager>,
175178
/// The allowed stores for this component instance.
176179
allowed_stores: HashSet<String>,
180+
observe_context: ObserveContext,
177181
}
178182

179183
impl FactorInstanceBuilder for InstanceBuilder {
@@ -183,11 +187,13 @@ impl FactorInstanceBuilder for InstanceBuilder {
183187
let Self {
184188
store_manager,
185189
allowed_stores,
190+
observe_context,
186191
} = self;
187192
Ok(KeyValueDispatch::new_with_capacity(
188193
allowed_stores,
189194
store_manager,
190195
u32::MAX,
196+
Some(observe_context),
191197
))
192198
}
193199
}

crates/factor-llm/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ llm-cublas = ["llm", "spin-llm-local/cublas"]
1717
anyhow = { workspace = true }
1818
async-trait = { workspace = true }
1919
serde = { workspace = true }
20+
spin-factor-observe = { path = "../factor-observe" }
2021
spin-factors = { path = "../factors" }
2122
spin-llm-local = { path = "../llm-local", optional = true }
2223
spin-llm-remote-http = { path = "../llm-remote-http" }

crates/factor-llm/src/host.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@ impl v2::Host for InstanceState {
1515
prompt: String,
1616
params: Option<v2::InferencingParams>,
1717
) -> Result<v2::InferencingResult, v2::Error> {
18+
self.observe_context.reparent_tracing_span();
19+
1820
if !self.allowed_models.contains(&model) {
1921
return Err(access_denied_error(&model));
2022
}
@@ -42,6 +44,8 @@ impl v2::Host for InstanceState {
4244
model: v1::EmbeddingModel,
4345
data: Vec<String>,
4446
) -> Result<v2::EmbeddingsResult, v2::Error> {
47+
self.observe_context.reparent_tracing_span();
48+
4549
if !self.allowed_models.contains(&model) {
4650
return Err(access_denied_error(&model));
4751
}

crates/factor-llm/src/lib.rs

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use std::collections::{HashMap, HashSet};
55
use std::sync::Arc;
66

77
use async_trait::async_trait;
8+
use spin_factor_observe::ObserveContext;
89
use spin_factors::{
910
ConfigureAppContext, Factor, PrepareContext, RuntimeFactors, SelfInstanceBuilder,
1011
};
@@ -76,7 +77,7 @@ impl Factor for LlmFactor {
7677

7778
fn prepare<T: RuntimeFactors>(
7879
&self,
79-
ctx: PrepareContext<T, Self>,
80+
mut ctx: PrepareContext<T, Self>,
8081
) -> anyhow::Result<Self::InstanceBuilder> {
8182
let allowed_models = ctx
8283
.app_state()
@@ -85,10 +86,12 @@ impl Factor for LlmFactor {
8586
.cloned()
8687
.unwrap_or_default();
8788
let engine = ctx.app_state().engine.clone();
89+
let observe_context = ObserveContext::from_prepare_context(&mut ctx)?;
8890

8991
Ok(InstanceState {
9092
engine,
9193
allowed_models,
94+
observe_context,
9295
})
9396
}
9497
}
@@ -103,6 +106,7 @@ pub struct AppState {
103106
pub struct InstanceState {
104107
engine: Arc<Mutex<dyn LlmEngine>>,
105108
pub allowed_models: Arc<HashSet<String>>,
109+
observe_context: ObserveContext,
106110
}
107111

108112
/// The runtime configuration for the LLM factor.

crates/factor-observe/Cargo.toml

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
[package]
2+
name = "spin-factor-observe"
3+
version = { workspace = true }
4+
authors = { workspace = true }
5+
edition = { workspace = true }
6+
7+
[dependencies]
8+
anyhow = { workspace = true }
9+
indexmap = "2.2.6"
10+
opentelemetry = { workspace = true }
11+
opentelemetry_sdk = { workspace = true }
12+
spin-core = { path = "../core" }
13+
spin-factors = { path = "../factors" }
14+
spin-resource-table = { path = "../table" }
15+
spin-world = { path = "../world" }
16+
tracing = { workspace = true }
17+
tracing-opentelemetry = { workspace = true }
18+
19+
[dev-dependencies]
20+
toml = "0.5"
21+
22+
[lints]
23+
workspace = true

0 commit comments

Comments
 (0)