Merge pull request #3423 from itowlson/kv-sqlite-async-who-will-rid-me-of-this-turbulent-git

Key-value and SQLite async

Author: itowlson

Cargo.lock

@@ -13,6 +13,7 @@ spin-factors = { path = "../factors" }
 spin-locked-app = { path = "../locked-app" }
 spin-resource-table = { path = "../table" }
 spin-telemetry = { path = "../telemetry" }
+spin-wasi-async = { path = "../wasi-async" }
 spin-world = { path = "../world" }
 thiserror = { workspace = true }
 tokio = { workspace = true, features = ["macros", "sync", "rt"] }

crates/factor-key-value/Cargo.toml

@@ -1,9 +1,13 @@
 use super::{Cas, SwapError};
 use anyhow::{Context, Result};
-use spin_core::{async_trait, wasmtime::component::Resource};
+use spin_core::{
+    async_trait,
+    wasmtime::component::{Accessor, FutureReader, Resource, StreamReader},
+};
 use spin_factor_otel::OtelFactorState;
 use spin_resource_table::Table;
 use spin_telemetry::traces::{self, Blame};
+use spin_world::spin::key_value::key_value as v3;
 use spin_world::v2::key_value;
 use spin_world::wasi::keyvalue as wasi_keyvalue;
 use spin_world::MAX_HOST_BUFFERED_BYTES;
@@ -43,6 +47,13 @@ pub trait Store: Sync + Send {
     async fn delete(&self, key: &str) -> Result<(), Error>;
     async fn exists(&self, key: &str) -> Result<bool, Error>;
     async fn get_keys(&self, max_result_bytes: usize) -> Result<Vec<String>, Error>;
+    async fn get_keys_async(
+        &self,
+        max_result_bytes: usize,
+    ) -> (
+        tokio::sync::mpsc::Receiver<String>,
+        tokio::sync::oneshot::Receiver<Result<(), v3::Error>>,
+    );
     async fn get_many(
         &self,
         keys: Vec<String>,
@@ -220,6 +231,140 @@ impl key_value::HostStore for KeyValueDispatch {
     }
 }
 
+impl spin_core::wasmtime::component::HasData for KeyValueDispatch {
+    type Data<'a> = &'a mut KeyValueDispatch;
+}
+
+impl v3::Host for KeyValueDispatch {}
+
+impl v3::HostStore for KeyValueDispatch {
+    async fn drop(&mut self, store: Resource<v3::Store>) -> Result<()> {
+        self.stores.remove(store.rep());
+        Ok(())
+    }
+}
+
+impl v3::HostStoreWithStore for crate::KeyValueFactorData {
+    async fn open<T>(
+        accessor: &Accessor<T, Self>,
+        label: String,
+    ) -> Result<Result<Resource<v3::Store>, v3::Error>> {
+        let (allowed, manager) = accessor.with(|mut access| {
+            let host = access.get();
+            host.otel.reparent_tracing_span();
+            (host.allowed_stores.contains(&label), host.manager.clone())
+        });
+
+        if !allowed {
+            return Ok(Err(v3::Error::AccessDenied));
+        }
+
+        let store = manager.get(&label).await?;
+        store.after_open().await?;
+
+        let rsrc = accessor.with(|mut access| {
+            let host = access.get();
+            host.stores
+                .push(store)
+                .map(Resource::new_own)
+                .map_err(|()| v3::Error::StoreTableFull)
+        });
+
+        Ok(rsrc)
+    }
+
+    async fn get<T>(
+        accessor: &Accessor<T, Self>,
+        store: Resource<v3::Store>,
+        key: String,
+    ) -> Result<Result<Option<Vec<u8>>, v3::Error>> {
+        let store = accessor.with(|mut access| {
+            let host = access.get();
+            host.otel.reparent_tracing_span();
+            host.get_store(store).cloned()
+        })?;
+        Ok(store
+            .get(&key, MAX_HOST_BUFFERED_BYTES)
+            .await
+            .map_err(to_v3_err)
+            .map_err(track_error_on_span_v3))
+    }
+
+    async fn set<T>(
+        accessor: &Accessor<T, Self>,
+        store: Resource<v3::Store>,
+        key: String,
+        value: Vec<u8>,
+    ) -> Result<Result<(), v3::Error>> {
+        let store = accessor.with(|mut access| {
+            let host = access.get();
+            host.otel.reparent_tracing_span();
+            host.get_store(store).cloned()
+        })?;
+        Ok(store
+            .set(&key, &value)
+            .await
+            .map_err(to_v3_err)
+            .map_err(track_error_on_span_v3))
+    }
+
+    async fn delete<T>(
+        accessor: &Accessor<T, Self>,
+        store: Resource<v3::Store>,
+        key: String,
+    ) -> Result<Result<(), v3::Error>> {
+        let store = accessor.with(|mut access| {
+            let host = access.get();
+            host.otel.reparent_tracing_span();
+            host.get_store(store).cloned()
+        })?;
+        Ok(store
+            .delete(&key)
+            .await
+            .map_err(to_v3_err)
+            .map_err(track_error_on_span_v3))
+    }
+
+    async fn exists<T>(
+        accessor: &Accessor<T, Self>,
+        store: Resource<v3::Store>,
+        key: String,
+    ) -> Result<Result<bool, v3::Error>> {
+        let store = accessor.with(|mut access| {
+            let host = access.get();
+            host.otel.reparent_tracing_span();
+            host.get_store(store).cloned()
+        })?;
+        Ok(store
+            .exists(&key)
+            .await
+            .map_err(to_v3_err)
+            .map_err(track_error_on_span_v3))
+    }
+
+    async fn get_keys<T>(
+        accessor: &Accessor<T, Self>,
+        store: Resource<v3::Store>,
+    ) -> Result<(StreamReader<String>, FutureReader<Result<(), v3::Error>>)> {
+        let store = accessor.with(|mut access| {
+            let host = access.get();
+            host.otel.reparent_tracing_span();
+            host.get_store(store).cloned()
+        })?;
+
+        let (keys_rx, err_rx) = store.get_keys_async(MAX_HOST_BUFFERED_BYTES).await;
+
+        let producer = spin_wasi_async::stream::producer(keys_rx);
+        let (ksr, efr) = accessor.with(|mut access| {
+            let ksr = StreamReader::new(&mut access, producer);
+            let efr = FutureReader::new(&mut access, err_rx);
+            (ksr, efr)
+        });
+
+        Ok((ksr, efr))
+    }
+}
+
 /// Make sure that infrastructure related errors are tracked in the current span.
 fn track_error_on_span(err: Error) -> Error {
     let blame = match err {
@@ -230,6 +375,16 @@ fn track_error_on_span(err: Error) -> Error {
     err
 }
 
+/// Make sure that infrastructure related errors are tracked in the current span.
+fn track_error_on_span_v3(err: v3::Error) -> v3::Error {
+    let blame = match err {
+        v3::Error::NoSuchStore | v3::Error::AccessDenied => Blame::Guest,
+        v3::Error::StoreTableFull | v3::Error::Other(_) => Blame::Host,
+    };
+    traces::mark_as_error(&err, Some(blame));
+    err
+}
+
 fn to_wasi_err(e: Error) -> wasi_keyvalue::store::Error {
     match track_error_on_span(e) {
         Error::AccessDenied => wasi_keyvalue::store::Error::AccessDenied,
@@ -239,6 +394,15 @@ fn to_wasi_err(e: Error) -> wasi_keyvalue::store::Error {
     }
 }
 
+pub fn to_v3_err(e: Error) -> v3::Error {
+    match track_error_on_span(e) {
+        Error::AccessDenied => v3::Error::AccessDenied,
+        Error::NoSuchStore => v3::Error::NoSuchStore,
+        Error::StoreTableFull => v3::Error::StoreTableFull,
+        Error::Other(msg) => v3::Error::Other(msg),
+    }
+}
+
 impl wasi_keyvalue::store::Host for KeyValueDispatch {
     #[instrument(name = "wasi_key_value.open", skip_all, fields(otel.kind = "client"))]
     async fn open(
@@ -484,6 +648,11 @@ pub fn log_error(err: impl std::fmt::Debug) -> Error {
     Error::Other(format!("{err:?}"))
 }
 
+pub fn log_error_v3(err: impl std::fmt::Debug) -> v3::Error {
+    tracing::warn!("key-value error: {err:?}");
+    v3::Error::Other(format!("{err:?}"))
+}
+
 pub fn log_cas_error(err: impl std::fmt::Debug) -> SwapError {
     tracing::warn!("key-value error: {err:?}");
     SwapError::Other(format!("{err:?}"))

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

@@ -17,9 +17,13 @@ use spin_locked_app::MetadataKey;
 
 /// Metadata key for key-value stores.
 pub const KEY_VALUE_STORES_KEY: MetadataKey<Vec<String>> = MetadataKey::new("key_value_stores");
-pub use host::{log_cas_error, log_error, Error, KeyValueDispatch, Store, StoreManager};
+pub use host::to_v3_err;
+pub use host::{
+    log_cas_error, log_error, log_error_v3, Error, KeyValueDispatch, Store, StoreManager,
+};
 pub use runtime_config::RuntimeConfig;
 use spin_core::async_trait;
+pub use spin_world::spin::key_value::key_value as v3;
 pub use util::DelegatingStoreManager;
 
 /// A factor that provides key-value storage.
@@ -43,6 +47,9 @@ impl Factor for KeyValueFactor {
     fn init(&mut self, ctx: &mut impl InitContext<Self>) -> anyhow::Result<()> {
         ctx.link_bindings(spin_world::v1::key_value::add_to_linker::<_, FactorData<Self>>)?;
         ctx.link_bindings(spin_world::v2::key_value::add_to_linker::<_, FactorData<Self>>)?;
+        ctx.link_bindings(
+            spin_world::spin::key_value::key_value::add_to_linker::<_, KeyValueFactorData>,
+        )?;
         ctx.link_bindings(spin_world::wasi::keyvalue::store::add_to_linker::<_, FactorData<Self>>)?;
         ctx.link_bindings(spin_world::wasi::keyvalue::batch::add_to_linker::<_, FactorData<Self>>)?;
         ctx.link_bindings(
@@ -200,3 +207,9 @@ impl FactorInstanceBuilder for InstanceBuilder {
         ))
     }
 }
+
+pub struct KeyValueFactorData(KeyValueFactor);
+
+impl spin_core::wasmtime::component::HasData for KeyValueFactorData {
+    type Data<'a> = &'a mut KeyValueDispatch;
+}

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

@@ -1,6 +1,6 @@
 use anyhow::bail;
 use spin_core::async_trait;
-use spin_factor_key_value::{Cas, KeyValueFactor, RuntimeConfig, Store, StoreManager};
+use spin_factor_key_value::{v3, Cas, KeyValueFactor, RuntimeConfig, Store, StoreManager};
 use spin_factors::RuntimeFactors;
 use spin_factors_test::{toml, TestEnvironment};
 use spin_world::v2::key_value::{Error, HostStore};
@@ -141,6 +141,15 @@ impl Store for MockStore {
         let _ = max_result_bytes;
         todo!()
     }
+    async fn get_keys_async(
+        &self,
+        _max_result_bytes: usize,
+    ) -> (
+        tokio::sync::mpsc::Receiver<String>,
+        tokio::sync::oneshot::Receiver<Result<(), v3::Error>>,
+    ) {
+        todo!()
+    }
 
     async fn get_many(
         &self,

crates/factor-key-value/tests/factor_test.rs

@@ -1,5 +1,4 @@
 use anyhow::Result;
-use spin_core::wasmtime;
 use spin_core::wasmtime::component::{Accessor, FutureReader, Resource, StreamReader};
 use spin_world::spin::postgres3_0_0::postgres::{self as v3};
 use spin_world::spin::postgres4_2_0::postgres::{self as v4};
@@ -243,8 +242,6 @@ impl<CF: ClientFactory> spin_world::spin::postgres4_2_0::postgres::HostConnectio
         ),
         v4::Error,
     > {
-        use wasmtime::AsContextMut;
-
         let client = accessor.with(|mut access| {
             let host = access.get();
             host.connections.get(connection.rep()).unwrap().clone()
@@ -261,8 +258,8 @@ impl<CF: ClientFactory> spin_world::spin::postgres4_2_0::postgres::HostConnectio
         let row_producer = spin_wasi_async::stream::producer(rows);
 
         let (sr, efr) = accessor.with(|mut access| {
-            let sr = StreamReader::new(access.as_context_mut(), row_producer);
-            let efr = FutureReader::new(access.as_context_mut(), error);
+            let sr = StreamReader::new(&mut access, row_producer);
+            let efr = FutureReader::new(&mut access, error);
             (sr, efr)
         });
 

crates/factor-outbound-pg/src/host.rs

@@ -10,10 +10,12 @@ rust-version.workspace = true
 
 [dependencies]
 async-trait = { 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" }
+spin-wasi-async = { path = "../wasi-async" }
 spin-world = { path = "../world" }
 tokio = { workspace = true }
 tracing = { workspace = true }