Add RedisMessenger::add_streams

rpcpool · Nov 17, 2023 · 32fac47 · 32fac47
1 parent a6508b4
commit 32fac47
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 0 deletions.
diff --git a/plerkle_messenger/src/plerkle_messenger.rs b/plerkle_messenger/src/plerkle_messenger.rs
@@ -55,6 +55,7 @@ pub trait Messenger: Sync + Send {
     where
         Self: Sized;
     fn messenger_type(&self) -> MessengerType;
+    async fn add_streams(&mut self, streams: &[(&'static str, usize)]) -> Result<(), MessengerError>;
     async fn add_stream(&mut self, stream_key: &'static str) -> Result<(), MessengerError>;
     async fn set_buffer_size(&mut self, stream_key: &'static str, max_buffer_size: usize);
     async fn send(&mut self, stream_key: &'static str, bytes: &[u8]) -> Result<(), MessengerError>;

diff --git a/plerkle_messenger/src/redis_messenger.rs b/plerkle_messenger/src/redis_messenger.rs
@@ -247,6 +247,36 @@ impl Messenger for RedisMessenger {
         }
     }
 
+    async fn add_streams(&mut self, streams: &[(&'static str, usize)]) -> Result<(), MessengerError> {
+        for (stream_key, max_buffer_size) in streams {
+            // Add to streams hashmap.
+            let _result = self.streams.insert(
+                stream_key,
+                RedisMessengerStream {
+                    max_len: Some(StreamMaxlen::Approx(*max_buffer_size)),
+                    local_buffer: LinkedList::new(),
+                    local_buffer_total: 0,
+                    local_buffer_last_flush: Instant::now(),
+                },
+            );
+
+            // Add stream to Redis.
+            let result: RedisResult<()> = self
+                .connection
+                .xgroup_create_mkstream(stream_key, self.consumer_group_name.as_str(), "$")
+                .await;
+
+            if let Err(error) = result {
+                if !(error.kind() == redis::ErrorKind::ExtensionError && error.code() == Some("BUSYGROUP")) {
+                    return Err(MessengerError::ConfigurationError {
+                        msg: format!("{:?}", error),
+                    });
+                }
+            }
+        }
+        Ok(())
+    }
+
     async fn add_stream(&mut self, stream_key: &'static str) -> Result<(), MessengerError> {
         // Add to streams hashmap.
         let _result = self.streams.insert(