Fix snapshot recovery

Author: kmatasfp

cli/golem-cli/src/command_handler/worker/stream.rs

@@ -328,6 +328,39 @@ impl WorkerConnection {
                                         .await;
                                 }
                             }
+                            AgentEvent::SnapshotRecoverySucceeded {
+                                timestamp,
+                                snapshot_index,
+                            } => {
+                                if matching {
+                                    output
+                                        .emit_log(
+                                            timestamp,
+                                            LogLevel::Info,
+                                            "snapshot-recovery".to_string(),
+                                            format!("loaded snapshot from oplog index {snapshot_index}"),
+                                        )
+                                        .await;
+                                }
+                            }
+                            AgentEvent::SnapshotRecoveryFailed {
+                                timestamp,
+                                snapshot_index,
+                                error,
+                            } => {
+                                if matching {
+                                    output
+                                        .emit_log(
+                                            timestamp,
+                                            LogLevel::Error,
+                                            "snapshot-recovery".to_string(),
+                                            format!(
+                                                "failed to load snapshot from oplog index {snapshot_index}: {error}"
+                                            ),
+                                        )
+                                        .await;
+                                }
+                            }
                             AgentEvent::ClientLagged {
                                 number_of_missed_messages,
                             } => {

golem-api-grpc/proto/golem/worker/log_event.proto

@@ -14,6 +14,8 @@ message LogEvent {
     InvocationFinished invocation_finished = 5;
     ClientLagged client_lagged = 6;
     PluginError plugin_error = 7;
+    SnapshotRecoverySucceeded snapshot_recovery_succeeded = 8;
+    SnapshotRecoveryFailed snapshot_recovery_failed = 9;
   }
 }
 
@@ -61,6 +63,17 @@ message PluginError {
   string message = 3;
 }
 
+message SnapshotRecoverySucceeded {
+  google.protobuf.Timestamp timestamp = 1;
+  uint64 snapshot_index = 2;
+}
+
+message SnapshotRecoveryFailed {
+  google.protobuf.Timestamp timestamp = 1;
+  uint64 snapshot_index = 2;
+  string error = 3;
+}
+
 /**
  * The client fell behind and the point it left of is no longer in our buffer.
  * {number_of_skipped_messages} is the number of messages between the client left of and the point it is now at.

golem-common/src/model/mod.rs

@@ -1264,6 +1264,15 @@ pub enum AgentEvent {
         plugin_name: String,
         message: String,
     },
+    SnapshotRecoverySucceeded {
+        timestamp: Timestamp,
+        snapshot_index: OplogIndex,
+    },
+    SnapshotRecoveryFailed {
+        timestamp: Timestamp,
+        snapshot_index: OplogIndex,
+        error: String,
+    },
     /// The client fell behind and the point it left of is no longer in our buffer.
     /// {number_of_skipped_messages} is the number of messages between the client left of and the point it is now at.
     ClientLagged { number_of_missed_messages: u64 },
@@ -1315,6 +1324,16 @@ impl Display for AgentEvent {
             } => {
                 write!(f, "<plugin-error> [{plugin_name}] {message}")
             }
+            AgentEvent::SnapshotRecoverySucceeded { snapshot_index, .. } => {
+                write!(f, "<snapshot-recovery-succeeded> {snapshot_index}")
+            }
+            AgentEvent::SnapshotRecoveryFailed {
+                snapshot_index,
+                error,
+                ..
+            } => {
+                write!(f, "<snapshot-recovery-failed> {snapshot_index}: {error}")
+            }
             AgentEvent::ClientLagged {
                 number_of_missed_messages,
             } => {

golem-common/src/model/protobuf.rs

@@ -545,6 +545,19 @@ impl TryFrom<golem::worker::LogEvent> for AgentEvent {
                         message: event.message,
                     })
                 }
+                golem::worker::log_event::Event::SnapshotRecoverySucceeded(event) => {
+                    Ok(AgentEvent::SnapshotRecoverySucceeded {
+                        timestamp: event.timestamp.ok_or("Missing timestamp")?.into(),
+                        snapshot_index: OplogIndex::from_u64(event.snapshot_index),
+                    })
+                }
+                golem::worker::log_event::Event::SnapshotRecoveryFailed(event) => {
+                    Ok(AgentEvent::SnapshotRecoveryFailed {
+                        timestamp: event.timestamp.ok_or("Missing timestamp")?.into(),
+                        snapshot_index: OplogIndex::from_u64(event.snapshot_index),
+                        error: event.error,
+                    })
+                }
             },
             None => Err("Missing event".to_string()),
         }
@@ -631,6 +644,30 @@ impl TryFrom<AgentEvent> for golem::worker::LogEvent {
                     },
                 )),
             }),
+            AgentEvent::SnapshotRecoverySucceeded {
+                timestamp,
+                snapshot_index,
+            } => Ok(golem::worker::LogEvent {
+                event: Some(golem::worker::log_event::Event::SnapshotRecoverySucceeded(
+                    golem::worker::SnapshotRecoverySucceeded {
+                        timestamp: Some(timestamp.into()),
+                        snapshot_index: snapshot_index.into(),
+                    },
+                )),
+            }),
+            AgentEvent::SnapshotRecoveryFailed {
+                timestamp,
+                snapshot_index,
+                error,
+            } => Ok(golem::worker::LogEvent {
+                event: Some(golem::worker::log_event::Event::SnapshotRecoveryFailed(
+                    golem::worker::SnapshotRecoveryFailed {
+                        timestamp: Some(timestamp.into()),
+                        snapshot_index: snapshot_index.into(),
+                        error,
+                    },
+                )),
+            }),
             AgentEvent::ClientLagged {
                 number_of_missed_messages,
             } => Ok(golem::worker::LogEvent {

golem-debugging-service/src/model/params.rs

@@ -120,6 +120,8 @@ impl LogNotification {
                 message,
             }),
             InternalWorkerEvent::PluginError { .. } => None,
+            InternalWorkerEvent::SnapshotRecoverySucceeded { .. } => None,
+            InternalWorkerEvent::SnapshotRecoveryFailed { .. } => None,
         }
     }
 }

golem-test-framework/src/dsl/mod.rs

@@ -1013,6 +1013,15 @@ pub fn log_event_to_string(event: &LogEvent) -> String {
         Some(log_event::Event::InvocationStarted(_)) => "".to_string(),
         Some(log_event::Event::ClientLagged { .. }) => "".to_string(),
         Some(log_event::Event::PluginError(err)) => err.message.clone(),
+        Some(log_event::Event::SnapshotRecoverySucceeded(event)) => {
+            format!("snapshot recovery loaded {}", event.snapshot_index)
+        }
+        Some(log_event::Event::SnapshotRecoveryFailed(event)) => {
+            format!(
+                "snapshot recovery failed {}: {}",
+                event.snapshot_index, event.error
+            )
+        }
         None => std::panic!("Unexpected event type"),
     }
 }

golem-worker-executor/src/durable_host/durability.rs

@@ -769,7 +769,7 @@ impl<Ctx: WorkerCtx> InFunctionRetryHost for DurableWorkerCtx<Ctx> {
 
     fn durable_execution_state(&self) -> DurableExecutionState {
         DurableExecutionState {
-            is_live: self.state.is_live(),
+            is_live: self.state.is_live() || self.state.snapshotting_mode.is_some(),
             persistence_level: self.state.persistence_level,
             snapshotting_mode: self.state.snapshotting_mode,
             assume_idempotence: self.state.assume_idempotence,
@@ -782,6 +782,10 @@ impl<Ctx: WorkerCtx> InFunctionRetryHost for DurableWorkerCtx<Ctx> {
         retry_from: OplogIndex,
         retry_policy_state: Option<RetryPolicyState>,
     ) {
+        if self.state.snapshotting_mode.is_some() {
+            return;
+        }
+
         use golem_common::model::oplog::AgentError;
         let inside_atomic_region = self.state.outermost_atomic_region_has_side_effects();
         let entry = OplogEntry::error(
@@ -816,13 +820,14 @@ impl<Ctx: WorkerCtx> DurabilityHost for DurableWorkerCtx<Ctx> {
         forced_commit: bool,
     ) -> Result<(), WorkerExecutorError> {
         self.end_function(function_type, begin_index).await?;
-        if function_type == &DurableFunctionType::WriteRemote
-            || matches!(function_type, DurableFunctionType::WriteRemoteBatched(_))
-            || matches!(
-                function_type,
-                DurableFunctionType::WriteRemoteTransaction(_)
-            )
-            || forced_commit
+        if self.state.snapshotting_mode.is_none()
+            && (function_type == &DurableFunctionType::WriteRemote
+                || matches!(function_type, DurableFunctionType::WriteRemoteBatched(_))
+                || matches!(
+                    function_type,
+                    DurableFunctionType::WriteRemoteTransaction(_)
+                )
+                || forced_commit)
         {
             self.public_state
                 .worker()

golem-worker-executor/src/durable_host/mod.rs

@@ -1183,6 +1183,12 @@ impl<Ctx: WorkerCtx> DurableWorkerCtx<Ctx> {
         &mut self,
         function_type: &DurableFunctionType,
     ) -> Result<OplogIndex, WorkerExecutorError> {
+        if self.state.snapshotting_mode.is_some() {
+            let begin_index = self.state.current_oplog_index().await;
+            self.state.current_retry_point = begin_index;
+            return Ok(begin_index);
+        }
+
         if (*function_type == DurableFunctionType::WriteRemote && !self.state.assume_idempotence)
             || matches!(
                 *function_type,
@@ -1325,6 +1331,10 @@ impl<Ctx: WorkerCtx> DurableWorkerCtx<Ctx> {
         function_type: &DurableFunctionType,
         begin_index: OplogIndex,
     ) -> Result<(), WorkerExecutorError> {
+        if self.state.snapshotting_mode.is_some() {
+            return Ok(());
+        }
+
         if (*function_type == DurableFunctionType::WriteRemote && !self.state.assume_idempotence)
             || matches!(
                 *function_type,
@@ -1717,11 +1727,6 @@ impl<Ctx: WorkerCtx> DurableWorkerCtx<Ctx> {
                             return Ok(Some(RetryDecision::Immediate));
                         }
 
-                        store
-                            .as_context_mut()
-                            .data_mut()
-                            .begin_call_snapshotting_function();
-
                         let load_result = invoke_observed_and_traced(
                             lowered,
                             store,
@@ -1733,7 +1738,8 @@ impl<Ctx: WorkerCtx> DurableWorkerCtx<Ctx> {
                         store
                             .as_context_mut()
                             .data_mut()
-                            .end_call_snapshotting_function();
+                            .durable_ctx_mut()
+                            .end_call_snapshotting_function_if_active();
 
                         for span_id in local_span_ids {
                             let _ = store
@@ -1877,9 +1883,11 @@ impl<Ctx: WorkerCtx> DurableWorkerCtx<Ctx> {
                 ..
             } => (payload, mime_type),
             _ => {
-                warn!(
+                let error = format!(
                     "Expected Snapshot entry at oplog index {snapshot_index}, found different entry; falling back to full replay"
                 );
+                warn!("{error}");
+                Self::emit_snapshot_recovery_event(store, snapshot_index, false, Some(error));
                 if let Err(err) = store
                     .as_context_mut()
                     .data_mut()
@@ -1906,7 +1914,11 @@ impl<Ctx: WorkerCtx> DurableWorkerCtx<Ctx> {
         {
             Ok(data) => data,
             Err(err) => {
-                warn!("Failed to download snapshot payload: {err}; falling back to full replay");
+                let error = format!(
+                    "Failed to download snapshot payload: {err}; falling back to full replay"
+                );
+                warn!("{error}");
+                Self::emit_snapshot_recovery_event(store, snapshot_index, false, Some(error));
                 if let Err(err) = store
                     .as_context_mut()
                     .data_mut()
@@ -1950,7 +1962,10 @@ impl<Ctx: WorkerCtx> DurableWorkerCtx<Ctx> {
         ) {
             Ok(lowered) => lowered,
             Err(err) => {
-                warn!("Snapshot recovery failed to lower load-snapshot invocation: {err}");
+                let error =
+                    format!("Snapshot recovery failed to lower load-snapshot invocation: {err}");
+                warn!("{error}");
+                Self::emit_snapshot_recovery_event(store, snapshot_index, false, Some(error));
                 return SnapshotRecoveryResult::Failed;
             }
         };
@@ -1964,13 +1979,16 @@ impl<Ctx: WorkerCtx> DurableWorkerCtx<Ctx> {
             .set_current_invocation_context(invocation_context)
             .await
         {
-            warn!("Snapshot recovery failed to install invocation context: {err}");
+            let error = format!("Snapshot recovery failed to install invocation context: {err}");
+            warn!("{error}");
+            Self::emit_snapshot_recovery_event(store, snapshot_index, false, Some(error));
             return SnapshotRecoveryResult::Failed;
         }
 
         store
             .as_context_mut()
             .data_mut()
+            .durable_ctx_mut()
             .begin_call_snapshotting_function();
 
         let load_result =
@@ -1979,7 +1997,8 @@ impl<Ctx: WorkerCtx> DurableWorkerCtx<Ctx> {
         store
             .as_context_mut()
             .data_mut()
-            .end_call_snapshotting_function();
+            .durable_ctx_mut()
+            .end_call_snapshotting_function_if_active();
 
         for span_id in local_span_ids {
             let _ = store
@@ -2024,12 +2043,38 @@ impl<Ctx: WorkerCtx> DurableWorkerCtx<Ctx> {
 
         if let Some(error) = failed {
             warn!("{error}; re-creating instance for full replay");
+            Self::emit_snapshot_recovery_event(store, snapshot_index, false, Some(error));
             SnapshotRecoveryResult::Failed
         } else {
             debug!("Snapshot loaded successfully from oplog index {snapshot_index}");
+            Self::emit_snapshot_recovery_event(store, snapshot_index, true, None);
             SnapshotRecoveryResult::Success
         }
     }
+
+    fn emit_snapshot_recovery_event(
+        store: &mut (impl AsContextMut<Data = Ctx> + Send),
+        snapshot_index: OplogIndex,
+        succeeded: bool,
+        error: Option<String>,
+    ) {
+        store
+            .as_context_mut()
+            .data_mut()
+            .get_public_state()
+            .event_service()
+            .emit_event(
+                if succeeded {
+                    InternalWorkerEvent::snapshot_recovery_succeeded(snapshot_index)
+                } else {
+                    InternalWorkerEvent::snapshot_recovery_failed(
+                        snapshot_index,
+                        error.unwrap_or_else(|| "unknown".to_string()),
+                    )
+                },
+                true,
+            );
+    }
 }
 
 enum SnapshotRecoveryResult {
@@ -2039,6 +2084,12 @@ enum SnapshotRecoveryResult {
 }
 
 impl<Ctx: WorkerCtx> DurableWorkerCtx<Ctx> {
+    pub(crate) fn end_call_snapshotting_function_if_active(&mut self) {
+        if self.state.snapshotting_mode.is_some() {
+            self.end_call_snapshotting_function();
+        }
+    }
+
     pub(crate) fn register_open_websocket(
         &mut self,
         rep: u32,
@@ -2601,6 +2652,14 @@ impl<Ctx: WorkerCtx> UpdateManagement for DurableWorkerCtx<Ctx> {
         // While calling a snapshotting function (load/save), we completely turn off persistence
         // In addition to the user-controllable persistence level we also skip writing the
         // oplog entries marking the exported function call.
+        if self.state.snapshotting_mode.is_some() {
+            warn!(
+                "begin_call_snapshotting_function called while snapshotting is already active; \
+                 leaving persistence level unchanged"
+            );
+            return;
+        }
+
         let previous_level = self.state.persistence_level;
         self.state.snapshotting_mode = Some(previous_level);
         self.state.persistence_level = PersistenceLevel::PersistNothing;