Merge pull request #83 from balena-io-modules/create-rollback-bug

Fix task error rollback

Author: flowzone-app[bot]

mahler-core/src/runtime/channel.rs

@@ -1,57 +1,30 @@
-use std::ops::Deref;
-use std::sync::Arc;
-
-use tokio::sync::Mutex;
-
 use crate::error::Error;
 use crate::json::{Patch, Value};
 use crate::result::Result;
 use crate::sync::Sender;
 
-#[derive(Clone)]
-struct Checkpoint<T>(Arc<Mutex<Option<T>>>);
-
-impl<T> Default for Checkpoint<T> {
-    fn default() -> Self {
-        Self(Arc::new(Mutex::new(None)))
-    }
-}
-
-impl<T> Deref for Checkpoint<T> {
-    type Target = Mutex<Option<T>>;
-
-    fn deref(&self) -> &Self::Target {
-        self.0.as_ref()
-    }
-}
-
 /// A channel to communicate state changes at runtime
 ///
 /// The `Channel` allows tasks to send state changes and rollback checkpoints back to the worker during execution,
 /// enabling real-time progress updates. This can be used by extractors to propagate changes
 /// during a long operation.
 #[derive(Clone)]
 pub struct Channel {
-    sender: Option<Sender<Patch>>,
-    checkpoint: Checkpoint<Value>,
+    sender: Option<Sender<(Patch, Option<Value>)>>,
 }
 
 impl std::fmt::Debug for Channel {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let mut dbg = f.debug_struct("Channel");
-        dbg.field(
-            "sender",
-            if self.sender.is_some() {
-                &"attached"
-            } else {
-                &"detached"
-            },
-        );
-        match self.checkpoint.try_lock() {
-            Ok(guard) => dbg.field("checkpoint", &*guard),
-            Err(_) => dbg.field("checkpoint", &"locked"),
-        };
-        dbg.finish()
+        f.debug_struct("Channel")
+            .field(
+                "sender",
+                if self.sender.is_some() {
+                    &"attached"
+                } else {
+                    &"detached"
+                },
+            )
+            .finish()
     }
 }
 
@@ -63,51 +36,40 @@ impl Channel {
 
     /// Create a detached channel. A detached channel is not connected to a worker
     pub fn detached() -> Self {
-        Self {
-            sender: None,
-            checkpoint: Checkpoint::default(),
-        }
+        Self { sender: None }
     }
 
     /// Communicate the changes to the global state
     pub async fn send(&self, changes: Patch) -> Result<()> {
         if let Some(sender) = self.sender.as_ref() {
-            sender.send(changes).await.map_err(Error::internal)?;
+            sender
+                .send((changes, None))
+                .await
+                .map_err(Error::internal)?;
         }
         Ok(())
     }
 
     /// Communicate changes and update the checkpoint
     ///
-    /// Like [`send`](Channel::send), but also records the given value as the
-    /// current checkpoint. The checkpoint can be read by the caller via
-    /// [`checkpoint`](Channel::checkpoint) after the task completes.
-    pub async fn send_and_commit(&self, changes: Patch, new_checkpoint: Value) -> Result<()> {
+    /// Like [`send`](Channel::send), but also sends the given value as the
+    /// current checkpoint alongside the patch.
+    pub async fn send_and_commit(&self, changes: Patch, checkpoint: Value) -> Result<()> {
         if let Some(sender) = self.sender.as_ref() {
-            sender.send(changes).await.map_err(Error::internal)?;
-            *self.checkpoint.lock().await = Some(new_checkpoint);
+            sender
+                .send((changes, Some(checkpoint)))
+                .await
+                .map_err(Error::internal)?;
         }
         Ok(())
     }
-
-    /// Return the last checkpointed value, if any.
-    ///
-    /// Returns `Some(value)` if [`send_and_commit`](Channel::send_and_commit) was called
-    /// at least once, `None` otherwise.
-    pub async fn checkpoint(&self) -> Option<Value> {
-        self.checkpoint.lock().await.clone()
-    }
 }
 
-impl From<Sender<Patch>> for Channel {
-    /// Create a channel with a checkpoint for tracking committed state.
-    ///
-    /// The checkpoint allows [`send_and_commit`](Channel::send_and_commit) to record
-    /// state snapshots, retrievable later via [`checkpoint`](Channel::checkpoint).
-    fn from(sender: Sender<Patch>) -> Self {
+impl From<Sender<(Patch, Option<Value>)>> for Channel {
+    /// Create an attached channel from a sender.
+    fn from(sender: Sender<(Patch, Option<Value>)>) -> Self {
         Self {
             sender: Some(sender),
-            checkpoint: Checkpoint::default(),
         }
     }
 }

mahler/src/extract/view.rs

@@ -816,14 +816,14 @@ mod tests {
         })
         .unwrap();
 
-        let (tx, mut rx) = sync::channel::<Patch>(1);
+        let (tx, mut rx) = sync::channel::<(Patch, Option<Value>)>(1);
         let channel = Channel::from(tx);
 
         let received = Arc::new(tokio::sync::Mutex::new(Vec::new()));
         let received_clone = received.clone();
         let ack_task = tokio::spawn(async move {
             while let Some(msg) = rx.recv().await {
-                received_clone.lock().await.push(msg.data.clone());
+                received_clone.lock().await.push(msg.data.0.clone());
                 msg.ack();
             }
         });
@@ -875,14 +875,14 @@ mod tests {
         })
         .unwrap();
 
-        let (tx, mut rx) = sync::channel::<Patch>(10);
+        let (tx, mut rx) = sync::channel::<(Patch, Option<Value>)>(10);
         let channel = Channel::from(tx);
 
         let received = Arc::new(tokio::sync::Mutex::new(Vec::new()));
         let received_clone = received.clone();
         let ack_task = tokio::spawn(async move {
             while let Some(msg) = rx.recv().await {
-                received_clone.lock().await.push(msg.data.clone());
+                received_clone.lock().await.push(msg.data.0.clone());
                 msg.ack();
             }
         });
@@ -930,15 +930,15 @@ mod tests {
         };
         let system = System::try_from(state).unwrap();
 
-        let (tx, mut rx) = sync::channel::<Patch>(1);
+        let (tx, mut rx) = sync::channel::<(Patch, Option<Value>)>(1);
         let channel = Channel::from(tx);
 
         // Spawn a task to receive and acknowledge messages
         let received_patches = Arc::new(tokio::sync::Mutex::new(Vec::new()));
         let received_patches_clone = received_patches.clone();
         let ack_task = tokio::spawn(async move {
             while let Some(msg) = rx.recv().await {
-                received_patches_clone.lock().await.push(msg.data.clone());
+                received_patches_clone.lock().await.push(msg.data.0.clone());
                 msg.ack();
             }
         });
@@ -987,15 +987,15 @@ mod tests {
         let state = MyState { numbers };
         let system = System::try_from(state).unwrap();
 
-        let (tx, mut rx) = sync::channel::<Patch>(1);
+        let (tx, mut rx) = sync::channel::<(Patch, Option<Value>)>(1);
         let channel = Channel::from(tx);
 
         // Spawn a task to receive and acknowledge messages
         let received_patches = Arc::new(tokio::sync::Mutex::new(Vec::new()));
         let received_patches_clone = received_patches.clone();
         let ack_task = tokio::spawn(async move {
             while let Some(msg) = rx.recv().await {
-                received_patches_clone.lock().await.push(msg.data.clone());
+                received_patches_clone.lock().await.push(msg.data.0.clone());
                 msg.ack();
             }
         });
@@ -1080,7 +1080,7 @@ mod tests {
         })
         .unwrap();
 
-        let (tx, mut rx) = sync::channel::<Patch>(1);
+        let (tx, mut rx) = sync::channel::<(Patch, Option<Value>)>(1);
         let channel = Channel::from(tx);
 
         let received = Arc::new(tokio::sync::Mutex::new(Vec::new()));
@@ -1101,18 +1101,19 @@ mod tests {
 
         tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
 
-        let patches = received.lock().await;
-        assert_eq!(patches.len(), 1);
+        let messages = received.lock().await;
+        assert_eq!(messages.len(), 1);
+        let (patch, checkpoint) = &messages[0];
         assert_eq!(
-            patches[0],
+            *patch,
             serde_json::from_value::<Patch>(json!([
               { "op": "add", "path": "/numbers/0", "value": 10 },
             ]))
             .unwrap()
         );
-        drop(patches);
-
-        assert_eq!(channel.checkpoint().await, Some(json!([10])));
+        // commit() sends the new state value as the checkpoint
+        assert_eq!(*checkpoint, Some(json!([10])));
+        drop(messages);
 
         drop(channel);
         ack_task.abort();
@@ -1138,7 +1139,7 @@ mod tests {
         };
         let system = System::try_from(state).unwrap();
 
-        let (tx, mut rx) = sync::channel::<Patch>(1);
+        let (tx, mut rx) = sync::channel::<(Patch, Option<Value>)>(1);
         let channel = Channel::from(tx);
 
         let received = Arc::new(tokio::sync::Mutex::new(Vec::new()));
@@ -1159,18 +1160,19 @@ mod tests {
 
         tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
 
-        let patches = received.lock().await;
-        assert_eq!(patches.len(), 1);
+        let messages = received.lock().await;
+        assert_eq!(messages.len(), 1);
+        let (patch, checkpoint) = &messages[0];
         assert_eq!(
-            patches[0],
+            *patch,
             serde_json::from_value::<Patch>(json!([
               { "op": "add", "path": "/numbers/new", "value": 42 },
             ]))
             .unwrap()
         );
-        drop(patches);
-
-        assert_eq!(channel.checkpoint().await, Some(json!(42)));
+        // commit() sends the committed value as the checkpoint
+        assert_eq!(*checkpoint, Some(json!(42)));
+        drop(messages);
 
         drop(channel);
         ack_task.abort();