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

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

@@ -1,5 +1,6 @@
 use std::{
-	fmt::{self, Debug},
+	borrow::Cow,
+	fmt::{self, Debug, Display},
 	marker::PhantomData,
 	sync::Arc,
 };
@@ -8,12 +9,11 @@ use rivet_pools::UpsPool;
 use rivet_util::Id;
 use tokio_util::sync::{CancellationToken, DropGuard};
 use tracing::Instrument;
-use universalpubsub::{NextOutput, Subscriber};
+use universalpubsub::{NextOutput, Subject, Subscriber};
 
 use crate::{
 	error::{WorkflowError, WorkflowResult},
 	message::{Message, PubsubMessage, PubsubMessageWrapper},
-	utils::{self, tags::AsTags},
 };
 
 #[derive(Clone)]
@@ -54,21 +54,19 @@ impl MessageCtx {
 	/// finish publishing. This is done since there are very few cases where a
 	/// service should need to wait or fail if a message does not publish
 	/// successfully.
-	#[tracing::instrument(skip_all, fields(message=M::NAME))]
-	pub async fn message<M>(
-		&self,
-		tags: impl AsTags + 'static,
-		message_body: M,
-	) -> WorkflowResult<()>
+	#[tracing::instrument(skip_all, fields(message=M::NAME, %topic))]
+	pub async fn message<M>(&self, topic: &str, message_body: M) -> WorkflowResult<()>
 	where
 		M: Message,
 	{
 		let client = self.clone();
-		let spawn_res = tokio::task::Builder::new()
+		let topic = topic.to_string();
+
+		tokio::task::Builder::new()
 			.name("gasoline::message_async")
 			.spawn(
 				async move {
-					match client.message_wait::<M>(tags, message_body).await {
+					match client.message_wait::<M>(&topic, message_body).await {
 						Ok(_) => {}
 						Err(err) => {
 							tracing::error!(?err, "failed to publish message");
@@ -76,12 +74,9 @@ impl MessageCtx {
 					}
 				}
 				.instrument(tracing::info_span!("message_bg")),
-			);
-		if let Err(err) = spawn_res {
-			tracing::error!(?err, "failed to spawn message_async task");
-		}
-
-		Ok(())
+			)
+			.map_err(|err| WorkflowError::PublishMessage(err.into()))
+			.map(|_| ())
 	}
 
 	/// Same as `message` but waits for the message to successfully publish.
@@ -89,12 +84,15 @@ impl MessageCtx {
 	/// This is useful in scenarios where we need to publish a large amount of
 	/// messages at once so we put the messages in a queue instead of submitting
 	/// a large number of tasks to Tokio at once.
-	#[tracing::instrument(skip_all, fields(message = M::NAME))]
-	pub async fn message_wait<M>(&self, tags: impl AsTags, message_body: M) -> WorkflowResult<()>
+	#[tracing::instrument(skip_all, fields(message = M::NAME, %topic))]
+	pub async fn message_wait<M>(&self, topic: &str, message_body: M) -> WorkflowResult<()>
 	where
 		M: Message,
 	{
-		let subject = M::subject();
+		let subject = MsgSubject {
+			topic,
+			msg_marker: PhantomData::<M>,
+		};
 		let duration_since_epoch = std::time::SystemTime::now()
 			.duration_since(std::time::UNIX_EPOCH)
 			.unwrap_or_else(|err| unreachable!("time is broken: {}", err));
@@ -112,7 +110,7 @@ impl MessageCtx {
 		let message = PubsubMessageWrapper {
 			req_id,
 			ray_id: self.ray_id,
-			tags: tags.as_tags()?,
+			topic: topic.to_string(),
 			ts,
 			body: &body_buf,
 		};
@@ -128,15 +126,14 @@ impl MessageCtx {
 		// It's important to write to the stream as fast as possible in order to
 		// ensure messages are handled quickly.
 		let message_buf = Arc::new(message_buf);
-		self.message_publish_pubsub::<M>(&subject, message_buf)
-			.await;
+		self.message_publish_pubsub::<M>(subject, message_buf).await;
 
 		Ok(())
 	}
 
 	/// Publishes the message to pubsub.
 	#[tracing::instrument(level = "debug", skip_all)]
-	async fn message_publish_pubsub<M>(&self, subject: &str, message_buf: Arc<Vec<u8>>)
+	async fn message_publish_pubsub<M>(&self, subject: MsgSubject<'_, M>, message_buf: Arc<Vec<u8>>)
 	where
 		M: Message,
 	{
@@ -146,8 +143,6 @@ impl MessageCtx {
 			// Ignore for infinite backoff
 			backoff.tick().await;
 
-			let subject = subject.to_owned();
-
 			tracing::trace!(
 				%subject,
 				message_len = message_buf.len(),
@@ -156,7 +151,7 @@ impl MessageCtx {
 			if let Err(err) = self
 				.pubsub
 				.publish(
-					&subject,
+					subject.clone(),
 					&(*message_buf),
 					universalpubsub::PublishOpts::broadcast(),
 				)
@@ -181,13 +176,13 @@ impl MessageCtx {
 // MARK: Subscriptions
 impl MessageCtx {
 	/// Listens for gasoline messages globally on pubsub.
-	#[tracing::instrument(skip_all, fields(message = M::NAME))]
-	pub async fn subscribe<M>(&self, tags: impl AsTags) -> WorkflowResult<SubscriptionHandle<M>>
+	#[tracing::instrument(skip_all, fields(message = M::NAME, %topic))]
+	pub async fn subscribe<M>(&self, topic: &str) -> WorkflowResult<SubscriptionHandle<M>>
 	where
 		M: Message,
 	{
 		self.subscribe_opt::<M>(SubscribeOpts {
-			tags: tags.as_tags()?,
+			topic: topic.to_string(),
 			flush: true,
 		})
 		.in_current_span()
@@ -203,10 +198,10 @@ impl MessageCtx {
 	where
 		M: Message,
 	{
-		let subject = M::subject();
+		let subject = format!("{}:{}", M::subject(), opts.topic);
 
 		// Create subscription and flush immediately.
-		tracing::debug!(%subject, tags = ?opts.tags, "creating subscription");
+		tracing::debug!(%subject, "creating subscription");
 		let subscription = self
 			.pubsub
 			.subscribe(&subject)
@@ -220,14 +215,14 @@ impl MessageCtx {
 		}
 
 		// Return handle
-		let subscription = SubscriptionHandle::new(subject, subscription, opts.tags.clone());
+		let subscription = SubscriptionHandle::new(subject, subscription);
 		Ok(subscription)
 	}
 }
 
 #[derive(Debug)]
 pub struct SubscribeOpts {
-	pub tags: serde_json::Value,
+	pub topic: String,
 	pub flush: bool,
 }
 
@@ -242,7 +237,6 @@ where
 	_guard: DropGuard,
 	subject: String,
 	subscription: Subscriber,
-	tags: serde_json::Value,
 }
 
 impl<M> Debug for SubscriptionHandle<M>
@@ -252,7 +246,6 @@ where
 	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 		f.debug_struct("SubscriptionHandle")
 			.field("subject", &self.subject)
-			.field("tags", &self.tags)
 			.finish()
 	}
 }
@@ -262,7 +255,7 @@ where
 	M: Message,
 {
 	#[tracing::instrument(level = "debug", skip_all)]
-	fn new(subject: String, subscription: Subscriber, tags: serde_json::Value) -> Self {
+	fn new(subject: String, subscription: Subscriber) -> Self {
 		let token = CancellationToken::new();
 
 		{
@@ -291,7 +284,6 @@ where
 			_guard: token.drop_guard(),
 			subject,
 			subscription,
-			tags,
 		}
 	}
 
@@ -302,35 +294,28 @@ where
 	pub async fn next(&mut self) -> WorkflowResult<PubsubMessage<M>> {
 		tracing::debug!("waiting for message");
 
-		loop {
-			// Poll the subscription.
-			//
-			// Use blocking threads instead of `try_next`, since I'm not sure
-			// try_next works as intended.
-			let message = match self.subscription.next().await {
-				Ok(NextOutput::Message(msg)) => msg,
-				Ok(NextOutput::Unsubscribed) => {
-					tracing::debug!("unsubscribed");
-					return Err(WorkflowError::SubscriptionUnsubscribed);
-				}
-				Err(err) => {
-					tracing::warn!(?err, "subscription error");
-					return Err(WorkflowError::CreateSubscription(err.into()));
-				}
-			};
-
-			let message_wrapper = PubsubMessage::<M>::deserialize_wrapper(&message.payload)?;
+		// Poll the subscription.
+		//
+		// Use blocking threads instead of `try_next`, since I'm not sure
+		// try_next works as intended.
+		let message = match self.subscription.next().await {
+			Ok(NextOutput::Message(msg)) => msg,
+			Ok(NextOutput::Unsubscribed) => {
+				tracing::debug!("unsubscribed");
+				return Err(WorkflowError::SubscriptionUnsubscribed);
+			}
+			Err(err) => {
+				tracing::warn!(?err, "subscription error");
+				return Err(WorkflowError::CreateSubscription(err.into()));
+			}
+		};
 
-			// Check if the subscription tags match a subset of the message tags
-			if utils::is_value_subset(&self.tags, &message_wrapper.tags) {
-				let message = PubsubMessage::<M>::deserialize_from_wrapper(message_wrapper)?;
-				tracing::debug!(?message, "received message");
+		let message_wrapper = PubsubMessage::<M>::deserialize_wrapper(&message.payload)?;
 
-				return Ok(message);
-			}
+		let message = PubsubMessage::<M>::deserialize_from_wrapper(message_wrapper)?;
+		tracing::debug!(?message, "received message");
 
-			// Message tags don't match, continue with loop
-		}
+		Ok(message)
 	}
 
 	/// Converts the subscription in to a stream.
@@ -344,3 +329,30 @@ where
 		})
 	}
 }
+
+// Helper struct
+struct MsgSubject<'a, M: Message> {
+	topic: &'a str,
+	msg_marker: PhantomData<M>,
+}
+
+impl<M: Message> Clone for MsgSubject<'_, M> {
+	fn clone(&self) -> Self {
+		MsgSubject {
+			topic: self.topic,
+			msg_marker: PhantomData::<M>,
+		}
+	}
+}
+
+impl<M: Message> Display for MsgSubject<'_, M> {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		write!(f, "{}:{}", M::subject(), self.topic)
+	}
+}
+
+impl<M: Message> Subject for MsgSubject<'_, M> {
+	fn root<'a>() -> Option<Cow<'a, str>> {
+		Some(Cow::Owned(M::subject()))
+	}
+}

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

@@ -13,7 +13,7 @@ use crate::{
 	message::Message,
 	operation::{Operation, OperationInput},
 	signal::Signal,
-	utils::tags::AsTags,
+	utils::{tags::AsTags, topic::AsTopic},
 	workflow::{Workflow, WorkflowInput},
 };
 
@@ -91,6 +91,17 @@ impl OperationCtx {
 			.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>> {
@@ -101,7 +112,6 @@ impl OperationCtx {
 
 	/// Creates a signal builder.
 	pub fn signal<T: Signal + Serialize>(&self, body: T) -> builder::signal::SignalBuilder<T> {
-		// TODO: Add check for from_workflow so you cant dispatch a signal
 		builder::signal::SignalBuilder::new(
 			self.db.clone(),
 			self.config.clone(),
@@ -138,13 +148,13 @@ impl OperationCtx {
 		builder::message::MessageBuilder::new(self.msg_ctx.clone(), body)
 	}
 
-	#[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/standalone.rs

@@ -13,7 +13,7 @@ use crate::{
 	message::Message,
 	operation::{Operation, OperationInput},
 	signal::Signal,
-	utils::tags::AsTags,
+	utils::{tags::AsTags, topic::AsTopic},
 	workflow::{Workflow, WorkflowInput},
 };
 
@@ -66,6 +66,23 @@ impl StandaloneCtx {
 		})
 	}
 
+	#[tracing::instrument(skip_all)]
+	pub fn with_ray(&self, ray_id: Id, req_id: Id) -> WorkflowResult<Self> {
+		let mut ctx = StandaloneCtx::new(
+			self.db.clone(),
+			self.config.clone(),
+			self.pools.clone(),
+			self.cache.clone(),
+			&self.name,
+			ray_id,
+			req_id,
+		)?;
+
+		ctx.from_workflow = self.from_workflow;
+
+		Ok(ctx)
+	}
+
 	#[tracing::instrument(skip_all)]
 	pub fn new_from_activity(ctx: &ActivityCtx, req_id: Id) -> WorkflowResult<Self> {
 		let mut ctx = StandaloneCtx::new(
@@ -129,6 +146,17 @@ impl StandaloneCtx {
 			.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>> {
@@ -175,13 +203,13 @@ impl StandaloneCtx {
 		builder::message::MessageBuilder::new(self.msg_ctx.clone(), body)
 	}
 
-	#[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/test.rs

@@ -14,7 +14,7 @@ use crate::{
 	operation::{Operation, OperationInput},
 	prelude::*,
 	signal::Signal,
-	utils::tags::AsTags,
+	utils::{tags::AsTags, topic::AsTopic},
 	workflow::{Workflow, WorkflowInput},
 };
 
@@ -55,8 +55,9 @@ impl TestCtx {
 		let cache = rivet_cache::CacheInner::from_env(&config, pools.clone())
 			.expect("failed to create cache");
 
-		let db = db::DatabaseKv::from_pools(pools.clone()).await?;
-		let debug_db = db::DatabaseKv::from_pools(pools.clone()).await? as Arc<dyn DatabaseDebug>;
+		let db = db::DatabaseKv::new(config.clone(), pools.clone()).await?;
+		let debug_db =
+			db::DatabaseKv::new(config.clone(), pools.clone()).await? as Arc<dyn DatabaseDebug>;
 
 		let service_name = format!("{}-test--{}", rivet_env::service_name(), "gasoline_test");
 		let ray_id = Id::new_v1(config.dc_label());
@@ -93,6 +94,19 @@ impl TestCtx {
 		&*self.debug_db
 	}
 
+	pub fn standalone(&self) -> Result<StandaloneCtx> {
+		StandaloneCtx::new(
+			self.db.clone(),
+			self.config.clone(),
+			self.pools.clone(),
+			self.cache.clone(),
+			&self.name,
+			self.ray_id,
+			Id::new_v1(self.config.dc_label()),
+		)
+		.map_err(Into::into)
+	}
+
 	pub async fn shutdown(&mut self) -> Result<()> {
 		if let Some(worker_handle) = self.worker_handle.take() {
 			tracing::info!("stopping workflow worker");
@@ -148,6 +162,17 @@ impl TestCtx {
 			.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>> {
@@ -193,13 +218,13 @@ impl TestCtx {
 		builder::message::MessageBuilder::new(self.msg_ctx.clone(), body)
 	}
 
-	#[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/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,14 +125,12 @@ impl<'a> VersionedWorkflowCtx<'a> {
 		})
 	}
 
-	/// Execute a custom listener.
+	/// Listens for 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>> {
 		wrap!(self, "listen", {
-			self.inner.custom_listener(listener).in_current_span().await
+			self.inner.listen_n::<T>(limit).in_current_span().await
 		})
 	}
 
@@ -181,12 +179,40 @@ impl<'a> VersionedWorkflowCtx<'a> {
 		})
 	}
 
+	#[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>> {
+		wrap!(self, "listen with timeout", {
+			self.inner
+				.listen_n_with_timeout::<T>(duration, limit)
+				.in_current_span()
+				.await
+		})
+	}
+
 	#[tracing::instrument(skip_all, fields(t=std::any::type_name::<T>()))]
 	pub async fn listen_until<T: Listen>(&mut self, time: impl TsToMillis) -> Result<Option<T>> {
 		wrap!(self, "listen until", {
 			self.inner.listen_until::<T>(time).in_current_span().await
 		})
 	}
+
+	#[tracing::instrument(skip_all, fields(t=std::any::type_name::<T>()))]
+	pub async fn listen_n_until<T: Listen>(
+		&mut self,
+		time: impl TsToMillis,
+		limit: usize,
+	) -> Result<Vec<T>> {
+		wrap!(self, "listen until", {
+			self.inner
+				.listen_n_until::<T>(time, limit)
+				.in_current_span()
+				.await
+		})
+	}
 }
 
 impl<'a> Deref for VersionedWorkflowCtx<'a> {