fix: made network operations cancellable

Author: bmuddha

magicblock-replicator/src/nats/consumer.rs

@@ -4,6 +4,7 @@ use async_nats::jetstream::consumer::{
     pull::{Config as PullConfig, Stream as MessageStream},
     AckPolicy, DeliverPolicy, PullConsumer,
 };
+use tokio_util::sync::CancellationToken;
 use tracing::warn;
 
 use super::cfg;
@@ -58,18 +59,27 @@ impl Consumer {
     ///
     /// Use this in a `tokio::select!` loop to process messages as they arrive.
     /// Messages are fetched in batches for efficiency.
-    pub async fn messages(&self) -> MessageStream {
+    pub async fn messages(
+        &self,
+        cancel: &CancellationToken,
+    ) -> Option<MessageStream> {
         loop {
-            let result = self
+            let messages = self
                 .inner
                 .stream()
                 .max_messages_per_batch(cfg::BATCH_SIZE)
-                .messages()
-                .await;
-            match result {
-                Ok(s) => break s,
-                Err(error) => {
-                    warn!(%error, "failed to create message stream")
+                .messages();
+            tokio::select! {
+                result = messages => {
+                    match result {
+                        Ok(s) => break Some(s),
+                        Err(error) => {
+                            warn!(%error, "failed to create message stream")
+                        }
+                    }
+                }
+                _ = cancel.cancelled() => {
+                    break None;
                 }
             }
         }

magicblock-replicator/src/nats/lock_watcher.rs

@@ -2,6 +2,7 @@
 
 use async_nats::jetstream::kv::{Operation, Watch};
 use futures::StreamExt;
+use tokio_util::sync::CancellationToken;
 use tracing::warn;
 
 use super::cfg;
@@ -17,8 +18,14 @@ pub struct LockWatcher {
 
 impl LockWatcher {
     /// Creates a new lock watcher.
-    pub(crate) async fn new(broker: &Broker) -> Self {
+    pub(crate) async fn new(
+        broker: &Broker,
+        cancel: &CancellationToken,
+    ) -> Option<Self> {
         let watch = loop {
+            if cancel.is_cancelled() {
+                return None;
+            }
             let store = match broker.ctx.get_key_value(cfg::PRODUCER_LOCK).await
             {
                 Ok(s) => s,
@@ -35,7 +42,7 @@ impl LockWatcher {
                 }
             }
         };
-        Self { watch }
+        Some(Self { watch })
     }
 
     /// Waits for the lock to be deleted or expire.

magicblock-replicator/src/service/context.rs

@@ -186,7 +186,10 @@ impl ReplicationContext {
         let Some(consumer) = self.create_consumer(reset).await else {
             return Ok(None);
         };
-        let watcher = LockWatcher::new(&self.broker).await;
+        let Some(watcher) = LockWatcher::new(&self.broker, &self.cancel).await
+        else {
+            return Ok(None);
+        };
         self.enter_replica_mode().await;
         Ok(Some(Standby::new(
             self,

magicblock-replicator/src/service/standby.rs

@@ -49,7 +49,10 @@ impl Standby {
     /// Returns `Some(Primary)` on promotion, `None` on shutdown.
     pub async fn run(mut self) -> Result<Option<Primary>> {
         let mut timeout_check = tokio::time::interval(Duration::from_secs(1));
-        let mut stream = self.consumer.messages().await;
+        let Some(mut stream) = self.consumer.messages(&self.ctx.cancel).await
+        else {
+            return Ok(None);
+        };
 
         loop {
             tokio::select! {
@@ -63,8 +66,12 @@ impl Standby {
                 }
                 result = stream.next() => {
                     let Some(result) = result else {
-                        stream = self.consumer.messages().await;
-                        continue;
+                        if let Some(s) = self.consumer.messages(&self.ctx.cancel).await {
+                            stream = s;
+                            continue;
+                        } else {
+                            return Ok(None);
+                        };
                     };
                     match result {
                         Ok(msg) => {