spin_factor_key_value/
lib.rs

1mod host;
2pub mod runtime_config;
3mod util;
4
5use std::{
6    collections::{HashMap, HashSet},
7    sync::Arc,
8};
9
10use anyhow::ensure;
11use spin_connection_semaphore::{ConnectionSemaphore, LimitedSemaphore};
12use spin_factor_otel::OtelFactorState;
13use spin_factors::{
14    ConfigureAppContext, Factor, FactorData, FactorInstanceBuilder, InitContext, PrepareContext,
15    RuntimeFactors,
16};
17use spin_locked_app::APP_NAME_KEY;
18use spin_locked_app::MetadataKey;
19
20/// Metadata key for key-value stores.
21pub const KEY_VALUE_STORES_KEY: MetadataKey<Vec<String>> = MetadataKey::new("key_value_stores");
22pub use host::to_v3_err;
23pub use host::{
24    Error, KeyValueDispatch, Store, StoreManager, log_cas_error, log_error, log_error_v3,
25};
26pub use runtime_config::RuntimeConfig;
27use spin_core::async_trait;
28pub use spin_world::spin::key_value::key_value as v3;
29pub use util::DelegatingStoreManager;
30
31/// A factor that provides key-value storage.
32#[derive(Default)]
33pub struct KeyValueFactor {
34    _priv: (),
35}
36
37impl KeyValueFactor {
38    /// Create a new KeyValueFactor.
39    pub fn new() -> Self {
40        Self { _priv: () }
41    }
42}
43
44impl Factor for KeyValueFactor {
45    type RuntimeConfig = RuntimeConfig;
46    type AppState = AppState;
47    type InstanceBuilder = InstanceBuilder;
48
49    fn init<T: InitContext<Self>>(&mut self, ctx: &mut T) -> anyhow::Result<()> {
50        ctx.link_bindings(spin_world::v1::key_value::add_to_linker::<_, FactorData<Self>>)?;
51        ctx.link_bindings(spin_world::v2::key_value::add_to_linker::<_, FactorData<Self>>)?;
52        ctx.link_bindings(
53            spin_world::spin::key_value::key_value::add_to_linker::<_, KeyValueFactorData>,
54        )?;
55        ctx.link_bindings(spin_world::wasi::keyvalue::store::add_to_linker::<_, FactorData<Self>>)?;
56        ctx.link_bindings(spin_world::wasi::keyvalue::batch::add_to_linker::<_, FactorData<Self>>)?;
57        ctx.link_bindings(
58            spin_world::wasi::keyvalue::atomics::add_to_linker::<_, FactorData<Self>>,
59        )?;
60        Ok(())
61    }
62
63    fn configure_app<T: RuntimeFactors>(
64        &self,
65        mut ctx: ConfigureAppContext<T, Self>,
66    ) -> anyhow::Result<Self::AppState> {
67        let runtime_config = ctx.take_runtime_config().unwrap_or_default();
68        let store_managers = runtime_config.clone();
69
70        let delegating_manager = DelegatingStoreManager::new(store_managers);
71        let store_manager = Arc::new(delegating_manager);
72
73        // Build component -> allowed stores map
74        let mut component_allowed_stores = HashMap::new();
75        for component in ctx.app().components() {
76            let component_id = component.id().to_string();
77            let key_value_stores = component
78                .get_metadata(KEY_VALUE_STORES_KEY)?
79                .unwrap_or_default()
80                .into_iter()
81                .collect::<HashSet<_>>();
82            for label in &key_value_stores {
83                // TODO: port nicer errors from KeyValueComponent (via error type?)
84                ensure!(
85                    store_manager.is_defined(label),
86                    "unknown key_value_stores label {label:?} for component {component_id:?}"
87                );
88            }
89            component_allowed_stores.insert(component_id, key_value_stores);
90            // TODO: warn (?) on unused store?
91        }
92
93        let app_id: Arc<str> = ctx
94            .app()
95            .get_metadata(APP_NAME_KEY)?
96            .unwrap_or_else(|| "<unnamed>".into())
97            .into();
98
99        // The global connection semaphore is not used here because the KV factor limits
100        // operations (not connections) and cannot access the underlying client through the
101        // `Store` trait abstraction.
102        let semaphore = ConnectionSemaphore::new(
103            None,
104            runtime_config
105                .max_concurrent_operations()
106                .map(LimitedSemaphore::new),
107            "key-value",
108            app_id,
109            runtime_config.wait_timeout(),
110        );
111
112        Ok(AppState {
113            store_manager,
114            component_allowed_stores,
115            semaphore,
116        })
117    }
118
119    fn prepare<T: RuntimeFactors>(
120        &self,
121        mut ctx: PrepareContext<T, Self>,
122    ) -> anyhow::Result<InstanceBuilder> {
123        let app_state = ctx.app_state();
124        let allowed_stores = app_state
125            .component_allowed_stores
126            .get(ctx.app_component().id())
127            .expect("component should be in component_stores")
128            .clone();
129        let otel = OtelFactorState::from_prepare_context(&mut ctx)?;
130        Ok(InstanceBuilder {
131            store_manager: app_state.store_manager.clone(),
132            allowed_stores,
133            semaphore: app_state.semaphore.clone(),
134            otel,
135        })
136    }
137}
138
139type AppStoreManager = DelegatingStoreManager;
140
141pub struct AppState {
142    /// The store manager for the app.
143    ///
144    /// This is a cache around a delegating store manager. For `get` requests,
145    /// first checks the cache before delegating to the underlying store
146    /// manager.
147    store_manager: Arc<AppStoreManager>,
148    /// The allowed stores for each component.
149    ///
150    /// This is a map from component ID to the set of store labels that the
151    /// component is allowed to use.
152    component_allowed_stores: HashMap<String, HashSet<String>>,
153    /// App-scoped semaphore used to limit in-flight key-value operations.
154    semaphore: ConnectionSemaphore,
155}
156
157impl AppState {
158    /// Returns the [`StoreManager::summary`] for the given store label.
159    pub fn store_summary(&self, label: &str) -> Option<String> {
160        self.store_manager.summary(label)
161    }
162
163    /// Returns true if the given store label is used by any component.
164    pub fn store_is_used(&self, label: &str) -> bool {
165        self.component_allowed_stores
166            .values()
167            .any(|stores| stores.contains(label))
168    }
169
170    /// Get a store by label.
171    pub async fn get_store(&self, label: &str) -> Option<Arc<dyn Store>> {
172        self.store_manager.get(label).await.ok()
173    }
174}
175
176/// `SwapError` are errors that occur during compare and swap operations
177#[derive(Debug, thiserror::Error)]
178pub enum SwapError {
179    #[error("{0}")]
180    CasFailed(String),
181
182    #[error("{0}")]
183    Other(String),
184}
185
186/// `Cas` trait describes the interface a key value compare and swap implementor must fulfill.
187///
188/// `current` is expected to get the current value for the key associated with the CAS operation
189/// while also starting what is needed to ensure the value to be replaced will not have mutated
190/// between the time of calling `current` and `swap`. For example, a get from a backend store
191/// may provide the caller with an etag (a version stamp), which can be used with an if-match
192/// header to ensure the version updated is the version that was read (optimistic concurrency).
193/// Rather than an etag, one could start a transaction, if supported by the backing store, which
194/// would provide atomicity.
195///
196/// `swap` is expected to replace the old value with the new value respecting the atomicity of the
197/// operation. If there was no key / value with the given key in the store, the `swap` operation
198/// should **insert** the key and value, disallowing an update.
199#[async_trait]
200pub trait Cas: Sync + Send {
201    async fn current(&self, max_result_bytes: usize) -> anyhow::Result<Option<Vec<u8>>, Error>;
202    async fn swap(&self, value: Vec<u8>) -> anyhow::Result<(), SwapError>;
203    async fn bucket_rep(&self) -> u32;
204    async fn key(&self) -> String;
205}
206
207pub struct InstanceBuilder {
208    /// The store manager for the app.
209    ///
210    /// This is a cache around a delegating store manager. For `get` requests,
211    /// first checks the cache before delegating to the underlying store
212    /// manager.
213    store_manager: Arc<AppStoreManager>,
214    /// The allowed stores for this component instance.
215    allowed_stores: HashSet<String>,
216    /// App-scoped semaphore shared by all component instances for this app.
217    semaphore: ConnectionSemaphore,
218    otel: OtelFactorState,
219}
220
221impl FactorInstanceBuilder for InstanceBuilder {
222    type InstanceState = KeyValueDispatch;
223
224    fn build(self) -> anyhow::Result<Self::InstanceState> {
225        let Self {
226            store_manager,
227            allowed_stores,
228            semaphore,
229            otel,
230        } = self;
231        Ok(KeyValueDispatch::new_with_capacity_and_semaphore(
232            allowed_stores,
233            store_manager,
234            u32::MAX,
235            semaphore,
236            otel,
237        ))
238    }
239}
240
241pub struct KeyValueFactorData(KeyValueFactor);
242
243impl spin_core::wasmtime::component::HasData for KeyValueFactorData {
244    type Data<'a> = &'a mut KeyValueDispatch;
245}
spin_factor_key_value/lib.rs

spin_factor_key_value/
lib.rs
