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

package/packages/gasoline/src/signal.rs

@@ -46,28 +46,31 @@ macro_rules! join_signal {
 		join_signal!(@ $vis $join [] [] $($tt)*);
 	};
 	(@
-	    $vis:vis $join:ident
-	    [$({ $names:tt } { $types:tt })*]
-	    [$({ $just_types:tt })*]
+		$vis:vis $join:ident
+		[$({ $names:tt } { $types:tt })*]
+		[$({ $just_types:tt })*]
 	) => {
-    	$vis enum $join {
+		$vis enum $join {
 			$( $names ($types) ),*
 		}
 
-    	#[async_trait::async_trait]
+		#[async_trait::async_trait]
 		impl Listen for $join {
 			#[tracing::instrument(skip_all, fields(t=std::any::type_name::<Self>()))]
-			async fn listen(ctx: &mut gas::prelude::ListenCtx) -> gas::prelude::WorkflowResult<Self> {
-				let row = ctx.listen_any(&[
-				    $(<$just_types 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(&[
+						$(<$just_types 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> {
 				$(
-				    if name == <$types as gas::signal::Signal>::NAME {
+					if name == <$types as gas::signal::Signal>::NAME {
 						std::result::Result::Ok(
 							Self::$names(
 								serde_json::from_str(body.get())
@@ -84,33 +87,33 @@ macro_rules! join_signal {
 				}
 			}
 		}
-    };
+	};
 	(@
-	    $vis:vis $join:ident
-	    [$({ $names:tt } { $types:tt })*]
-	    [$({ $just_types:tt })*]
-	    $name:ident,
-	    $($tail:tt)*
+		$vis:vis $join:ident
+		[$({ $names:tt } { $types:tt })*]
+		[$({ $just_types:tt })*]
+		$name:ident,
+		$($tail:tt)*
 	) => {
 	   join_signal!(@
-	       $vis $join
-	       [$( { $names } { $types } )* { $name } { $name }]
-	       [$( { $just_types } )* { $name }]
-	       $($tail)*
+	   	$vis $join
+	   	[$( { $names } { $types } )* { $name } { $name }]
+	   	[$( { $just_types } )* { $name }]
+	   	$($tail)*
 	   );
 	};
 	(@
-	    $vis:vis $join:ident
-	    [$({ $names:tt } { $types:tt })*]
-	    [$({ $just_types:tt })*]
-	    $name:ident($ty:ty),
-	    $($tail:tt)*
+		$vis:vis $join:ident
+		[$({ $names:tt } { $types:tt })*]
+		[$({ $just_types:tt })*]
+		$name:ident($ty:ty),
+		$($tail:tt)*
 	) => {
 	   join_signal!(@
-	       $vis $join
-	       [$( { $names } { $types } )* { $name } { $ty }]
-	       [$( { $just_types } )* { $ty }]
-	       $($tail)*
+	   	$vis $join
+	   	[$( { $names } { $types } )* { $name } { $ty }]
+	   	[$( { $just_types } )* { $ty }]
+	   	$($tail)*
 	   );
 	};
 }

package/packages/gasoline/src/utils/mod.rs

@@ -1,20 +1,3 @@
 pub mod tags;
 pub mod time;
-
-/// Returns true if `subset` is a subset of `superset`.
-pub fn is_value_subset(subset: &serde_json::Value, superset: &serde_json::Value) -> bool {
-	match (subset, superset) {
-		(serde_json::Value::Object(sub_obj), serde_json::Value::Object(super_obj)) => {
-			sub_obj.iter().all(|(k, sub_val)| {
-				super_obj
-					.get(k)
-					.map_or(false, |super_val| is_value_subset(sub_val, super_val))
-			})
-		}
-		(serde_json::Value::Array(sub_arr), serde_json::Value::Array(super_arr)) => sub_arr
-			.iter()
-			.zip(super_arr)
-			.all(|(sub_val, super_val)| is_value_subset(sub_val, super_val)),
-		_ => subset == superset,
-	}
-}
+pub mod topic;

package/packages/gasoline/src/utils/topic.rs

@@ -0,0 +1,35 @@
+use std::fmt::Display;
+
+pub trait AsTopic: Send + Sync {
+	fn as_topic(&self) -> String;
+}
+
+impl AsTopic for String {
+	fn as_topic(&self) -> String {
+		self.clone()
+	}
+}
+
+impl AsTopic for &str {
+	fn as_topic(&self) -> String {
+		self.to_string()
+	}
+}
+
+impl<T: Display + Send + Sync, U: Display + Send + Sync> AsTopic for (T, U) {
+	fn as_topic(&self) -> String {
+		format!("{}:{}", self.0, self.1)
+	}
+}
+
+impl AsTopic for () {
+	fn as_topic(&self) -> String {
+		String::new()
+	}
+}
+
+impl<T: AsTopic> AsTopic for &T {
+	fn as_topic(&self) -> String {
+		(*self).as_topic()
+	}
+}

package/packages/gasoline/src/worker.rs

@@ -16,6 +16,7 @@ use crate::{
 	ctx::WorkflowCtx,
 	db::{BumpSubSubject, DatabaseHandle},
 	error::WorkflowError,
+	metrics,
 	registry::RegistryHandle,
 };
 
@@ -25,6 +26,7 @@ pub(crate) const PING_INTERVAL: Duration = Duration::from_secs(10);
 const METRICS_INTERVAL: Duration = Duration::from_secs(20);
 // How long the pull workflows function can take before shutting down the runtime.
 const PULL_WORKFLOWS_TIMEOUT: Duration = Duration::from_secs(10);
+const SHUTDOWN_PROGRESS_INTERVAL: Duration = Duration::from_secs(7);
 
 /// Used to spawn a new thread that indefinitely polls the database for new workflows. Only pulls workflows
 /// that are registered in its registry. After pulling, the workflows are ran and their state is written to
@@ -72,16 +74,21 @@ impl Worker {
 
 		let cache = rivet_cache::CacheInner::from_env(&self.config, self.pools.clone())?;
 
-		let mut bump_sub = { self.db.bump_sub(BumpSubSubject::Worker).await? };
+		// We use ready_chunks because multiple bumps in a row should be processed as 1 bump
+		let mut bump_sub = self
+			.db
+			.bump_sub(BumpSubSubject::Worker)
+			.await?
+			.ready_chunks(1024);
 
 		let mut tick_interval = tokio::time::interval(self.db.worker_poll_interval());
 		tick_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
 
-		let mut term_signal = TermSignal::new().await;
+		let mut term_signal = TermSignal::get();
 
 		// Update ping at least once before doing anything else
 		self.db
-			.update_worker_ping(self.worker_id)
+			.update_worker_ping(self.worker_id, true)
 			.await
 			.context("failed updating worker ping")?;
 
@@ -101,8 +108,13 @@ impl Worker {
 			tokio::select! {
 				_ = tick_interval.tick() => {},
 				res = bump_sub.next() => {
-					if res.is_none() {
-						break Err(WorkflowError::SubscriptionUnsubscribed.into());
+					match res {
+						Some(bumps) => {
+							metrics::WORKER_BUMPS_PER_TICK
+								.with_label_values(&[self.worker_id.to_string().as_str()])
+								.observe(bumps.len() as f64);
+						}
+						None => break Err(WorkflowError::SubscriptionUnsubscribed.into()),
 					}
 
 					tick_interval.reset();
@@ -128,17 +140,13 @@ impl Worker {
 			}
 
 			if let Err(err) = self.tick(&cache).await {
-				// Cancel background tasks. We abort because these are not critical tasks.
-				gc_handle.abort();
-				metrics_handle.abort();
-
 				break Err(err);
 			}
 		};
 
 		// Cancel background tasks
-		gc_handle.abort();
 		metrics_handle.abort();
+		gc_handle.abort();
 
 		if let Err(err) = &res {
 			tracing::error!(?err, "worker errored, attempting graceful shutdown");
@@ -177,11 +185,12 @@ impl Worker {
 			let workflow_id = workflow.workflow_id;
 
 			if self.running_workflows.contains_key(&workflow_id) {
-				tracing::error!(?workflow_id, "workflow already running");
+				tracing::error!(?workflow_id, "workflow already running on this worker");
 				continue;
 			}
 
 			let (stop_tx, stop_rx) = watch::channel(());
+			let name = workflow.workflow_name.clone();
 
 			let ctx = WorkflowCtx::new(
 				self.registry.clone(),
@@ -211,12 +220,20 @@ impl Worker {
 			self.running_workflows.insert(
 				workflow_id,
 				WorkflowHandle {
+					name,
 					stop: stop_tx,
 					handle,
 				},
 			);
 		}
 
+		metrics::WORKER_WORKFLOW_ACTIVE.reset();
+		for (_, wf) in &self.running_workflows {
+			metrics::WORKER_WORKFLOW_ACTIVE
+				.with_label_values(&[self.worker_id.to_string().as_str(), wf.name.as_str()])
+				.inc();
+		}
+
 		Ok(())
 	}
 
@@ -232,7 +249,7 @@ impl Worker {
 				loop {
 					ping_interval.tick().await;
 
-					if let Err(err) = db.update_worker_ping(worker_id).await {
+					if let Err(err) = db.update_worker_ping(worker_id, true).await {
 						tracing::error!(?err, "unhandled update ping error");
 					}
 
@@ -266,10 +283,34 @@ impl Worker {
 		)
 	}
 
+	fn shutdown_ping(&self) -> JoinHandle<()> {
+		let db = self.db.clone();
+		let worker_id = self.worker_id;
+
+		tokio::spawn(
+			async move {
+				let mut ping_interval = tokio::time::interval(PING_INTERVAL);
+				ping_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
+
+				loop {
+					ping_interval.tick().await;
+
+					if let Err(err) = db.update_worker_ping(worker_id, false).await {
+						tracing::error!(?err, "unhandled update ping error");
+					}
+				}
+			}
+			.instrument(tracing::info_span!("worker_shutdown_ping_task")),
+		)
+	}
+
 	#[tracing::instrument(skip_all)]
 	async fn shutdown(mut self, mut term_signal: TermSignal) {
 		let shutdown_duration = self.config.runtime.worker_shutdown_duration();
 
+		// Spawn a special ping task to continue updating ping but not the active worker idx
+		let ping_handle = self.shutdown_ping();
+
 		tracing::info!(
 			duration=?shutdown_duration,
 			remaining_workflows=?self.running_workflows.len(),
@@ -297,15 +338,21 @@ impl Worker {
 			.map(|(_, wf)| &mut wf.handle)
 			.collect::<FuturesUnordered<_>>();
 
+		let mut progress_interval = tokio::time::interval(SHUTDOWN_PROGRESS_INTERVAL);
+		progress_interval.tick().await;
+
 		let shutdown_start = Instant::now();
 		loop {
 			// Future will resolve once all workflow tasks complete
-			let join_fut = async { while let Some(_) = wf_futs.next().await {} };
+			let complete_fut = async { while let Some(_) = wf_futs.next().await {} };
 
 			tokio::select! {
-				_ = join_fut => {
+				_ = complete_fut => {
 					break;
 				}
+				_ = progress_interval.tick() => {
+					tracing::info!(remaining_workflows=%wf_futs.len(), "worker still shutting down");
+				}
 				abort = term_signal.recv() => {
 					if abort {
 						tracing::warn!("aborting worker shutdown");
@@ -319,11 +366,20 @@ impl Worker {
 			}
 		}
 
+		// Stop updating ping
+		ping_handle.abort();
+
+		metrics::WORKER_WORKFLOW_ACTIVE.reset();
+
 		let remaining_workflows = wf_futs.into_iter().count();
 		if remaining_workflows == 0 {
 			tracing::info!("all workflows evicted");
 		} else {
 			tracing::warn!(?remaining_workflows, "not all workflows evicted");
+
+			for (workflow_id, handle) in self.running_workflows.iter() {
+				tracing::warn!(?workflow_id, name=?handle.name, "not evicted");
+			}
 		}
 
 		tracing::info!("worker shutdown complete");
@@ -331,6 +387,7 @@ impl Worker {
 }
 
 struct WorkflowHandle {
+	name: String,
 	stop: watch::Sender<()>,
 	handle: JoinHandle<()>,
 }

package/packages/gasoline/src/workflow.rs

@@ -13,6 +13,7 @@ pub trait Workflow {
 	type Output: Serialize + DeserializeOwned + Debug + Send;
 
 	const NAME: &'static str;
+	const PRUNE_VARIANT: PruneVariant;
 
 	async fn run(ctx: &mut WorkflowCtx, input: &Self::Input) -> Result<Self::Output>;
 }
@@ -62,3 +63,15 @@ impl<'a, T: DeserializeOwned + Serialize> Drop for StateGuard<'a, T> {
 		self.guard.0 = serde_json::value::to_raw_value(&self.inner).expect("bad state");
 	}
 }
+
+/// Determines how the data related to a workflow is handled after the workflow completes.
+#[derive(Debug, Default, Clone, Copy, strum::FromRepr)]
+pub enum PruneVariant {
+	/// Delete all of the workflow's data after it completes and is pruned (state + history).
+	#[default]
+	All = 0,
+	/// Delete just the workflow's history after it completes and is pruned (state persists).
+	History = 1,
+	/// Don't delete any workflow data after it completes.
+	None = 2,
+}

package/packages/gasoline/tests/workflows/eviction_test.rs

@@ -15,7 +15,7 @@ pub async fn eviction_test_workflow(
 	ctx.msg(RunningMessage {
 		value: "foo".to_string(),
 	})
-	.tag("test", input.test_id)
+	.topic(("test", input.test_id))
 	.send()
 	.await?;
 
@@ -24,7 +24,7 @@ pub async fn eviction_test_workflow(
 	ctx.msg(RunningMessage {
 		value: "bar".to_string(),
 	})
-	.tag("test", input.test_id)
+	.topic(("test", input.test_id))
 	.send()
 	.await?;
 

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

@@ -8,12 +8,25 @@ use syn::{
 };
 
 #[derive(Default)]
-struct Config {
+enum PruneVariant {
+	#[default]
+	All,
+	History,
+	None,
+}
+
+#[derive(Default)]
+struct WorkflowConfig {
+	prune: PruneVariant,
+}
+
+#[derive(Default)]
+struct OperationConfig {
 	max_retries: Option<syn::Expr>,
 	timeout: Option<syn::Expr>,
 }
 
-impl Config {
+impl OperationConfig {
 	fn max_retries(&self) -> proc_macro2::TokenStream {
 		self.max_retries
 			.as_ref()
@@ -47,9 +60,10 @@ pub fn workflow(attr: TokenStream, item: TokenStream) -> TokenStream {
 		.unwrap_or_else(|| "Workflow".to_string());
 	let item_fn = parse_macro_input!(item as ItemFn);
 
-	if let Err(err) = parse_empty_config(&item_fn.attrs) {
-		return err.into_compile_error().into();
-	}
+	let config = match parse_wf_config(&item_fn.attrs) {
+		Ok(x) => x,
+		Err(err) => return err.into_compile_error().into(),
+	};
 
 	let ctx_ty = syn::parse_str("&mut WorkflowCtx").unwrap();
 	let TraitFnOutput {
@@ -74,6 +88,12 @@ pub fn workflow(attr: TokenStream, item: TokenStream) -> TokenStream {
 		);
 	}
 
+	let prune_variant = match config.prune {
+		PruneVariant::All => quote! { gas::workflow::PruneVariant::All },
+		PruneVariant::History => quote! { gas::workflow::PruneVariant::History },
+		PruneVariant::None => quote! { gas::workflow::PruneVariant::None },
+	};
+
 	let fn_body = item_fn.block;
 	let vis = item_fn.vis;
 
@@ -90,6 +110,7 @@ pub fn workflow(attr: TokenStream, item: TokenStream) -> TokenStream {
 			type Output = #output_type;
 
 			const NAME: &'static str = #fn_name;
+			const PRUNE_VARIANT: gas::workflow::PruneVariant = #prune_variant;
 
 			async fn run(#ctx_ident: #ctx_ty, #input_ident: &Self::Input) -> Result<Self::Output> {
 				#fn_body
@@ -107,7 +128,7 @@ pub fn activity(attr: TokenStream, item: TokenStream) -> TokenStream {
 	let name = parse_macro_input!(attr as Ident).to_string();
 	let item_fn = parse_macro_input!(item as ItemFn);
 
-	let config = match parse_config(&item_fn.attrs) {
+	let config = match parse_op_config(&item_fn.attrs) {
 		Ok(x) => x,
 		Err(err) => return err.into_compile_error().into(),
 	};
@@ -177,7 +198,7 @@ pub fn operation(attr: TokenStream, item: TokenStream) -> TokenStream {
 
 	let item_fn = parse_macro_input!(item as ItemFn);
 
-	let config = match parse_config(&item_fn.attrs) {
+	let config = match parse_op_config(&item_fn.attrs) {
 		Ok(x) => x,
 		Err(err) => return err.into_compile_error().into(),
 	};
@@ -372,9 +393,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> {
@@ -437,8 +462,44 @@ fn error(span: proc_macro2::Span, msg: &str) -> TokenStream {
 	syn::Error::new(span, msg).to_compile_error().into()
 }
 
-fn parse_config(attrs: &[syn::Attribute]) -> syn::Result<Config> {
-	let mut config = Config::default();
+fn parse_wf_config(attrs: &[syn::Attribute]) -> syn::Result<WorkflowConfig> {
+	let mut config = WorkflowConfig::default();
+
+	for attr in attrs {
+		let syn::Meta::NameValue(name_value) = &attr.meta else {
+			continue;
+		};
+
+		let ident = name_value.path.require_ident()?;
+
+		// Verify config property
+		if ident == "prune" {
+			let variant = syn::parse::<syn::Ident>(name_value.value.to_token_stream().into())?;
+
+			config.prune = match variant.to_string().as_str() {
+				"all" => PruneVariant::All,
+				"history" => PruneVariant::History,
+				"none" => PruneVariant::None,
+				_ => {
+					return Err(syn::Error::new(
+						name_value.value.span(),
+						format!("Invalid prune variant. Expected: all, history, none"),
+					));
+				}
+			};
+		} else if ident != "doc" {
+			return Err(syn::Error::new(
+				name_value.span(),
+				format!("Unknown config property `{ident}`"),
+			));
+		}
+	}
+
+	Ok(config)
+}
+
+fn parse_op_config(attrs: &[syn::Attribute]) -> syn::Result<OperationConfig> {
+	let mut config = OperationConfig::default();
 
 	for attr in attrs {
 		let syn::Meta::NameValue(name_value) = &attr.meta else {
@@ -492,6 +553,7 @@ fn parse_msg_config(attrs: &[syn::Attribute]) -> syn::Result<MessageConfig> {
 	Ok(config)
 }
 
+#[allow(dead_code)]
 fn parse_empty_config(attrs: &[syn::Attribute]) -> syn::Result<()> {
 	for attr in attrs {
 		let syn::Meta::NameValue(name_value) = &attr.meta else {

package/packages/gasoline-runtime/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "gasoline-runtime"
+version.workspace = true
+authors.workspace = true
+license.workspace = true
+edition.workspace = true
+
+[dependencies]
+anyhow.workspace = true
+epoxy.workspace = true
+gas.workspace = true
+namespace.workspace = true
+pegboard.workspace = true
+rivet-config.workspace = true
+rivet-types.workspace = true
+serde.workspace = true
+tracing.workspace = true
+universaldb.workspace = true

package/packages/gasoline-runtime/src/lib.rs

@@ -0,0 +1,12 @@
+use gas::prelude::*;
+
+pub mod workflows;
+
+pub fn registry() -> WorkflowResult<Registry> {
+	use workflows::*;
+
+	let mut registry = Registry::new();
+	registry.register_workflow::<pruner::Workflow>()?;
+
+	Ok(registry)
+}

package/packages/gasoline-runtime/src/workflows/mod.rs

@@ -0,0 +1 @@
+pub mod pruner;

package/packages/gasoline-runtime/src/workflows/pruner.rs

@@ -0,0 +1,55 @@
+use std::sync::Arc;
+
+use anyhow::Result;
+use futures_util::FutureExt;
+use gas::{db::debug::DatabaseDebug, prelude::*};
+
+#[derive(Debug, Deserialize, Serialize)]
+pub struct Input {}
+
+#[workflow]
+pub async fn gasoline_pruner(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> {
+	ctx.repeat(|ctx| {
+		async move {
+			ctx.activity(PruneInput {}).await?;
+
+			ctx.sleep(ctx.config().runtime.gasoline_prune_interval_duration())
+				.await?;
+
+			Ok(Loop::<()>::Continue)
+		}
+		.boxed()
+	})
+	.await?;
+
+	Ok(())
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, Hash)]
+struct PruneInput {}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct PruneOutput {
+	prune_count: usize,
+}
+
+#[activity(Prune)]
+async fn prune(ctx: &ActivityCtx, _input: &PruneInput) -> Result<PruneOutput> {
+	// Create new db instance with debug trait
+	let db = db::DatabaseKv::new(ctx.config().clone(), ctx.pools().clone()).await?
+		as Arc<dyn DatabaseDebug + Send + Sync>;
+
+	// Check if pruning is enabled
+	let Some(prune_eligibility_duration) =
+		ctx.config().runtime.gasoline_prune_eligibility_duration()
+	else {
+		return Ok(PruneOutput { prune_count: 0 });
+	};
+
+	let before_ts = util::timestamp::now() - prune_eligibility_duration.as_millis() as i64;
+	let prune_count = db.prune_workflows(before_ts).await?;
+
+	tracing::debug!(%prune_count, "pruned workflows");
+
+	Ok(PruneOutput { prune_count })
+}