Rewrote Redis input_provider to not spawn. Hopefully stabilizes it

Author: mortenhaahr

src/io/builders/input_provider_builder.rs

@@ -186,17 +186,18 @@ impl InputProviderBuilder {
                         .map(|topic| (topic.clone(), format!("{}", topic)))
                         .collect(),
                 };
-                let redis_input_provider = tc::io::redis::RedisInputProvider::new(
-                    self.executor.unwrap().clone(),
+                let mut redis_input_provider = tc::io::redis::RedisInputProvider::new(
                     REDIS_HOSTNAME,
                     self.redis_port,
                     var_topics,
                 )
                 .expect("Redis input provider could not be created");
+
                 redis_input_provider
-                    .ready()
+                    .connect()
                     .await
-                    .expect("Redis input provider failed to start");
+                    .expect("Redis input provider failed to connect");
+
                 Box::new(redis_input_provider) as Box<dyn InputProvider<Val = Value>>
             }
             InputProviderSpec::Manual => {

src/io/redis/input_provider.rs

@@ -8,27 +8,21 @@ use futures::{
     StreamExt,
     future::{self, LocalBoxFuture},
 };
-use smol::{LocalExecutor, future::yield_now};
-use tracing::info;
+use tracing::{info, warn};
 
 use crate::{InputProvider, OutputStream, Value, VarName};
 
-pub struct VarData {
-    pub variable: VarName,
-    pub channel_name: String,
-    stream: Option<OutputStream<Value>>,
-}
-
 pub struct RedisInputProvider {
-    pub host: String,
-    pub var_data: BTreeMap<VarName, VarData>,
-    pub result: Rc<AsyncCell<anyhow::Result<()>>>,
     pub started: Rc<AsyncCell<bool>>,
+    client: redis::Client,
+    redis_stream: Option<redis::aio::PubSubStream>,
+    var_topics: BTreeMap<VarName, String>,
+    var_streams: BTreeMap<VarName, OutputStream<Value>>,
+    senders: BTreeMap<VarName, Sender<Value>>,
 }
 
 impl RedisInputProvider {
     pub fn new(
-        ex: Rc<LocalExecutor<'static>>,
         hostname: &str,
         port: Option<u16>,
         var_topics: BTreeMap<VarName, String>,
@@ -46,90 +40,54 @@ impl RedisInputProvider {
                     ((v.clone(), tx), (v.clone(), rx))
                 })
                 .unzip();
-        let topic_vars = var_topics
-            .iter()
-            .map(|(k, v)| (v.clone(), k.clone()))
-            .collect::<BTreeMap<_, _>>();
 
-        let started = AsyncCell::shared();
-
-        let client = redis::Client::open(url.clone())?;
-
-        let result = AsyncCell::shared();
-
-        // TODO: Should not be spawning a task inside new. Should fix the potential race conditions
-        // instead of circumventing them like this.
-        info!("Spawning RedisInputProvider on url: {:?}", url);
-        ex.spawn(RedisInputProvider::input_monitor(
-            result.clone(),
-            client,
-            var_topics.clone(),
-            topic_vars,
-            senders,
-            started.clone(),
-        ))
-        .detach();
-
-        let var_data = var_topics
+        let var_streams: BTreeMap<VarName, OutputStream<Value>> = receivers
             .into_iter()
-            .zip(receivers.into_values())
-            .map(|((v, topic), mut rx)| {
-                let stream = stream! {
+            .map(|(v, mut rx)| {
+                let stream: OutputStream<Value> = Box::pin(stream! {
                     while let Some(x) = rx.recv().await {
                         yield x
                     }
-                };
-                (
-                    v.clone(),
-                    VarData {
-                        variable: v,
-                        channel_name: topic,
-                        stream: Some(Box::pin(stream)),
-                    },
-                )
+                });
+                (v.clone(), stream)
             })
             .collect();
 
+        let started = AsyncCell::shared();
+
+        let client = redis::Client::open(url.clone())?;
+
         Ok(RedisInputProvider {
-            host: url,
-            result,
-            var_data,
             started,
+            client: client,
+            redis_stream: None,
+            var_topics: var_topics,
+            senders: senders,
+            var_streams: var_streams,
         })
     }
 
-    async fn input_monitor(
-        result: Rc<AsyncCell<anyhow::Result<()>>>,
-        client: redis::Client,
-        var_topics: BTreeMap<VarName, String>,
-        topic_vars: BTreeMap<String, VarName>,
-        senders: BTreeMap<VarName, Sender<Value>>,
-        started: Rc<AsyncCell<bool>>,
-    ) {
-        let result = result.guard_shared(Err(anyhow!("RedisInputProvider crashed")));
-        result.set(
-            RedisInputProvider::input_monitor_with_result(
-                client, var_topics, topic_vars, senders, started,
-            )
-            .await,
-        );
+    pub async fn connect(&mut self) -> anyhow::Result<()> {
+        let mut pubsub = self.client.get_async_pubsub().await?;
+        let channel_names = self.var_topics.values().collect::<Vec<_>>();
+        info!("Subscribing to Redis channel_names: {:?}", channel_names);
+        pubsub.subscribe(channel_names).await?;
+        let stream = pubsub.into_on_message();
+        self.redis_stream = Some(stream);
+        self.started.set(true);
+        Ok(())
     }
 
-    async fn input_monitor_with_result(
-        client: redis::Client,
+    async fn run_logic(
         var_topics: BTreeMap<VarName, String>,
-        topic_vars: BTreeMap<String, VarName>,
         mut senders: BTreeMap<VarName, Sender<Value>>,
-        started: Rc<AsyncCell<bool>>,
+        mut redis_stream: redis::aio::PubSubStream,
     ) -> anyhow::Result<()> {
-        let mut pubsub = client.get_async_pubsub().await?;
-        let channel_names = var_topics.values().collect::<Vec<_>>();
-        info!("Subscribing to Redis channel_names: {:?}", channel_names);
-        pubsub.subscribe(channel_names).await?;
-        started.set(true);
-        let mut stream = pubsub.on_message();
-
-        while let Some(msg) = stream.next().await {
+        let topic_vars = var_topics
+            .iter()
+            .map(|(k, v)| (v.clone(), k.clone()))
+            .collect::<BTreeMap<_, _>>();
+        while let Some(msg) = redis_stream.next().await {
             let var_name = topic_vars
                 .get(msg.get_channel_name())
                 .ok_or_else(|| anyhow!("Unknown channel name"))?;
@@ -163,26 +121,48 @@ impl InputProvider for RedisInputProvider {
     type Val = Value;
 
     fn input_stream(&mut self, var: &VarName) -> Option<OutputStream<Value>> {
-        let var_data = self.var_data.get_mut(var)?;
-        let stream = var_data.stream.take()?;
-        Some(stream)
+        self.var_streams.remove(var)
     }
 
     fn run(&mut self) -> LocalBoxFuture<'static, anyhow::Result<()>> {
-        Box::pin(self.result.take_shared())
+        let stream = self.redis_stream.take();
+
+        match stream {
+            Some(stream) => Box::pin(Self::run_logic(
+                self.var_topics.clone(),
+                std::mem::take(&mut self.senders),
+                stream,
+            )),
+            None => Box::pin(future::ready(Err(anyhow!("Not connected to Redis yet")))),
+        }
     }
 
     fn ready(&self) -> LocalBoxFuture<'static, Result<(), anyhow::Error>> {
         let started = self.started.clone();
         Box::pin(async move {
+            info!("Checking if Redis input provider is ready");
+            let mut attempts = 0;
             while !started.get().await {
-                yield_now().await;
+                attempts += 1;
+                info!(
+                    "Redis input provider not ready yet, checking again (attempt #{})",
+                    attempts
+                );
+                smol::Timer::after(std::time::Duration::from_millis(100)).await;
+
+                if attempts > 50 {
+                    warn!(
+                        "Redis input provider still not ready after 5 seconds, continuing to wait"
+                    );
+                    attempts = 0;
+                }
             }
+            info!("Redis input provider is ready");
             Ok(())
         })
     }
 
     fn vars(&self) -> Vec<VarName> {
-        self.var_data.keys().cloned().collect()
+        self.var_topics.keys().cloned().collect()
     }
 }

tests/test_redis_io.rs

@@ -160,13 +160,10 @@ mod integration_tests {
         ]);
 
         // Create the MQTT input provider
-        let mut input_provider = RedisInputProvider::new(
-            executor.clone(),
-            REDIS_HOSTNAME,
-            Some(redis_port),
-            var_topics,
-        )
-        .map_err(|e| anyhow::anyhow!("Failed to create Redis input provider: {}", e))?;
+        let mut input_provider =
+            RedisInputProvider::new(REDIS_HOSTNAME, Some(redis_port), var_topics)
+                .map_err(|e| anyhow::anyhow!("Failed to create Redis input provider: {}", e))?;
+        input_provider.connect().await?;
 
         let x_stream = input_provider
             .input_stream(&"x".into())
@@ -228,13 +225,10 @@ mod integration_tests {
         ]);
 
         // Create the MQTT input provider
-        let mut input_provider = RedisInputProvider::new(
-            executor.clone(),
-            REDIS_HOSTNAME,
-            Some(redis_port),
-            var_topics,
-        )
-        .map_err(|e| anyhow::anyhow!("Failed to create Redis input provider: {}", e))?;
+        let mut input_provider =
+            RedisInputProvider::new(REDIS_HOSTNAME, Some(redis_port), var_topics)
+                .map_err(|e| anyhow::anyhow!("Failed to create Redis input provider: {}", e))?;
+        input_provider.connect().await?;
 
         let x_stream = input_provider
             .input_stream(&"x".into())
@@ -293,13 +287,10 @@ mod integration_tests {
         ]);
 
         // Create the MQTT input provider
-        let mut input_provider = RedisInputProvider::new(
-            executor.clone(),
-            REDIS_HOSTNAME,
-            Some(redis_port),
-            var_topics,
-        )
-        .map_err(|e| anyhow::anyhow!("Failed to create Redis input provider: {}", e))?;
+        let mut input_provider =
+            RedisInputProvider::new(REDIS_HOSTNAME, Some(redis_port), var_topics)
+                .map_err(|e| anyhow::anyhow!("Failed to create Redis input provider: {}", e))?;
+        input_provider.connect().await?;
 
         let x_stream = input_provider
             .input_stream(&"x".into())