feat(gas): add batch listen

engine/packages/engine/src/util/wf/mod.rs

@@ -467,6 +467,35 @@ pub async fn print_history(
 					}
 				}
 			}
+			EventData::Signals(data) => {
+				// Indent
+				print!("{}{c} ", "  ".repeat(indent));
+
+				for ((signal_id, name), body) in
+					data.signal_ids.iter().zip(&data.names).zip(&data.bodies)
+				{
+					// Indent
+					print!("{}{c}   - ", "  ".repeat(indent));
+					println!("{}", event_style.apply_to(name));
+
+					print!("{}{c}   ", "  ".repeat(indent));
+					println!("id {}", style(signal_id).green());
+
+					if !exclude_json {
+						// Indent
+						print!("{}{c}   ", "  ".repeat(indent));
+
+						println!(
+							"body {}",
+							indent_string(
+								&colored_json(body)?,
+								format!("{}{c}   ", "  ".repeat(indent)),
+								true
+							)
+						);
+					}
+				}
+			}
 			_ => {}
 		}
 	}
@@ -543,7 +572,7 @@ pub fn event_style(event: &Event) -> Style {
 		EventData::Removed(_) => Style::new().red(),
 		EventData::VersionCheck => Style::new().red(),
 		EventData::Branch => Style::new(),
-		EventData::Empty => Style::new(),
+		EventData::Signals(_) => Style::new().cyan(),
 	}
 }
 
@@ -595,7 +624,11 @@ pub fn print_event_name(event: &Event) {
 		}
 		EventData::VersionCheck => print!("{}", style.apply_to("version check").bold()),
 		EventData::Branch => print!("{}", style.apply_to("branch").bold()),
-		EventData::Empty => print!("{}", style.apply_to("empty").bold()),
+		EventData::Signals(signal) => print!(
+			"{} {}",
+			style.apply_to("signal receive").bold(),
+			style.apply_to(&signal.names.len())
+		),
 	}
 }
 

engine/packages/gasoline-macros/src/lib.rs

@@ -372,9 +372,13 @@ pub fn signal(attr: TokenStream, item: TokenStream) -> TokenStream {
 
 		#[async_trait::async_trait]
 		impl gas::listen::Listen for #ident {
-			async fn listen(ctx: &mut gas::prelude::ListenCtx) -> gas::prelude::WorkflowResult<Self> {
-				let row = ctx.listen_any(&[<Self as gas::signal::Signal>::NAME]).await?;
-				Self::parse(&row.signal_name, &row.body)
+			async fn listen(ctx: &mut gas::prelude::ListenCtx, limit: usize) -> gas::prelude::WorkflowResult<Vec<Self>> {
+				ctx
+					.listen_any(&[<Self as gas::signal::Signal>::NAME], limit)
+					.await?
+					.into_iter()
+					.map(|signal| Self::parse(&signal.signal_name, &signal.body))
+					.collect()
 			}
 
 			fn parse(_name: &str, body: &serde_json::value::RawValue) -> gas::prelude::WorkflowResult<Self> {

engine/packages/gasoline/src/ctx/listen.rs

@@ -1,5 +1,5 @@
 use rivet_metrics::KeyValue;
-use std::{ops::Deref, time::Instant};
+use std::ops::Deref;
 
 use crate::{
 	ctx::WorkflowCtx,
@@ -14,7 +14,7 @@ pub struct ListenCtx<'a> {
 	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,
+	last_attempt: bool,
 	// HACK: Prevent `ListenCtx::listen_any` from being called more than once
 	used: bool,
 }
@@ -24,14 +24,14 @@ impl<'a> ListenCtx<'a> {
 		ListenCtx {
 			ctx,
 			location,
-			last_try: false,
+			last_attempt: false,
 			used: false,
 		}
 	}
 
-	pub(crate) fn reset(&mut self, last_try: bool) {
+	pub(crate) fn reset(&mut self, last_attempt: bool) {
 		self.used = false;
-		self.last_try = last_try;
+		self.last_attempt = last_attempt;
 	}
 
 	/// Checks for a signal to this workflow with any of the given signal names.
@@ -40,75 +40,63 @@ impl<'a> ListenCtx<'a> {
 	pub async fn listen_any(
 		&mut self,
 		signal_names: &[&'static str],
-	) -> WorkflowResult<SignalData> {
+		limit: usize,
+	) -> WorkflowResult<Vec<SignalData>> {
 		if self.used {
 			return Err(WorkflowError::ListenCtxUsed);
 		} else {
 			self.used = true;
 		}
 
-		let start_instant = Instant::now();
-
-		// Fetch new pending signal
-		let signal = self
+		// Fetch new pending signals
+		let signals = self
 			.ctx
 			.db()
-			.pull_next_signal(
+			.pull_next_signals(
 				self.ctx.workflow_id(),
 				self.ctx.name(),
 				signal_names,
 				self.location,
 				self.ctx.version(),
 				self.ctx.loop_location(),
-				self.last_try,
+				limit,
+				self.last_attempt,
 			)
 			.await?;
 
-		let dt = start_instant.elapsed().as_secs_f64();
-		metrics::SIGNAL_PULL_DURATION.record(
-			dt,
-			&[
-				KeyValue::new("workflow_name", self.ctx.name().to_string()),
-				KeyValue::new(
-					"signal_name",
-					signal
-						.as_ref()
-						.map(|signal| signal.signal_name.clone())
-						.unwrap_or("<none>".into()),
-				),
-			],
-		);
-
-		let Some(signal) = signal else {
+		if signals.is_empty() {
 			return Err(WorkflowError::NoSignalFound(Box::from(signal_names)));
-		};
+		}
 
-		let recv_lag = (rivet_util::timestamp::now() as f64 - signal.create_ts as f64) / 1000.;
-		crate::metrics::SIGNAL_RECV_LAG.record(
-			recv_lag,
-			&[
-				KeyValue::new("workflow_name", self.ctx.name().to_string()),
-				KeyValue::new("signal_name", signal.signal_name.clone()),
-			],
-		);
+		let now = rivet_util::timestamp::now();
+		for signal in &signals {
+			let recv_lag = (now as f64 - signal.create_ts as f64) / 1000.0;
+			metrics::SIGNAL_RECV_LAG.record(
+				recv_lag,
+				&[
+					KeyValue::new("workflow_name", self.ctx.name().to_string()),
+					KeyValue::new("signal_name", signal.signal_name.clone()),
+				],
+			);
+
+			if recv_lag > 3.0 {
+				// We print an error here so the trace of this workflow does not get dropped
+				tracing::error!(
+					?recv_lag,
+					signal_id=%signal.signal_id,
+					signal_name=%signal.signal_name,
+					"long signal recv time",
+				);
+			}
 
-		if recv_lag > 3.0 {
-			// We print an error here so the trace of this workflow does not get dropped
-			tracing::error!(
-				?recv_lag,
+			tracing::debug!(
 				signal_id=%signal.signal_id,
 				signal_name=%signal.signal_name,
-				"long signal recv time",
+				"signal received",
 			);
 		}
 
-		tracing::debug!(
-			signal_id=%signal.signal_id,
-			signal_name=%signal.signal_name,
-			"signal received",
-		);
-
-		Ok(signal)
+		Ok(signals)
 	}
 }
 

engine/packages/gasoline/src/ctx/versioned_workflow.rs

@@ -9,7 +9,7 @@ use crate::{
 	builder::{WorkflowRepr, workflow as builder},
 	ctx::{WorkflowCtx, workflow::Loop},
 	executable::{AsyncResult, Executable},
-	listen::{CustomListener, Listen},
+	listen::Listen,
 	message::Message,
 	signal::Signal,
 	utils::time::{DurationToMillis, TsToMillis},
@@ -125,17 +125,6 @@ impl<'a> VersionedWorkflowCtx<'a> {
 		})
 	}
 
-	/// Execute a custom listener.
-	#[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> {
-		wrap!(self, "listen", {
-			self.inner.custom_listener(listener).in_current_span().await
-		})
-	}
-
 	/// Creates a message builder.
 	pub fn msg<M: Message>(&mut self, body: M) -> builder::message::MessageBuilder<'_, M> {
 		builder::message::MessageBuilder::new(self.inner, self.version(), body)