Merge pull request #2175 from golemcloud/retry-improvements

Retry logic improvements

Author: vigoo

Cargo.lock

@@ -481,10 +481,11 @@ dependencies = [
 [tasks.integration-tests-group1]
 description = "Runs integration tests only"
 dependencies = ["build-bins"]
-env = { "RUST_LOG" = "info", "RUST_BACKTRACE" = "1" }
+env = { "RUST_LOG" = "debug", "RUST_BACKTRACE" = "1" }
 script = '''
-cargo test --package integration-tests --test integration -- --nocapture --report-time $JUNIT_OPTS
+cargo test --package integration-tests --test integration -- --nocapture --test-threads=1 --report-time
 '''
+# cargo test --package integration-tests --test integration -- --nocapture --report-time $JUNIT_OPTS
 
 [tasks.integration-tests-group2]
 description = "Runs component service integration tests only"

Makefile.toml

@@ -119,6 +119,7 @@ message ExportedFunctionCompletedParameters {
 message ErrorParameters {
   google.protobuf.Timestamp timestamp = 1;
   string error = 2;
+  uint64 retry_from = 3;
 }
 
 message JumpParameters {

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

@@ -305,3 +305,40 @@ impl From<String> for TransactionId {
         TransactionId(value)
     }
 }
+
+#[cfg(feature = "sql")]
+mod sql {
+    use crate::model::TransactionId;
+    use sqlx::encode::IsNull;
+    use sqlx::error::BoxDynError;
+    use sqlx::postgres::PgTypeInfo;
+    use sqlx::{Database, Postgres, Type};
+    use std::io::Write;
+
+    impl sqlx::Decode<'_, Postgres> for TransactionId {
+        fn decode(value: <Postgres as Database>::ValueRef<'_>) -> Result<Self, BoxDynError> {
+            let bytes = value.as_bytes()?;
+            Ok(TransactionId(
+                u64::from_be_bytes(bytes.try_into()?).to_string(),
+            ))
+        }
+    }
+
+    impl sqlx::Encode<'_, Postgres> for TransactionId {
+        fn encode_by_ref(
+            &self,
+            buf: &mut <Postgres as Database>::ArgumentBuffer<'_>,
+        ) -> Result<IsNull, BoxDynError> {
+            let u64 = self.0.parse::<u64>()?;
+            let bytes = u64.to_be_bytes();
+            buf.write_all(&bytes)?;
+            Ok(IsNull::No)
+        }
+    }
+
+    impl Type<Postgres> for TransactionId {
+        fn type_info() -> PgTypeInfo {
+            PgTypeInfo::with_name("xid8")
+        }
+    }
+}

golem-common/src/base_model.rs

@@ -625,7 +625,9 @@ pub struct WorkerStatusRecord {
     /// The component version at the starting point of the replay. Will be the version of the Create oplog entry
     /// if only automatic updates were used or the version of the latest snapshot-based update
     pub component_version_for_replay: ComponentVersion,
-    pub current_retry_count: u32,
+    /// The number of encountered error entries grouped by their 'retry_from' index, calculated from
+    /// the last invocation boundary.
+    pub current_retry_count: HashMap<OplogIndex, u32>,
 }
 
 impl<Context> bincode::Decode<Context> for WorkerStatusRecord {
@@ -699,7 +701,7 @@ impl Default for WorkerStatusRecord {
             active_plugins: HashSet::new(),
             deleted_regions: DeletedRegions::new(),
             component_version_for_replay: 0,
-            current_retry_count: 0,
+            current_retry_count: HashMap::new(),
         }
     }
 }

golem-common/src/model/mod.rs

@@ -287,6 +287,11 @@ pub enum OplogEntry {
     Error {
         timestamp: Timestamp,
         error: WorkerError,
+        /// Points to the oplog index where the retry should start from. Normally this can be just the
+        /// current oplog index (after the last persisted side-effect). When failing in an atomic region
+        /// or batched remote writes, this should point to the start of the region.
+        /// When counting the number of retries for a specific error, the error entries are grouped by this index.
+        retry_from: OplogIndex,
     },
     /// Marker entry added when get-oplog-index is called from the worker, to make the jumping behavior
     /// more predictable.
@@ -455,22 +460,31 @@ pub enum OplogEntry {
         timestamp: Timestamp,
         level: PersistenceLevel,
     },
+    /// Marks the beginning of a remote transaction
     BeginRemoteTransaction {
         timestamp: Timestamp,
         transaction_id: TransactionId,
+        /// BeginRemoteTransaction entries need to be repeated on retries, because they may need a new
+        /// transaction_id. The `begin_index` field always points to the original, first entry. This makes
+        /// error grouping work. When None, this is the original begin entry.
+        original_begin_index: Option<OplogIndex>,
     },
+    /// Marks the point before a remote transaction is committed
     PreCommitRemoteTransaction {
         timestamp: Timestamp,
         begin_index: OplogIndex,
     },
+    /// Marks the point before a remote transaction is rolled back
     PreRollbackRemoteTransaction {
         timestamp: Timestamp,
         begin_index: OplogIndex,
     },
+    /// Marks the point after a remote transaction is committed
     CommittedRemoteTransaction {
         timestamp: Timestamp,
         begin_index: OplogIndex,
     },
+    /// Marks the point after a remote transaction is rolled back
     RolledBackRemoteTransaction {
         timestamp: Timestamp,
         begin_index: OplogIndex,
@@ -526,10 +540,11 @@ impl OplogEntry {
         }
     }
 
-    pub fn error(error: WorkerError) -> OplogEntry {
+    pub fn error(error: WorkerError, retry_from: OplogIndex) -> OplogEntry {
         OplogEntry::Error {
             timestamp: Timestamp::now_utc(),
             error,
+            retry_from,
         }
     }
 
@@ -718,10 +733,14 @@ impl OplogEntry {
         }
     }
 
-    pub fn begin_remote_transaction(transaction_id: TransactionId) -> OplogEntry {
+    pub fn begin_remote_transaction(
+        transaction_id: TransactionId,
+        original_begin_index: Option<OplogIndex>,
+    ) -> OplogEntry {
         OplogEntry::BeginRemoteTransaction {
             timestamp: Timestamp::now_utc(),
             transaction_id,
+            original_begin_index,
         }
     }
 
@@ -761,6 +780,10 @@ impl OplogEntry {
         matches!(self, OplogEntry::EndRemoteWrite { begin_index, .. } if *begin_index == idx)
     }
 
+    pub fn is_end_remote_write_s<S>(&self, idx: OplogIndex, _: &S) -> bool {
+        matches!(self, OplogEntry::EndRemoteWrite { begin_index, .. } if *begin_index == idx)
+    }
+
     pub fn is_pre_commit_remote_transaction(&self, idx: OplogIndex) -> bool {
         matches!(self, OplogEntry::PreCommitRemoteTransaction { begin_index, .. } if *begin_index == idx)
     }
@@ -773,43 +796,79 @@ impl OplogEntry {
         self.is_pre_commit_remote_transaction(idx) || self.is_pre_rollback_remote_transaction(idx)
     }
 
+    pub fn is_pre_remote_transaction_s<S>(&self, idx: OplogIndex, _: &S) -> bool {
+        self.is_pre_commit_remote_transaction(idx) || self.is_pre_rollback_remote_transaction(idx)
+    }
+
     pub fn is_committed_remote_transaction(&self, idx: OplogIndex) -> bool {
         matches!(self, OplogEntry::CommittedRemoteTransaction { begin_index, .. } if *begin_index == idx)
     }
 
+    pub fn is_committed_remote_transaction_s<S>(&self, idx: OplogIndex, _: &S) -> bool {
+        matches!(self, OplogEntry::CommittedRemoteTransaction { begin_index, .. } if *begin_index == idx)
+    }
+
     pub fn is_rolled_back_remote_transaction(&self, idx: OplogIndex) -> bool {
         matches!(self, OplogEntry::RolledBackRemoteTransaction { begin_index, .. } if *begin_index == idx)
     }
 
+    pub fn is_rolled_back_remote_transaction_s<S>(&self, idx: OplogIndex, _: &S) -> bool {
+        matches!(self, OplogEntry::RolledBackRemoteTransaction { begin_index, .. } if *begin_index == idx)
+    }
+
     pub fn is_end_remote_transaction(&self, idx: OplogIndex) -> bool {
         self.is_committed_remote_transaction(idx) || self.is_rolled_back_remote_transaction(idx)
     }
 
+    pub fn is_end_remote_transaction_s<S>(&self, idx: OplogIndex, s: &S) -> bool {
+        self.is_committed_remote_transaction_s(idx, s)
+            || self.is_rolled_back_remote_transaction_s(idx, s)
+    }
+
     /// Checks that an "intermediate oplog entry" between a `BeginRemoteWrite` and an `EndRemoteWrite`
     /// is not a RemoteWrite entry which does not belong to the batched remote write started at `idx`.
-    pub fn no_concurrent_side_effect(&self, idx: OplogIndex) -> bool {
-        match self {
-            OplogEntry::ImportedFunctionInvoked {
-                durable_function_type,
-                ..
-            } => match durable_function_type {
-                DurableFunctionType::WriteRemoteBatched(Some(begin_index))
-                    if *begin_index == idx =>
-                {
-                    true
-                }
-                DurableFunctionType::WriteRemoteTransaction(Some(begin_index))
-                    if *begin_index == idx =>
-                {
-                    true
-                }
-                DurableFunctionType::ReadLocal => true,
-                DurableFunctionType::WriteLocal => true,
-                DurableFunctionType::ReadRemote => true,
-                _ => false,
-            },
-            OplogEntry::ExportedFunctionCompleted { .. } => false,
-            _ => true,
+    /// Side effects in a PersistenceLevel::PersistNothing region are ignored.
+    pub fn no_concurrent_side_effect(
+        &self,
+        idx: OplogIndex,
+        persistence_level: &PersistenceLevel,
+    ) -> bool {
+        if persistence_level == &PersistenceLevel::PersistNothing {
+            true
+        } else {
+            match self {
+                OplogEntry::ImportedFunctionInvoked {
+                    durable_function_type,
+                    ..
+                } => match durable_function_type {
+                    DurableFunctionType::WriteRemoteBatched(Some(begin_index))
+                        if *begin_index == idx =>
+                    {
+                        true
+                    }
+                    DurableFunctionType::WriteRemoteTransaction(Some(begin_index))
+                        if *begin_index == idx =>
+                    {
+                        true
+                    }
+                    DurableFunctionType::ReadLocal => true,
+                    DurableFunctionType::WriteLocal => true,
+                    DurableFunctionType::ReadRemote => true,
+                    _ => false,
+                },
+                OplogEntry::ExportedFunctionCompleted { .. } => false,
+                _ => true,
+            }
+        }
+    }
+
+    pub fn track_persistence_level(
+        &self,
+        _idx: OplogIndex,
+        persistence_level: &mut PersistenceLevel,
+    ) {
+        if let OplogEntry::ChangePersistenceLevel { level, .. } = self {
+            *persistence_level = *level
         }
     }
 

golem-common/src/model/oplog.rs

@@ -390,6 +390,8 @@ pub struct TimestampParameter {
 pub struct ErrorParameters {
     pub timestamp: Timestamp,
     pub error: String,
+    #[wit_field(skip)]
+    pub retry_from: OplogIndex,
 }
 
 #[derive(Clone, Debug, Serialize, PartialEq, Deserialize)]

golem-common/src/model/public_oplog/mod.rs

@@ -226,6 +226,7 @@ impl TryFrom<golem_api_grpc::proto::golem::worker::OplogEntry> for PublicOplogEn
             oplog_entry::Entry::Error(error) => Ok(PublicOplogEntry::Error(ErrorParameters {
                 timestamp: error.timestamp.ok_or("Missing timestamp field")?.into(),
                 error: error.error,
+                retry_from: OplogIndex::from_u64(error.retry_from),
             })),
             oplog_entry::Entry::NoOp(no_op) => Ok(PublicOplogEntry::NoOp(TimestampParameter {
                 timestamp: no_op.timestamp.ok_or("Missing timestamp field")?.into(),
@@ -602,6 +603,7 @@ impl TryFrom<PublicOplogEntry> for golem_api_grpc::proto::golem::worker::OplogEn
                     golem_api_grpc::proto::golem::worker::ErrorParameters {
                         timestamp: Some(error.timestamp.into()),
                         error: error.error,
+                        retry_from: error.retry_from.0,
                     },
                 )),
             },

golem-common/src/model/public_oplog/protobuf.rs

@@ -180,6 +180,7 @@ fn error_serialization_poem_serde_equivalence() {
     let entry = PublicOplogEntry::Error(ErrorParameters {
         timestamp: rounded_ts(Timestamp::now_utc()),
         error: "test".to_string(),
+        retry_from: OplogIndex::INITIAL,
     });
     let serialized = entry.to_json_string();
     let deserialized: PublicOplogEntry = serde_json::from_str(&serialized).unwrap();

golem-common/src/model/public_oplog/tests.rs

@@ -15,6 +15,7 @@
 use crate::additional_deps::AdditionalDeps;
 use anyhow::Error;
 use async_trait::async_trait;
+use golem_common::base_model::OplogIndex;
 use golem_common::model::agent::AgentId;
 use golem_common::model::invocation_context::{
     self, AttributeValue, InvocationContextStack, SpanId,
@@ -237,6 +238,10 @@ impl InvocationHooks for DebugContext {
             .on_invocation_success(full_function_name, function_input, consumed_fuel, output)
             .await
     }
+
+    async fn get_current_retry_point(&self) -> OplogIndex {
+        self.durable_ctx.get_current_retry_point().await
+    }
 }
 
 #[async_trait]