feat(workflows): implement retry backoff for activity errors

Author: MasterPtato

lib/chirp-workflow/core/src/ctx/workflow.rs

@@ -21,8 +21,8 @@ use crate::{
 	workflow::{Workflow, WorkflowInput},
 };
 
-// Time to delay a worker from retrying after an error
-const RETRY_TIMEOUT: Duration = Duration::from_millis(2000);
+// Time to delay a workflow from retrying after an error
+pub const RETRY_TIMEOUT_MS: usize = 2000;
 // Poll interval when polling for signals in-process
 const SIGNAL_RETRY: Duration = Duration::from_millis(100);
 // Most in-process signal poll tries
@@ -200,8 +200,10 @@ impl WorkflowCtx {
 			}
 			Err(err) => {
 				// Retry the workflow if its recoverable
-				let deadline = if err.is_recoverable() {
-					Some(rivet_util::timestamp::now() + RETRY_TIMEOUT.as_millis() as i64)
+				let deadline_ts = if let Some(deadline_ts) = err.backoff() {
+					Some(deadline_ts)
+				} else if err.is_recoverable() {
+					Some(rivet_util::timestamp::now() + RETRY_TIMEOUT_MS as i64)
 				} else {
 					None
 				};
@@ -214,7 +216,7 @@ impl WorkflowCtx {
 				// finish. This workflow will be retried when the sub workflow completes
 				let wake_sub_workflow = err.sub_workflow();
 
-				if deadline.is_some() || !wake_signals.is_empty() || wake_sub_workflow.is_some() {
+				if deadline_ts.is_some() || !wake_signals.is_empty() || wake_sub_workflow.is_some() {
 					tracing::info!(name=%self.name, id=%self.workflow_id, ?err, "workflow sleeping");
 				} else {
 					tracing::error!(name=%self.name, id=%self.workflow_id, ?err, "workflow error");
@@ -236,7 +238,7 @@ impl WorkflowCtx {
 						.fail_workflow(
 							self.workflow_id,
 							false,
-							deadline,
+							deadline_ts,
 							wake_signals,
 							wake_sub_workflow,
 							&err_str,
@@ -319,7 +321,7 @@ impl WorkflowCtx {
 					)
 					.await?;
 
-				Err(WorkflowError::ActivityFailure(err))
+				Err(WorkflowError::ActivityFailure(err, 0))
 			}
 			Err(err) => {
 				tracing::debug!("activity timeout");
@@ -643,11 +645,12 @@ impl WorkflowCtx {
 						// Convert error in the case of max retries exceeded. This will only act on retryable
 						// errors
 						let err = match err {
-							WorkflowError::ActivityFailure(err) => {
+							WorkflowError::ActivityFailure(err, _) => {
 								if error_count + 1 >= I::Activity::MAX_RETRIES {
 									WorkflowError::ActivityMaxFailuresReached(err)
 								} else {
-									WorkflowError::ActivityFailure(err)
+									// Add error count to the error
+									WorkflowError::ActivityFailure(err, error_count)
 								}
 							}
 							WorkflowError::ActivityTimeout | WorkflowError::OperationTimeout => {

lib/chirp-workflow/core/src/error.rs

@@ -1,6 +1,11 @@
+use std::time::{SystemTime, UNIX_EPOCH};
+
+use tokio::time::Instant;
 use global_error::GlobalError;
 use uuid::Uuid;
 
+use crate::ctx::workflow::RETRY_TIMEOUT_MS;
+
 pub type WorkflowResult<T> = Result<T, WorkflowError>;
 
 /// Throwing this will eject from the workflow scope back in to the engine.
@@ -12,8 +17,9 @@ pub enum WorkflowError {
 	#[error("workflow failure: {0:?}")]
 	WorkflowFailure(GlobalError),
 
+	// Includes error count
 	#[error("activity failure: {0:?}")]
-	ActivityFailure(GlobalError),
+	ActivityFailure(GlobalError, u32),
 
 	#[error("activity failure, max retries reached: {0:?}")]
 	ActivityMaxFailuresReached(GlobalError),
@@ -119,9 +125,31 @@ pub enum WorkflowError {
 }
 
 impl WorkflowError {
+	pub fn backoff(&self) -> Option<i64> {
+		if let WorkflowError::ActivityFailure(_, error_count) = self {
+			// NOTE: Max retry is handled in `WorkflowCtx::activity`
+			let mut backoff =
+				rivet_util::Backoff::new_at(8, None, RETRY_TIMEOUT_MS, 500, *error_count as usize);
+			let next = backoff.step().expect("should not have max retry");
+
+			// Calculate timestamp based on the backoff
+			let duration_until = next.duration_since(Instant::now());
+			let deadline_ts = (SystemTime::now() + duration_until)
+			.duration_since(UNIX_EPOCH)
+			.unwrap_or_else(|err| unreachable!("time is broken: {}", err))
+			.as_millis()
+			.try_into()
+			.expect("doesn't fit in i64");
+
+			Some(deadline_ts)
+		} else {
+			None
+		}
+	}
+	
 	pub fn is_recoverable(&self) -> bool {
 		match self {
-			WorkflowError::ActivityFailure(_) => true,
+			WorkflowError::ActivityFailure(_, _) => true,
 			WorkflowError::ActivityTimeout => true,
 			WorkflowError::OperationTimeout => true,
 			_ => false,

lib/chirp/client/src/client.rs

@@ -411,17 +411,17 @@ impl Client {
 				{
 					tracing::warn!(tick_index = %service_unavailable_backoff.tick_index(), "service unavailable");
 					if service_unavailable_backoff.tick().await {
-						return Err(ClientError::NatsResponseStatus(status));
-					} else {
 						continue 'req;
+					} else {
+						return Err(ClientError::NatsResponseStatus(status));
 					}
 				}
 				Err(ClientError::RpcAckTimedOut) => {
 					tracing::warn!(tick_index = %service_unavailable_backoff.tick_index(), "rpc ack timed out");
 					if service_unavailable_backoff.tick().await {
-						return Err(ClientError::RpcAckTimedOut);
-					} else {
 						continue 'req;
+					} else {
+						return Err(ClientError::RpcAckTimedOut);
 					}
 				}
 				Err(err) => {

lib/chirp/worker/src/manager.rs

@@ -890,10 +890,10 @@ where
 				tracing::info!("ticking request retry backoff");
 
 				if backoff.tick().await {
+					tracing::info!("retrying request");
+				} else {
 					tracing::warn!("retry request failed too many times");
 					return res;
-				} else {
-					tracing::info!("retrying request");
 				}
 			} else {
 				// Return result immediately
@@ -1040,7 +1040,7 @@ where
 	async fn consumer_ack(self: Arc<Self>, msg_meta: RedisMessageMeta) {
 		// let mut backoff = rivet_util::Backoff::default();
 		// loop {
-		// 	if backoff.tick().await {
+		// 	if !backoff.tick().await {
 		// 		tracing::error!("acking stream message failed too many times, aborting");
 		// 		return;
 		// 	}

lib/util/core/src/lib.rs

@@ -150,17 +150,34 @@ impl Backoff {
 			sleep_until: Instant::now(),
 		}
 	}
+	
+	pub fn new_at(
+		max_exponent: usize,
+		max_retries: Option<usize>,
+		wait: usize,
+		randomness: usize,
+		i: usize,
+	) -> Backoff {
+		Backoff {
+			max_exponent,
+			max_retries,
+			wait,
+			randomness,
+			i,
+			sleep_until: Instant::now(),
+		}
+	}
 
 	pub fn tick_index(&self) -> usize {
 		self.i
 	}
 
 	/// Waits for the next backoff tick.
 	///
-	/// Returns true if the index is greater than `max_retries`.
+	/// Returns false if the index is greater than `max_retries`.
 	pub async fn tick(&mut self) -> bool {
 		if self.max_retries.map_or(false, |x| self.i > x) {
-			return true;
+			return false;
 		}
 
 		tokio::time::sleep_until(self.sleep_until.into()).await;
@@ -171,7 +188,24 @@ impl Backoff {
 
 		self.i += 1;
 
-		false
+		true
+	}
+
+	/// Returns the instant of the next backoff tick. Does not wait.
+	///
+	/// Returns None if the index is greater than `max_retries`.
+	pub fn step(&mut self) -> Option<Instant> {
+		if self.max_retries.map_or(false, |x| self.i > x) {
+			return None;
+		}
+
+		let next_wait = self.wait * 2usize.pow(self.i.min(self.max_exponent) as u32)
+			+ rand::thread_rng().gen_range(0..self.randomness);
+		self.sleep_until += Duration::from_millis(next_wait as u64);
+
+		self.i += 1;
+
+		Some(self.sleep_until)
 	}
 
 	pub fn default_infinite() -> Backoff {
@@ -202,7 +236,7 @@ mod tests {
 			last_tick = now;
 			println!("tick: {}", dt.as_secs_f64());
 
-			if backoff.tick().await {
+			if !backoff.tick().await {
 				println!("cancelling");
 				break;
 			}