Merge pull request #687 from FrankYang0529/reuse-redis-connection

feat(outbound-redis): reuse connecton

Author: fibonacci1729

Cargo.lock

@@ -12,3 +12,6 @@ redis = { version = "0.21", features = [ "tokio-comp" ] }
 spin-engine = { path = "../engine" }
 spin-manifest = { path = "../manifest" }
 wit-bindgen-wasmtime = { git = "https://github.com/bytecodealliance/wit-bindgen", rev = "dde4694aaa6acf9370206527a798ac4ba6a8c5b8" }
+tracing = { version = "0.1", features = [ "log" ] }
+tracing-futures = "0.2"
+owning_ref = "0.4.1"

crates/outbound-redis/Cargo.toml

@@ -1,5 +1,10 @@
 use outbound_redis::*;
+use owning_ref::RwLockReadGuardRef;
 use redis::Commands;
+use std::{
+    collections::HashMap,
+    sync::{Arc, Mutex, RwLock},
+};
 
 pub use outbound_redis::add_to_linker;
 use spin_engine::{
@@ -11,8 +16,9 @@ use wit_bindgen_wasmtime::wasmtime::Linker;
 wit_bindgen_wasmtime::export!("../../wit/ephemeral/outbound-redis.wit");
 
 /// A simple implementation to support outbound Redis commands.
-#[derive(Default, Clone)]
-pub struct OutboundRedis;
+pub struct OutboundRedis {
+    pub connections: Arc<RwLock<HashMap<String, Mutex<redis::Connection>>>>,
+}
 
 impl HostComponent for OutboundRedis {
     type State = Self;
@@ -24,42 +30,87 @@ impl HostComponent for OutboundRedis {
         add_to_linker(linker, move |ctx| state_handle.get_mut(ctx))
     }
 
-    fn build_state(
-        &self,
-        _component: &spin_manifest::CoreComponent,
-    ) -> anyhow::Result<Self::State> {
-        Ok(Self)
+    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)),
+        })
     }
 }
 
 impl outbound_redis::OutboundRedis for OutboundRedis {
     fn publish(&mut self, address: &str, channel: &str, payload: &[u8]) -> Result<(), Error> {
-        let client = redis::Client::open(address).map_err(|_| Error::Error)?;
-        let mut pubsub_conn = client.get_connection().map_err(|_| Error::Error)?;
-        pubsub_conn
-            .publish(channel, payload)
+        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)?;
         Ok(())
     }
 
     fn get(&mut self, address: &str, key: &str) -> Result<Vec<u8>, Error> {
-        let client = redis::Client::open(address).map_err(|_| Error::Error)?;
-        let mut conn = client.get_connection().map_err(|_| Error::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)?;
         Ok(value)
     }
 
     fn set(&mut self, address: &str, key: &str, value: &[u8]) -> Result<(), Error> {
-        let client = redis::Client::open(address).map_err(|_| Error::Error)?;
-        let mut conn = client.get_connection().map_err(|_| Error::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)?;
         Ok(())
     }
 
     fn incr(&mut self, address: &str, key: &str) -> Result<i64, Error> {
-        let client = redis::Client::open(address).map_err(|_| Error::Error)?;
-        let mut conn = client.get_connection().map_err(|_| Error::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)?;
         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))
+    }
+}

crates/outbound-redis/src/lib.rs

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

crates/trigger/src/lib.rs

@@ -790,7 +790,7 @@ mod integration_tests {
     async fn wait_tcp(url: &str, process: &mut Child, target: &str) -> Result<()> {
         let mut wait_count = 0;
         loop {
-            if wait_count >= 180 {
+            if wait_count >= 240 {
                 panic!(
                     "Ran out of retries waiting for {} to start on URL {}",
                     target, url