@rivetkit/engine 1.0.0 → 2.2.1-pr.4600.b74ff3b

package/packages/gasoline/src/ctx/workflow.rs

@@ -4,7 +4,7 @@ use std::{
 	time::{Duration, Instant},
 };
 
-use anyhow::Result;
+use anyhow::{Context, Result};
 use futures_util::StreamExt;
 use opentelemetry::trace::SpanContext;
 use rivet_util::Id;
@@ -13,8 +13,6 @@ use tokio::sync::{Mutex, watch};
 use tracing::Instrument;
 use tracing_opentelemetry::OpenTelemetrySpanExt;
 
-use rivet_metrics::KeyValue;
-
 use crate::{
 	activity::{Activity, ActivityInput},
 	builder::{WorkflowRepr, workflow as builder},
@@ -24,12 +22,12 @@ use crate::{
 	executable::{AsyncResult, Executable},
 	history::{
 		History,
-		cursor::{Cursor, HistoryResult},
+		cursor::{CheckVersionHistoryResult, Cursor, HistoryResult, RemovedHistoryResult},
 		event::SleepState,
-		location::{Coordinate, Location},
+		location::Location,
 		removed::Removed,
 	},
-	listen::{CustomListener, Listen},
+	listen::Listen,
 	message::Message,
 	metrics,
 	registry::RegistryHandle,
@@ -158,6 +156,7 @@ impl WorkflowCtx {
 
 		// Lookup workflow
 		let workflow = self.registry.get_workflow(&self.name)?;
+		let prune_variant = workflow.prune_variant;
 
 		// Run workflow
 		let mut res = (workflow.run)(&mut self).await;
@@ -184,7 +183,7 @@ impl WorkflowCtx {
 					// Write output
 					if let Err(err) = self
 						.db
-						.complete_workflow(self.workflow_id, &self.name, &output)
+						.complete_workflow(self.workflow_id, &self.name, &output, prune_variant)
 						.await
 					{
 						if retries > MAX_DB_ACTION_RETRIES {
@@ -214,22 +213,18 @@ impl WorkflowCtx {
 				// finish. This workflow will be retried when the sub workflow completes
 				let wake_sub_workflow = err.sub_workflow();
 
+				let err_str = err.to_string();
+
 				if err.is_recoverable() && !err.is_retryable() {
 					tracing::debug!(?err, "workflow sleeping");
 				} else {
 					tracing::error!(?err, "workflow error");
 
-					metrics::WORKFLOW_ERRORS.add(
-						1,
-						&[
-							KeyValue::new("workflow_name", self.name.clone()),
-							KeyValue::new("error", err.to_string()),
-						],
-					);
+					metrics::WORKFLOW_ERRORS
+						.with_label_values(&[self.name.as_str(), err_str.as_str()])
+						.inc();
 				}
 
-				let err_str = err.to_string();
-
 				let mut retries = 0;
 				let mut interval = tokio::time::interval(DB_ACTION_RETRY);
 				interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
@@ -299,7 +294,7 @@ impl WorkflowCtx {
 
 		let res = tokio::time::timeout(A::TIMEOUT, A::run(&ctx, input).in_current_span())
 			.await
-			.map_err(|_| WorkflowError::ActivityTimeout(0));
+			.map_err(|_| WorkflowError::ActivityTimeout(A::NAME, 0));
 
 		let dt = start_instant.elapsed().as_secs_f64();
 
@@ -342,14 +337,9 @@ impl WorkflowCtx {
 					},
 				)?;
 
-				metrics::ACTIVITY_DURATION.record(
-					dt,
-					&[
-						KeyValue::new("workflow_name", self.name.clone()),
-						KeyValue::new("activity_name", A::NAME),
-						KeyValue::new("error_code", ""),
-					],
-				);
+				metrics::ACTIVITY_DURATION
+					.with_label_values(&[self.name.as_str(), A::NAME, ""])
+					.observe(dt);
 
 				Ok(output)
 			}
@@ -381,25 +371,15 @@ impl WorkflowCtx {
 					.unwrap_or_default();
 
 				if !is_recoverable {
-					metrics::ACTIVITY_ERRORS.add(
-						1,
-						&[
-							KeyValue::new("workflow_name", self.name.clone()),
-							KeyValue::new("activity_name", A::NAME),
-							KeyValue::new("error_code", err_str.clone()),
-						],
-					);
+					metrics::ACTIVITY_ERRORS
+						.with_label_values(&[self.name.as_str(), A::NAME, err_str.as_str()])
+						.inc();
 				}
-				metrics::ACTIVITY_DURATION.record(
-					dt,
-					&[
-						KeyValue::new("workflow_name", self.name.clone()),
-						KeyValue::new("activity_name", A::NAME),
-						KeyValue::new("error_code", err_str.clone()),
-					],
-				);
+				metrics::ACTIVITY_DURATION
+					.with_label_values(&[self.name.as_str(), A::NAME, err_str.as_str()])
+					.observe(dt);
 
-				Err(WorkflowError::ActivityFailure(err, 0))
+				Err(WorkflowError::ActivityFailure(A::NAME, err, 0))
 			}
 			Err(err) => {
 				tracing::debug!("activity timeout");
@@ -421,22 +401,12 @@ impl WorkflowCtx {
 					)
 					.await?;
 
-				metrics::ACTIVITY_ERRORS.add(
-					1,
-					&[
-						KeyValue::new("workflow_name", self.name.clone()),
-						KeyValue::new("activity_name", A::NAME),
-						KeyValue::new("error_code", err_str.clone()),
-					],
-				);
-				metrics::ACTIVITY_DURATION.record(
-					dt,
-					&[
-						KeyValue::new("workflow_name", self.name.clone()),
-						KeyValue::new("activity_name", A::NAME),
-						KeyValue::new("error_code", err_str.clone()),
-					],
-				);
+				metrics::ACTIVITY_ERRORS
+					.with_label_values(&[self.name.as_str(), A::NAME, err_str.as_str()])
+					.inc();
+				metrics::ACTIVITY_DURATION
+					.with_label_values(&[self.name.as_str(), A::NAME, err_str.as_str()])
+					.observe(dt);
 
 				Err(err)
 			}
@@ -602,28 +572,31 @@ 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(name, err, _) => {
 								if error_count.saturating_add(1) >= I::Activity::MAX_RETRIES {
-									WorkflowError::ActivityMaxFailuresReached(err)
+									WorkflowError::ActivityMaxFailuresReached(name, err)
 								} else {
 									// Add error count to the error for backoff calculation
-									WorkflowError::ActivityFailure(err, error_count)
+									WorkflowError::ActivityFailure(name, err, error_count)
 								}
 							}
-							WorkflowError::ActivityTimeout(_) => {
+							WorkflowError::ActivityTimeout(name, _) => {
 								if error_count.saturating_add(1) >= I::Activity::MAX_RETRIES {
-									WorkflowError::ActivityMaxFailuresReached(err.into())
+									WorkflowError::ActivityMaxFailuresReached(name, err.into())
 								} else {
 									// Add error count to the error for backoff calculation
-									WorkflowError::ActivityTimeout(error_count)
+									WorkflowError::ActivityTimeout(name, error_count)
 								}
 							}
-							WorkflowError::OperationTimeout(_) => {
+							WorkflowError::OperationTimeout(op_name, _) => {
 								if error_count.saturating_add(1) >= I::Activity::MAX_RETRIES {
-									WorkflowError::ActivityMaxFailuresReached(err.into())
+									WorkflowError::ActivityMaxFailuresReached(
+										I::Activity::NAME,
+										err.into(),
+									)
 								} else {
 									// Add error count to the error for backoff calculation
-									WorkflowError::OperationTimeout(error_count)
+									WorkflowError::OperationTimeout(op_name, error_count)
 								}
 							}
 							_ => err,
@@ -698,91 +671,40 @@ impl WorkflowCtx {
 	/// received, the workflow will be woken up and continue.
 	#[tracing::instrument(skip_all, fields(t=std::any::type_name::<T>()))]
 	pub async fn listen<T: Listen>(&mut self) -> Result<T> {
-		self.check_stop()?;
-
-		let history_res = self.cursor.compare_signal(self.version)?;
-		let location = self.cursor.current_location_for(&history_res);
-
-		// Signal received before
-		let signal = if let HistoryResult::Event(signal) = history_res {
-			tracing::debug!(
-				signal_name=%signal.name,
-				"replaying signal"
-			);
-
-			T::parse(&signal.name, &signal.body)?
-		}
-		// Listen for new signal
-		else {
-			tracing::debug!("listening for signal");
-
-			let mut bump_sub = self
-				.db
-				.bump_sub(BumpSubSubject::SignalPublish {
-					to_workflow_id: self.workflow_id,
-				})
-				.await?;
-			let mut retries = self.db.max_signal_poll_retries();
-			let mut interval = tokio::time::interval(self.db.signal_poll_interval());
-			interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
-
-			// Skip first tick, we wait after the db call instead of before
-			interval.tick().await;
-
-			let mut ctx = ListenCtx::new(self, &location);
-
-			loop {
-				ctx.reset(retries == 0);
-
-				match T::listen(&mut ctx).in_current_span().await {
-					Ok(res) => break res,
-					Err(err) if matches!(err, WorkflowError::NoSignalFound(_)) => {
-						if retries == 0 {
-							return Err(err.into());
-						}
-						retries -= 1;
-					}
-					Err(err) => return Err(err.into()),
-				}
+		let signals = self.listen_n::<T>(1).in_current_span().await?;
 
-				// Poll and wait for a wake at the same time
-				tokio::select! {
-					_ = bump_sub.next() => {},
-					_ = interval.tick() => {},
-					res = self.wait_stop() => res?,
-				}
-			}
-		};
-
-		// Move to next event
-		self.cursor.update(&location);
-
-		Ok(signal)
+		signals
+			.into_iter()
+			.next()
+			.context("must return at least 1 signal")
 	}
 
-	/// Execute a custom listener.
+	/// Listens for a N signals for a short time before setting the workflow to sleep. Once signals are
+	/// received, the workflow will be woken up and continue. Never returns an empty vec.
 	#[tracing::instrument(skip_all, fields(t=std::any::type_name::<T>()))]
-	pub async fn custom_listener<T: CustomListener>(
-		&mut self,
-		listener: &T,
-	) -> Result<<T as CustomListener>::Output> {
+	pub async fn listen_n<T: Listen>(&mut self, limit: usize) -> Result<Vec<T>> {
 		self.check_stop()?;
 
-		let history_res = self.cursor.compare_signal(self.version)?;
+		let history_res = self.cursor.compare_signals(self.version)?;
 		let location = self.cursor.current_location_for(&history_res);
 
-		// Signal received before
-		let signal = if let HistoryResult::Event(signal) = history_res {
+		// Signals received before
+		let signals = if let HistoryResult::Event(signals) = history_res {
 			tracing::debug!(
-				signal_name=%signal.name,
-				"replaying signal",
+				count=%signals.names.len(),
+				"replaying signals"
 			);
 
-			T::parse(&signal.name, &signal.body)?
+			signals
+				.names
+				.iter()
+				.zip(&signals.bodies)
+				.map(|(name, body)| T::parse(name, &body))
+				.collect::<std::result::Result<Vec<_>, _>>()?
 		}
-		// Listen for new signal
+		// Listen for new signals
 		else {
-			tracing::debug!("listening for signal");
+			tracing::debug!("listening for signals");
 
 			let mut bump_sub = self
 				.db
@@ -802,7 +724,7 @@ impl WorkflowCtx {
 			loop {
 				ctx.reset(retries == 0);
 
-				match listener.listen(&mut ctx).in_current_span().await {
+				match T::listen(&mut ctx, limit).in_current_span().await {
 					Ok(res) => break res,
 					Err(err) if matches!(err, WorkflowError::NoSignalFound(_)) => {
 						if retries == 0 {
@@ -825,7 +747,7 @@ impl WorkflowCtx {
 		// Move to next event
 		self.cursor.update(&location);
 
-		Ok(signal)
+		Ok(signals)
 	}
 
 	/// Creates a message builder.
@@ -840,8 +762,7 @@ impl WorkflowCtx {
 		F: for<'a> FnMut(&'a mut WorkflowCtx) -> AsyncResult<'a, Loop<T>>,
 		T: Serialize + DeserializeOwned,
 	{
-		self.loop_inner((), |ctx, _| cb(ctx))
-			.in_current_span()
+		builder::lupe::LoopBuilder::run(builder::lupe::LoopBuilder::new(self, ()), |ctx, _| cb(ctx))
 			.await
 	}
 
@@ -853,170 +774,11 @@ impl WorkflowCtx {
 		F: for<'a> FnMut(&'a mut WorkflowCtx, &'a mut S) -> AsyncResult<'a, Loop<T>>,
 		T: Serialize + DeserializeOwned,
 	{
-		self.loop_inner(state, cb).in_current_span().await
+		builder::lupe::LoopBuilder::new(self, state).run(cb).await
 	}
 
-	async fn loop_inner<S, F, T>(&mut self, state: S, mut cb: F) -> Result<T>
-	where
-		S: Serialize + DeserializeOwned,
-		F: for<'a> FnMut(&'a mut WorkflowCtx, &'a mut S) -> AsyncResult<'a, Loop<T>>,
-		T: Serialize + DeserializeOwned,
-	{
-		self.check_stop()?;
-
-		let history_res = self.cursor.compare_loop(self.version)?;
-		let loop_location = self.cursor.current_location_for(&history_res);
-
-		// Loop existed before
-		let (mut iteration, mut state, output) =
-			if let HistoryResult::Event(loop_event) = history_res {
-				let state = loop_event.parse_state()?;
-				let output = loop_event.parse_output()?;
-
-				(loop_event.iteration, state, output)
-			} else {
-				let state_val = serde_json::value::to_raw_value(&state)
-					.map_err(WorkflowError::SerializeLoopOutput)?;
-
-				// Insert event before loop is run so the history is consistent
-				self.db
-					.upsert_workflow_loop_event(
-						self.workflow_id,
-						&self.name,
-						&loop_location,
-						self.version,
-						0,
-						&state_val,
-						None,
-						self.loop_location(),
-					)
-					.await?;
-
-				(0, state, None)
-			};
-
-		// Create a branch but no branch event (loop event takes its place)
-		let mut loop_branch =
-			self.branch_inner(self.input.clone(), self.version, loop_location.clone());
-
-		// Loop complete
-		let output = if let Some(output) = output {
-			tracing::debug!("replaying loop output");
-
-			output
-		}
-		// Run loop
-		else {
-			tracing::debug!("running loop");
-
-			loop {
-				self.check_stop()?;
-
-				let start_instant = Instant::now();
-
-				// Create a new branch for each iteration of the loop at location {...loop location, iteration idx}
-				let mut iteration_branch = loop_branch.branch_inner(
-					self.input.clone(),
-					self.version,
-					loop_branch
-						.cursor
-						.root()
-						.join(Coordinate::simple(iteration + 1)),
-				);
-
-				// Set branch loop location to the current loop
-				iteration_branch.loop_location = Some(loop_location.clone());
-
-				let i = iteration;
-
-				// Async block for instrumentation purposes
-				let (dt2, res) = async {
-					// Insert event if iteration is not a replay
-					if !loop_branch.cursor.compare_loop_branch(iteration)? {
-						self.db
-							.commit_workflow_branch_event(
-								self.workflow_id,
-								iteration_branch.cursor.root(),
-								self.version,
-								Some(&loop_location),
-							)
-							.await?;
-					}
-
-					let start_instant2 = Instant::now();
-
-					// Run loop
-					match cb(&mut iteration_branch, &mut state).await? {
-						Loop::Continue => {
-							let dt2 = start_instant2.elapsed().as_secs_f64();
-							iteration += 1;
-
-							let state_val = serde_json::value::to_raw_value(&state)
-								.map_err(WorkflowError::SerializeLoopOutput)?;
-
-							self.db
-								.upsert_workflow_loop_event(
-									self.workflow_id,
-									&self.name,
-									&loop_location,
-									self.version,
-									iteration,
-									&state_val,
-									None,
-									self.loop_location(),
-								)
-								.await?;
-
-							anyhow::Ok((dt2, None))
-						}
-						Loop::Break(res) => {
-							let dt2 = start_instant2.elapsed().as_secs_f64();
-							iteration += 1;
-
-							let state_val = serde_json::value::to_raw_value(&state)
-								.map_err(WorkflowError::SerializeLoopOutput)?;
-							let output_val = serde_json::value::to_raw_value(&res)
-								.map_err(WorkflowError::SerializeLoopOutput)?;
-
-							self.db
-								.upsert_workflow_loop_event(
-									self.workflow_id,
-									&self.name,
-									&loop_location,
-									self.version,
-									iteration,
-									&state_val,
-									Some(&output_val),
-									self.loop_location(),
-								)
-								.await?;
-
-							Ok((dt2, Some(res)))
-						}
-					}
-				}
-				.instrument(tracing::info_span!("iteration", iteration=%i))
-				.await?;
-
-				// Validate no leftover events
-				iteration_branch.cursor.check_clear()?;
-
-				let dt = start_instant.elapsed().as_secs_f64();
-				metrics::LOOP_ITERATION_DURATION.record(
-					dt - dt2,
-					&[KeyValue::new("workflow_name", self.name.clone())],
-				);
-
-				if let Some(res) = res {
-					break res;
-				}
-			}
-		};
-
-		// Move to next event
-		self.cursor.update(&loop_location);
-
-		Ok(output)
+	pub fn lupe(&mut self) -> builder::lupe::LoopBuilder<'_, ()> {
+		builder::lupe::LoopBuilder::new(self, ())
 	}
 
 	#[tracing::instrument(skip_all)]
@@ -1095,19 +857,47 @@ impl WorkflowCtx {
 		&mut self,
 		duration: impl DurationToMillis,
 	) -> Result<Option<T>> {
+		let signals = self.listen_n_with_timeout(duration, 1).await?;
+
+		Ok(signals.into_iter().next())
+	}
+
+	/// Listens for signals with a timeout. Returns an empty vec if the timeout is reached.
+	///
+	/// Internally this is a sleep event and a signals event.
+	#[tracing::instrument(skip_all, fields(t=std::any::type_name::<T>()))]
+	pub async fn listen_n_with_timeout<T: Listen>(
+		&mut self,
+		duration: impl DurationToMillis,
+		limit: usize,
+	) -> Result<Vec<T>> {
 		let time = (rivet_util::timestamp::now() as u64 + duration.to_millis()?) as i64;
 
-		self.listen_until(time).await
+		self.listen_n_until(time, limit).await
+	}
+
+	/// Listens for a signal until the given timestamp. Returns `None` if the timestamp is reached.
+	///
+	/// Internally this is a sleep event and a signals event.
+	#[tracing::instrument(skip_all, fields(t=std::any::type_name::<T>(), duration))]
+	pub async fn listen_until<T: Listen>(&mut self, time: impl TsToMillis) -> Result<Option<T>> {
+		let signals = self.listen_n_until(time, 1).await?;
+
+		Ok(signals.into_iter().next())
 	}
 
 	// TODO: Potential bad transaction: if the signal gets pulled and saved in history but an error occurs
 	// before the sleep event state is set to "interrupted", the next time this workflow is run it will error
 	// because it tries to pull a signal again
-	/// Listens for a signal until the given timestamp. Returns `None` if the timestamp is reached.
+	/// Listens for signals until the given timestamp. Returns an empty vec if the timestamp is reached.
 	///
 	/// Internally this is a sleep event and a signal event.
 	#[tracing::instrument(skip_all, fields(t=std::any::type_name::<T>(), duration))]
-	pub async fn listen_until<T: Listen>(&mut self, time: impl TsToMillis) -> Result<Option<T>> {
+	pub async fn listen_n_until<T: Listen>(
+		&mut self,
+		time: impl TsToMillis,
+		limit: usize,
+	) -> Result<Vec<T>> {
 		self.check_stop()?;
 
 		let history_res = self.cursor.compare_sleep(self.version)?;
@@ -1140,51 +930,56 @@ impl WorkflowCtx {
 		// Move to next event
 		self.cursor.update(&sleep_location);
 
-		// Signal received before
+		// Signals received before
 		if matches!(state, SleepState::Interrupted) {
-			let history_res = self.cursor.compare_signal(self.version)?;
-			let signal_location = self.cursor.current_location_for(&history_res);
+			let history_res = self.cursor.compare_signals(self.version)?;
+			let signals_location = self.cursor.current_location_for(&history_res);
 
-			if let HistoryResult::Event(signal) = history_res {
+			if let HistoryResult::Event(signals) = history_res {
 				tracing::debug!(
-					signal_name=%signal.name,
-					"replaying signal",
+					count=?signals.names.len(),
+					"replaying signals",
 				);
 
-				let signal = T::parse(&signal.name, &signal.body)?;
+				let signals = signals
+					.names
+					.iter()
+					.zip(&signals.bodies)
+					.map(|(name, body)| T::parse(name, &body))
+					.collect::<std::result::Result<Vec<_>, _>>()?;
 
 				// Move to next event
-				self.cursor.update(&signal_location);
+				self.cursor.update(&signals_location);
 
 				// Short circuit
-				return Ok(Some(signal));
+				return Ok(signals);
 			} else {
 				return Err(WorkflowError::HistoryDiverged(format!(
-					"expected signal at {}, found nothing",
-					signal_location,
+					"expected signals at {}, found nothing",
+					signals_location,
 				))
 				.into());
 			}
 		}
 
-		// Location of the signal event (comes after the sleep event)
-		let signal_location = self.cursor.current_location_for(&history_res2);
+		// Location of the signals event (comes after the sleep event)
+		let signals_location = self.cursor.current_location_for(&history_res2);
 		let duration = deadline_ts.saturating_sub(rivet_util::timestamp::now());
 		tracing::Span::current().record("duration", &duration);
 
 		// Duration is now 0, timeout is over
-		let signal = if duration <= 0 {
-			// After timeout is over, check once for signal
+		let signals = if duration <= 0 {
+			// After timeout is over, check once for signals
 			if matches!(state, SleepState::Normal) {
-				let mut ctx = ListenCtx::new(self, &signal_location);
+				let mut ctx = ListenCtx::new(self, &signals_location);
 
-				match T::listen(&mut ctx).in_current_span().await {
-					Ok(x) => Some(x),
-					Err(WorkflowError::NoSignalFound(_)) => None,
+				match T::listen(&mut ctx, limit).in_current_span().await {
+					Ok(x) => x,
+					Err(WorkflowError::NoSignalFound(_)) => Vec::new(),
 					Err(err) => return Err(err.into()),
 				}
 			} else {
-				None
+				Vec::new()
 			}
 		}
 		// Sleep in memory if duration is shorter than the worker tick
@@ -1194,7 +989,7 @@ impl WorkflowCtx {
 			let res = tokio::time::timeout(
 				Duration::from_millis(duration.try_into()?),
 				(async {
-					tracing::debug!("listening for signal with timeout");
+					tracing::debug!("listening for signals with timeout");
 
 					let mut bump_sub = self
 						.db
@@ -1208,12 +1003,12 @@ impl WorkflowCtx {
 					// Skip first tick, we wait after the db call instead of before
 					interval.tick().await;
 
-					let mut ctx = ListenCtx::new(self, &signal_location);
+					let mut ctx = ListenCtx::new(self, &signals_location);
 
 					loop {
 						ctx.reset(false);
 
-						match T::listen(&mut ctx).in_current_span().await {
+						match T::listen(&mut ctx, limit).in_current_span().await {
 							// Retry
 							Err(WorkflowError::NoSignalFound(_)) => {}
 							x => return x,
@@ -1232,17 +1027,17 @@ impl WorkflowCtx {
 			.await;
 
 			match res {
-				Ok(res) => Some(res?),
+				Ok(res) => res?,
 				Err(_) => {
-					tracing::debug!("timed out listening for signal");
+					tracing::debug!("timed out listening for signals");
 
-					None
+					Vec::new()
 				}
 			}
 		}
 		// Workflow sleep for long durations
 		else {
-			tracing::debug!("listening for signal with timeout");
+			tracing::debug!("listening for signals with timeout");
 
 			let mut bump_sub = self
 				.db
@@ -1257,13 +1052,13 @@ impl WorkflowCtx {
 			// Skip first tick, we wait after the db call instead of before
 			interval.tick().await;
 
-			let mut ctx = ListenCtx::new(self, &signal_location);
+			let mut ctx = ListenCtx::new(self, &signals_location);
 
 			loop {
 				ctx.reset(retries == 0);
 
-				match T::listen(&mut ctx).in_current_span().await {
-					Ok(res) => break Some(res),
+				match T::listen(&mut ctx, limit).in_current_span().await {
+					Ok(res) => break res,
 					Err(WorkflowError::NoSignalFound(signals)) => {
 						if retries == 0 {
 							return Err(
@@ -1285,7 +1080,7 @@ impl WorkflowCtx {
 		};
 
 		// Update sleep state
-		if signal.is_some() {
+		if !signals.is_empty() {
 			self.db
 				.update_workflow_sleep_event_state(
 					self.workflow_id,
@@ -1295,7 +1090,7 @@ impl WorkflowCtx {
 				.await?;
 
 			// Move to next event
-			self.cursor.update(&signal_location);
+			self.cursor.update(&signals_location);
 		} else if matches!(state, SleepState::Normal) {
 			self.db
 				.update_workflow_sleep_event_state(
@@ -1306,7 +1101,7 @@ impl WorkflowCtx {
 				.await?;
 		}
 
-		Ok(signal)
+		Ok(signals)
 	}
 
 	/// Represents a removed workflow step.
@@ -1314,70 +1109,81 @@ impl WorkflowCtx {
 	pub async fn removed<T: Removed>(&mut self) -> Result<()> {
 		self.check_stop()?;
 
-		// Existing event
-		if self.cursor.compare_removed::<T>()? {
-			tracing::debug!("skipping removed step",);
-		}
-		// New "removed" event
-		else {
-			tracing::debug!("inserting removed step");
+		match self.cursor.compare_removed::<T>() {
+			RemovedHistoryResult::New => {
+				tracing::debug!("inserting removed step");
 
-			self.db
-				.commit_workflow_removed_event(
-					self.workflow_id,
-					&self.cursor.current_location(),
-					T::event_type(),
-					T::name(),
-					self.loop_location(),
-				)
-				.await?;
-		};
+				self.db
+					.commit_workflow_removed_event(
+						self.workflow_id,
+						&self.cursor.current_location(),
+						T::event_type(),
+						T::name(),
+						self.loop_location(),
+					)
+					.await?;
 
-		// Move to next event
-		self.cursor.inc();
+				// Move to next event
+				self.cursor.inc();
+			}
+			RemovedHistoryResult::Skip => {
+				tracing::debug!("skipping removed step");
+
+				// Move to next event
+				self.cursor.inc();
+			}
+			RemovedHistoryResult::Ignore(msg) => {
+				tracing::debug!(
+					"removed event filter doesn't match existing event, ignoring: {msg}"
+				);
+			}
+		}
 
 		Ok(())
 	}
 
 	/// Returns the version of the current event in history. If no event exists, returns `current_version` and
 	/// inserts a version check event.
-	#[tracing::instrument(skip_all, fields(current_version))]
-	pub async fn check_version(&mut self, current_version: usize) -> Result<usize> {
+	#[tracing::instrument(skip_all, fields(latest_version))]
+	pub async fn check_version(&mut self, latest_version: usize) -> Result<usize> {
 		self.check_stop()?;
 
-		if current_version == 0 {
+		if latest_version == 0 {
 			return Err(WorkflowError::InvalidVersion(
 				"version for `check_version` must be greater than 0".into(),
 			)
 			.into());
 		}
 
-		let (is_version_check, version) =
-			if let Some((is_version_check, step_version)) = self.cursor.compare_version_check()? {
-				tracing::debug!("checking existing version");
-
-				(is_version_check, step_version)
-			} else {
-				tracing::debug!("inserting version check");
+		let history_res = self.cursor.compare_version_check()?;
+		let event_location = self.cursor.current_location_for(&history_res.equivalent());
 
-				self.db
-					.commit_workflow_version_check_event(
-						self.workflow_id,
-						&self.cursor.current_location(),
-						current_version + self.version - 1,
-						self.loop_location(),
-					)
-					.await?;
+		let (event_version, checked_version) = match history_res {
+			CheckVersionHistoryResult::New => (Some(self.version), latest_version),
+			CheckVersionHistoryResult::Event(checked_version) => (None, checked_version),
+			CheckVersionHistoryResult::Insertion(next_event_version) => {
+				(Some(next_event_version + 1), next_event_version)
+			}
+		};
 
-				(true, current_version)
-			};
+		if let Some(event_version) = event_version {
+			tracing::debug!("inserting version check");
 
-		if is_version_check {
-			// Move to next event
-			self.cursor.inc();
+			self.db
+				.commit_workflow_version_check_event(
+					self.workflow_id,
+					&event_location,
+					event_version + self.version - 1,
+					checked_version,
+					self.loop_location(),
+				)
+				.await?;
 		}
 
-		Ok(version + 1 - self.version)
+		// Move to next event
+		self.cursor.update(&event_location);
+
+		Ok(checked_version + 1 - self.version)
 	}
 }
 
@@ -1390,6 +1196,10 @@ impl WorkflowCtx {
 		self.loop_location.as_ref()
 	}
 
+	pub(crate) fn set_loop_location(&mut self, loop_location: Location) {
+		self.loop_location = Some(loop_location);
+	}
+
 	pub(crate) fn db(&self) -> &DatabaseHandle {
 		&self.db
 	}