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

package/packages/error/tests/basic.rs

@@ -121,8 +121,8 @@ fn test_error_chaining() {
 // Test struct without formatted description (from derive.rs)
 #[derive(RivetError, Serialize, Deserialize)]
 #[error(
-	"namespace",
-	"invalid_name",
+	"test",
+	"namespace_invalid_name",
 	"
 	Invalid namespace name.
 	
@@ -144,8 +144,8 @@ fn test_struct_without_formatted_description() {
 	.build();
 	let rivet_error = RivetError::extract(&error);
 
-	assert_eq!(rivet_error.group(), "namespace");
-	assert_eq!(rivet_error.code(), "invalid_name");
+	assert_eq!(rivet_error.group(), "test");
+	assert_eq!(rivet_error.code(), "namespace_invalid_name");
 	assert!(rivet_error.message().contains("Invalid namespace name"));
 	assert!(rivet_error.meta.is_some());
 
@@ -159,8 +159,8 @@ fn test_struct_without_formatted_description() {
 // Test multiline description formatting (from derive.rs)
 #[derive(RivetError, Serialize, Deserialize)]
 #[error(
-	"api",
-	"rate_limited",
+	"test",
+	"api_rate_limited",
 	"
 	Rate limit exceeded.
 	
@@ -184,8 +184,8 @@ fn test_multiline_descriptions() {
 	.build();
 	let rivet_error = RivetError::extract(&error);
 
-	assert_eq!(rivet_error.group(), "api");
-	assert_eq!(rivet_error.code(), "rate_limited");
+	assert_eq!(rivet_error.group(), "test");
+	assert_eq!(rivet_error.code(), "api_rate_limited");
 	assert_eq!(
 		rivet_error.message(),
 		"Rate limit exceeded. Limit: 100, resets at: 1234567890"

package/packages/gasoline/Cargo.toml

@@ -11,6 +11,7 @@ async-stream.workspace = true
 async-trait.workspace = true
 gasoline-macros.workspace = true
 cjson.workspace = true
+clickhouse.workspace = true
 dirs.workspace = true
 futures-util.workspace = true
 hex.workspace = true

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

@@ -1,15 +1,15 @@
-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::MessageCtx, message::Message, metrics};
+use crate::{
+	builder::BuilderError, ctx::MessageCtx, message::Message, metrics, utils::topic::AsTopic,
+};
 
 pub struct MessageBuilder<M: Message> {
 	msg_ctx: MessageCtx,
 	body: M,
-	tags: serde_json::Map<String, serde_json::Value>,
+	topic: Option<String>,
 	wait: bool,
 	error: Option<BuilderError>,
 }
@@ -19,38 +19,18 @@ impl<M: Message> MessageBuilder<M> {
 		MessageBuilder {
 			msg_ctx,
 			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
 	}
@@ -71,33 +51,25 @@ impl<M: Message> MessageBuilder<M> {
 			return Err(err.into());
 		}
 
-		tracing::debug!(tags=?self.tags, "dispatching message");
+		tracing::debug!(topic=?self.topic, "dispatching message");
 
-		let start_instant = Instant::now();
+		let topic = self.topic.unwrap_or_else(|| "*".to_string());
 
-		let tags = serde_json::Value::Object(self.tags);
+		let start_instant = Instant::now();
 
 		if self.wait {
-			self.msg_ctx.message_wait(tags, self.body).await?;
+			self.msg_ctx.message_wait(&topic, self.body).await?;
 		} else {
-			self.msg_ctx.message(tags, self.body).await?;
+			self.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", ""),
-				KeyValue::new("message_name", M::NAME),
-			],
-		);
-		metrics::MESSAGE_PUBLISHED.add(
-			1,
-			&[
-				KeyValue::new("workflow_name", ""),
-				KeyValue::new("message_name", M::NAME),
-			],
-		);
+		metrics::MESSAGE_SEND_DURATION
+			.with_label_values(&["", M::NAME])
+			.observe(dt);
+		metrics::MESSAGE_PUBLISHED
+			.with_label_values(&["", M::NAME])
+			.inc();
 
 		Ok(())
 	}

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

@@ -1,7 +1,6 @@
 use std::{fmt::Display, time::Instant};
 
 use anyhow::Result;
-use rivet_metrics::KeyValue;
 use rivet_util::Id;
 use serde::Serialize;
 
@@ -18,6 +17,7 @@ pub struct SignalBuilder<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>,
 }
 
@@ -37,11 +37,12 @@ impl<T: Signal + Serialize> SignalBuilder<T> {
 			to_workflow_name: None,
 			to_workflow_id: None,
 			tags: serde_json::Map::new(),
+			graceful_not_found: false,
 			error: from_workflow.then_some(BuilderError::CannotDispatchFromOpInWorkflow),
 		}
 	}
 
-	// TODO: Get rid of this
+	// TODO: Get rid of this (RVT-5281)
 	// NOTE: This is a bad implementation because it disregards other errors that may have happened earlier
 	#[allow(non_snake_case)]
 	pub fn bypass_signal_from_workflow_I_KNOW_WHAT_IM_DOING(mut self) -> Self {
@@ -102,8 +103,21 @@ impl<T: Signal + Serialize> SignalBuilder<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>> {
 		if let Some(err) = self.error {
 			return Err(err.into());
 		}
@@ -132,8 +146,18 @@ impl<T: Signal + Serialize> SignalBuilder<T> {
 				let workflow_id = self
 					.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");
+
+						return Ok(None);
+					} else {
+						return Err(WorkflowError::WorkflowNotFound.into());
+					}
+				};
 
 				self.db
 					.publish_signal(self.ray_id, workflow_id, signal_id, T::NAME, &input_val)
@@ -173,21 +197,13 @@ impl<T: Signal + Serialize> SignalBuilder<T> {
 		}
 
 		let dt = start_instant.elapsed().as_secs_f64();
-		metrics::SIGNAL_SEND_DURATION.record(
-			dt,
-			&[
-				KeyValue::new("workflow_name", ""),
-				KeyValue::new("signal_name", T::NAME),
-			],
-		);
-		metrics::SIGNAL_PUBLISHED.add(
-			1,
-			&[
-				KeyValue::new("workflow_name", ""),
-				KeyValue::new("signal_name", T::NAME),
-			],
-		);
-
-		Ok(signal_id)
+		metrics::SIGNAL_SEND_DURATION
+			.with_label_values(&["", T::NAME])
+			.observe(dt);
+		metrics::SIGNAL_PUBLISHED
+			.with_label_values(&["", T::NAME])
+			.inc();
+
+		Ok(Some(signal_id))
 	}
 }

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

@@ -1,7 +1,6 @@
 use std::{fmt::Display, marker::PhantomData, time::Instant};
 
 use anyhow::Result;
-use rivet_metrics::KeyValue;
 use rivet_util::Id;
 use serde::Serialize;
 use tracing::Instrument;
@@ -51,6 +50,17 @@ where
 		}
 	}
 
+	// TODO: Get rid of this (RVT-5281)
+	// NOTE: This is a bad implementation because it disregards other errors that may have happened earlier
+	#[allow(non_snake_case)]
+	pub fn bypass_signal_from_workflow_I_KNOW_WHAT_IM_DOING(mut self) -> Self {
+		if let Some(BuilderError::CannotDispatchFromOpInWorkflow) = &self.error {
+			self.error = None;
+		}
+
+		self
+	}
+
 	pub fn tags(mut self, tags: serde_json::Value) -> Self {
 		if self.error.is_some() {
 			return self;
@@ -138,24 +148,18 @@ where
 			} else {
 				tracing::debug!(?tags, "unique workflow already exists");
 			}
+		} else {
+			tracing::debug!(?actual_workflow_id, "dispatched workflow");
 		}
 
 		if workflow_id == actual_workflow_id {
 			let dt = start_instant.elapsed().as_secs_f64();
-			metrics::WORKFLOW_DISPATCH_DURATION.record(
-				dt,
-				&[
-					KeyValue::new("workflow_name", ""),
-					KeyValue::new("sub_workflow_name", workflow_name),
-				],
-			);
-			metrics::WORKFLOW_DISPATCHED.add(
-				1,
-				&[
-					KeyValue::new("workflow_name", ""),
-					KeyValue::new("sub_workflow_name", workflow_name),
-				],
-			);
+			metrics::WORKFLOW_DISPATCH_DURATION
+				.with_label_values(&["", workflow_name])
+				.observe(dt);
+			metrics::WORKFLOW_DISPATCHED
+				.with_label_values(&["", workflow_name])
+				.inc();
 		}
 
 		Ok(actual_workflow_id)

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

@@ -0,0 +1,295 @@
+use std::time::{Duration, Instant};
+
+use anyhow::Result;
+use serde::{Serialize, de::DeserializeOwned};
+use tracing::Instrument;
+
+use crate::{
+	ctx::WorkflowCtx,
+	ctx::workflow::Loop,
+	error::WorkflowError,
+	executable::AsyncResult,
+	history::{cursor::HistoryResult, location::Coordinate},
+	metrics,
+};
+
+/// How often to commit loop event data to db and mark previous loop history to forgotten
+const DEFAULT_LOOP_COMMIT_INTERVAL: usize = 20;
+
+pub struct LoopBuilder<'a, S> {
+	ctx: &'a mut WorkflowCtx,
+	state: S,
+	commit_interval: Option<usize>,
+}
+
+impl<'a, S: Serialize + DeserializeOwned> LoopBuilder<'a, S> {
+	pub(crate) fn new(ctx: &'a mut WorkflowCtx, state: S) -> Self {
+		LoopBuilder {
+			ctx,
+			state,
+			commit_interval: None,
+		}
+	}
+
+	pub fn with_state<S2: Serialize + DeserializeOwned>(self, state: S2) -> LoopBuilder<'a, S2> {
+		LoopBuilder {
+			ctx: self.ctx,
+			state,
+			commit_interval: self.commit_interval,
+		}
+	}
+
+	pub fn commit_interval(self, commit_interval: usize) -> Self {
+		LoopBuilder {
+			ctx: self.ctx,
+			state: self.state,
+			commit_interval: Some(commit_interval),
+		}
+	}
+
+	#[tracing::instrument(skip_all)]
+	pub async fn run<F, T>(self, mut cb: F) -> Result<T>
+	where
+		F: for<'b> FnMut(&'b mut WorkflowCtx, &'b mut S) -> AsyncResult<'b, Loop<T>>,
+		T: Serialize + DeserializeOwned,
+	{
+		let LoopBuilder {
+			ctx,
+			state,
+			commit_interval,
+		} = self;
+
+		ctx.check_stop()?;
+
+		let history_res = ctx.cursor().compare_loop(ctx.version())?;
+		let loop_location = ctx.cursor().current_location_for(&history_res);
+
+		// Loop existed before
+		let (mut iteration, mut state, output, mut loop_event_init_fut) =
+			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, None)
+			} else {
+				let state_val = serde_json::value::to_raw_value(&state)
+					.map_err(WorkflowError::SerializeLoopOutput)?;
+
+				// Clone data to move into future
+				let loop_location = loop_location.clone();
+				let db2 = ctx.db().clone();
+				let workflow_id = ctx.workflow_id();
+				let name = ctx.name().to_string();
+				let version = ctx.version();
+				let nested_loop_location = ctx.loop_location().cloned();
+
+				// This future is deferred until later for parallelization
+				let loop_event_init_fut = async move {
+					db2.upsert_workflow_loop_event(
+						workflow_id,
+						&name,
+						&loop_location,
+						version,
+						0,
+						&state_val,
+						None,
+						nested_loop_location.as_ref(),
+					)
+					.await
+				};
+
+				(0, state, None, Some(loop_event_init_fut))
+			};
+
+		// Create a branch for the loop event
+		let mut loop_branch =
+			ctx.branch_inner(ctx.input().clone(), ctx.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");
+
+			// Used to defer loop upsertion for parallelization
+			let mut loop_event_upsert_fut = None;
+			let mut iteration_dt = Duration::ZERO;
+
+			loop {
+				ctx.check_stop()?;
+
+				// Create a new branch for each iteration of the loop at location {...loop location, iteration idx}
+				let mut iteration_branch = loop_branch.branch_inner(
+					ctx.input().clone(),
+					ctx.version(),
+					loop_branch
+						.cursor()
+						.root()
+						.join(Coordinate::simple(iteration + 1)),
+				);
+				let iteration_branch_root = iteration_branch.cursor().root().clone();
+
+				// Set branch loop location to the current loop
+				iteration_branch.set_loop_location(loop_location.clone());
+
+				let previous_iteration = iteration;
+
+				// Async block for instrumentation purposes
+				let res = async {
+					let db2 = ctx.db().clone();
+
+					// NOTE: Great care has been taken to optimize this function. This join allows multiple
+					// txns to run simultaneously instead of in series but is hard to read.
+					//
+					// 1. First (but not necessarily chronologically first because its parallelized), we
+					//    commit the loop event. This only happens on the first iteration of the loop
+					// 2. Second, we commit the branch event for the current iteration
+					// 3. Third, we run the user's loop code
+					// 4. Last, if we have to upsert the loop event, we save the future and poll it in the
+					//    next iteration of the loop as part of this join
+					let (
+						loop_event_commit_res,
+						loop_event_upsert_res,
+						branch_commit_res,
+						(loop_res, cb_dt),
+					) = tokio::join!(
+						async {
+							if let Some(loop_event_init_fut) = loop_event_init_fut.take() {
+								loop_event_init_fut.await
+							} else {
+								Ok(())
+							}
+						},
+						async {
+							if let Some(loop_event_upsert_fut) = loop_event_upsert_fut.take() {
+								let start_instant = Instant::now();
+								loop_event_upsert_fut.await?;
+								metrics::LOOP_COMMIT_DURATION
+									.with_label_values(&[&ctx.name().to_string()])
+									.observe(start_instant.elapsed().as_secs_f64());
+							}
+
+							anyhow::Ok(())
+						},
+						async {
+							// Insert event if iteration is not a replay
+							if !loop_branch.cursor().compare_loop_branch(iteration)? {
+								db2.commit_workflow_branch_event(
+									ctx.workflow_id(),
+									&iteration_branch_root,
+									ctx.version(),
+									Some(&loop_location),
+								)
+								.await?;
+
+								// Only record iteration duration if its not a replay
+								metrics::LOOP_ITERATION_DURATION
+									.with_label_values(&[&ctx.name().to_string()])
+									.observe(iteration_dt.as_secs_f64());
+							}
+
+							anyhow::Ok(())
+						},
+						async {
+							let iteration_start_instant = Instant::now();
+
+							(
+								cb(&mut iteration_branch, &mut state).await,
+								iteration_start_instant.elapsed(),
+							)
+						}
+					);
+
+					loop_event_commit_res?;
+					loop_event_upsert_res?;
+					branch_commit_res?;
+
+					iteration_dt = cb_dt;
+
+					// Run loop
+					match loop_res? {
+						Loop::Continue => {
+							iteration += 1;
+
+							// Commit workflow state to db
+							if iteration % commit_interval.unwrap_or(DEFAULT_LOOP_COMMIT_INTERVAL)
+								== 0
+							{
+								let state_val = serde_json::value::to_raw_value(&state)
+									.map_err(WorkflowError::SerializeLoopOutput)?;
+
+								// Clone data to move into future
+								let loop_location = loop_location.clone();
+								let db2 = ctx.db().clone();
+								let workflow_id = ctx.workflow_id();
+								let name = ctx.name();
+								let version = ctx.version();
+								let nested_loop_location = ctx.loop_location().cloned();
+
+								// Defer upsertion to next iteration so it runs in parallel
+								loop_event_upsert_fut = Some(async move {
+									db2.upsert_workflow_loop_event(
+										workflow_id,
+										&name,
+										&loop_location,
+										version,
+										iteration,
+										&state_val,
+										None,
+										nested_loop_location.as_ref(),
+									)
+									.await
+								});
+							}
+
+							anyhow::Ok(None)
+						}
+						Loop::Break(res) => {
+							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)?;
+
+							// Commit loop output and final state to db. Note that we don't defer this because
+							// there will be no more loop iterations afterwards.
+							ctx.db()
+								.upsert_workflow_loop_event(
+									ctx.workflow_id(),
+									&ctx.name(),
+									&loop_location,
+									ctx.version(),
+									iteration,
+									&state_val,
+									Some(&output_val),
+									ctx.loop_location(),
+								)
+								.await?;
+
+							Ok(Some(res))
+						}
+					}
+				}
+				.instrument(tracing::info_span!("iteration", iteration=%previous_iteration))
+				.await?;
+
+				// Validate no leftover events
+				iteration_branch.cursor().check_clear()?;
+
+				if let Some(res) = res {
+					break res;
+				}
+			}
+		};
+
+		// Move to next event
+		ctx.cursor_mut().update(&loop_location);
+
+		Ok(output)
+	}
+}