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

package/packages/gasoline/src/db/kv/system.rs

@@ -1,33 +1,170 @@
-use std::time::Instant;
+use std::fs;
+use std::sync::{Arc, OnceLock};
+use std::time::{Duration, Instant};
+use tokio::sync::Mutex;
 
-use sysinfo::{CpuRefreshKind, MINIMUM_CPU_UPDATE_INTERVAL, RefreshKind, System};
+use sysinfo::{CpuRefreshKind, Pid, ProcessRefreshKind, ProcessesToUpdate, RefreshKind, System};
+
+const CPU_UPDATE_INTERVAL: Duration = Duration::from_millis(150);
+
+static SYSTEM_INFO_CELL: OnceLock<Arc<Mutex<SystemInfo>>> = OnceLock::new();
 
 pub struct SystemInfo {
 	system: System,
+	pid: Pid,
 	last_cpu_usage_read: Instant,
+	cgroup_cpu_max: Option<CgroupCpuMax>,
+	last_cgroup_usage_usec: Option<u64>,
+	last_cgroup_cores: f32,
 }
 
 impl SystemInfo {
-	pub fn new() -> Self {
-		SystemInfo {
-			system: System::new_with_specifics(
-				RefreshKind::nothing().with_cpu(CpuRefreshKind::nothing().with_cpu_usage()),
-			),
-			last_cpu_usage_read: Instant::now(),
+	pub fn get() -> Arc<Mutex<Self>> {
+		SYSTEM_INFO_CELL
+			.get_or_init(|| {
+				Arc::new(Mutex::new(SystemInfo {
+					system: System::new_with_specifics(
+						RefreshKind::nothing()
+							.with_cpu(CpuRefreshKind::nothing().with_cpu_usage())
+							.with_processes(ProcessRefreshKind::nothing().with_cpu()),
+					),
+					pid: Pid::from_u32(std::process::id()),
+					last_cpu_usage_read: Instant::now(),
+					cgroup_cpu_max: CgroupCpuMax::read(),
+					last_cgroup_usage_usec: CgroupCpuUsage::read(),
+					last_cgroup_cores: 0.0,
+				}))
+			})
+			.clone()
+	}
+
+	/// Returns a float 0.0-1.0 of the avg cpu usage in the current container (if cgroups are configured) or
+	/// otherwise for the current process.
+	pub fn cpu_usage_ratio(&mut self, cpu_max: Option<usize>) -> f32 {
+		// 1 = 1 core
+		let cpu_max = if let Some(cpu_max) = cpu_max {
+			cpu_max as f32 / 1000.0
+		} else {
+			if let Some(CgroupCpuMax { quota, period }) = self.cgroup_cpu_max {
+				if quota > 0 {
+					quota as f32 / period as f32
+				} else {
+					// Negative quota means unlimited, use cpu count
+					self.system.cpus().len() as f32
+				}
+			} else {
+				self.system.cpus().len() as f32
+			}
+		};
+
+		let total = if let Some(last_usage_usec) = self.last_cgroup_usage_usec {
+			// Use cgroup cpu.stat for usage (cumulative counter)
+			if self.last_cpu_usage_read.elapsed() > CPU_UPDATE_INTERVAL {
+				if let Some(current_usage_usec) = CgroupCpuUsage::read() {
+					let elapsed_usec = self.last_cpu_usage_read.elapsed().as_micros() as u64;
+					let usage_delta_usec = current_usage_usec.saturating_sub(last_usage_usec);
+
+					// Calculate cores used: (usage_delta / elapsed_time)
+					let cores_used = if elapsed_usec > 0 {
+						usage_delta_usec as f32 / elapsed_usec as f32
+					} else {
+						0.0
+					};
+
+					self.last_cgroup_usage_usec = Some(current_usage_usec);
+					self.last_cpu_usage_read = Instant::now();
+					self.last_cgroup_cores = cores_used;
+
+					cores_used
+				} else {
+					// Failed to read cgroup, disable cgroup usage tracking
+					self.last_cgroup_usage_usec = None;
+					0.0
+				}
+			} else {
+				// Not time to update yet, return last calculated value
+				self.last_cgroup_cores
+			}
+		} else {
+			// Use per-process CPU metrics
+			if self.last_cpu_usage_read.elapsed() > CPU_UPDATE_INTERVAL {
+				self.system.refresh_processes_specifics(
+					ProcessesToUpdate::Some(&[self.pid]),
+					true,
+					ProcessRefreshKind::nothing().with_cpu(),
+				);
+				self.last_cpu_usage_read = Instant::now();
+			}
+
+			// Get CPU usage for current process (returns percentage 0-100 per core)
+			self.system
+				.process(self.pid)
+				.map(|p| p.cpu_usage() / 100.0)
+				.unwrap_or(0.0)
+		};
+
+		crate::metrics::CPU_USAGE.observe(total as f64);
+
+		total / cpu_max
+	}
+}
+
+struct CgroupCpuMax {
+	quota: i64,
+	period: u64,
+}
+
+impl CgroupCpuMax {
+	fn read() -> Option<Self> {
+		// cgroups v2
+		if let Ok(content) = fs::read_to_string("/sys/fs/cgroup/cpu.max") {
+			let parts = content.trim().split_whitespace().collect::<Vec<&str>>();
+			if parts.len() == 2 {
+				return Some(CgroupCpuMax {
+					quota: parts[0].parse::<i64>().ok()?,
+					period: parts[1].parse::<u64>().ok()?,
+				});
+			}
 		}
+
+		// cgroups v1
+		let quota = fs::read_to_string("/sys/fs/cgroup/cpu/cpu.cfs_quota_us")
+			.ok()?
+			.trim()
+			.parse()
+			.ok()?;
+		let period = fs::read_to_string("/sys/fs/cgroup/cpu/cpu.cfs_period_us")
+			.ok()?
+			.trim()
+			.parse()
+			.ok()?;
+
+		Some(CgroupCpuMax { quota, period })
 	}
+}
+
+struct CgroupCpuUsage;
+
+impl CgroupCpuUsage {
+	/// Reads CPU usage from cgroup cpu.stat
+	/// Returns usage in microseconds
+	fn read() -> Option<u64> {
+		// cgroups v2
+		if let Ok(content) = fs::read_to_string("/sys/fs/cgroup/cpu.stat") {
+			for line in content.lines() {
+				if let Some(usage) = line.strip_prefix("usage_usec ") {
+					return usage.trim().parse().ok();
+				}
+			}
+		}
 
-	/// Returns a float 0.0-100.0 of the avg cpu usage over the entire system.
-	pub fn cpu_usage(&mut self) -> f32 {
-		if self.last_cpu_usage_read.elapsed() > MINIMUM_CPU_UPDATE_INTERVAL {
-			self.system.refresh_cpu_usage();
-			self.last_cpu_usage_read = Instant::now();
+		// cgroups v1
+		if let Ok(content) = fs::read_to_string("/sys/fs/cgroup/cpuacct/cpuacct.usage") {
+			// cpuacct.usage is in nanoseconds, convert to microseconds
+			let usage_nsec: u64 = content.trim().parse().ok()?;
+			return Some(usage_nsec / 1000);
 		}
 
-		self.system
-			.cpus()
-			.iter()
-			.fold(0.0, |s, cpu| s + cpu.cpu_usage())
-			/ self.system.cpus().len() as f32
+		None
 	}
 }

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

@@ -11,7 +11,7 @@ use crate::{
 		event::{Event, EventType, SleepState},
 		location::Location,
 	},
-	workflow::Workflow,
+	workflow::{PruneVariant, Workflow},
 };
 
 pub mod debug;
@@ -24,7 +24,7 @@ pub type DatabaseHandle = Arc<dyn Database + Sync>;
 #[async_trait::async_trait]
 pub trait Database: Send {
 	/// Create a new DB instance.
-	async fn from_pools(pools: rivet_pools::Pools) -> Result<Arc<Self>>
+	async fn new(config: rivet_config::Config, pools: rivet_pools::Pools) -> Result<Arc<Self>>
 	where
 		Self: Sized;
 
@@ -65,7 +65,11 @@ pub trait Database: Send {
 	) -> WorkflowResult<BoxStream<'b, ()>>;
 
 	/// Updates the last ping ts for this worker.
-	async fn update_worker_ping(&self, worker_id: Id) -> WorkflowResult<()>;
+	async fn update_worker_ping(
+		&self,
+		worker_id: Id,
+		update_active_idx: bool,
+	) -> WorkflowResult<()>;
 
 	/// Removes the worker from consideration for `pull_workflows` delegation.
 	async fn mark_worker_inactive(&self, worker_id: Id) -> WorkflowResult<()>;
@@ -101,6 +105,12 @@ pub trait Database: Send {
 		tags: &serde_json::Value,
 	) -> WorkflowResult<Option<Id>>;
 
+	/// Retrieves the first incomplete workflow for each (name, tags) pair in a single batch transaction.
+	async fn find_workflows(
+		&self,
+		queries: &[(&str, serde_json::Value)],
+	) -> WorkflowResult<Vec<Option<Id>>>;
+
 	/// Pulls workflows for processing by the worker. Will only pull workflows with names matching the filter.
 	/// Should also update the ping of this worker.
 	async fn pull_workflows(
@@ -115,6 +125,7 @@ pub trait Database: Send {
 		workflow_id: Id,
 		workflow_name: &str,
 		output: &serde_json::value::RawValue,
+		prune_variant: PruneVariant,
 	) -> WorkflowResult<()>;
 
 	/// Write a workflow sleep/failure to the database.
@@ -129,8 +140,8 @@ pub trait Database: Send {
 		error: &str,
 	) -> WorkflowResult<()>;
 
-	/// Pulls the oldest signal with the given filter.
-	async fn pull_next_signal(
+	/// Pulls signals in order from oldest to newest with the given filter.
+	async fn pull_next_signals(
 		&self,
 		workflow_id: Id,
 		workflow_name: &str,
@@ -138,8 +149,9 @@ pub trait Database: Send {
 		location: &Location,
 		version: usize,
 		loop_location: Option<&Location>,
-		last_try: bool,
-	) -> WorkflowResult<Option<SignalData>>;
+		limit: usize,
+		last_attempt: bool,
+	) -> WorkflowResult<Vec<SignalData>>;
 
 	/// Retrieves a workflow with the given ID. Can only be called from a workflow context.
 	async fn get_sub_workflow(
@@ -286,6 +298,7 @@ pub trait Database: Send {
 		from_workflow_id: Id,
 		location: &Location,
 		version: usize,
+		inner_version: usize,
 		loop_location: Option<&Location>,
 	) -> WorkflowResult<()>;
 }
@@ -293,6 +306,7 @@ pub trait Database: Send {
 #[derive(Debug)]
 pub struct WorkflowData {
 	pub workflow_id: Id,
+	pub name: String,
 	input: Box<serde_json::value::RawValue>,
 	state: Box<serde_json::value::RawValue>,
 	output: Option<Box<serde_json::value::RawValue>>,
@@ -315,6 +329,14 @@ impl WorkflowData {
 			.transpose()
 			.map_err(WorkflowError::DeserializeWorkflowOutput)
 	}
+
+	pub fn is_complete(&self) -> bool {
+		self.output.is_some()
+	}
+
+	pub fn is_dead(&self) -> bool {
+		!self.is_complete() && !self.has_wake_condition
+	}
 }
 
 #[derive(Debug)]

package/packages/gasoline/src/error.rs

@@ -9,18 +9,18 @@ pub type WorkflowResult<T> = Result<T, WorkflowError>;
 
 #[derive(thiserror::Error, Debug)]
 pub enum WorkflowError {
-	#[error("workflow failure: {0:?}")]
-	WorkflowFailure(#[source] anyhow::Error),
+	#[error("workflow {0} failed: {1:?}")]
+	WorkflowFailure(&'static str, #[source] anyhow::Error),
 
 	// Includes error count
-	#[error("activity failure: {0:?}")]
-	ActivityFailure(#[source] anyhow::Error, usize),
+	#[error("activity {0} failed: {1:?}")]
+	ActivityFailure(&'static str, #[source] anyhow::Error, usize),
 
-	#[error("activity failure, max retries reached: {0:?}")]
-	ActivityMaxFailuresReached(#[source] anyhow::Error),
+	#[error("activity {0} failed, max retries reached: {1:?}")]
+	ActivityMaxFailuresReached(&'static str, #[source] anyhow::Error),
 
-	#[error("operation failure: {0:?}")]
-	OperationFailure(#[source] anyhow::Error),
+	#[error("operation {0} failed: {1:?}")]
+	OperationFailure(&'static str, #[source] anyhow::Error),
 
 	#[error("workflow missing from registry: {0}")]
 	WorkflowMissingFromRegistry(String),
@@ -112,6 +112,9 @@ pub enum WorkflowError {
 	#[error("failed to flush pubsub: {0}")]
 	FlushPubsub(#[source] anyhow::Error),
 
+	#[error("failed to publish message: {0}")]
+	PublishMessage(#[source] anyhow::Error),
+
 	#[error("subscription unsubscribed")]
 	SubscriptionUnsubscribed,
 
@@ -146,12 +149,12 @@ pub enum WorkflowError {
 	Config(#[source] anyhow::Error),
 
 	// Includes error count
-	#[error("activity timed out")]
-	ActivityTimeout(usize),
+	#[error("activity {0} timed out")]
+	ActivityTimeout(&'static str, usize),
 
 	// Includes error count
-	#[error("operation timed out")]
-	OperationTimeout(usize),
+	#[error("operation {0} timed out")]
+	OperationTimeout(&'static str, usize),
 
 	#[error("duplicate registered workflow: {0}")]
 	DuplicateRegisteredWorkflow(String),
@@ -186,9 +189,9 @@ impl WorkflowError {
 	/// Returns the next deadline for a workflow to be woken up again based on the error.
 	pub(crate) fn deadline_ts(&self) -> Option<i64> {
 		match self {
-			WorkflowError::ActivityFailure(_, error_count)
-			| WorkflowError::ActivityTimeout(error_count)
-			| WorkflowError::OperationTimeout(error_count) => {
+			WorkflowError::ActivityFailure(_, _, error_count)
+			| WorkflowError::ActivityTimeout(_, error_count)
+			| WorkflowError::OperationTimeout(_, error_count) => {
 				// NOTE: Max retry is handled in `WorkflowCtx::activity`
 				let mut backoff = rivet_util::backoff::Backoff::new_at(
 					8,
@@ -211,6 +214,8 @@ impl WorkflowError {
 				Some(deadline_ts)
 			}
 			WorkflowError::Sleep(ts) | WorkflowError::NoSignalFoundAndSleep(_, ts) => Some(*ts),
+			// Always retry udb errors, they're usually caused by ephemeral outages and can be retried
+			WorkflowError::Udb(_) => Some(rivet_util::timestamp::now() + 30_000),
 			_ => None,
 		}
 	}
@@ -218,9 +223,9 @@ impl WorkflowError {
 	/// Any error that the workflow can continue on with its execution from.
 	pub(crate) fn is_recoverable(&self) -> bool {
 		match self {
-			WorkflowError::ActivityFailure(_, _)
-			| WorkflowError::ActivityTimeout(_)
-			| WorkflowError::OperationTimeout(_)
+			WorkflowError::ActivityFailure(_, _, _)
+			| WorkflowError::ActivityTimeout(_, _)
+			| WorkflowError::OperationTimeout(_, _)
 			| WorkflowError::NoSignalFound(_)
 			| WorkflowError::NoSignalFoundAndSleep(_, _)
 			| WorkflowError::SubWorkflowIncomplete(_)
@@ -233,9 +238,9 @@ impl WorkflowError {
 	/// Any error that the workflow can try again on a fixed number of times. Only used for printing.
 	pub(crate) fn is_retryable(&self) -> bool {
 		match self {
-			WorkflowError::ActivityFailure(_, _)
-			| WorkflowError::ActivityTimeout(_)
-			| WorkflowError::OperationTimeout(_) => true,
+			WorkflowError::ActivityFailure(_, _, _)
+			| WorkflowError::ActivityTimeout(_, _)
+			| WorkflowError::OperationTimeout(_, _) => true,
 			_ => false,
 		}
 	}

package/packages/gasoline/src/executable.rs

@@ -71,7 +71,9 @@ impl<T: Executable> Executable for Option<T> {
 	}
 
 	fn shift_cursor(&self, ctx: &mut WorkflowCtx) -> WorkflowResult<()> {
-		ctx.cursor_mut().inc();
+		if self.is_some() {
+			ctx.cursor_mut().inc();
+		}
 		Ok(())
 	}
 }