feat: graceful worker shutdown (#2274)

<!-- Please make sure there is an issue that this PR is correlated to. -->
Fixes RVT-4594
## Changes

<!-- If there are frontend changes, please include screenshots. -->

Author: MasterPtato

examples/system-test/tests/client.ts

@@ -28,6 +28,7 @@ const client = new RivetClient({
 async function run() {
 	let actorId: string | undefined;
 	try {
+		console.time(`create actor`);
 		console.log("Creating actor", { region });
 		const { actor } = await client.actor.create({
 			project: RIVET_PROJECT,
@@ -54,16 +55,18 @@ async function run() {
 				},
 				...(BUILD_NAME === "ws-container"
 					? {
-							resources: {
-								cpu: 100,
-								memory: 100,
-							},
-						}
+						resources: {
+							cpu: 100,
+							memory: 100,
+						},
+					}
 					: {}),
 			},
 		});
 		actorId = actor.id;
 
+		console.timeEnd(`create actor`);
+
 		const port = actor.network.ports.http;
 
 		const actorOrigin = `${port.protocol}://${port.hostname}:${port.port}${port.path ?? ""}`;
@@ -84,13 +87,13 @@ async function run() {
 		//}
 
 		// Retry loop for HTTP health check
-		console.time(`ready-${actorId}`);
+		console.time(`ready ${actorId}`);
 		while (true) {
 			try {
 				const response = await fetch(`${actorOrigin}/health`);
 				if (response.ok) {
 					console.log("Health check passed");
-					console.timeEnd(`ready-${actorId}`);
+					console.timeEnd(`ready ${actorId}`);
 					break;
 				} else {
 					console.error(

packages/common/chirp-workflow/core/src/builder/workflow/message.rs

@@ -73,6 +73,8 @@ impl<'a, M: Message> MessageBuilder<'a, M> {
 
 	#[tracing::instrument(skip_all)]
 	pub async fn send(self) -> GlobalResult<()> {
+		self.ctx.check_stop().map_err(GlobalError::raw)?;
+
 		if let Some(err) = self.error {
 			return Err(err.into());
 		}

packages/common/chirp-workflow/core/src/builder/workflow/signal.rs

@@ -86,6 +86,8 @@ impl<'a, T: Signal + Serialize> SignalBuilder<'a, T> {
 
 	#[tracing::instrument(skip_all)]
 	pub async fn send(self) -> GlobalResult<Uuid> {
+		self.ctx.check_stop().map_err(GlobalError::raw)?;
+
 		if let Some(err) = self.error {
 			return Err(err.into());
 		}

packages/common/chirp-workflow/core/src/builder/workflow/sub_workflow.rs

@@ -88,6 +88,8 @@ where
 
 	#[tracing::instrument(skip_all)]
 	pub async fn dispatch(self) -> GlobalResult<Uuid> {
+		self.ctx.check_stop().map_err(GlobalError::raw)?;
+
 		if let Some(err) = self.error {
 			return Err(err.into());
 		}
@@ -237,6 +239,8 @@ where
 	pub async fn output(
 		self,
 	) -> GlobalResult<<<I as WorkflowInput>::Workflow as Workflow>::Output> {
+		self.ctx.check_stop().map_err(GlobalError::raw)?;
+
 		if let Some(err) = self.error {
 			return Err(err.into());
 		}
@@ -288,6 +292,8 @@ where
 		&self,
 		sub_workflow_id: Uuid,
 	) -> GlobalResult<<<I as WorkflowInput>::Workflow as Workflow>::Output> {
+		self.ctx.check_stop().map_err(GlobalError::raw)?;
+
 		tracing::debug!(name=%self.ctx.name(), id=%self.ctx.workflow_id(), sub_workflow_name=%I::Workflow::NAME, ?sub_workflow_id, "waiting for sub workflow");
 
 		let mut wake_sub = self.ctx.db().wake_sub().await?;

packages/common/chirp-workflow/core/src/ctx/listen.rs

@@ -10,7 +10,7 @@ use crate::{
 
 /// Indirection struct to prevent invalid implementations of listen traits.
 pub struct ListenCtx<'a> {
-	ctx: &'a mut WorkflowCtx,
+	ctx: &'a WorkflowCtx,
 	location: &'a Location,
 	// Used by certain db drivers to know when to update internal indexes for signal wake conditions
 	last_try: bool,
@@ -19,7 +19,7 @@ pub struct ListenCtx<'a> {
 }
 
 impl<'a> ListenCtx<'a> {
-	pub(crate) fn new(ctx: &'a mut WorkflowCtx, location: &'a Location) -> Self {
+	pub(crate) fn new(ctx: &'a WorkflowCtx, location: &'a Location) -> Self {
 		ListenCtx {
 			ctx,
 			location,

packages/common/chirp-workflow/core/src/ctx/workflow.rs

@@ -6,6 +6,7 @@ use std::{
 use futures_util::StreamExt;
 use global_error::{GlobalError, GlobalResult};
 use serde::{de::DeserializeOwned, Serialize};
+use tokio::sync::watch;
 use tracing::Instrument as _;
 use uuid::Uuid;
 
@@ -69,6 +70,8 @@ pub struct WorkflowCtx {
 	loop_location: Option<Location>,
 
 	msg_ctx: MessageCtx,
+	/// Used to stop workflow execution by the worker.
+	stop: watch::Receiver<()>,
 }
 
 impl WorkflowCtx {
@@ -79,6 +82,7 @@ impl WorkflowCtx {
 		config: rivet_config::Config,
 		conn: rivet_connection::Connection,
 		data: PulledWorkflowData,
+		stop: watch::Receiver<()>,
 	) -> GlobalResult<Self> {
 		let msg_ctx = MessageCtx::new(&conn, data.ray_id).await?;
 		let event_history = Arc::new(data.events);
@@ -105,6 +109,7 @@ impl WorkflowCtx {
 			loop_location: None,
 
 			msg_ctx,
+			stop,
 		})
 	}
 
@@ -135,6 +140,9 @@ impl WorkflowCtx {
 	pub(crate) async fn run(mut self) -> WorkflowResult<()> {
 		tracing::debug!(name=%self.name, id=%self.workflow_id, "running workflow");
 
+		// Check for stop before running
+		self.check_stop()?;
+
 		// Lookup workflow
 		let workflow = self.registry.get_workflow(&self.name)?;
 
@@ -176,8 +184,10 @@ impl WorkflowCtx {
 				}
 			}
 			Err(err) => {
+				let wake_immediate = err.wake_immediate();
+
 				// Retry the workflow if its recoverable
-				let deadline_ts = if let Some(deadline_ts) = err.deadline_ts() {
+				let wake_deadline_ts = if let Some(deadline_ts) = err.deadline_ts() {
 					Some(deadline_ts)
 				} else {
 					None
@@ -217,8 +227,8 @@ impl WorkflowCtx {
 						.commit_workflow(
 							self.workflow_id,
 							&self.name,
-							false,
-							deadline_ts,
+							wake_immediate,
+							wake_deadline_ts,
 							wake_signals,
 							wake_sub_workflow,
 							&err_str,
@@ -423,6 +433,7 @@ impl WorkflowCtx {
 			loop_location: self.loop_location.clone(),
 
 			msg_ctx: self.msg_ctx.clone(),
+			stop: self.stop.clone(),
 		}
 	}
 
@@ -434,6 +445,21 @@ impl WorkflowCtx {
 
 		branch
 	}
+
+	pub(crate) fn check_stop(&self) -> WorkflowResult<()> {
+		if self.stop.has_changed().unwrap_or(true) {
+			Err(WorkflowError::WorkflowStopped)
+		} else {
+			Ok(())
+		}
+	}
+
+	pub(crate) async fn wait_stop(&self) -> WorkflowResult<()> {
+		// We have to clone here because this function can't have a mutable reference to self. The state of
+		// the stop channel doesn't matter because it only ever receives one message
+		let _ = self.stop.clone().changed().await;
+		Err(WorkflowError::WorkflowStopped)
+	}
 }
 
 impl WorkflowCtx {
@@ -459,6 +485,8 @@ impl WorkflowCtx {
 		I: ActivityInput,
 		<I as ActivityInput>::Activity: Activity<Input = I>,
 	{
+		self.check_stop().map_err(GlobalError::raw)?;
+
 		let event_id = EventId::new(I::Activity::NAME, &input);
 
 		let history_res = self
@@ -556,6 +584,8 @@ impl WorkflowCtx {
 	/// short circuit in the event of an error to make sure activity side effects are recorded.
 	#[tracing::instrument(skip_all)]
 	pub async fn join<T: Executable>(&mut self, exec: T) -> GlobalResult<T::Output> {
+		self.check_stop().map_err(GlobalError::raw)?;
+
 		exec.execute(self).await
 	}
 
@@ -571,7 +601,7 @@ impl WorkflowCtx {
 					Ok(inner_err) => {
 						// Despite "history diverged" errors being unrecoverable, they should not have be returned
 						// by this function because the state of the history is already messed up and no new
-						// workflow items can be run.
+						// workflow items should be run.
 						if !inner_err.is_recoverable()
 							&& !matches!(*inner_err, WorkflowError::HistoryDiverged(_))
 						{
@@ -599,6 +629,8 @@ impl WorkflowCtx {
 	/// received, the workflow will be woken up and continue.
 	#[tracing::instrument(skip_all)]
 	pub async fn listen<T: Listen>(&mut self) -> GlobalResult<T> {
+		self.check_stop().map_err(GlobalError::raw)?;
+
 		let history_res = self
 			.cursor
 			.compare_signal(self.version)
@@ -648,6 +680,7 @@ impl WorkflowCtx {
 				tokio::select! {
 					_ = wake_sub.next() => {},
 					_ = interval.tick() => {},
+					res = self.wait_stop() => res.map_err(GlobalError::raw)?,
 				}
 			}
 		};
@@ -664,6 +697,8 @@ impl WorkflowCtx {
 		&mut self,
 		listener: &T,
 	) -> GlobalResult<<T as CustomListener>::Output> {
+		self.check_stop().map_err(GlobalError::raw)?;
+
 		let history_res = self
 			.cursor
 			.compare_signal(self.version)
@@ -713,6 +748,7 @@ impl WorkflowCtx {
 				tokio::select! {
 					_ = wake_sub.next() => {},
 					_ = interval.tick() => {},
+					res = self.wait_stop() => res.map_err(GlobalError::raw)?,
 				}
 			}
 		};
@@ -756,6 +792,8 @@ impl WorkflowCtx {
 		F: for<'a> FnMut(&'a mut WorkflowCtx, &'a mut S) -> AsyncResult<'a, Loop<T>>,
 		T: Serialize + DeserializeOwned,
 	{
+		self.check_stop().map_err(GlobalError::raw)?;
+
 		let history_res = self
 			.cursor
 			.compare_loop(self.version)
@@ -806,6 +844,8 @@ impl WorkflowCtx {
 			tracing::debug!(name=%self.name, id=%self.workflow_id, "running loop");
 
 			loop {
+				self.check_stop().map_err(GlobalError::raw)?;
+
 				let start_instant = Instant::now();
 
 				// Create a new branch for each iteration of the loop at location {...loop location, iteration idx}
@@ -928,6 +968,8 @@ impl WorkflowCtx {
 
 	#[tracing::instrument(skip_all)]
 	pub async fn sleep_until(&mut self, time: impl TsToMillis) -> GlobalResult<()> {
+		self.check_stop().map_err(GlobalError::raw)?;
+
 		let history_res = self
 			.cursor
 			.compare_sleep(self.version)
@@ -969,7 +1011,10 @@ impl WorkflowCtx {
 		else if duration < self.db.worker_poll_interval().as_millis() as i64 + 1 {
 			tracing::debug!(name=%self.name, id=%self.workflow_id, %deadline_ts, "sleeping in memory");
 
-			tokio::time::sleep(Duration::from_millis(duration.try_into()?)).await;
+			tokio::select! {
+				_ = tokio::time::sleep(Duration::from_millis(duration.try_into()?)) => {},
+				res = self.wait_stop() => res?,
+			}
 		}
 		// Workflow sleep
 		else {
@@ -1008,6 +1053,8 @@ impl WorkflowCtx {
 		&mut self,
 		time: impl TsToMillis,
 	) -> GlobalResult<Option<T>> {
+		self.check_stop().map_err(GlobalError::raw)?;
+
 		let history_res = self
 			.cursor
 			.compare_sleep(self.version)
@@ -1122,6 +1169,7 @@ impl WorkflowCtx {
 						tokio::select! {
 							_ = wake_sub.next() => {},
 							_ = interval.tick() => {},
+							res = self.wait_stop() => res?,
 						}
 					}
 				})
@@ -1173,6 +1221,7 @@ impl WorkflowCtx {
 				tokio::select! {
 					_ = wake_sub.next() => {},
 					_ = interval.tick() => {},
+					res = self.wait_stop() => res.map_err(GlobalError::raw)?,
 				}
 			}
 		};
@@ -1205,6 +1254,8 @@ impl WorkflowCtx {
 	/// Represents a removed workflow step.
 	#[tracing::instrument(skip_all)]
 	pub async fn removed<T: Removed>(&mut self) -> GlobalResult<()> {
+		self.check_stop().map_err(GlobalError::raw)?;
+
 		// Existing event
 		if self
 			.cursor
@@ -1242,6 +1293,8 @@ impl WorkflowCtx {
 	/// inserts a version check event.
 	#[tracing::instrument(skip_all)]
 	pub async fn check_version(&mut self, current_version: usize) -> GlobalResult<usize> {
+		self.check_stop().map_err(GlobalError::raw)?;
+
 		if current_version == 0 {
 			return Err(GlobalError::raw(WorkflowError::InvalidVersion(
 				"version for `check_version` must be greater than 0".into(),

packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/mod.rs

@@ -214,7 +214,7 @@ impl Database for DatabaseFdbSqliteNats {
 				);
 
 				// Shut down entire runtime
-				rivet_runtime::shutdown();
+				rivet_runtime::shutdown().await;
 			}
 		});
 

packages/common/chirp-workflow/core/src/error.rs

@@ -34,6 +34,9 @@ pub enum WorkflowError {
 	#[error("workflow not found")]
 	WorkflowNotFound,
 
+	#[error("workflow stopped")]
+	WorkflowStopped,
+
 	#[error("history diverged: {0}")]
 	HistoryDiverged(String),
 
@@ -195,6 +198,10 @@ pub enum WorkflowError {
 }
 
 impl WorkflowError {
+	pub(crate) fn wake_immediate(&self) -> bool {
+		matches!(self, WorkflowError::WorkflowStopped)
+	}
+
 	/// Returns the next deadline for a workflow to be woken up again based on the error.
 	pub(crate) fn deadline_ts(&self) -> Option<i64> {
 		match self {
@@ -231,12 +238,13 @@ impl WorkflowError {
 			| WorkflowError::NoSignalFound(_)
 			| WorkflowError::NoSignalFoundAndSleep(_, _)
 			| WorkflowError::SubWorkflowIncomplete(_)
-			| WorkflowError::Sleep(_) => true,
+			| WorkflowError::Sleep(_)
+			| WorkflowError::WorkflowStopped => true,
 			_ => false,
 		}
 	}
 
-	/// Any error that the workflow can try again on. Only used for printing.
+	/// Any error that the workflow can try again on a fixed number of times. Only used for printing.
 	pub(crate) fn is_retryable(&self) -> bool {
 		match self {
 			WorkflowError::ActivityFailure(_, _)