Merge pull request #2242 from karthik2804/allow_multiple_redis_server

allow redis trigger to connect to multiple servers

Author: itowlson

Cargo.lock

@@ -19,6 +19,7 @@ spin-trigger = { path = "../trigger" }
 spin-world = { path = "../world" }
 redis = { version = "0.21", features = ["tokio-comp"] }
 tracing = { workspace = true }
+tokio = { version = "1.23", features = ["full"] }
 
 [dev-dependencies]
 spin-testing = { path = "../testing" }

crates/trigger-redis/Cargo.toml

@@ -2,27 +2,27 @@
 
 mod spin;
 
-use std::collections::HashMap;
-
 use anyhow::{anyhow, Context, Result};
 use futures::{future::join_all, StreamExt};
 use redis::{Client, ConnectionLike};
 use serde::{de::IgnoredAny, Deserialize, Serialize};
 use spin_core::async_trait;
 use spin_trigger::{cli::NoArgs, TriggerAppEngine, TriggerExecutor};
+use std::collections::HashMap;
+use std::sync::Arc;
 
 use crate::spin::SpinRedisExecutor;
 
 pub(crate) type RuntimeData = ();
 pub(crate) type Store = spin_core::Store<RuntimeData>;
 
+type ChannelComponents = HashMap<String, Vec<String>>;
 /// The Spin Redis trigger.
+#[derive(Clone)]
 pub struct RedisTrigger {
-    engine: TriggerAppEngine<Self>,
-    // Redis address to connect to
-    address: String,
-    // Mapping of subscription channels to component IDs
-    channel_components: HashMap<String, Vec<String>>,
+    engine: Arc<TriggerAppEngine<Self>>,
+    // Mapping of server url with subscription channel and associated component IDs
+    server_channels: HashMap<String, ChannelComponents>,
 }
 
 /// Redis trigger configuration.
@@ -33,6 +33,8 @@ pub struct RedisTriggerConfig {
     pub component: String,
     /// Channel to subscribe to
     pub channel: String,
+    /// optional overide address for trigger
+    pub address: Option<String>,
     /// Trigger executor (currently unused)
     #[serde(default, skip_serializing)]
     pub executor: IgnoredAny,
@@ -52,32 +54,90 @@ impl TriggerExecutor for RedisTrigger {
     type RunConfig = NoArgs;
 
     async fn new(engine: TriggerAppEngine<Self>) -> Result<Self> {
-        let address = engine
+        let default_address: String = engine
             .trigger_metadata::<TriggerMetadata>()?
             .unwrap_or_default()
             .address;
-        let address_expr = spin_expressions::Template::new(address)?;
-        let address = engine.resolve_template(&address_expr)?;
+        let default_address_expr = spin_expressions::Template::new(default_address)?;
+        let default_address = engine.resolve_template(&default_address_expr)?;
 
-        let mut channel_components: HashMap<String, Vec<String>> = HashMap::new();
+        let mut server_channels: HashMap<String, ChannelComponents> = HashMap::new();
 
         for (_, config) in engine.trigger_configs() {
-            channel_components
+            let address = config.address.clone().unwrap_or(default_address.clone());
+            let address_expr = spin_expressions::Template::new(address)?;
+            let address = engine.resolve_template(&address_expr)?;
+            let server = server_channels.entry(address).or_default();
+            server
                 .entry(config.channel.clone())
                 .or_default()
                 .push(config.component.clone());
         }
         Ok(Self {
-            engine,
-            address,
-            channel_components,
+            engine: Arc::new(engine),
+            server_channels,
         })
     }
 
     /// Run the Redis trigger indefinitely.
     async fn run(self, _config: Self::RunConfig) -> Result<()> {
-        let address = &self.address;
+        let tasks: Vec<_> = self
+            .server_channels
+            .clone()
+            .into_iter()
+            .map(|(server_address, channel_components)| {
+                let trigger = self.clone();
+                tokio::spawn(async move {
+                    trigger
+                        .run_listener(server_address.clone(), channel_components.clone())
+                        .await
+                })
+            })
+            .collect();
+
+        // wait for the first handle to be returned and drop the rest
+        let (_, _, rest) = futures::future::select_all(tasks).await;
+        drop(rest);
+
+        Ok(())
+    }
+}
 
+impl RedisTrigger {
+    // Handle the message.
+    async fn handle(
+        &self,
+        address: &str,
+        channel_components: &ChannelComponents,
+        msg: redis::Msg,
+    ) -> Result<()> {
+        let channel = msg.get_channel_name();
+        tracing::info!("Received message on channel {address}:{:?}", channel);
+
+        if let Some(component_ids) = channel_components.get(channel) {
+            let futures = component_ids.iter().map(|id| {
+                tracing::trace!("Executing Redis component {id:?}");
+                SpinRedisExecutor.execute(&self.engine, id, channel, msg.get_payload_bytes())
+            });
+            let results: Vec<_> = join_all(futures).await.into_iter().collect();
+            let errors = results
+                .into_iter()
+                .filter_map(|r| r.err())
+                .collect::<Vec<_>>();
+            if !errors.is_empty() {
+                return Err(anyhow!("{errors:#?}"));
+            }
+        } else {
+            tracing::debug!("No subscription found for {:?}", channel);
+        }
+        Ok(())
+    }
+
+    async fn run_listener(
+        &self,
+        address: String,
+        channel_components: ChannelComponents,
+    ) -> Result<()> {
         tracing::info!("Connecting to Redis server at {}", address);
         let mut client = Client::open(address.to_string())?;
         let mut pubsub = client
@@ -88,55 +148,30 @@ impl TriggerExecutor for RedisTrigger {
 
         println!("Active Channels on {address}:");
         // Subscribe to channels
-        for (channel, component) in self.channel_components.iter() {
+        for (channel, component) in channel_components.iter() {
             tracing::info!("Subscribing component {component:?} to channel {channel:?}");
             pubsub.subscribe(channel).await?;
-            println!("\t{channel}: [{}]", component.join(","));
+            println!("\t{address}:{channel}: [{}]", component.join(","));
         }
 
         let mut stream = pubsub.on_message();
         loop {
             match stream.next().await {
                 Some(msg) => {
-                    if let Err(err) = self.handle(msg).await {
+                    if let Err(err) = self.handle(&address, &channel_components, msg).await {
                         tracing::warn!("Error handling message: {err}");
                     }
                 }
                 None => {
                     tracing::trace!("Empty message");
                     if !client.check_connection() {
                         tracing::info!("No Redis connection available");
-                        break Ok(());
+                        println!("Disconnected from {address}");
+                        break;
                     }
                 }
             };
         }
-    }
-}
-
-impl RedisTrigger {
-    // Handle the message.
-    async fn handle(&self, msg: redis::Msg) -> Result<()> {
-        let channel = msg.get_channel_name();
-        tracing::info!("Received message on channel {:?}", channel);
-
-        if let Some(component_ids) = self.channel_components.get(channel) {
-            let futures = component_ids.iter().map(|id| {
-                tracing::trace!("Executing Redis component {id:?}");
-                SpinRedisExecutor.execute(&self.engine, id, channel, msg.get_payload_bytes())
-            });
-            let results: Vec<_> = join_all(futures).await.into_iter().collect();
-            let errors = results
-                .into_iter()
-                .filter_map(|r| r.err())
-                .collect::<Vec<_>>();
-            if !errors.is_empty() {
-                return Err(anyhow!("{errors:#?}"));
-            }
-        } else {
-            tracing::debug!("No subscription found for {:?}", channel);
-        }
-
         Ok(())
     }
 }

crates/trigger-redis/src/lib.rs

@@ -18,9 +18,9 @@ async fn test_pubsub() -> Result<()> {
         .test_program("redis-rust.wasm")
         .build_trigger("messages")
         .await;
-
+    let test = HashMap::new();
     let msg = create_trigger_event("messages", "hello");
-    trigger.handle(msg).await?;
+    trigger.handle("", &test, msg).await?;
 
     Ok(())
 }