@rivetkit/engine 1.0.0

package/packages/universaldb/src/utils/keys.rs

@@ -0,0 +1,134 @@
+macro_rules! define_keys {
+	($(($value:literal, $const_name:ident, $str_name:literal)),* $(,)?) => {
+		$(
+			pub const $const_name: usize = $value;
+		)*
+
+		pub fn key_from_str(key: &str) -> Option<usize> {
+			match key {
+				$(
+					$str_name => Some($const_name),
+				)*
+				_ => None,
+			}
+		}
+
+		pub fn str_from_key(key: usize) -> Option<&'static str> {
+			match key {
+				$(
+					$const_name => Some($str_name),
+				)*
+				_ => None,
+			}
+		}
+	};
+}
+
+define_keys! {
+	(0, RIVET, "rivet"),
+	(1, GASOLINE, "gasoline"),
+	(2, KV, "kv"),
+	(3, PEGBOARD, "pegboard"),
+	(4, DATA, "data"),
+	(5, WORKFLOW, "workflow"),
+	(6, SIGNAL, "signal"),
+	(7, BODY, "body"),
+	(8, LEASE, "lease"),
+	(9, TAG, "tag"),
+	(10, INPUT, "input"),
+	(11, OUTPUT, "output"),
+	(12, WAKE_SIGNAL, "wake_signal"),
+	(13, WAKE_DEADLINE, "wake_deadline"),
+	(14, ACK_TS, "ack_ts"),
+	(15, CREATE_TS, "create_ts"),
+	(16, SILENCE_TS, "silence_ts"),
+	(17, PENDING, "pending"),
+	(18, RAY_ID, "ray_id"),
+	(19, NAME, "name"),
+	(20, WORKFLOW_ID, "workflow_id"),
+	(21, WAKE, "wake"),
+	(22, SUB_WORKFLOW, "sub_workflow"),
+	(23, WORKER, "worker"),
+	(24, LAST_PING_TS, "last_ping_ts"),
+	(25, METRICS_LOCK, "metrics_lock"),
+	(26, ERROR, "error"),
+	(27, WAKE_SUB_WORKFLOW_ID, "wake_sub_workflow_id"),
+	(28, BY_NAME_AND_TAG, "by_name_and_tag"),
+	(29, HAS_WAKE_CONDITION, "has_wake_condition"),
+	(30, WORKER_ID, "worker_id"),
+	(31, DBS, "dbs"),
+	(32, ACTOR, "actor"),
+	(33, BY_NAME, "by_name"),
+	(34, DEFAULT, "default"),
+	(35, REMAINING_MEMORY, "remaining_memory"),
+	(36, REMAINING_CPU, "remaining_cpu"),
+	(37, TOTAL_MEMORY, "total_memory"),
+	(38, TOTAL_CPU, "total_cpu"),
+	(39, NAMESPACE, "namespace"),
+	(40, HIBERNATING_REQUEST, "hibernating_request"),
+	(41, DISPLAY_NAME, "display_name"),
+	(42, CONNECTABLE, "connectable"),
+	(43, SLEEP_TS, "sleep_ts"),
+	(44, METADATA, "metadata"),
+	(45, COMPRESSED_DATA, "compressed_data"),
+	(46, RUNNER, "runner"),
+	(47, RUNNER_ALLOC_IDX, "runner_alloc_idx"),
+	(48, REMAINING_SLOTS, "remaining_slots"),
+	(49, TOTAL_SLOTS, "total_slots"),
+	(50, STATE, "state"),
+	(51, RUNNER_ID, "runner_id"),
+	(52, NAMESPACE_ID, "namespace_id"),
+	(53, PENDING_ACTOR_BY_RUNNER_NAME_SELECTOR, "pending_actor_by_runner_name_selector"),
+	(54, CONTAINER, "container"),
+	(55, HISTORY, "history"),
+	(56, ACTIVE, "active"),
+	(57, FORGOTTEN, "forgotten"),
+	(58, EVENT_TYPE, "event_type"),
+	(59, VERSION, "version"),
+	// 60
+	(61, SIGNAL_ID, "signal_id"),
+	(62, SUB_WORKFLOW_ID, "sub_workflow_id"),
+	(63, ITERATION, "iteration"),
+	(64, DEADLINE_TS, "deadline_ts"),
+	(65, SLEEP_STATE, "sleep_state"),
+	(66, INNER_EVENT_TYPE, "inner_event_type"),
+	(67, ALL, "all"),
+	(68, BY_NAME_AND_KEY, "by_name_and_key"),
+	(69, DESTROY_TS, "destroy_ts"),
+	(70, EPOXY, "epoxy"),
+	(71, PEER_STATE, "peer_state"),
+	(72, ACTOR_KV, "actor_kv"),
+	(73, STOP_TS, "stop_ts"),
+	(74, DRAIN_TS, "drain_ts"),
+	(75, LAST_RTT, "last_rtt"),
+	(76, CONNECTED_TS, "connected_ts"),
+	(77, EXPIRED_TS, "expired_ts"),
+	(78, KEY, "key"),
+	(79, VALUE, "value"),
+	(80, COMMITTED_VALUE, "committed_value"),
+	(81, OPTIMISTIC_CACHE, "optimistic_cache"),
+	(82, OPTIMISTIC_CACHED_VALUE, "optimistic_cached_value"),
+	(83, RESERVATION, "reservation"),
+	(84, INSTANCE, "instance"),
+	(85, LOG, "log"),
+	(86, ENTRY, "entry"),
+	(87, KEY_INSTANCE, "key_instance"),
+	(88, INSTANCE_NUMBER, "instance_number"),
+	(89, REPLICA, "replica"),
+	(90, CONFIG, "config"),
+	(91, METRIC, "metric"),
+	(92, CURRENT_BALLOT, "current_ballot"),
+	(93, INSTANCE_BALLOT, "instance_ballot"),
+	(94, SERVERLESS, "serverless"),
+	(95, DESIRED_SLOTS, "desired_slots"),
+	(96, BY_VARIANT, "by_variant"),
+	(97, ACL, "acl"),
+	(98, TOKEN, "token"),
+	(99, SECRET, "secret"),
+	// 100 - RESERVED BY EE
+	// 101 - RESERVED BY EE
+	// 102 - RESERVED BY EE
+	// 103 - RESERVED BY EE
+	// 104 - RESERVED BY EE
+	// 105 - RESERVED BY EE
+}

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

@@ -0,0 +1,92 @@
+use crate::tuple::{PackError, PackResult};
+
+mod cherry_pick;
+pub mod codes;
+mod ext;
+mod formal_key;
+pub mod keys;
+mod subspace;
+
+pub use cherry_pick::*;
+pub use ext::*;
+pub use formal_key::*;
+pub use subspace::Subspace;
+
+pub const CHUNK_SIZE: usize = 10_000; // 10 KB, not KiB, see https://apple.github.io/foundationdb/blob.html
+
+#[derive(Debug, Clone, Copy)]
+pub enum IsolationLevel {
+	Serializable,
+	Snapshot,
+}
+
+/// Indicates the transaction might have committed
+#[derive(Debug, Clone, Copy)]
+pub struct MaybeCommitted(pub bool);
+
+/// Calculate exponential backoff based on attempt.
+///
+/// Ours:
+///   0 -> 10ms + 0-1ms jitter
+///   1 -> 20ms + 0-2ms jitter
+///   2 -> 40ms + 0-4ms jitter
+///   ...
+///   7 (max) -> 1280ms + 0-128ms jitter
+/// FDB (see https://github.com/apple/foundationdb/blob/b1fbbd87a794b7c6c2f456925c45d8af339a8ae0/fdbclient/NativeAPI.actor.cpp#L4333 and https://github.com/apple/foundationdb/blob/b1fbbd87a794b7c6c2f456925c45d8af339a8ae0/fdbclient/ClientKnobs.cpp#L74-L76):
+///   0 -> 10ms
+///   1 -> 20ms
+///   2 -> 40ms
+///   ...
+///   X -> max 1s
+pub fn calculate_tx_retry_backoff(attempt: usize) -> u64 {
+	let base = 2_u64.pow((attempt as u32).min(7));
+	let base_backoff_ms = base * 10;
+
+	// Jitter is 0-10% of backoff ms
+	let jitter_ms = rand::random::<u64>() % base;
+
+	base_backoff_ms + jitter_ms
+}
+
+/// When using `add_conflict_range` to add a conflict for a single key, you cannot set both the start and end
+/// keys to the same key. Instead, the end key must be the start key + a 0 byte.
+/// See Python bindings: https://github.com/apple/foundationdb/blob/ec714791df4a6e4dafb5a926130d5789ce0c497a/bindings/python/fdb/impl.py#L633-L635
+pub fn end_of_key_range(key: &[u8]) -> Vec<u8> {
+	let mut end_key = Vec::with_capacity(key.len() + 1);
+	end_key.extend_from_slice(key);
+	end_key.push(0);
+	end_key
+}
+
+// Copied from foundationdb crate
+#[inline]
+pub fn parse_bytes(input: &[u8], num: usize) -> PackResult<(&[u8], &[u8])> {
+	if input.len() < num {
+		Err(PackError::MissingBytes)
+	} else {
+		Ok((&input[num..], &input[..num]))
+	}
+}
+
+// Copied from foundationdb crate
+#[inline]
+pub fn parse_byte(input: &[u8]) -> PackResult<(&[u8], u8)> {
+	if input.is_empty() {
+		Err(PackError::MissingBytes)
+	} else {
+		Ok((&input[1..], input[0]))
+	}
+}
+
+// Copied from foundationdb crate
+pub fn parse_code(input: &[u8], expected: u8) -> PackResult<&[u8]> {
+	let (input, found) = parse_byte(input)?;
+	if found == expected {
+		Ok(input)
+	} else {
+		Err(PackError::BadCode {
+			found,
+			expected: Some(expected),
+		})
+	}
+}

package/packages/universaldb/src/utils/subspace.rs

@@ -0,0 +1,95 @@
+use std::{borrow::Cow, ops::Deref};
+
+use crate::{
+	key_selector::KeySelector,
+	range_option::RangeOption,
+	tuple::{self, PackResult, TuplePack, TupleUnpack},
+};
+use rivet_metrics::KeyValue;
+
+use crate::metrics;
+
+/// Wrapper type around `universaldb::tuple::Subspace` that records metrics.
+#[derive(Clone)]
+pub struct Subspace {
+	inner: tuple::Subspace,
+}
+
+impl Subspace {
+	/// Creates a subspace with the given tuple.
+	pub fn new<T: TuplePack>(t: &T) -> Self {
+		Self {
+			inner: tuple::Subspace::all().subspace(t),
+		}
+	}
+
+	pub fn all() -> Self {
+		Self {
+			inner: tuple::Subspace::all(),
+		}
+	}
+
+	/// Returns a new Subspace whose prefix extends this Subspace with a given tuple encodable.
+	pub fn subspace<T: TuplePack>(&self, t: &T) -> Self {
+		Self {
+			inner: self.inner.subspace(t),
+		}
+	}
+
+	pub fn join(&self, s: &Subspace) -> Self {
+		Self {
+			inner: tuple::Subspace::from_bytes([self.inner.bytes(), s.bytes()].concat()),
+		}
+	}
+
+	/// Returns the key encoding the specified Tuple with the prefix of this Subspace
+	/// prepended.
+	pub fn pack<T: TuplePack>(&self, t: &T) -> Vec<u8> {
+		metrics::KEY_PACK_COUNT.add(1, &[KeyValue::new("type", std::any::type_name::<T>())]);
+
+		self.inner.pack(t)
+	}
+
+	/// Returns the key encoding the specified Tuple with the prefix of this Subspace
+	/// prepended, with a versionstamp.
+	pub fn pack_with_versionstamp<T: TuplePack>(&self, t: &T) -> Vec<u8> {
+		metrics::KEY_PACK_COUNT.add(1, &[KeyValue::new("type", std::any::type_name::<T>())]);
+
+		self.inner.pack_with_versionstamp(t)
+	}
+
+	/// `unpack` returns the Tuple encoded by the given key with the prefix of this Subspace
+	/// removed.  `unpack` will return an error if the key is not in this Subspace or does not
+	/// encode a well-formed Tuple.
+	pub fn unpack<'de, T: TupleUnpack<'de>>(&self, key: &'de [u8]) -> PackResult<T> {
+		metrics::KEY_UNPACK_COUNT.add(1, &[KeyValue::new("type", std::any::type_name::<T>())]);
+
+		self.inner.unpack(key)
+	}
+}
+
+impl Deref for Subspace {
+	type Target = tuple::Subspace;
+
+	fn deref(&self) -> &Self::Target {
+		&self.inner
+	}
+}
+
+impl From<tuple::Subspace> for Subspace {
+	fn from(value: tuple::Subspace) -> Self {
+		Subspace { inner: value }
+	}
+}
+
+impl<'a> From<&'a Subspace> for RangeOption<'static> {
+	fn from(subspace: &Subspace) -> Self {
+		let (begin, end) = subspace.range();
+
+		Self {
+			begin: KeySelector::first_greater_or_equal(Cow::Owned(begin)),
+			end: KeySelector::first_greater_or_equal(Cow::Owned(end)),
+			..Self::default()
+		}
+	}
+}

package/packages/universaldb/src/value.rs

@@ -0,0 +1,159 @@
+use std::{
+	ops::{Deref, DerefMut},
+	pin::Pin,
+};
+
+use anyhow::Result;
+
+#[derive(Debug, PartialEq, Eq)]
+pub struct Slice(Vec<u8>);
+
+impl Slice {
+	pub fn new() -> Self {
+		Slice(Vec::new())
+	}
+}
+
+impl Deref for Slice {
+	type Target = Vec<u8>;
+
+	fn deref(&self) -> &Self::Target {
+		&self.0
+	}
+}
+
+impl DerefMut for Slice {
+	fn deref_mut(&mut self) -> &mut Self::Target {
+		&mut self.0
+	}
+}
+
+impl From<Vec<u8>> for Slice {
+	fn from(value: Vec<u8>) -> Self {
+		Slice(value)
+	}
+}
+
+impl From<Slice> for Vec<u8> {
+	fn from(value: Slice) -> Self {
+		value.0
+	}
+}
+
+#[derive(Debug, Clone)]
+pub struct Value(KeyValue);
+
+impl Value {
+	pub fn new(key: Vec<u8>, value: Vec<u8>) -> Self {
+		Value(KeyValue::new(key, value))
+	}
+
+	pub fn from_keyvalue(kv: KeyValue) -> Self {
+		Value(kv)
+	}
+
+	pub fn key(&self) -> &[u8] {
+		self.0.key()
+	}
+
+	pub fn value(&self) -> &[u8] {
+		self.0.value()
+	}
+
+	pub fn into_parts(self) -> (Vec<u8>, Vec<u8>) {
+		self.0.into_parts()
+	}
+}
+
+// Values wraps a Vec<KeyValue> to match FoundationDB API
+#[derive(Debug, Clone)]
+pub struct Values {
+	values: Vec<KeyValue>,
+	more: bool,
+}
+
+impl Values {
+	pub fn new(values: Vec<KeyValue>) -> Self {
+		Values {
+			values,
+			more: false,
+		}
+	}
+
+	pub fn with_more(values: Vec<KeyValue>, more: bool) -> Self {
+		Values { values, more }
+	}
+
+	pub fn more(&self) -> bool {
+		self.more
+	}
+
+	pub fn into_vec(self) -> Vec<KeyValue> {
+		self.values
+	}
+
+	pub fn len(&self) -> usize {
+		self.values.len()
+	}
+
+	pub fn is_empty(&self) -> bool {
+		self.values.is_empty()
+	}
+
+	pub fn iter(&self) -> std::slice::Iter<'_, KeyValue> {
+		self.values.iter()
+	}
+
+	pub fn into_iter(self) -> std::vec::IntoIter<KeyValue> {
+		self.values.into_iter()
+	}
+}
+
+// impl Deref for Values {
+// 	type Target = [KeyValue];
+// 	fn deref(&self) -> &Self::Target {
+// 		&self.values
+// 	}
+// }
+// impl AsRef<[KeyValue]> for Values {
+// 	fn as_ref(&self) -> &[KeyValue] {
+// 		self.deref()
+// 	}
+// }
+
+// KeyValue type with key() and value() methods
+#[derive(Debug, Clone)]
+pub struct KeyValue {
+	key: Vec<u8>,
+	value: Vec<u8>,
+}
+
+impl KeyValue {
+	pub fn new(key: Vec<u8>, value: Vec<u8>) -> Self {
+		KeyValue { key, value }
+	}
+
+	pub fn key(&self) -> &[u8] {
+		&self.key
+	}
+
+	pub fn value(&self) -> &[u8] {
+		&self.value
+	}
+
+	pub fn into_parts(self) -> (Vec<u8>, Vec<u8>) {
+		(self.key, self.value)
+	}
+
+	pub fn to_value(self) -> Value {
+		Value::from_keyvalue(self)
+	}
+
+	pub fn value_ref(&self) -> Value {
+		Value::from_keyvalue(self.clone())
+	}
+}
+
+// Stream type for range queries - generic over item type
+pub type Stream<'a, T = KeyValue> =
+	Pin<Box<dyn futures_util::Stream<Item = Result<T>> + Send + 'a>>;

package/packages/universaldb/src/versionstamp.rs

@@ -0,0 +1,173 @@
+use std::{
+	sync::atomic::{AtomicU16, AtomicU64, Ordering},
+	time::{SystemTime, UNIX_EPOCH},
+};
+
+use crate::tuple::{TuplePack, Versionstamp, pack_with_versionstamp};
+
+static TRANSACTION_COUNTER: AtomicU16 = AtomicU16::new(0);
+static LAST_TIMESTAMP: AtomicU64 = AtomicU64::new(0);
+
+pub fn generate_versionstamp(user_version: u16) -> Versionstamp {
+	// HACK: Using SystemTime::now() for versionstamp generation is problematic because:
+	// (a) System time can go backwards due to NTP adjustments, daylight savings, etc.
+	// (b) System time is not synchronized across machines, so versionstamps generated
+	//     on different machines may not be correctly ordered
+	//
+	// This implementation tries to mitigate issue (a) by using max(current_time, last_time)
+	// but cannot solve issue (b)
+	let current_timestamp = SystemTime::now()
+		.duration_since(UNIX_EPOCH)
+		.unwrap()
+		.as_micros() as u64;
+
+	// Check if we've moved to a new microsecond to reset the transaction counter
+	//
+	// NOTE: Time can go backwards, so we use the max of current and last timestamp
+	let last_ts = LAST_TIMESTAMP.load(Ordering::Acquire);
+	let timestamp = current_timestamp.max(last_ts);
+
+	if timestamp > last_ts {
+		// Reset counter for new microsecond
+		LAST_TIMESTAMP.store(timestamp, Ordering::Release);
+		TRANSACTION_COUNTER.store(0, Ordering::Release);
+	}
+
+	let counter = TRANSACTION_COUNTER.fetch_add(1, Ordering::SeqCst);
+
+	// Handle counter overflow: if we've generated 65,536 versionstamps in the same
+	// microsecond, increment the timestamp and reset the counter
+	let final_timestamp = if counter == u16::MAX {
+		// Counter overflowed, increment timestamp
+		let new_timestamp = timestamp + 1;
+		LAST_TIMESTAMP.store(new_timestamp, Ordering::Release);
+		TRANSACTION_COUNTER.store(0, Ordering::Release);
+		new_timestamp
+	} else {
+		timestamp
+	};
+
+	let mut bytes = [0u8; 12];
+
+	bytes[0..8].copy_from_slice(&final_timestamp.to_be_bytes());
+	bytes[8..10].copy_from_slice(&counter.to_be_bytes());
+
+	bytes[10..12].copy_from_slice(&user_version.to_be_bytes());
+
+	Versionstamp::from(bytes)
+}
+
+pub fn substitute_versionstamp(
+	packed_data: &mut Vec<u8>,
+	versionstamp: Versionstamp,
+) -> Result<(), String> {
+	const VERSIONSTAMP_MARKER: u8 = 0x33;
+	const VERSIONSTAMP_SIZE: usize = 12;
+
+	if packed_data.len() < 4 {
+		return Err("Packed data too short to contain versionstamp offset".to_string());
+	}
+
+	let offset_bytes = packed_data.split_off(packed_data.len() - 4);
+	let offset = u32::from_le_bytes([
+		offset_bytes[0],
+		offset_bytes[1],
+		offset_bytes[2],
+		offset_bytes[3],
+	]) as usize;
+
+	if offset >= packed_data.len() {
+		return Err(format!(
+			"Invalid versionstamp offset: {} exceeds data length {}",
+			offset,
+			packed_data.len()
+		));
+	}
+
+	// The offset might point to the marker or the first byte of the versionstamp
+	let versionstamp_start = if packed_data.get(offset) == Some(&VERSIONSTAMP_MARKER) {
+		offset + 1
+	} else if offset > 0 && packed_data.get(offset - 1) == Some(&VERSIONSTAMP_MARKER) {
+		// The offset points to the first byte of the versionstamp data
+		offset
+	} else {
+		return Err(format!(
+			"No versionstamp marker (0x33) found at or before offset {}",
+			offset
+		));
+	};
+
+	let versionstamp_end = versionstamp_start + VERSIONSTAMP_SIZE;
+
+	if versionstamp_end > packed_data.len() {
+		return Err("Versionstamp extends beyond data bounds".to_string());
+	}
+
+	let existing_bytes = &packed_data[versionstamp_start..versionstamp_end];
+	if existing_bytes[0..10] != [0xff; 10] {
+		// Versionstamp is already complete, nothing to do
+		return Ok(());
+	}
+
+	let versionstamp_bytes = versionstamp.as_bytes();
+	packed_data[versionstamp_start..versionstamp_end].copy_from_slice(versionstamp_bytes);
+
+	Ok(())
+}
+
+pub fn pack_and_substitute_versionstamp<T: TuplePack>(
+	value: &T,
+	user_version: u16,
+) -> Result<Vec<u8>, String> {
+	let mut packed_data = pack_with_versionstamp(value);
+
+	let versionstamp = generate_versionstamp(user_version);
+
+	substitute_versionstamp(&mut packed_data, versionstamp)?;
+
+	Ok(packed_data)
+}
+
+/// Checks if a value might contain an incomplete versionstamp and attempts to substitute it
+///
+/// This is a helper function for database drivers that want to support versionstamp substitution.
+/// It detects if a value contains an incomplete versionstamp by checking for the offset marker
+/// at the end, and substitutes it with a generated versionstamp.
+///
+/// Returns the potentially modified value. If the value doesn't contain a versionstamp marker
+/// or substitution fails, returns the original value.
+pub fn substitute_versionstamp_if_incomplete(mut value: Vec<u8>, user_version: u16) -> Vec<u8> {
+	// Check if the value contains an incomplete versionstamp
+	// An incomplete versionstamp has a 4-byte offset at the end
+	if value.len() >= 4 {
+		// Try to read the offset from the last 4 bytes
+		let offset_bytes = &value[value.len() - 4..];
+		let offset = u32::from_le_bytes([
+			offset_bytes[0],
+			offset_bytes[1],
+			offset_bytes[2],
+			offset_bytes[3],
+		]) as usize;
+
+		// Check if this could be a valid versionstamp offset
+		// The offset should point within the value (excluding the 4 offset bytes)
+		if offset < value.len() - 4 {
+			// Check for versionstamp marker at or near the offset
+			const VERSIONSTAMP_MARKER: u8 = 0x33;
+			let has_marker = value.get(offset) == Some(&VERSIONSTAMP_MARKER)
+				|| (offset > 0 && value.get(offset - 1) == Some(&VERSIONSTAMP_MARKER));
+
+			if has_marker {
+				// This looks like it contains an incomplete versionstamp
+				// Generate a versionstamp and substitute it
+				let versionstamp = generate_versionstamp(user_version);
+
+				// Substitute the versionstamp
+				// This will do nothing if the versionstamp is already complete
+				let _ = substitute_versionstamp(&mut value, versionstamp);
+			}
+		}
+	}
+
+	value
+}