spin_factor_key_value/
host.rs

1use super::{Cas, SwapError};
2use anyhow::{Context, Result};
3use spin_connection_semaphore::{ConnectionPermit, ConnectionSemaphore};
4use spin_core::{
5    async_trait,
6    wasmtime::component::{Accessor, FutureReader, Resource, StreamReader},
7};
8use spin_factor_otel::OtelFactorState;
9use spin_resource_table::Table;
10use spin_telemetry::traces::{self, Blame};
11use spin_world::MAX_HOST_BUFFERED_BYTES;
12use spin_world::spin::key_value::key_value as v3;
13use spin_world::v2::key_value;
14use spin_world::wasi::keyvalue as wasi_keyvalue;
15use std::{any::Any, collections::HashSet, sync::Arc};
16use tracing::instrument;
17
18const DEFAULT_STORE_TABLE_CAPACITY: u32 = 256;
19
20pub use key_value::Error;
21
22#[async_trait]
23pub trait StoreManager: Sync + Send {
24    async fn get(&self, name: &str) -> Result<Arc<dyn Store>, Error>;
25    fn is_defined(&self, store_name: &str) -> bool;
26
27    /// A human-readable summary of the given store's configuration
28    ///
29    /// Example: "Redis at localhost:1234"
30    fn summary(&self, store_name: &str) -> Option<String> {
31        let _ = store_name;
32        None
33    }
34
35    /// Metadata about the store manager that can be accessed for data collection or debugging purposes. This is opaque and can be in any format the store manager chooses.
36    fn metadata(&self) -> Arc<dyn Any> {
37        Arc::new(())
38    }
39}
40
41#[async_trait]
42pub trait Store: Sync + Send {
43    async fn after_open(&self) -> Result<(), Error> {
44        Ok(())
45    }
46    async fn get(&self, key: &str, max_result_bytes: usize) -> Result<Option<Vec<u8>>, Error>;
47    async fn set(&self, key: &str, value: &[u8]) -> Result<(), Error>;
48    async fn delete(&self, key: &str) -> Result<(), Error>;
49    async fn exists(&self, key: &str) -> Result<bool, Error>;
50    async fn get_keys(&self, max_result_bytes: usize) -> Result<Vec<String>, Error>;
51    async fn get_keys_async(
52        &self,
53        max_result_bytes: usize,
54    ) -> (
55        tokio::sync::mpsc::Receiver<String>,
56        tokio::sync::oneshot::Receiver<Result<(), v3::Error>>,
57    );
58    async fn get_many(
59        &self,
60        keys: Vec<String>,
61        max_result_bytes: usize,
62    ) -> Result<Vec<(String, Option<Vec<u8>>)>, Error>;
63    async fn set_many(&self, key_values: Vec<(String, Vec<u8>)>) -> Result<(), Error>;
64    async fn delete_many(&self, keys: Vec<String>) -> Result<(), Error>;
65    async fn increment(&self, key: String, delta: i64) -> Result<i64, Error>;
66    async fn new_compare_and_swap(&self, bucket_rep: u32, key: &str)
67    -> Result<Arc<dyn Cas>, Error>;
68}
69
70pub struct KeyValueDispatch {
71    allowed_stores: HashSet<String>,
72    manager: Arc<dyn StoreManager>,
73    stores: Table<Arc<dyn Store>>,
74    compare_and_swaps: Table<Arc<dyn Cas>>,
75    semaphore: ConnectionSemaphore,
76    otel: OtelFactorState,
77}
78
79impl KeyValueDispatch {
80    pub fn new(
81        allowed_stores: HashSet<String>,
82        manager: Arc<dyn StoreManager>,
83        app_id: Arc<str>,
84    ) -> Self {
85        Self::new_with_capacity_and_semaphore(
86            allowed_stores,
87            manager,
88            DEFAULT_STORE_TABLE_CAPACITY,
89            ConnectionSemaphore::new(None, None, "key-value", app_id, None),
90            Default::default(),
91        )
92    }
93
94    pub fn new_with_capacity(
95        allowed_stores: HashSet<String>,
96        manager: Arc<dyn StoreManager>,
97        capacity: u32,
98        app_id: Arc<str>,
99        otel: OtelFactorState,
100    ) -> Self {
101        Self::new_with_capacity_and_semaphore(
102            allowed_stores,
103            manager,
104            capacity,
105            ConnectionSemaphore::new(None, None, "key-value", app_id, None),
106            otel,
107        )
108    }
109
110    pub fn new_with_capacity_and_semaphore(
111        allowed_stores: HashSet<String>,
112        manager: Arc<dyn StoreManager>,
113        capacity: u32,
114        semaphore: ConnectionSemaphore,
115        otel: OtelFactorState,
116    ) -> Self {
117        Self {
118            allowed_stores,
119            manager,
120            stores: Table::new(capacity),
121            compare_and_swaps: Table::new(capacity),
122            semaphore,
123            otel,
124        }
125    }
126
127    async fn acquire_permit(&self) -> std::result::Result<ConnectionPermit, Error> {
128        self.semaphore.acquire().await.map_err(|err| {
129            tracing::warn!("key-value error: {err:?}");
130            Error::Other("too many requests".into())
131        })
132    }
133
134    async fn acquire_permit_wasi(
135        &self,
136    ) -> std::result::Result<ConnectionPermit, wasi_keyvalue::store::Error> {
137        self.semaphore.acquire().await.map_err(|err| {
138            tracing::warn!("key-value error: {err:?}");
139            wasi_keyvalue::store::Error::Other("too many requests".into())
140        })
141    }
142
143    pub fn get_store<T: 'static>(&self, store: Resource<T>) -> anyhow::Result<&Arc<dyn Store>> {
144        let res = self.stores.get(store.rep()).context("invalid store");
145        if let Err(err) = &res {
146            traces::mark_as_error(err, Some(Blame::Host));
147        }
148        res
149    }
150
151    pub fn get_cas<T: 'static>(&self, cas: Resource<T>) -> Result<&Arc<dyn Cas>> {
152        self.compare_and_swaps
153            .get(cas.rep())
154            .context("invalid compare and swap")
155    }
156
157    pub fn allowed_stores(&self) -> &HashSet<String> {
158        &self.allowed_stores
159    }
160
161    pub fn get_store_wasi<T: 'static>(
162        &self,
163        store: Resource<T>,
164    ) -> Result<&Arc<dyn Store>, wasi_keyvalue::store::Error> {
165        self.stores
166            .get(store.rep())
167            .ok_or(wasi_keyvalue::store::Error::NoSuchStore)
168    }
169
170    pub fn get_cas_wasi<T: 'static>(
171        &self,
172        cas: Resource<T>,
173    ) -> Result<&Arc<dyn Cas>, wasi_keyvalue::atomics::Error> {
174        self.compare_and_swaps
175            .get(cas.rep())
176            .ok_or(wasi_keyvalue::atomics::Error::Other(
177                "compare and swap not found".to_string(),
178            ))
179    }
180
181    pub fn manager_metadata(&self) -> Arc<dyn Any> {
182        self.manager.metadata()
183    }
184}
185
186impl key_value::Host for KeyValueDispatch {}
187
188impl key_value::HostStore for KeyValueDispatch {
189    #[instrument(name = "spin_key_value.open", skip(self), err, fields(otel.kind = "client", kv.backend=self.manager.summary(&name).unwrap_or("unknown".to_string())))]
190    async fn open(&mut self, name: String) -> Result<Result<Resource<key_value::Store>, Error>> {
191        self.otel.reparent_tracing_span();
192        Ok(async {
193            if self.allowed_stores.contains(&name) {
194                let store = self.manager.get(&name).await?;
195                store.after_open().await?;
196                let store_idx = self
197                    .stores
198                    .push(store)
199                    .map_err(|()| Error::StoreTableFull)?;
200                Ok(Resource::new_own(store_idx))
201            } else {
202                Err(Error::AccessDenied)
203            }
204        }
205        .await)
206    }
207
208    #[instrument(name = "spin_key_value.get", skip_all, fields(otel.kind = "client"))]
209    async fn get(
210        &mut self,
211        store: Resource<key_value::Store>,
212        key: String,
213    ) -> Result<Result<Option<Vec<u8>>, Error>> {
214        self.otel.reparent_tracing_span();
215        let store = self.get_store(store)?;
216        let _permit = self.acquire_permit().await.map_err(track_error_on_span)?;
217        Ok(store
218            .get(&key, MAX_HOST_BUFFERED_BYTES)
219            .await
220            .map_err(track_error_on_span))
221    }
222
223    #[instrument(name = "spin_key_value.set", skip_all, fields(otel.kind = "client"))]
224    async fn set(
225        &mut self,
226        store: Resource<key_value::Store>,
227        key: String,
228        value: Vec<u8>,
229    ) -> Result<Result<(), Error>> {
230        self.otel.reparent_tracing_span();
231        let store = self.get_store(store)?;
232        let _permit = self.acquire_permit().await.map_err(track_error_on_span)?;
233        Ok(store.set(&key, &value).await.map_err(track_error_on_span))
234    }
235
236    #[instrument(name = "spin_key_value.delete", skip_all, fields(otel.kind = "client"))]
237    async fn delete(
238        &mut self,
239        store: Resource<key_value::Store>,
240        key: String,
241    ) -> Result<Result<(), Error>> {
242        self.otel.reparent_tracing_span();
243        let store = self.get_store(store)?;
244        let _permit = self.acquire_permit().await.map_err(track_error_on_span)?;
245        Ok(store.delete(&key).await.map_err(track_error_on_span))
246    }
247
248    #[instrument(name = "spin_key_value.exists", skip_all, fields(otel.kind = "client"))]
249    async fn exists(
250        &mut self,
251        store: Resource<key_value::Store>,
252        key: String,
253    ) -> Result<Result<bool, Error>> {
254        self.otel.reparent_tracing_span();
255        let store = self.get_store(store)?;
256        let _permit = self.acquire_permit().await.map_err(track_error_on_span)?;
257        Ok(store.exists(&key).await.map_err(track_error_on_span))
258    }
259
260    #[instrument(name = "spin_key_value.get_keys", skip_all, fields(otel.kind = "client"))]
261    async fn get_keys(
262        &mut self,
263        store: Resource<key_value::Store>,
264    ) -> Result<Result<Vec<String>, Error>> {
265        self.otel.reparent_tracing_span();
266        let store = self.get_store(store)?;
267        let _permit = self.acquire_permit().await.map_err(track_error_on_span)?;
268        Ok(store
269            .get_keys(MAX_HOST_BUFFERED_BYTES)
270            .await
271            .map_err(track_error_on_span))
272    }
273
274    async fn drop(&mut self, store: Resource<key_value::Store>) -> Result<()> {
275        self.stores.remove(store.rep());
276        Ok(())
277    }
278}
279
280impl spin_core::wasmtime::component::HasData for KeyValueDispatch {
281    type Data<'a> = &'a mut KeyValueDispatch;
282}
283
284impl v3::Host for KeyValueDispatch {
285    fn convert_error(&mut self, err: v3::Error) -> anyhow::Result<v3::Error> {
286        Ok(err)
287    }
288}
289
290impl v3::HostStore for KeyValueDispatch {
291    async fn drop(&mut self, store: Resource<v3::Store>) -> Result<()> {
292        self.stores.remove(store.rep());
293        Ok(())
294    }
295}
296
297impl<T> v3::HostStoreWithStore<T> for crate::KeyValueFactorData {
298    async fn open(
299        accessor: &Accessor<T, Self>,
300        label: String,
301    ) -> Result<Resource<v3::Store>, v3::Error> {
302        let (allowed, manager) = accessor.with(|mut access| {
303            let host = access.get();
304            host.otel.reparent_tracing_span();
305            (host.allowed_stores.contains(&label), host.manager.clone())
306        });
307
308        if !allowed {
309            return Err(v3::Error::AccessDenied);
310        }
311
312        let store = manager.get(&label).await.map_err(to_v3_err)?;
313        store.after_open().await.map_err(to_v3_err)?;
314
315        accessor.with(|mut access| {
316            let host = access.get();
317            host.stores
318                .push(store)
319                .map(Resource::new_own)
320                .map_err(|()| v3::Error::StoreTableFull)
321        })
322    }
323
324    async fn get(
325        accessor: &Accessor<T, Self>,
326        store: Resource<v3::Store>,
327        key: String,
328    ) -> Result<Option<Vec<u8>>, v3::Error> {
329        let (store_result, permit_fut) = accessor.with(|mut access| {
330            let host = access.get();
331            host.otel.reparent_tracing_span();
332            (host.get_store(store).cloned(), host.semaphore.clone())
333        });
334        let store = store_result.map_err(|_| v3::Error::NoSuchStore)?;
335        let _permit = acquire_permit_v3(&permit_fut)
336            .await
337            .map_err(track_error_on_span_v3)?;
338        store
339            .get(&key, MAX_HOST_BUFFERED_BYTES)
340            .await
341            .map_err(to_v3_err)
342            .map_err(track_error_on_span_v3)
343    }
344
345    async fn set(
346        accessor: &Accessor<T, Self>,
347        store: Resource<v3::Store>,
348        key: String,
349        value: Vec<u8>,
350    ) -> Result<(), v3::Error> {
351        let (store_result, semaphore) = accessor.with(|mut access| {
352            let host = access.get();
353            host.otel.reparent_tracing_span();
354            (host.get_store(store).cloned(), host.semaphore.clone())
355        });
356        let store = store_result.map_err(|_| v3::Error::NoSuchStore)?;
357        let _permit = acquire_permit_v3(&semaphore)
358            .await
359            .map_err(track_error_on_span_v3)?;
360        store
361            .set(&key, &value)
362            .await
363            .map_err(to_v3_err)
364            .map_err(track_error_on_span_v3)
365    }
366
367    async fn delete(
368        accessor: &Accessor<T, Self>,
369        store: Resource<v3::Store>,
370        key: String,
371    ) -> Result<(), v3::Error> {
372        let (store_result, semaphore) = accessor.with(|mut access| {
373            let host = access.get();
374            host.otel.reparent_tracing_span();
375            (host.get_store(store).cloned(), host.semaphore.clone())
376        });
377        let store = store_result.map_err(|_| v3::Error::NoSuchStore)?;
378        let _permit = acquire_permit_v3(&semaphore)
379            .await
380            .map_err(track_error_on_span_v3)?;
381        store
382            .delete(&key)
383            .await
384            .map_err(to_v3_err)
385            .map_err(track_error_on_span_v3)
386    }
387
388    async fn exists(
389        accessor: &Accessor<T, Self>,
390        store: Resource<v3::Store>,
391        key: String,
392    ) -> Result<bool, v3::Error> {
393        let (store_result, semaphore) = accessor.with(|mut access| {
394            let host = access.get();
395            host.otel.reparent_tracing_span();
396            (host.get_store(store).cloned(), host.semaphore.clone())
397        });
398        let store = store_result.map_err(|_| v3::Error::NoSuchStore)?;
399        let _permit = acquire_permit_v3(&semaphore)
400            .await
401            .map_err(track_error_on_span_v3)?;
402        store
403            .exists(&key)
404            .await
405            .map_err(to_v3_err)
406            .map_err(track_error_on_span_v3)
407    }
408
409    async fn get_keys(
410        accessor: &Accessor<T, Self>,
411        store: Resource<v3::Store>,
412    ) -> Result<(StreamReader<String>, FutureReader<Result<(), v3::Error>>)> {
413        let (store_result, semaphore) = accessor.with(|mut access| {
414            let host = access.get();
415            host.otel.reparent_tracing_span();
416            (host.get_store(store).cloned(), host.semaphore.clone())
417        });
418        let store = store_result.map_err(|_| v3::Error::NoSuchStore)?;
419
420        let _permit = acquire_permit_v3(&semaphore)
421            .await
422            .map_err(track_error_on_span_v3)?;
423
424        let (keys_rx, err_rx) = store.get_keys_async(MAX_HOST_BUFFERED_BYTES).await;
425
426        let producer = spin_wasi_async::stream::producer(keys_rx);
427        let (ksr, efr) = accessor.with(|mut access| {
428            let ksr = StreamReader::new(&mut access, producer)?;
429            let efr = FutureReader::new(&mut access, err_rx)?;
430            anyhow::Ok((ksr, efr))
431        })?;
432
433        Ok((ksr, efr))
434    }
435}
436
437/// Make sure that infrastructure related errors are tracked in the current span.
438fn track_error_on_span(err: Error) -> Error {
439    let blame = match &err {
440        Error::NoSuchStore | Error::AccessDenied => Blame::Guest,
441        Error::Other(msg) if msg.contains("too many requests") => Blame::Guest,
442        Error::StoreTableFull | Error::Other(_) => Blame::Host,
443    };
444    traces::mark_as_error(&err, Some(blame));
445    err
446}
447
448/// Make sure that infrastructure related errors are tracked in the current span.
449fn track_error_on_span_v3(err: v3::Error) -> v3::Error {
450    let blame = match &err {
451        v3::Error::NoSuchStore | v3::Error::AccessDenied => Blame::Guest,
452        v3::Error::Other(msg) if msg.contains("too many requests") => Blame::Guest,
453        v3::Error::StoreTableFull | v3::Error::Other(_) => Blame::Host,
454    };
455    traces::mark_as_error(&err, Some(blame));
456    err
457}
458
459/// Maps a semaphore acquisition failure to a v3 error with a consistent "too many requests"
460/// message so that `track_error_on_span_v3` correctly attributes the blame to the guest.
461async fn acquire_permit_v3(
462    semaphore: &ConnectionSemaphore,
463) -> std::result::Result<ConnectionPermit, v3::Error> {
464    semaphore.acquire().await.map_err(|err| {
465        tracing::warn!("key-value error: {err:?}");
466        v3::Error::Other("too many requests".into())
467    })
468}
469
470fn to_wasi_err(e: Error) -> wasi_keyvalue::store::Error {
471    match track_error_on_span(e) {
472        Error::AccessDenied => wasi_keyvalue::store::Error::AccessDenied,
473        Error::NoSuchStore => wasi_keyvalue::store::Error::NoSuchStore,
474        Error::StoreTableFull => wasi_keyvalue::store::Error::Other("store table full".to_string()),
475        Error::Other(msg) => wasi_keyvalue::store::Error::Other(msg),
476    }
477}
478
479pub fn to_v3_err(e: Error) -> v3::Error {
480    match track_error_on_span(e) {
481        Error::AccessDenied => v3::Error::AccessDenied,
482        Error::NoSuchStore => v3::Error::NoSuchStore,
483        Error::StoreTableFull => v3::Error::StoreTableFull,
484        Error::Other(msg) => v3::Error::Other(msg),
485    }
486}
487
488impl wasi_keyvalue::store::Host for KeyValueDispatch {
489    #[instrument(name = "wasi_key_value.open", skip_all, fields(otel.kind = "client"))]
490    async fn open(
491        &mut self,
492        identifier: String,
493    ) -> Result<Resource<wasi_keyvalue::store::Bucket>, wasi_keyvalue::store::Error> {
494        if self.allowed_stores.contains(&identifier) {
495            let store = self.manager.get(&identifier).await.map_err(to_wasi_err)?;
496            store.after_open().await.map_err(to_wasi_err)?;
497            let store_idx = self
498                .stores
499                .push(store)
500                .map_err(|()| wasi_keyvalue::store::Error::Other("store table full".to_string()))?;
501            Ok(Resource::new_own(store_idx))
502        } else {
503            Err(wasi_keyvalue::store::Error::AccessDenied)
504        }
505    }
506
507    fn convert_error(
508        &mut self,
509        error: spin_world::wasi::keyvalue::store::Error,
510    ) -> std::result::Result<spin_world::wasi::keyvalue::store::Error, anyhow::Error> {
511        Ok(error)
512    }
513}
514
515use wasi_keyvalue::store::Bucket;
516impl wasi_keyvalue::store::HostBucket for KeyValueDispatch {
517    #[instrument(name = "wasi_key_value.get", skip_all, fields(otel.kind = "client"))]
518    async fn get(
519        &mut self,
520        self_: Resource<Bucket>,
521        key: String,
522    ) -> Result<Option<Vec<u8>>, wasi_keyvalue::store::Error> {
523        let store = self.get_store_wasi(self_)?;
524        let _permit = self.acquire_permit_wasi().await?;
525        store
526            .get(&key, MAX_HOST_BUFFERED_BYTES)
527            .await
528            .map_err(to_wasi_err)
529    }
530
531    #[instrument(name = "wasi_key_value.set", skip_all, fields(otel.kind = "client"))]
532    async fn set(
533        &mut self,
534        self_: Resource<Bucket>,
535        key: String,
536        value: Vec<u8>,
537    ) -> Result<(), wasi_keyvalue::store::Error> {
538        let store = self.get_store_wasi(self_)?;
539        let _permit = self.acquire_permit_wasi().await?;
540        store.set(&key, &value).await.map_err(to_wasi_err)
541    }
542
543    #[instrument(name = "wasi_key_value.delete", skip_all, fields(otel.kind = "client"))]
544    async fn delete(
545        &mut self,
546        self_: Resource<Bucket>,
547        key: String,
548    ) -> Result<(), wasi_keyvalue::store::Error> {
549        let store = self.get_store_wasi(self_)?;
550        let _permit = self.acquire_permit_wasi().await?;
551        store.delete(&key).await.map_err(to_wasi_err)
552    }
553
554    #[instrument(name = "wasi_key_value.exists", skip_all, fields(otel.kind = "client"))]
555    async fn exists(
556        &mut self,
557        self_: Resource<Bucket>,
558        key: String,
559    ) -> Result<bool, wasi_keyvalue::store::Error> {
560        let store = self.get_store_wasi(self_)?;
561        let _permit = self.acquire_permit_wasi().await?;
562        store.exists(&key).await.map_err(to_wasi_err)
563    }
564
565    #[instrument(name = "wasi_key_value.list_keys", skip_all, fields(otel.kind = "client"))]
566    async fn list_keys(
567        &mut self,
568        self_: Resource<Bucket>,
569        cursor: Option<String>,
570    ) -> Result<wasi_keyvalue::store::KeyResponse, wasi_keyvalue::store::Error> {
571        match cursor {
572            Some(_) => Err(wasi_keyvalue::store::Error::Other(
573                "list_keys: cursor not supported".to_owned(),
574            )),
575            None => {
576                let store = self.get_store_wasi(self_)?;
577                let _permit = self.acquire_permit_wasi().await?;
578                let keys = store
579                    .get_keys(MAX_HOST_BUFFERED_BYTES)
580                    .await
581                    .map_err(to_wasi_err)?;
582                Ok(wasi_keyvalue::store::KeyResponse { keys, cursor: None })
583            }
584        }
585    }
586
587    async fn drop(&mut self, rep: Resource<Bucket>) -> anyhow::Result<()> {
588        self.stores.remove(rep.rep());
589        Ok(())
590    }
591}
592
593impl wasi_keyvalue::batch::Host for KeyValueDispatch {
594    #[instrument(name = "spin_key_value.get_many", skip_all, fields(otel.kind = "client"))]
595    #[allow(clippy::type_complexity)]
596    async fn get_many(
597        &mut self,
598        bucket: Resource<wasi_keyvalue::batch::Bucket>,
599        keys: Vec<String>,
600    ) -> std::result::Result<Vec<(String, Option<Vec<u8>>)>, wasi_keyvalue::store::Error> {
601        let store = self.get_store_wasi(bucket)?;
602        if keys.is_empty() {
603            return Ok(vec![]);
604        }
605        let _permit = self.acquire_permit_wasi().await?;
606        store
607            .get_many(keys, MAX_HOST_BUFFERED_BYTES)
608            .await
609            .map_err(to_wasi_err)
610    }
611
612    #[instrument(name = "spin_key_value.set_many", skip_all, fields(otel.kind = "client"))]
613    async fn set_many(
614        &mut self,
615        bucket: Resource<wasi_keyvalue::batch::Bucket>,
616        key_values: Vec<(String, Vec<u8>)>,
617    ) -> std::result::Result<(), wasi_keyvalue::store::Error> {
618        let store = self.get_store_wasi(bucket)?;
619        if key_values.is_empty() {
620            return Ok(());
621        }
622        let _permit = self.acquire_permit_wasi().await?;
623        store.set_many(key_values).await.map_err(to_wasi_err)
624    }
625
626    #[instrument(name = "spin_key_value.delete_many", skip_all, fields(otel.kind = "client"))]
627    async fn delete_many(
628        &mut self,
629        bucket: Resource<wasi_keyvalue::batch::Bucket>,
630        keys: Vec<String>,
631    ) -> std::result::Result<(), wasi_keyvalue::store::Error> {
632        let store = self.get_store_wasi(bucket)?;
633        if keys.is_empty() {
634            return Ok(());
635        }
636        let _permit = self.acquire_permit_wasi().await?;
637        store.delete_many(keys).await.map_err(to_wasi_err)
638    }
639}
640
641impl wasi_keyvalue::atomics::HostCas for KeyValueDispatch {
642    #[instrument(name = "wasi_key_value_cas.new", skip_all, fields(otel.kind = "client"))]
643    async fn new(
644        &mut self,
645        bucket: Resource<wasi_keyvalue::atomics::Bucket>,
646        key: String,
647    ) -> Result<Resource<wasi_keyvalue::atomics::Cas>, wasi_keyvalue::store::Error> {
648        let bucket_rep = bucket.rep();
649        let bucket: Resource<Bucket> = Resource::new_own(bucket_rep);
650        let store = self.get_store_wasi(bucket)?;
651        let cas = store
652            .new_compare_and_swap(bucket_rep, &key)
653            .await
654            .map_err(to_wasi_err)?;
655        self.compare_and_swaps
656            .push(cas)
657            .map_err(|()| {
658                spin_world::wasi::keyvalue::store::Error::Other(
659                    "too many compare_and_swaps opened".to_string(),
660                )
661            })
662            .map(Resource::new_own)
663    }
664
665    #[instrument(name = "wasi_key_value_cas.current", skip_all, fields(otel.kind = "client"))]
666    async fn current(
667        &mut self,
668        cas: Resource<wasi_keyvalue::atomics::Cas>,
669    ) -> Result<Option<Vec<u8>>, wasi_keyvalue::store::Error> {
670        let cas = self
671            .get_cas(cas)
672            .map_err(|e| wasi_keyvalue::store::Error::Other(e.to_string()))?;
673        let _permit = self.acquire_permit_wasi().await?;
674        cas.current(MAX_HOST_BUFFERED_BYTES)
675            .await
676            .map_err(to_wasi_err)
677    }
678
679    async fn drop(&mut self, rep: Resource<wasi_keyvalue::atomics::Cas>) -> Result<()> {
680        self.compare_and_swaps.remove(rep.rep());
681        Ok(())
682    }
683}
684
685impl wasi_keyvalue::atomics::Host for KeyValueDispatch {
686    fn convert_cas_error(
687        &mut self,
688        error: spin_world::wasi::keyvalue::atomics::CasError,
689    ) -> std::result::Result<spin_world::wasi::keyvalue::atomics::CasError, anyhow::Error> {
690        Ok(error)
691    }
692
693    #[instrument(name = "spin_key_value.increment", skip_all, fields(otel.kind = "client"))]
694    async fn increment(
695        &mut self,
696        bucket: Resource<wasi_keyvalue::atomics::Bucket>,
697        key: String,
698        delta: i64,
699    ) -> Result<i64, wasi_keyvalue::store::Error> {
700        let store = self.get_store_wasi(bucket)?;
701        let _permit = self.acquire_permit_wasi().await?;
702        store.increment(key, delta).await.map_err(to_wasi_err)
703    }
704
705    #[instrument(name = "spin_key_value.swap", skip_all, fields(otel.kind = "client"))]
706    async fn swap(
707        &mut self,
708        cas_res: Resource<atomics::Cas>,
709        value: Vec<u8>,
710    ) -> Result<(), CasError> {
711        let cas_rep = cas_res.rep();
712        let cas = self
713            .get_cas(Resource::<Bucket>::new_own(cas_rep))
714            .map_err(|e| CasError::StoreError(atomics::Error::Other(e.to_string())))?;
715        let _permit = self
716            .acquire_permit_wasi()
717            .await
718            .map_err(CasError::StoreError)?;
719
720        match cas.swap(value).await {
721            Ok(_) => Ok(()),
722            Err(err) => match err {
723                SwapError::CasFailed(_) => {
724                    let bucket = Resource::new_own(cas.bucket_rep().await);
725                    let new_cas = self
726                        .new(bucket, cas.key().await)
727                        .await
728                        .map_err(CasError::StoreError)?;
729                    let new_cas_rep = new_cas.rep();
730                    self.current(Resource::new_own(new_cas_rep))
731                        .await
732                        .map_err(CasError::StoreError)?;
733                    let res = Resource::new_own(new_cas_rep);
734                    Err(CasError::CasFailed(res))
735                }
736                SwapError::Other(msg) => Err(CasError::StoreError(atomics::Error::Other(msg))),
737            },
738        }
739    }
740}
741
742pub fn log_error(err: impl std::fmt::Debug) -> Error {
743    tracing::warn!("key-value error: {err:?}");
744    Error::Other(format!("{err:?}"))
745}
746
747pub fn log_error_v3(err: impl std::fmt::Debug) -> v3::Error {
748    tracing::warn!("key-value error: {err:?}");
749    v3::Error::Other(format!("{err:?}"))
750}
751
752pub fn log_cas_error(err: impl std::fmt::Debug) -> SwapError {
753    tracing::warn!("key-value error: {err:?}");
754    SwapError::Other(format!("{err:?}"))
755}
756
757use spin_world::v1::key_value::Error as LegacyError;
758use spin_world::wasi::keyvalue::atomics;
759use spin_world::wasi::keyvalue::atomics::{CasError, HostCas};
760
761fn to_legacy_error(err: Error) -> LegacyError {
762    match err {
763        Error::StoreTableFull => LegacyError::StoreTableFull,
764        Error::NoSuchStore => LegacyError::NoSuchStore,
765        Error::AccessDenied => LegacyError::AccessDenied,
766        Error::Other(s) => LegacyError::Io(s),
767    }
768}
769
770impl spin_world::v1::key_value::Host for KeyValueDispatch {
771    async fn open(&mut self, name: String) -> Result<Result<u32, LegacyError>> {
772        let result = <Self as key_value::HostStore>::open(self, name).await?;
773        Ok(result.map_err(to_legacy_error).map(|s| s.rep()))
774    }
775
776    async fn get(&mut self, store: u32, key: String) -> Result<Result<Vec<u8>, LegacyError>> {
777        let this = Resource::new_borrow(store);
778        let result = <Self as key_value::HostStore>::get(self, this, key).await?;
779        Ok(result
780            .map_err(to_legacy_error)
781            .and_then(|v| v.ok_or(LegacyError::NoSuchKey)))
782    }
783
784    async fn set(
785        &mut self,
786        store: u32,
787        key: String,
788        value: Vec<u8>,
789    ) -> Result<Result<(), LegacyError>> {
790        let this = Resource::new_borrow(store);
791        let result = <Self as key_value::HostStore>::set(self, this, key, value).await?;
792        Ok(result.map_err(to_legacy_error))
793    }
794
795    async fn delete(&mut self, store: u32, key: String) -> Result<Result<(), LegacyError>> {
796        let this = Resource::new_borrow(store);
797        let result = <Self as key_value::HostStore>::delete(self, this, key).await?;
798        Ok(result.map_err(to_legacy_error))
799    }
800
801    async fn exists(&mut self, store: u32, key: String) -> Result<Result<bool, LegacyError>> {
802        let this = Resource::new_borrow(store);
803        let result = <Self as key_value::HostStore>::exists(self, this, key).await?;
804        Ok(result.map_err(to_legacy_error))
805    }
806
807    async fn get_keys(&mut self, store: u32) -> Result<Result<Vec<String>, LegacyError>> {
808        let this = Resource::new_borrow(store);
809        let result = <Self as key_value::HostStore>::get_keys(self, this).await?;
810        Ok(result.map_err(to_legacy_error))
811    }
812
813    async fn close(&mut self, store: u32) -> Result<()> {
814        let this = Resource::new_borrow(store);
815        <Self as key_value::HostStore>::drop(self, this).await
816    }
817}
spin_factor_key_value/host.rs

spin_factor_key_value/
host.rs
