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

package/packages/gasoline/src/db/kv/keys/metric.rs

@@ -2,57 +2,68 @@ use anyhow::*;
 use universaldb::prelude::*;
 
 #[derive(Debug, PartialEq, Eq)]
-pub enum GaugeMetric {
+pub enum Metric {
+	// Count (workflow name)
 	WorkflowActive(String),
+	/// Count (workflow name)
 	WorkflowSleeping(String),
+	/// Count (workflow name, error)
 	WorkflowDead(String, String),
+	/// Count (workflow name)
 	WorkflowComplete(String),
+	/// Deprecated
 	SignalPending(String),
+	/// Count (signal name)
+	SignalPending2(String),
 }
 
-impl GaugeMetric {
-	fn variant(&self) -> GaugeMetricVariant {
+impl Metric {
+	fn variant(&self) -> MetricVariant {
 		match self {
-			GaugeMetric::WorkflowActive(_) => GaugeMetricVariant::WorkflowActive,
-			GaugeMetric::WorkflowSleeping(_) => GaugeMetricVariant::WorkflowSleeping,
-			GaugeMetric::WorkflowDead(_, _) => GaugeMetricVariant::WorkflowDead,
-			GaugeMetric::WorkflowComplete(_) => GaugeMetricVariant::WorkflowComplete,
-			GaugeMetric::SignalPending(_) => GaugeMetricVariant::SignalPending,
+			Metric::WorkflowActive(_) => MetricVariant::WorkflowActive,
+			Metric::WorkflowSleeping(_) => MetricVariant::WorkflowSleeping,
+			Metric::WorkflowDead(_, _) => MetricVariant::WorkflowDead,
+			Metric::WorkflowComplete(_) => MetricVariant::WorkflowComplete,
+			Metric::SignalPending(_) => MetricVariant::SignalPending,
+			Metric::SignalPending2(_) => MetricVariant::SignalPending2,
 		}
 	}
 }
 
 #[derive(strum::FromRepr)]
-enum GaugeMetricVariant {
+enum MetricVariant {
 	WorkflowActive = 0,
 	WorkflowSleeping = 1,
 	WorkflowDead = 2,
 	WorkflowComplete = 3,
+	// Deprecated
 	SignalPending = 4,
+	SignalPending2 = 5,
 }
 
+/// Stores gauge metrics for global database usage.
 #[derive(Debug)]
-pub struct GaugeMetricKey {
-	pub metric: GaugeMetric,
+pub struct MetricKey {
+	pub metric: Metric,
 }
 
-impl GaugeMetricKey {
-	pub fn new(metric: GaugeMetric) -> Self {
-		GaugeMetricKey { metric }
+impl MetricKey {
+	pub fn new(metric: Metric) -> Self {
+		MetricKey { metric }
 	}
 
-	pub fn subspace() -> GaugeMetricSubspaceKey {
-		GaugeMetricSubspaceKey::new()
+	pub fn subspace() -> MetricSubspaceKey {
+		MetricSubspaceKey::new()
 	}
 }
 
-impl FormalKey for GaugeMetricKey {
+impl FormalKey for MetricKey {
 	// IMPORTANT: Uses LE bytes, not BE
 	/// Count.
-	type Value = usize;
+	type Value = i64;
 
 	fn deserialize(&self, raw: &[u8]) -> Result<Self::Value> {
-		Ok(usize::from_le_bytes(raw.try_into()?))
+		Ok(i64::from_le_bytes(raw.try_into()?))
 	}
 
 	fn serialize(&self, value: Self::Value) -> Result<Vec<u8>> {
@@ -60,7 +71,7 @@ impl FormalKey for GaugeMetricKey {
 	}
 }
 
-impl TuplePack for GaugeMetricKey {
+impl TuplePack for MetricKey {
 	fn pack<W: std::io::Write>(
 		&self,
 		w: &mut W,
@@ -68,79 +79,90 @@ impl TuplePack for GaugeMetricKey {
 	) -> std::io::Result<VersionstampOffset> {
 		let mut offset = VersionstampOffset::None { size: 0 };
 
-		let t = (METRIC, self.metric.variant() as usize);
+		let t = (METRIC, self.metric.variant() as i64);
 		offset += t.pack(w, tuple_depth)?;
 
 		offset += match &self.metric {
-			GaugeMetric::WorkflowActive(workflow_name) => workflow_name.pack(w, tuple_depth)?,
-			GaugeMetric::WorkflowSleeping(workflow_name) => workflow_name.pack(w, tuple_depth)?,
-			GaugeMetric::WorkflowDead(workflow_name, error) => {
+			Metric::WorkflowActive(workflow_name) => workflow_name.pack(w, tuple_depth)?,
+			Metric::WorkflowSleeping(workflow_name) => workflow_name.pack(w, tuple_depth)?,
+			Metric::WorkflowDead(workflow_name, error) => {
 				(workflow_name, error).pack(w, tuple_depth)?
 			}
-			GaugeMetric::WorkflowComplete(workflow_name) => workflow_name.pack(w, tuple_depth)?,
-			GaugeMetric::SignalPending(signal_name) => signal_name.pack(w, tuple_depth)?,
+			Metric::WorkflowComplete(workflow_name) => workflow_name.pack(w, tuple_depth)?,
+			Metric::SignalPending(signal_name) => signal_name.pack(w, tuple_depth)?,
+			Metric::SignalPending2(signal_name) => signal_name.pack(w, tuple_depth)?,
 		};
 
 		std::result::Result::Ok(offset)
 	}
 }
 
-impl<'de> TupleUnpack<'de> for GaugeMetricKey {
+impl<'de> TupleUnpack<'de> for MetricKey {
 	fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> {
 		let (input, (_, variant)) = <(usize, usize)>::unpack(input, tuple_depth)?;
-		let variant = GaugeMetricVariant::from_repr(variant).ok_or_else(|| {
+		let variant = MetricVariant::from_repr(variant).ok_or_else(|| {
 			PackError::Message(format!("invalid metric variant `{variant}` in key").into())
 		})?;
 
 		let (input, v) = match variant {
-			GaugeMetricVariant::WorkflowActive => {
+			MetricVariant::WorkflowActive => {
 				let (input, workflow_name) = String::unpack(input, tuple_depth)?;
 
 				(
 					input,
-					GaugeMetricKey {
-						metric: GaugeMetric::WorkflowActive(workflow_name),
+					MetricKey {
+						metric: Metric::WorkflowActive(workflow_name),
 					},
 				)
 			}
-			GaugeMetricVariant::WorkflowSleeping => {
+			MetricVariant::WorkflowSleeping => {
 				let (input, workflow_name) = String::unpack(input, tuple_depth)?;
 
 				(
 					input,
-					GaugeMetricKey {
-						metric: GaugeMetric::WorkflowSleeping(workflow_name),
+					MetricKey {
+						metric: Metric::WorkflowSleeping(workflow_name),
 					},
 				)
 			}
-			GaugeMetricVariant::WorkflowDead => {
+			MetricVariant::WorkflowDead => {
 				let (input, (workflow_name, error)) =
 					<(String, String)>::unpack(input, tuple_depth)?;
 
 				(
 					input,
-					GaugeMetricKey {
-						metric: GaugeMetric::WorkflowDead(workflow_name, error),
+					MetricKey {
+						metric: Metric::WorkflowDead(workflow_name, error),
 					},
 				)
 			}
-			GaugeMetricVariant::WorkflowComplete => {
+			MetricVariant::WorkflowComplete => {
 				let (input, workflow_name) = String::unpack(input, tuple_depth)?;
 
 				(
 					input,
-					GaugeMetricKey {
-						metric: GaugeMetric::WorkflowComplete(workflow_name),
+					MetricKey {
+						metric: Metric::WorkflowComplete(workflow_name),
 					},
 				)
 			}
-			GaugeMetricVariant::SignalPending => {
+			MetricVariant::SignalPending => {
 				let (input, signal_name) = String::unpack(input, tuple_depth)?;
 
 				(
 					input,
-					GaugeMetricKey {
-						metric: GaugeMetric::SignalPending(signal_name),
+					MetricKey {
+						metric: Metric::SignalPending(signal_name),
+					},
+				)
+			}
+			MetricVariant::SignalPending2 => {
+				let (input, signal_name) = String::unpack(input, tuple_depth)?;
+
+				(
+					input,
+					MetricKey {
+						metric: Metric::SignalPending2(signal_name),
 					},
 				)
 			}
@@ -150,15 +172,16 @@ impl<'de> TupleUnpack<'de> for GaugeMetricKey {
 	}
 }
 
-pub struct GaugeMetricSubspaceKey {}
+/// Used to list all global gauge metrics.
+pub struct MetricSubspaceKey {}
 
-impl GaugeMetricSubspaceKey {
+impl MetricSubspaceKey {
 	pub fn new() -> Self {
-		GaugeMetricSubspaceKey {}
+		MetricSubspaceKey {}
 	}
 }
 
-impl TuplePack for GaugeMetricSubspaceKey {
+impl TuplePack for MetricSubspaceKey {
 	fn pack<W: std::io::Write>(
 		&self,
 		w: &mut W,

package/packages/gasoline/src/db/kv/keys/signal.rs

@@ -332,11 +332,19 @@ impl<'de> TupleUnpack<'de> for WorkflowIdKey {
 	}
 }
 
-pub struct DataSubspaceKey {}
+pub struct DataSubspaceKey {
+	signal_id: Option<Id>,
+}
 
 impl DataSubspaceKey {
 	pub fn new() -> Self {
-		DataSubspaceKey {}
+		DataSubspaceKey { signal_id: None }
+	}
+
+	pub fn new_with_signal_id(signal_id: Id) -> Self {
+		DataSubspaceKey {
+			signal_id: Some(signal_id),
+		}
 	}
 }
 
@@ -346,8 +354,16 @@ impl TuplePack for DataSubspaceKey {
 		w: &mut W,
 		tuple_depth: TupleDepth,
 	) -> std::io::Result<VersionstampOffset> {
+		let mut offset = VersionstampOffset::None { size: 0 };
+
 		let t = (SIGNAL, DATA);
-		t.pack(w, tuple_depth)
+		offset += t.pack(w, tuple_depth)?;
+
+		if let Some(signal_id) = &self.signal_id {
+			offset += signal_id.pack(w, tuple_depth)?;
+		}
+
+		Ok(offset)
 	}
 }
 

package/packages/gasoline/src/db/kv/keys/workflow.rs

@@ -4,6 +4,8 @@ use anyhow::*;
 use rivet_util::Id;
 use universaldb::prelude::*;
 
+use crate::workflow::PruneVariant;
+
 #[derive(Debug)]
 pub struct LeaseKey {
 	pub workflow_id: Id,
@@ -1106,11 +1108,19 @@ impl<'de> TupleUnpack<'de> for WorkerIdKey {
 	}
 }
 
-pub struct DataSubspaceKey {}
+pub struct DataSubspaceKey {
+	workflow_id: Option<Id>,
+}
 
 impl DataSubspaceKey {
 	pub fn new() -> Self {
-		DataSubspaceKey {}
+		DataSubspaceKey { workflow_id: None }
+	}
+
+	pub fn new_with_workflow_id(workflow_id: Id) -> Self {
+		DataSubspaceKey {
+			workflow_id: Some(workflow_id),
+		}
 	}
 }
 
@@ -1120,8 +1130,16 @@ impl TuplePack for DataSubspaceKey {
 		w: &mut W,
 		tuple_depth: TupleDepth,
 	) -> std::io::Result<VersionstampOffset> {
+		let mut offset = VersionstampOffset::None { size: 0 };
+
 		let t = (WORKFLOW, DATA);
-		t.pack(w, tuple_depth)
+		offset += t.pack(w, tuple_depth)?;
+
+		if let Some(workflow_id) = &self.workflow_id {
+			offset += workflow_id.pack(w, tuple_depth)?;
+		}
+
+		Ok(offset)
 	}
 }
 
@@ -1173,3 +1191,331 @@ impl<'de> TupleUnpack<'de> for SilenceTsKey {
 		Ok((input, v))
 	}
 }
+
+#[derive(Debug, PartialEq, Eq)]
+pub enum Metric {
+	/// Count (signal name)
+	SignalPending(String),
+}
+
+impl Metric {
+	fn variant(&self) -> MetricVariant {
+		match self {
+			Metric::SignalPending(_) => MetricVariant::SignalPending,
+		}
+	}
+}
+
+#[derive(strum::FromRepr)]
+enum MetricVariant {
+	SignalPending = 0,
+}
+
+/// Stores gauge metrics for a single workflow.
+#[derive(Debug)]
+pub struct MetricKey {
+	pub workflow_id: Id,
+	pub metric: Metric,
+}
+
+impl MetricKey {
+	pub fn new(workflow_id: Id, metric: Metric) -> Self {
+		MetricKey {
+			workflow_id,
+			metric,
+		}
+	}
+
+	pub fn subspace(workflow_id: Id) -> MetricSubspaceKey {
+		MetricSubspaceKey::new(workflow_id)
+	}
+}
+
+impl FormalKey for MetricKey {
+	// IMPORTANT: Uses LE bytes, not BE
+	/// Count.
+	type Value = i64;
+
+	fn deserialize(&self, raw: &[u8]) -> Result<Self::Value> {
+		Ok(i64::from_le_bytes(raw.try_into()?))
+	}
+
+	fn serialize(&self, value: Self::Value) -> Result<Vec<u8>> {
+		Ok(value.to_le_bytes().to_vec())
+	}
+}
+
+impl TuplePack for MetricKey {
+	fn pack<W: std::io::Write>(
+		&self,
+		w: &mut W,
+		tuple_depth: TupleDepth,
+	) -> std::io::Result<VersionstampOffset> {
+		let mut offset = VersionstampOffset::None { size: 0 };
+
+		let t = (
+			WORKFLOW,
+			METRIC,
+			self.workflow_id,
+			self.metric.variant() as usize,
+		);
+		offset += t.pack(w, tuple_depth)?;
+
+		offset += match &self.metric {
+			Metric::SignalPending(signal_name) => signal_name.pack(w, tuple_depth)?,
+		};
+
+		Ok(offset)
+	}
+}
+
+impl<'de> TupleUnpack<'de> for MetricKey {
+	fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> {
+		let (input, (_, _, workflow_id, variant)) =
+			<(usize, usize, Id, usize)>::unpack(input, tuple_depth)?;
+		let variant = MetricVariant::from_repr(variant).ok_or_else(|| {
+			PackError::Message(format!("invalid metric variant `{variant}` in key").into())
+		})?;
+
+		let (input, v) = match variant {
+			MetricVariant::SignalPending => {
+				let (input, signal_name) = String::unpack(input, tuple_depth)?;
+
+				(
+					input,
+					MetricKey {
+						workflow_id,
+						metric: Metric::SignalPending(signal_name),
+					},
+				)
+			}
+		};
+
+		Ok((input, v))
+	}
+}
+
+/// Used to list all gauge metrics for a workflow.
+pub struct MetricSubspaceKey {
+	workflow_id: Id,
+}
+
+impl MetricSubspaceKey {
+	pub fn new(workflow_id: Id) -> Self {
+		MetricSubspaceKey { workflow_id }
+	}
+}
+
+impl TuplePack for MetricSubspaceKey {
+	fn pack<W: std::io::Write>(
+		&self,
+		w: &mut W,
+		tuple_depth: TupleDepth,
+	) -> std::io::Result<VersionstampOffset> {
+		let t = (WORKFLOW, METRIC, self.workflow_id);
+		t.pack(w, tuple_depth)
+	}
+}
+
+#[derive(Debug)]
+pub struct CompleteTsKey {
+	workflow_id: Id,
+}
+
+impl CompleteTsKey {
+	pub fn new(workflow_id: Id) -> Self {
+		CompleteTsKey { workflow_id }
+	}
+}
+
+impl FormalKey for CompleteTsKey {
+	// Timestamp.
+	type Value = i64;
+
+	fn deserialize(&self, raw: &[u8]) -> Result<Self::Value> {
+		Ok(i64::from_be_bytes(raw.try_into()?))
+	}
+
+	fn serialize(&self, value: Self::Value) -> Result<Vec<u8>> {
+		Ok(value.to_be_bytes().to_vec())
+	}
+}
+
+impl TuplePack for CompleteTsKey {
+	fn pack<W: std::io::Write>(
+		&self,
+		w: &mut W,
+		tuple_depth: TupleDepth,
+	) -> std::io::Result<VersionstampOffset> {
+		let t = (WORKFLOW, DATA, self.workflow_id, COMPLETE_TS);
+		t.pack(w, tuple_depth)
+	}
+}
+
+impl<'de> TupleUnpack<'de> for CompleteTsKey {
+	fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> {
+		let (input, (_, _, workflow_id, data)) =
+			<(usize, usize, Id, usize)>::unpack(input, tuple_depth)?;
+		if data != COMPLETE_TS {
+			return Err(PackError::Message("expected COMPLETE_TS data".into()));
+		}
+
+		let v = CompleteTsKey { workflow_id };
+
+		Ok((input, v))
+	}
+}
+
+#[derive(Debug)]
+pub struct PruneTsKey {
+	workflow_id: Id,
+}
+
+impl PruneTsKey {
+	pub fn new(workflow_id: Id) -> Self {
+		PruneTsKey { workflow_id }
+	}
+}
+
+impl FormalKey for PruneTsKey {
+	// Timestamp.
+	type Value = i64;
+
+	fn deserialize(&self, raw: &[u8]) -> Result<Self::Value> {
+		Ok(i64::from_be_bytes(raw.try_into()?))
+	}
+
+	fn serialize(&self, value: Self::Value) -> Result<Vec<u8>> {
+		Ok(value.to_be_bytes().to_vec())
+	}
+}
+
+impl TuplePack for PruneTsKey {
+	fn pack<W: std::io::Write>(
+		&self,
+		w: &mut W,
+		tuple_depth: TupleDepth,
+	) -> std::io::Result<VersionstampOffset> {
+		let t = (WORKFLOW, DATA, self.workflow_id, PRUNE_TS);
+		t.pack(w, tuple_depth)
+	}
+}
+
+impl<'de> TupleUnpack<'de> for PruneTsKey {
+	fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> {
+		let (input, (_, _, workflow_id, data)) =
+			<(usize, usize, Id, usize)>::unpack(input, tuple_depth)?;
+		if data != PRUNE_TS {
+			return Err(PackError::Message("expected PRUNE_TS data".into()));
+		}
+
+		let v = PruneTsKey { workflow_id };
+
+		Ok((input, v))
+	}
+}
+
+#[derive(Debug)]
+pub struct PruneIdxKey {
+	pub ts: i64,
+	pub workflow_id: Id,
+	pub variant: PruneVariant,
+}
+
+impl PruneIdxKey {
+	pub fn new(workflow_id: Id, variant: PruneVariant) -> Self {
+		PruneIdxKey {
+			ts: rivet_util::timestamp::now(),
+			workflow_id,
+			variant,
+		}
+	}
+
+	pub fn subspace(ts: i64) -> PruneIdxSubspaceKey {
+		PruneIdxSubspaceKey::new(ts)
+	}
+
+	pub fn entire_subspace() -> PruneIdxSubspaceKey {
+		PruneIdxSubspaceKey::entire()
+	}
+}
+
+impl FormalKey for PruneIdxKey {
+	type Value = ();
+
+	fn deserialize(&self, _raw: &[u8]) -> Result<Self::Value> {
+		Ok(())
+	}
+
+	fn serialize(&self, _value: Self::Value) -> Result<Vec<u8>> {
+		Ok(Vec::new())
+	}
+}
+
+impl TuplePack for PruneIdxKey {
+	fn pack<W: std::io::Write>(
+		&self,
+		w: &mut W,
+		tuple_depth: TupleDepth,
+	) -> std::io::Result<VersionstampOffset> {
+		let t = (
+			WORKFLOW,
+			PRUNE_IDX,
+			self.ts,
+			self.workflow_id,
+			self.variant as usize,
+		);
+		t.pack(w, tuple_depth)
+	}
+}
+
+impl<'de> TupleUnpack<'de> for PruneIdxKey {
+	fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> {
+		let (input, (_, _, ts, workflow_id, prune_variant)) =
+			<(usize, usize, i64, Id, usize)>::unpack(input, tuple_depth)?;
+		let prune_variant = PruneVariant::from_repr(prune_variant).ok_or_else(|| {
+			PackError::Message(format!("invalid prune variant `{prune_variant}` in key").into())
+		})?;
+
+		let v = PruneIdxKey {
+			ts,
+			workflow_id,
+			variant: prune_variant,
+		};
+
+		Ok((input, v))
+	}
+}
+
+pub struct PruneIdxSubspaceKey {
+	ts: Option<i64>,
+}
+
+impl PruneIdxSubspaceKey {
+	pub fn new(ts: i64) -> Self {
+		PruneIdxSubspaceKey { ts: Some(ts) }
+	}
+
+	pub fn entire() -> Self {
+		PruneIdxSubspaceKey { ts: None }
+	}
+}
+
+impl TuplePack for PruneIdxSubspaceKey {
+	fn pack<W: std::io::Write>(
+		&self,
+		w: &mut W,
+		tuple_depth: TupleDepth,
+	) -> std::io::Result<VersionstampOffset> {
+		let mut offset = VersionstampOffset::None { size: 0 };
+
+		let t = (WORKFLOW, PRUNE_IDX);
+		offset += t.pack(w, tuple_depth)?;
+
+		if let Some(ts) = &self.ts {
+			offset += ts.pack(w, tuple_depth)?;
+		}
+
+		Ok(offset)
+	}
+}