r: Save underlying pending state inside PendingStatePaused

Author: tomasol

crates/concepts/src/storage.rs

@@ -1654,14 +1654,11 @@ pub enum PendingState {
     #[display("BlockedByJoinSet({_0})")]
     BlockedByJoinSet(PendingStateBlockedByJoinSet),
 
-    #[display("Paused")]
-    Paused,
+    #[display("Paused({_0})")]
+    Paused(PendingStatePaused),
 
-    #[display("Finished({finished})")]
-    Finished {
-        #[serde(flatten)]
-        finished: PendingStateFinished,
-    },
+    #[display("Finished({_0})")]
+    Finished(PendingStateFinished),
 }
 
 #[derive(Debug, Clone, derive_more::Display, PartialEq, Eq, Serialize)]
@@ -1671,15 +1668,15 @@ pub struct PendingStateLocked {
     pub lock_expires_at: DateTime<Utc>,
 }
 
-#[derive(Debug, Clone, derive_more::Display, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, derive_more::Display, PartialEq, Eq, Serialize)]
 #[display("`{scheduled_at}`, last_lock={last_lock:?}")]
 pub struct PendingStatePendingAt {
     pub scheduled_at: DateTime<Utc>,
     /// `last_lock` is needed for lock extension.
     pub last_lock: Option<LockedBy>,
 }
 
-#[derive(Debug, Clone, derive_more::Display, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, derive_more::Display, PartialEq, Eq, Serialize)]
 #[display("{join_set_id}, `{lock_expires_at}`, closing={closing}")]
 pub struct PendingStateBlockedByJoinSet {
     pub join_set_id: JoinSetId,
@@ -1689,6 +1686,17 @@ pub struct PendingStateBlockedByJoinSet {
     pub closing: bool,
 }
 
+/// State of execution before it was paused.
+#[derive(Debug, Clone, derive_more::Display, PartialEq, Eq, Serialize)]
+pub enum PendingStatePaused {
+    #[display("Locked({_0})")]
+    Locked(PendingStateLocked),
+    #[display("PendingAt({_0})")]
+    PendingAt(PendingStatePendingAt),
+    #[display("BlockedByJoinSet({_0})")]
+    BlockedByJoinSet(PendingStateBlockedByJoinSet),
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct LockedBy {
     pub executor_id: ExecutorId,
@@ -1703,7 +1711,8 @@ impl From<&Locked> for LockedBy {
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
+#[cfg_attr(any(test, feature = "test"), derive(Deserialize))]
 pub struct PendingStateFinished {
     pub version: VersionType, // not Version since it must be Copy
     pub finished_at: DateTime<Utc>,
@@ -1820,7 +1829,7 @@ impl PendingState {
                 source: None,
                 loc: Location::caller(),
             }),
-            PendingState::Paused => Err(DbErrorWriteNonRetriable::IllegalState {
+            PendingState::Paused(..) => Err(DbErrorWriteNonRetriable::IllegalState {
                 reason: "cannot lock, execution is paused".into(),
                 context: SpanTrace::capture(),
                 source: None,
@@ -1833,6 +1842,11 @@ impl PendingState {
     pub fn is_finished(&self) -> bool {
         matches!(self, PendingState::Finished { .. })
     }
+
+    #[must_use]
+    pub fn is_paused(&self) -> bool {
+        matches!(self, PendingState::Paused(_))
+    }
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]

crates/db-mem/src/inmemory_dao.rs

@@ -574,7 +574,7 @@ mod index {
                 }
                 PendingState::BlockedByJoinSet(..)
                 | PendingState::Finished { .. }
-                | PendingState::Paused => {}
+                | PendingState::Paused(..) => {}
             }
             // Add all open async timers
             let mut delay_req_resp = journal

crates/db-mem/src/journal.rs

@@ -4,8 +4,8 @@ use concepts::storage::{
     CreateRequest, DbErrorWrite, DbErrorWriteNonRetriable, ExecutionEvent, ExecutionRequest,
     HistoryEvent, JoinSetRequest, JoinSetResponse, JoinSetResponseEvent, JoinSetResponseEventOuter,
     Locked, LockedBy, PendingStateBlockedByJoinSet, PendingStateFinished,
-    PendingStateFinishedResultKind, PendingStateLocked, PendingStatePendingAt, ResponseCursor,
-    ResponseWithCursor, VersionType,
+    PendingStateFinishedResultKind, PendingStateLocked, PendingStatePaused, PendingStatePendingAt,
+    ResponseCursor, ResponseWithCursor, VersionType,
 };
 use concepts::storage::{ExecutionLog, PendingState, Version};
 use concepts::{ComponentId, JoinSetId};
@@ -297,41 +297,47 @@ impl ExecutionJournal {
         }
     }
 
-    fn update_pending_state(&mut self) {
+    fn find_current_pending_state(&self) -> PendingState {
+        if let Some(last_event) = self.execution_events.last()
+            && let ExecutionRequest::Finished { result, .. } = &last_event.event
+        {
+            let idx = self.execution_events.len() - 1;
+            return PendingState::Finished(PendingStateFinished {
+                version: VersionType::try_from(idx).expect("version limit reached"),
+                finished_at: last_event.created_at,
+                result_kind: PendingStateFinishedResultKind::from(result),
+            });
+        }
+
         let mut unpause_encountered = false;
+        let mut is_paused = false;
 
-        let pending_state = self
+        // Find the underlying state (ignoring Paused/Unpaused for now), store it in
+        // `PendingStatePaused` independent of whether the execution is actually paused.
+        let underlying_state: PendingStatePaused = self
             .execution_events
             .iter()
             .enumerate()
             .rev()
-            .find_map(|(idx, event)| match &event.event {
+            .find_map(|(_idx, event)| match &event.event {
+                ExecutionRequest::Finished { .. } => {
+                    unreachable!("finished state was already handled above")
+                }
                 ExecutionRequest::Created { scheduled_at, .. } => {
-                    Some(PendingState::PendingAt(PendingStatePendingAt {
+                    Some(PendingStatePaused::PendingAt(PendingStatePendingAt {
                         scheduled_at: *scheduled_at,
                         last_lock: None,
                     }))
                 }
 
-                ExecutionRequest::Finished { result, .. } => {
-                    assert_eq!(self.execution_events.len() - 1, idx);
-                    Some(PendingState::Finished {
-                        finished: PendingStateFinished {
-                            version: VersionType::try_from(idx).expect("version limit reached"),
-                            finished_at: event.created_at,
-                            result_kind: PendingStateFinishedResultKind::from(result),
-                        },
-                    })
-                }
-
                 ExecutionRequest::Locked(Locked {
                     executor_id,
                     lock_expires_at,
                     run_id,
                     component_id: _,
                     deployment_id: _,
                     retry_config: _,
-                }) => Some(PendingState::Locked(PendingStateLocked {
+                }) => Some(PendingStatePaused::Locked(PendingStateLocked {
                     locked_by: LockedBy {
                         executor_id: *executor_id,
                         run_id: *run_id,
@@ -350,7 +356,7 @@ impl ExecutionJournal {
                 | ExecutionRequest::Unlocked {
                     backoff_expires_at: expires_at,
                     ..
-                } => Some(PendingState::PendingAt(PendingStatePendingAt {
+                } => Some(PendingStatePaused::PendingAt(PendingStatePendingAt {
                     scheduled_at: *expires_at,
                     last_lock: self.find_last_lock().map(LockedBy::from),
                 })),
@@ -393,13 +399,13 @@ impl ExecutionJournal {
                     if let Some(nth_created_at) = resp {
                         // Original executor has a chance to continue, but after expiry any executor can pick up the execution.
                         let scheduled_at = max(*lock_expires_at, *nth_created_at);
-                        Some(PendingState::PendingAt(PendingStatePendingAt {
+                        Some(PendingStatePaused::PendingAt(PendingStatePendingAt {
                             scheduled_at,
                             last_lock: self.find_last_lock().map(LockedBy::from),
                         }))
                     } else {
                         // Still waiting for response
-                        Some(PendingState::BlockedByJoinSet(
+                        Some(PendingStatePaused::BlockedByJoinSet(
                             PendingStateBlockedByJoinSet {
                                 join_set_id: expected_join_set_id.clone(),
                                 lock_expires_at: *lock_expires_at,
@@ -419,9 +425,9 @@ impl ExecutionJournal {
                         unpause_encountered = false;
                         None
                     } else {
-                        // No unpauses were found with bigger version
-                        unpause_encountered = false; // For sanity check below.
-                        Some(PendingState::Paused)
+                        // No unpauses were found in the later events - execution is paused
+                        is_paused = true;
+                        None // Continue looking for underlying state
                     }
                 }
                 // No pending state change for following events:
@@ -444,7 +450,29 @@ impl ExecutionJournal {
             .expect("journal must begin with Created event");
 
         assert!(!unpause_encountered, "unpause must be preceeded with pause");
-        self.pending_state = pending_state;
+
+        // Check if execution is finished (overrides paused state)
+
+        {
+            if is_paused {
+                PendingState::Paused(underlying_state)
+            } else {
+                // Convert PendingStatePaused to PendingState
+                match underlying_state {
+                    PendingStatePaused::Locked(locked) => PendingState::Locked(locked),
+                    PendingStatePaused::PendingAt(pending_at) => {
+                        PendingState::PendingAt(pending_at)
+                    }
+                    PendingStatePaused::BlockedByJoinSet(blocked) => {
+                        PendingState::BlockedByJoinSet(blocked)
+                    }
+                }
+            }
+        }
+    }
+
+    fn update_pending_state(&mut self) {
+        self.pending_state = self.find_current_pending_state();
         self.component_id = self
             .find_last_lock()
             .map(|locked| locked.component_id.clone())