Migrate outbound-redis to new core

Signed-off-by: Lann Martin <lann.martin@fermyon.com>

Author: lann

Cargo.lock

@@ -8,11 +8,10 @@ doctest = false
 
 [dependencies]
 anyhow = "1.0"
-owning_ref = "0.4.1"
-redis = { version = "0.21", features = [ "tokio-comp" ] }
-spin-engine = { path = "../engine" }
-spin-manifest = { path = "../manifest" }
-tracing = { version = "0.1", features = [ "log" ] }
+redis = { version = "0.21", features = ["tokio-comp"] }
+spin-core = { path = "../core" }
+tokio = { version = "1", features = ["sync"] }
+tracing = { version = "0.1", features = ["log"] }
 
 [dependencies.wit-bindgen-wasmtime]
 git = "https://github.com/bytecodealliance/wit-bindgen"

crates/outbound-redis/Cargo.toml

@@ -1,118 +1,84 @@
-use outbound_redis::*;
-use owning_ref::RwLockReadGuardRef;
-use redis::Commands;
-use std::{
-    collections::HashMap,
-    sync::{Arc, Mutex, RwLock},
-};
+use std::{collections::HashMap, sync::Arc};
 
-pub use outbound_redis::add_to_linker;
-use spin_engine::{
-    host_component::{HostComponent, HostComponentsStateHandle},
-    RuntimeContext,
-};
-use wit_bindgen_wasmtime::{async_trait, wasmtime::Linker};
+use anyhow::Result;
+use redis::{aio::Connection, AsyncCommands};
+use spin_core::{HostComponent, Linker};
+use tokio::sync::{Mutex, RwLock};
+use wit_bindgen_wasmtime::async_trait;
 
 wit_bindgen_wasmtime::export!({paths: ["../../wit/ephemeral/outbound-redis.wit"], async: *});
+use outbound_redis::Error;
 
-/// A simple implementation to support outbound Redis commands.
+#[derive(Clone, Default)]
 pub struct OutboundRedis {
-    pub connections: Arc<RwLock<HashMap<String, Mutex<redis::Connection>>>>,
+    connections: Arc<RwLock<HashMap<String, Arc<Mutex<Connection>>>>>,
 }
 
 impl HostComponent for OutboundRedis {
-    type State = Self;
+    type Data = Self;
 
     fn add_to_linker<T: Send>(
-        linker: &mut Linker<RuntimeContext<T>>,
-        state_handle: HostComponentsStateHandle<Self::State>,
+        linker: &mut Linker<T>,
+        get: impl Fn(&mut spin_core::Data<T>) -> &mut Self::Data + Send + Sync + Copy + 'static,
     ) -> anyhow::Result<()> {
-        add_to_linker(linker, move |ctx| state_handle.get_mut(ctx))
+        crate::outbound_redis::add_to_linker(linker, get)
     }
 
-    fn build_state(&self, component: &spin_manifest::CoreComponent) -> anyhow::Result<Self::State> {
-        let mut conn_map = HashMap::new();
-        if let Some(address) = component.wasm.environment.get("REDIS_ADDRESS") {
-            let client = redis::Client::open(address.to_string())?;
-            let conn = client.get_connection()?;
-            conn_map.insert(address.to_owned(), Mutex::new(conn));
-        }
-        Ok(Self {
-            connections: Arc::new(RwLock::new(conn_map)),
-        })
+    fn build_data(&self) -> Self::Data {
+        self.clone()
     }
 }
 
-// TODO: use spawn_blocking or async client methods (redis::aio)
 #[async_trait]
 impl outbound_redis::OutboundRedis for OutboundRedis {
     async fn publish(&mut self, address: &str, channel: &str, payload: &[u8]) -> Result<(), Error> {
-        let conn_map = self.get_reused_conn_map(address)?;
-        let mut conn = conn_map
-            .get(address)
-            .unwrap()
-            .lock()
-            .map_err(|_| Error::Error)?;
-        conn.publish(channel, payload).map_err(|_| Error::Error)?;
+        let conn = self.get_conn(address).await.map_err(log_error)?;
+        conn.lock()
+            .await
+            .publish(channel, payload)
+            .await
+            .map_err(log_error)?;
         Ok(())
     }
 
     async fn get(&mut self, address: &str, key: &str) -> Result<Vec<u8>, Error> {
-        let conn_map = self.get_reused_conn_map(address)?;
-        let mut conn = conn_map
-            .get(address)
-            .unwrap()
-            .lock()
-            .map_err(|_| Error::Error)?;
-        let value = conn.get(key).map_err(|_| Error::Error)?;
+        let conn = self.get_conn(address).await.map_err(log_error)?;
+        let value = conn.lock().await.get(key).await.map_err(log_error)?;
         Ok(value)
     }
 
     async fn set(&mut self, address: &str, key: &str, value: &[u8]) -> Result<(), Error> {
-        let conn_map = self.get_reused_conn_map(address)?;
-        let mut conn = conn_map
-            .get(address)
-            .unwrap()
-            .lock()
-            .map_err(|_| Error::Error)?;
-        conn.set(key, value).map_err(|_| Error::Error)?;
+        let conn = self.get_conn(address).await.map_err(log_error)?;
+        conn.lock().await.set(key, value).await.map_err(log_error)?;
         Ok(())
     }
 
     async fn incr(&mut self, address: &str, key: &str) -> Result<i64, Error> {
-        let conn_map = self.get_reused_conn_map(address)?;
-        let mut conn = conn_map
-            .get(address)
-            .unwrap()
-            .lock()
-            .map_err(|_| Error::Error)?;
-        let value = conn.incr(key, 1).map_err(|_| Error::Error)?;
+        let conn = self.get_conn(address).await.map_err(log_error)?;
+        let value = conn.lock().await.incr(key, 1).await.map_err(log_error)?;
         Ok(value)
     }
 }
 
 impl OutboundRedis {
-    fn get_reused_conn_map<'ret, 'me: 'ret, 'c>(
-        &'me mut self,
-        address: &'c str,
-    ) -> Result<RwLockReadGuardRef<'ret, HashMap<String, Mutex<redis::Connection>>>, Error> {
-        let conn_map = self.connections.read().map_err(|_| Error::Error)?;
-        if conn_map.get(address).is_some() {
-            tracing::debug!("Reuse connection: {:?}", address);
-            return Ok(RwLockReadGuardRef::new(conn_map));
-        }
-        // Get rid of our read lock
-        drop(conn_map);
-
-        let mut conn_map = self.connections.write().map_err(|_| Error::Error)?;
-        let client = redis::Client::open(address).map_err(|_| Error::Error)?;
-        let conn = client.get_connection().map_err(|_| Error::Error)?;
-        tracing::debug!("Build new connection: {:?}", address);
-        conn_map.insert(address.to_string(), Mutex::new(conn));
-        // Get rid of our write lock
-        drop(conn_map);
-
-        let conn_map = self.connections.read().map_err(|_| Error::Error)?;
-        Ok(RwLockReadGuardRef::new(conn_map))
+    async fn get_conn(&self, address: &str) -> Result<Arc<Mutex<Connection>>> {
+        let conn_map = self.connections.read().await;
+        let conn = if let Some(conn) = conn_map.get(address) {
+            conn.clone()
+        } else {
+            let conn = redis::Client::open(address)?.get_async_connection().await?;
+            let conn = Arc::new(Mutex::new(conn));
+            self.connections
+                .write()
+                .await
+                .insert(address.to_string(), conn.clone());
+            conn
+        };
+        Ok(conn)
     }
 }
+
+fn log_error(err: impl std::fmt::Debug) -> Error {
+    tracing::warn!("Outbound Redis error: {err:?}");
+    Error::Error
+}

crates/outbound-redis/src/lib.rs

@@ -93,8 +93,7 @@ impl<Executor: TriggerExecutor> TriggerExecutorBuilder<Executor> {
             let mut builder = Engine::builder(&self.config)?;
 
             if !self.disable_default_host_components {
-                // FIXME(lann): migrate host components from prototype
-                // builder.add_host_component(outbound_redis::OutboundRedis::default())?;
+                builder.add_host_component(outbound_redis::OutboundRedis::default())?;
                 builder.add_host_component(outbound_pg::OutboundPg)?;
                 self.loader.add_dynamic_host_component(
                     &mut builder,

crates/trigger-new/src/lib.rs

@@ -1,10 +1,4 @@
-use std::{
-    collections::HashMap,
-    error::Error,
-    marker::PhantomData,
-    path::PathBuf,
-    sync::{Arc, RwLock},
-};
+use std::{error::Error, marker::PhantomData, path::PathBuf};
 
 use anyhow::Result;
 use async_trait::async_trait;
@@ -124,10 +118,7 @@ impl<Executor: TriggerExecutor> TriggerExecutorBuilder<Executor> {
 
 /// Add the default set of host components to the given builder.
 pub fn add_default_host_components<T: Default + Send + 'static>(
-    builder: &mut Builder<T>,
+    _builder: &mut Builder<T>,
 ) -> Result<()> {
-    builder.add_host_component(outbound_redis::OutboundRedis {
-        connections: Arc::new(RwLock::new(HashMap::new())),
-    })?;
     Ok(())
 }