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

package/packages/gasoline/src/builder/workflow/message.rs

@@ -1,12 +1,10 @@
-use std::{fmt::Display, time::Instant};
+use std::time::Instant;
 
 use anyhow::Result;
-use rivet_metrics::KeyValue;
-use serde::Serialize;
 
 use crate::{
 	builder::BuilderError, ctx::WorkflowCtx, error::WorkflowError, history::cursor::HistoryResult,
-	message::Message, metrics,
+	message::Message, metrics, utils::topic::AsTopic,
 };
 
 pub struct MessageBuilder<'a, M: Message> {
@@ -14,7 +12,7 @@ pub struct MessageBuilder<'a, M: Message> {
 	version: usize,
 
 	body: M,
-	tags: serde_json::Map<String, serde_json::Value>,
+	topic: Option<String>,
 	wait: bool,
 	error: Option<BuilderError>,
 }
@@ -26,38 +24,18 @@ impl<'a, M: Message> MessageBuilder<'a, M> {
 			version,
 
 			body,
-			tags: serde_json::Map::new(),
+			topic: None,
 			wait: false,
 			error: None,
 		}
 	}
 
-	pub fn tags(mut self, tags: serde_json::Value) -> Self {
+	pub fn topic(mut self, topic: impl AsTopic) -> Self {
 		if self.error.is_some() {
 			return self;
 		}
 
-		match tags {
-			serde_json::Value::Object(map) => {
-				self.tags.extend(map);
-			}
-			_ => self.error = Some(BuilderError::TagsNotMap),
-		}
-
-		self
-	}
-
-	pub fn tag(mut self, k: impl Display, v: impl Serialize) -> Self {
-		if self.error.is_some() {
-			return self;
-		}
-
-		match serde_json::to_value(&v) {
-			Ok(v) => {
-				self.tags.insert(k.to_string(), v);
-			}
-			Err(err) => self.error = Some(err.into()),
-		}
+		self.topic = Some(topic.as_topic());
 
 		self
 	}
@@ -93,15 +71,22 @@ impl<'a, M: Message> MessageBuilder<'a, M> {
 		}
 		// Send message
 		else {
-			tracing::debug!(tags=?self.tags, "dispatching message");
+			tracing::debug!(topic=?self.topic, "dispatching message");
 
 			let start_instant = Instant::now();
 
 			// Serialize body
 			let body_val = serde_json::value::to_raw_value(&self.body)
 				.map_err(WorkflowError::SerializeMessageBody)?;
-			let tags = serde_json::Value::Object(self.tags);
-			let tags2 = tags.clone();
+			let topic = self.topic.unwrap_or_else(|| "*".to_string());
+			let tags = serde_json::Value::Object(
+				[(
+					"topic".to_string(),
+					serde_json::Value::String(topic.clone()),
+				)]
+				.into_iter()
+				.collect(),
+			);
 
 			self.ctx
 				.db()
@@ -117,26 +102,18 @@ impl<'a, M: Message> MessageBuilder<'a, M> {
 				.await?;
 
 			if self.wait {
-				self.ctx.msg_ctx().message_wait(tags2, self.body).await?;
+				self.ctx.msg_ctx().message_wait(&topic, self.body).await?;
 			} else {
-				self.ctx.msg_ctx().message(tags2, self.body).await?;
+				self.ctx.msg_ctx().message(&topic, self.body).await?;
 			}
 
 			let dt = start_instant.elapsed().as_secs_f64();
-			metrics::MESSAGE_SEND_DURATION.record(
-				dt,
-				&[
-					KeyValue::new("workflow_name", self.ctx.name().to_string()),
-					KeyValue::new("message_name", M::NAME),
-				],
-			);
-			metrics::MESSAGE_PUBLISHED.add(
-				1,
-				&[
-					KeyValue::new("workflow_name", self.ctx.name().to_string()),
-					KeyValue::new("message_name", M::NAME),
-				],
-			);
+			metrics::MESSAGE_SEND_DURATION
+				.with_label_values(&[self.ctx.name(), M::NAME])
+				.observe(dt);
+			metrics::MESSAGE_PUBLISHED
+				.with_label_values(&[self.ctx.name(), M::NAME])
+				.inc();
 		}
 
 		// Move to next event

package/packages/gasoline/src/builder/workflow/mod.rs

@@ -1,5 +1,6 @@
 //! This module contains builders used specifically by the workflow ctx.
 
+pub mod lupe;
 pub mod message;
 pub mod signal;
 pub mod sub_workflow;

package/packages/gasoline/src/builder/workflow/signal.rs

@@ -1,13 +1,17 @@
 use std::{fmt::Display, time::Instant};
 
 use anyhow::Result;
-use rivet_metrics::KeyValue;
 use rivet_util::Id;
 use serde::Serialize;
 
 use crate::{
-	builder::BuilderError, ctx::WorkflowCtx, error::WorkflowError, history::cursor::HistoryResult,
-	metrics, signal::Signal, workflow::Workflow,
+	builder::BuilderError,
+	ctx::WorkflowCtx,
+	error::WorkflowError,
+	history::{cursor::HistoryResult, event::EventType},
+	metrics,
+	signal::Signal,
+	workflow::Workflow,
 };
 
 pub struct SignalBuilder<'a, T: Signal + Serialize> {
@@ -18,6 +22,7 @@ pub struct SignalBuilder<'a, T: Signal + Serialize> {
 	to_workflow_name: Option<&'static str>,
 	to_workflow_id: Option<Id>,
 	tags: serde_json::Map<String, serde_json::Value>,
+	graceful_not_found: bool,
 	error: Option<BuilderError>,
 }
 
@@ -31,6 +36,7 @@ impl<'a, T: Signal + Serialize> SignalBuilder<'a, T> {
 			to_workflow_name: None,
 			to_workflow_id: None,
 			tags: serde_json::Map::new(),
+			graceful_not_found: false,
 			error: None,
 		}
 	}
@@ -85,14 +91,37 @@ impl<'a, T: Signal + Serialize> SignalBuilder<'a, T> {
 		self
 	}
 
+	/// Does not throw an error when the signal target is not found and instead returns `Ok(None)`.
+	pub fn graceful_not_found(mut self) -> Self {
+		if self.error.is_some() {
+			return self;
+		}
+
+		self.graceful_not_found = true;
+
+		self
+	}
+
+	/// Returns the signal id that was just sent. Unless `graceful_not_found` is set and the workflow does not
+	/// exist, will always return `Some`.
 	#[tracing::instrument(skip_all, fields(signal_name=T::NAME, signal_id))]
-	pub async fn send(self) -> Result<Id> {
+	pub async fn send(self) -> Result<Option<Id>> {
 		self.ctx.check_stop()?;
 
 		if let Some(err) = self.error {
 			return Err(err.into());
 		}
 
+		// Check if this signal is being replayed and previously had no target (will have a removed event)
+		if self.graceful_not_found && self.ctx.cursor().current_event_is_removed() {
+			tracing::debug!("replaying gracefully not found signal dispatch");
+
+			// Move to next event
+			self.ctx.cursor_mut().inc();
+
+			return Ok(None);
+		}
+
 		// Error for version mismatch. This is done in the builder instead of in `VersionedWorkflowCtx` to
 		// defer the error.
 		self.ctx.compare_version("signal", self.version)?;
@@ -105,7 +134,7 @@ impl<'a, T: Signal + Serialize> SignalBuilder<'a, T> {
 
 		// Signal sent before
 		let signal_id = if let HistoryResult::Event(signal) = history_res {
-			tracing::debug!("replaying signal dispatch",);
+			tracing::debug!("replaying signal dispatch");
 
 			signal.signal_id
 		}
@@ -133,8 +162,33 @@ impl<'a, T: Signal + Serialize> SignalBuilder<'a, T> {
 						.ctx
 						.db()
 						.find_workflow(workflow_name, &serde_json::Value::Object(self.tags))
-						.await?
-						.ok_or(WorkflowError::WorkflowNotFound)?;
+						.await?;
+
+					let Some(workflow_id) = workflow_id else {
+						// Handle signal target not found gracefully
+						if self.graceful_not_found {
+							tracing::debug!("signal target not found");
+
+							// Insert removed event
+							self.ctx
+								.db()
+								.commit_workflow_removed_event(
+									self.ctx.workflow_id(),
+									&location,
+									EventType::SignalSend,
+									Some(T::NAME),
+									self.ctx.loop_location(),
+								)
+								.await?;
+
+							// Move to next event
+							self.ctx.cursor_mut().update(&location);
+
+							return Ok(None);
+						} else {
+							return Err(WorkflowError::WorkflowNotFound.into());
+						}
+					};
 
 					self.ctx
 						.db()
@@ -199,20 +253,12 @@ impl<'a, T: Signal + Serialize> SignalBuilder<'a, T> {
 			}
 
 			let dt = start_instant.elapsed().as_secs_f64();
-			metrics::SIGNAL_SEND_DURATION.record(
-				dt,
-				&[
-					KeyValue::new("workflow_name", self.ctx.name().to_string()),
-					KeyValue::new("signal_name", T::NAME),
-				],
-			);
-			metrics::SIGNAL_PUBLISHED.add(
-				1,
-				&[
-					KeyValue::new("workflow_name", self.ctx.name().to_string()),
-					KeyValue::new("signal_name", T::NAME),
-				],
-			);
+			metrics::SIGNAL_SEND_DURATION
+				.with_label_values(&[self.ctx.name(), T::NAME])
+				.observe(dt);
+			metrics::SIGNAL_PUBLISHED
+				.with_label_values(&[self.ctx.name(), T::NAME])
+				.inc();
 
 			signal_id
 		};
@@ -222,6 +268,6 @@ impl<'a, T: Signal + Serialize> SignalBuilder<'a, T> {
 		// Move to next event
 		self.ctx.cursor_mut().update(&location);
 
-		Ok(signal_id)
+		Ok(Some(signal_id))
 	}
 }

package/packages/gasoline/src/builder/workflow/sub_workflow.rs

@@ -2,7 +2,6 @@ use std::{fmt::Display, marker::PhantomData, sync::Arc, time::Instant};
 
 use anyhow::Result;
 use futures_util::StreamExt;
-use rivet_metrics::KeyValue;
 use rivet_util::Id;
 use serde::Serialize;
 use tracing::Instrument;
@@ -185,20 +184,12 @@ where
 
 			if sub_workflow_id == actual_sub_workflow_id {
 				let dt = start_instant.elapsed().as_secs_f64();
-				metrics::WORKFLOW_DISPATCH_DURATION.record(
-					dt,
-					&[
-						KeyValue::new("workflow_name", ctx.name().to_string()),
-						KeyValue::new("sub_workflow_name", sub_workflow_name),
-					],
-				);
-				metrics::WORKFLOW_DISPATCHED.add(
-					1,
-					&[
-						KeyValue::new("workflow_name", ctx.name().to_string()),
-						KeyValue::new("sub_workflow_name", sub_workflow_name),
-					],
-				);
+				metrics::WORKFLOW_DISPATCH_DURATION
+					.with_label_values(&[ctx.name(), sub_workflow_name])
+					.observe(dt);
+				metrics::WORKFLOW_DISPATCHED
+					.with_label_values(&[ctx.name(), sub_workflow_name])
+					.inc();
 			}
 
 			actual_sub_workflow_id

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

@@ -7,16 +7,18 @@ use tokio::sync::Mutex;
 use tracing::Instrument;
 
 use crate::{
+	builder::common as builder,
 	ctx::{
 		common,
 		message::{MessageCtx, SubscriptionHandle},
 	},
-	db::DatabaseHandle,
+	db::{DatabaseHandle, WorkflowData},
 	error::{WorkflowError, WorkflowResult},
 	message::Message,
 	operation::{Operation, OperationInput},
-	utils::tags::AsTags,
-	workflow::StateGuard,
+	signal::Signal,
+	utils::{tags::AsTags, topic::AsTopic},
+	workflow::{StateGuard, Workflow},
 };
 
 pub struct ActivityCtx {
@@ -77,6 +79,44 @@ impl ActivityCtx {
 }
 
 impl ActivityCtx {
+	/// Finds the first incomplete workflow with the given tags.
+	#[tracing::instrument(skip_all, ret(Debug), fields(workflow_name=W::NAME))]
+	pub async fn find_workflow<W: Workflow>(&self, tags: impl AsTags) -> Result<Option<Id>> {
+		common::find_workflow::<W>(&self.db, tags)
+			.in_current_span()
+			.await
+	}
+
+	/// Finds the first incomplete workflow for each (name, tags) pair in a single batch transaction.
+	#[tracing::instrument(skip_all)]
+	pub async fn find_workflows(
+		&self,
+		queries: &[(&str, serde_json::Value)],
+	) -> Result<Vec<Option<Id>>> {
+		common::find_workflows(&self.db, queries)
+			.in_current_span()
+			.await
+	}
+
+	/// Finds the first incomplete workflow with the given tags.
+	#[tracing::instrument(skip_all)]
+	pub async fn get_workflows(&self, workflow_ids: Vec<Id>) -> Result<Vec<WorkflowData>> {
+		common::get_workflows(&self.db, workflow_ids)
+			.in_current_span()
+			.await
+	}
+
+	/// Creates a signal builder.
+	pub fn signal<T: Signal + Serialize>(&self, body: T) -> builder::signal::SignalBuilder<T> {
+		builder::signal::SignalBuilder::new(
+			self.db.clone(),
+			self.config.clone(),
+			self.ray_id,
+			body,
+			true,
+		)
+	}
+
 	#[tracing::instrument(skip_all)]
 	pub fn state<T: Serialize + DeserializeOwned>(&self) -> Result<StateGuard<'_, T>> {
 		if self.parallelized {
@@ -124,13 +164,13 @@ impl ActivityCtx {
 
 	/// IMPORTANT: This is intended for ephemeral realtime events and should be used carefully. Use
 	/// signals if you need this to be durable.
-	#[tracing::instrument(skip_all, fields(message=M::NAME))]
-	pub async fn subscribe<M>(&self, tags: impl AsTags) -> Result<SubscriptionHandle<M>>
+	#[tracing::instrument(skip_all, fields(message=M::NAME, topic=%topic.as_topic()))]
+	pub async fn subscribe<M>(&self, topic: impl AsTopic) -> Result<SubscriptionHandle<M>>
 	where
 		M: Message,
 	{
 		self.msg_ctx
-			.subscribe::<M>(tags)
+			.subscribe::<M>(&topic.as_topic())
 			.in_current_span()
 			.await
 			.map_err(Into::into)

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

@@ -2,7 +2,6 @@ use std::time::{Duration, Instant};
 
 use anyhow::Result;
 use futures_util::StreamExt;
-use rivet_metrics::KeyValue;
 use rivet_util::Id;
 
 /// Time to delay a workflow from retrying after an error
@@ -67,6 +66,14 @@ pub async fn find_workflow<W: Workflow>(
 		.map_err(Into::into)
 }
 
+/// Finds the first incomplete workflow for each (name, tags) pair in a single batch transaction.
+pub async fn find_workflows(
+	db: &DatabaseHandle,
+	queries: &[(&str, serde_json::Value)],
+) -> Result<Vec<Option<Id>>> {
+	db.find_workflows(queries).await.map_err(Into::into)
+}
+
 /// Finds the first incomplete workflow with the given tags.
 pub async fn get_workflows(
 	db: &DatabaseHandle,
@@ -92,8 +99,11 @@ where
 
 	// Record metrics
 	crate::metrics::OPERATION_PENDING
-		.add(1, &[KeyValue::new("operation_name", I::Operation::NAME)]);
-	crate::metrics::OPERATION_TOTAL.add(1, &[KeyValue::new("operation_name", I::Operation::NAME)]);
+		.with_label_values(&[I::Operation::NAME])
+		.inc();
+	crate::metrics::OPERATION_TOTAL
+		.with_label_values(&[I::Operation::NAME])
+		.inc();
 
 	let start_instant = Instant::now();
 
@@ -109,24 +119,20 @@ where
 
 	let res = tokio::time::timeout(I::Operation::TIMEOUT, I::Operation::run(&ctx, &input))
 		.await
-		.map_err(|_| WorkflowError::OperationTimeout(0))
-		.map(|res| res.map_err(WorkflowError::OperationFailure));
+		.map_err(|_| WorkflowError::OperationTimeout(I::Operation::NAME, 0))
+		.map(|res| res.map_err(|err| WorkflowError::OperationFailure(I::Operation::NAME, err)));
 
 	// Record metrics
 	{
-		let error_code_str = match &res {
+		let error_str = match &res {
 			Ok(Err(err)) => {
-				let error_code_str = err.to_string();
+				let error_str = err.to_string();
 
-				crate::metrics::OPERATION_ERRORS.add(
-					1,
-					&[
-						KeyValue::new("operation_name", I::Operation::NAME),
-						KeyValue::new("error", error_code_str.clone()),
-					],
-				);
+				crate::metrics::OPERATION_ERRORS
+					.with_label_values(&[I::Operation::NAME, error_str.as_str()])
+					.inc();
 
-				error_code_str
+				error_str
 			}
 			Ok(_) => String::new(),
 			Err(_) => "timeout".to_string(),
@@ -135,14 +141,11 @@ where
 		// Other request metrics
 		let dt = start_instant.elapsed().as_secs_f64();
 		crate::metrics::OPERATION_PENDING
-			.add(-1, &[KeyValue::new("operation_name", I::Operation::NAME)]);
-		crate::metrics::OPERATION_DURATION.record(
-			dt,
-			&[
-				KeyValue::new("operation_name", I::Operation::NAME),
-				KeyValue::new("error", error_code_str.clone()),
-			],
-		);
+			.with_label_values(&[I::Operation::NAME])
+			.dec();
+		crate::metrics::OPERATION_DURATION
+			.with_label_values(&[I::Operation::NAME, error_str.as_str()])
+			.observe(dt);
 	}
 
 	let res = res?;

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

@@ -1,5 +1,4 @@
-use rivet_metrics::KeyValue;
-use std::{ops::Deref, time::Instant};
+use std::ops::Deref;
 
 use crate::{
 	ctx::WorkflowCtx,
@@ -14,7 +13,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 +23,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 +39,59 @@ 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
+				.with_label_values(&[self.ctx.name(), signal.signal_name.as_str()])
+				.observe(recv_lag);
 
-		if recv_lag > 3.0 {
-			// We print an error here so the trace of this workflow does not get dropped
-			tracing::error!(
-				?recv_lag,
+			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",
+				);
+			}
+
+			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)
 	}
 }