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

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

@@ -21,60 +21,60 @@ pub trait CherryPick {
 
 // Implements `CherryPick` for any tuple size
 macro_rules! impl_tuple {
-    ($($args:ident),*) => {
-        #[async_trait::async_trait]
-        impl<$($args: FormalKey + for<'de> TupleUnpack<'de>),*> CherryPick for ($($args),*)
-        where
-            $($args::Value: Send),*
-        {
-            type Output = ($($args::Value),*);
+	($($args:ident),*) => {
+		#[async_trait::async_trait]
+		impl<$($args: FormalKey + for<'de> TupleUnpack<'de>),*> CherryPick for ($($args),*)
+		where
+	    	$($args::Value: Send),*
+		{
+	    	type Output = ($($args::Value),*);
 
-            async fn cherry_pick<S: TuplePack + Send>(
-                tx: &Transaction,
-                subspace: S,
-                isolation_level: IsolationLevel,
-            ) -> Result<Self::Output> {
-                let tx = tx.with_subspace(Subspace::new(&subspace));
+	    	async fn cherry_pick<S: TuplePack + Send>(
+	        	tx: &Transaction,
+	        	subspace: S,
+	        	isolation_level: IsolationLevel,
+	    	) -> Result<Self::Output> {
+	        	let tx = tx.with_subspace(Subspace::new(&subspace));
 
-                let mut stream = tx.read_range(
-                    $crate::range_option::RangeOption {
-                        mode: StreamingMode::WantAll,
-                        ..(&Subspace::all()).into()
-                    },
-                    isolation_level,
-                );
+	        	let mut stream = tx.read_range(
+	            	$crate::range_option::RangeOption {
+	                	mode: StreamingMode::WantAll,
+	                	..(&Subspace::all()).into()
+	            	},
+	            	isolation_level,
+	        	);
 
-                $(
+	        	$(
 					#[allow(non_snake_case)]
-                    let mut $args = None;
-                )*
+	            	let mut $args = None;
+	        	)*
 
-                loop {
-                    let Some(entry) = stream.try_next().await? else {
-                        break;
-                    };
+	        	loop {
+	            	let Some(entry) = stream.try_next().await? else {
+	                	break;
+	            	};
 
-                    $(
-                        if let Ok(key) = tx.unpack::<$args>(entry.key()) {
-                            ensure!($args.is_none(), "{} already picked", std::any::type_name::<$args>());
+	            	$(
+	                	if let Ok(key) = tx.unpack::<$args>(entry.key()) {
+	                    	ensure!($args.is_none(), "{} already picked", std::any::type_name::<$args>());
 
-                            let value = key.deserialize(entry.value())?;
-                            $args = Some(value);
-                            continue;
-                        }
-                    )*
-                }
+	                    	let value = key.deserialize(entry.value())?;
+	                    	$args = Some(value);
+	                    	continue;
+	                	}
+	            	)*
+	        	}
 
-                Ok((
-                    $(
-                        $args.with_context(|| {
-                            format!("key not found in cherry pick: {}", std::any::type_name::<$args>())
-                        })?,
-                    )*
-                ))
-            }
-        }
-    }
+	        	Ok((
+	            	$(
+	                	$args.with_context(|| {
+	                    	format!("key not found in cherry pick: {}", std::any::type_name::<$args>())
+	                	})?,
+	            	)*
+	        	))
+	    	}
+		}
+	}
 }
 
 impl_tuple!(A, B);

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

@@ -85,7 +85,7 @@ define_keys! {
 	(57, FORGOTTEN, "forgotten"),
 	(58, EVENT_TYPE, "event_type"),
 	(59, VERSION, "version"),
-	// 60
+	(60, INNER_VERSION, "inner_version"),
 	(61, SIGNAL_ID, "signal_id"),
 	(62, SUB_WORKFLOW_ID, "sub_workflow_id"),
 	(63, ITERATION, "iteration"),
@@ -95,7 +95,7 @@ define_keys! {
 	(67, ALL, "all"),
 	(68, BY_NAME_AND_KEY, "by_name_and_key"),
 	(69, DESTROY_TS, "destroy_ts"),
-	(70, EPOXY, "epoxy"),
+	(70, EPOXY_V1, "epoxy"),
 	(71, PEER_STATE, "peer_state"),
 	(72, ACTOR_KV, "actor_kv"),
 	(73, STOP_TS, "stop_ts"),
@@ -131,4 +131,23 @@ define_keys! {
 	// 103 - RESERVED BY EE
 	// 104 - RESERVED BY EE
 	// 105 - RESERVED BY EE
+	(106, PROTOCOL_VERSION, "protocol_version"),
+	(107, LAST_COMMAND_IDX, "last_command_idx"),
+	(108, COMMAND, "command"),
+	(109, RUNNER_NAME_SELECTOR, "runner_name_selector"),
+	(110, BACKFILL, "backfill"),
+	(111, COMPLETE, "complete"),
+	(112, COMPLETE_TS, "complete_ts"),
+	(113, PRUNE_TS, "prune_ts"),
+	(114, PRUNE_IDX, "prune_idx"),
+	(115, ENVOY, "envoy"),
+	(116, ENVOY_KEY, "envoy_key"),
+	(117, POOL_NAME, "pool_name"),
+	(118, SLOTS, "slots"),
+	(119, ENVOY_LOAD_BALANCER_IDX, "envoy_load_balancer_idx"),
+	(120, BALLOT, "ballot"),
+	(121, ACCEPTED, "accepted"),
+	(122, CACHE, "cache"),
+	(123, CHANGELOG, "changelog"),
+	(124, EPOXY_V2, "epoxy_v2"),
 }

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

@@ -24,6 +24,14 @@ pub enum IsolationLevel {
 #[derive(Debug, Clone, Copy)]
 pub struct MaybeCommitted(pub bool);
 
+impl std::ops::Deref for MaybeCommitted {
+	type Target = bool;
+
+	fn deref(&self) -> &Self::Target {
+		&self.0
+	}
+}
+
 /// Calculate exponential backoff based on attempt.
 ///
 /// Ours:

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

@@ -5,7 +5,6 @@ use crate::{
 	range_option::RangeOption,
 	tuple::{self, PackResult, TuplePack, TupleUnpack},
 };
-use rivet_metrics::KeyValue;
 
 use crate::metrics;
 
@@ -45,7 +44,9 @@ impl Subspace {
 	/// 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>())]);
+		metrics::KEY_PACK_COUNT
+			.with_label_values(&[std::any::type_name::<T>()])
+			.inc();
 
 		self.inner.pack(t)
 	}
@@ -53,7 +54,9 @@ impl Subspace {
 	/// 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>())]);
+		metrics::KEY_PACK_COUNT
+			.with_label_values(&[std::any::type_name::<T>()])
+			.inc();
 
 		self.inner.pack_with_versionstamp(t)
 	}
@@ -62,7 +65,9 @@ impl 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>())]);
+		metrics::KEY_UNPACK_COUNT
+			.with_label_values(&[std::any::type_name::<T>()])
+			.inc();
 
 		self.inner.unpack(key)
 	}

package/packages/universaldb/tests/integration.rs

@@ -37,9 +37,11 @@ async fn test_postgres_driver() {
 	// Get the connection string from the secret
 	let connection_string = postgres_config.url.read().clone();
 
-	let driver = universaldb::driver::PostgresDatabaseDriver::new(connection_string)
-		.await
-		.unwrap();
+	let driver = universaldb::driver::PostgresDatabaseDriver::new_with_config(
+		universaldb::driver::postgres::PostgresConfig::new(connection_string),
+	)
+	.await
+	.unwrap();
 	let db = Database::new(Arc::new(driver));
 
 	run_all_tests(db).await

package/packages/universaldb/tests/integration_gas.rs

@@ -26,9 +26,11 @@ async fn test_postgres_gasoline() {
 	// Get the connection string from the secret
 	let connection_string = postgres_config.url.read().clone();
 
-	let driver = universaldb::driver::PostgresDatabaseDriver::new(connection_string)
-		.await
-		.unwrap();
+	let driver = universaldb::driver::PostgresDatabaseDriver::new_with_config(
+		universaldb::driver::postgres::PostgresConfig::new(connection_string),
+	)
+	.await
+	.unwrap();
 	let db = Database::new(Arc::new(driver));
 
 	test_gasoline_operations(&db).await;

package/packages/universaldb/tests/rocksdb.rs

@@ -1,12 +1,21 @@
-use std::{path::Path, sync::Arc};
+use std::{
+	path::Path,
+	sync::{
+		Arc,
+		atomic::{AtomicUsize, Ordering},
+	},
+	time::{Duration, Instant},
+};
 
 use anyhow::Context;
-use futures_util::StreamExt;
+use futures_util::{StreamExt, TryStreamExt};
+use rand::{distributions::WeightedError, seq::SliceRandom};
 use rivet_test_deps_docker::TestDatabase;
 use rocksdb::{OptimisticTransactionDB, Options, WriteOptions};
 use universaldb::{
 	Database,
-	utils::{IsolationLevel::*, calculate_tx_retry_backoff},
+	prelude::*,
+	utils::{calculate_tx_retry_backoff, end_of_key_range},
 };
 use uuid::Uuid;
 
@@ -106,36 +115,160 @@ async fn rocksdb_udb() {
 		.unwrap();
 	let db = Database::new(Arc::new(driver));
 
+	db.run(|tx| async move {
+		for i in 0..=255 {
+			for j in 0..=0 {
+				let key = vec![1, 2, 3, i, j];
+				let value = b"";
+
+				tx.set(&key, value);
+			}
+		}
+
+		Ok(())
+	})
+	.await
+	.unwrap();
+
 	let mut handles = Vec::new();
 
-	for _ in 0..64 {
+	for i in 0..64 {
 		let db = db.clone();
 
 		let handle = tokio::spawn(async move {
-			db.run(|tx| async move {
-				let key = vec![1, 2, 3];
-				let value = vec![4, 5, 6];
+			let tries = AtomicUsize::new(0);
+			let start = Instant::now();
+
+			let alloc = db
+				.run(|tx| {
+					let tries = &tries;
+
+					async move {
+						tries.fetch_add(1, Ordering::SeqCst);
+
+						let mut stream = tx.get_ranges_keyvalues(
+							universaldb::RangeOption {
+								begin: KeySelector::first_greater_or_equal(vec![1, 2, 3]),
+								end: KeySelector::last_less_or_equal(vec![1, 2, 3, 255, 255]),
+								mode: StreamingMode::Iterator,
+								..RangeOption::default()
+							},
+							Snapshot,
+						);
+
+						let mut chunk = Vec::with_capacity(100);
+
+						loop {
+							let empty = if chunk.len() >= 100 {
+								false
+							} else if let Some(entry) = stream.try_next().await? {
+								chunk.push(entry);
+								continue;
+							} else {
+								true
+							};
+
+							let entry = match chunk.choose_weighted(&mut rand::thread_rng(), |_| 1)
+							{
+								Ok(entry) => entry,
+								Err(WeightedError::NoItem) => break,
+								Err(err) => return Err(err.into()),
+							};
+
+							tx.add_conflict_range(
+								entry.key(),
+								&end_of_key_range(entry.key()),
+								ConflictRangeType::Read,
+							)?;
+
+							tx.clear(entry.key());
+
+							return Ok(Some(entry.key().to_vec()));
+
+							if empty {
+								break;
+							} else {
+								chunk.clear();
+							}
+						}
+
+						Ok(None)
+					}
+				})
+				.await?;
 
-				tx.get(&key, Serializable).await.unwrap();
-				tx.set(&key, &value);
+			let tries = tries.load(Ordering::SeqCst);
+			let duration = start.elapsed();
 
-				Ok(())
-			})
-			.await
-			.unwrap();
+			tracing::info!(%i, %tries, ?alloc, ?duration, "success");
 
-			tracing::info!("success");
+			anyhow::Ok(alloc.map(|_| (tries, duration)))
 		});
 
 		handles.push(handle);
 	}
 
-	futures_util::stream::iter(handles)
-		.buffer_unordered(1024)
-		.for_each(|result| async {
-			if let Err(err) = result {
-				tracing::error!(?err, "task failed");
+	let res = futures_util::stream::iter(handles)
+		.map(|result| async {
+			match result.await {
+				Ok(Ok(result)) => result,
+				Ok(Err(err)) => {
+					tracing::error!(?err, "task failed");
+					None
+				}
+				Err(err) => {
+					tracing::error!(?err, "task failed");
+					None
+				}
 			}
 		})
+		.buffer_unordered(1024)
+		.filter_map(|x| std::future::ready(x))
+		.collect::<Vec<_>>()
 		.await;
+
+	let mut retry_distribution = Vec::<(usize, usize)>::new();
+	let mut duration_buckets = vec![
+		(Duration::from_micros(0), Duration::from_micros(1000), 0),
+		(Duration::from_micros(1000), Duration::from_micros(2500), 0),
+		(Duration::from_micros(2500), Duration::from_micros(5000), 0),
+		(Duration::from_micros(5000), Duration::from_micros(10000), 0),
+		(Duration::from_millis(10), Duration::from_millis(25), 0),
+		(Duration::from_millis(25), Duration::from_millis(50), 0),
+		(Duration::from_millis(50), Duration::from_millis(100), 0),
+		(Duration::from_millis(100), Duration::from_millis(250), 0),
+		(Duration::from_millis(250), Duration::from_millis(500), 0),
+		(Duration::from_millis(500), Duration::from_millis(1000), 0),
+		(Duration::from_millis(1000), Duration::from_millis(2500), 0),
+		(Duration::from_millis(2500), Duration::MAX, 0),
+	];
+
+	for (retries, duration) in res {
+		if let Some((_, count)) = retry_distribution.iter_mut().find(|(n, _)| n == &retries) {
+			*count += 1;
+		} else {
+			retry_distribution.push((retries, 1));
+		}
+
+		for (_, end, count) in &mut duration_buckets {
+			if &duration < end {
+				*count += 1;
+				break;
+			}
+		}
+	}
+
+	for (retries, count) in retry_distribution {
+		println!("{retries}: {count}");
+	}
+
+	println!();
+
+	for (start, end, count) in duration_buckets {
+		if end == Duration::MAX {
+			println!("{start:?}-Inf: {count}");
+		} else {
+			println!("{start:?}-{end:?}: {count}");
+		}
+	}
 }

package/packages/universalpubsub/Cargo.toml

@@ -12,18 +12,24 @@ async-trait.workspace = true
 base64.workspace = true
 deadpool-postgres.workspace = true
 futures-util.workspace = true
+lazy_static.workspace = true
+rand.workspace = true
 rivet-error.workspace = true
+rivet-metrics.workspace = true
+rivet-postgres-util.workspace = true
 rivet-ups-protocol.workspace = true
 rivet-util.workspace = true
+scc.workspace = true
 serde_json.workspace = true
-vbare.workspace = true
 serde.workspace = true
 sha2.workspace = true
+tokio-postgres-rustls.workspace = true
 tokio-postgres.workspace = true
-tokio.workspace = true
 tokio-util.workspace = true
+tokio.workspace = true
 tracing.workspace = true
 uuid.workspace = true
+vbare.workspace = true
 
 [dev-dependencies]
 rivet-config.workspace = true

package/packages/universalpubsub/benches/simple.rs

@@ -665,9 +665,17 @@ async fn setup_pg_pair() -> Result<(PubSub, PubSub)> {
 		unreachable!()
 	};
 	let url = pg.url.read().clone();
-	let driver_pub =
-		universalpubsub::driver::postgres::PostgresDriver::connect(url.clone(), false).await?;
-	let driver_sub = universalpubsub::driver::postgres::PostgresDriver::connect(url, false).await?;
+	let driver_pub = universalpubsub::driver::postgres::PostgresDriver::connect(
+		url.clone(),
+		false,
+		None,
+		None,
+		None,
+	)
+	.await?;
+	let driver_sub =
+		universalpubsub::driver::postgres::PostgresDriver::connect(url, false, None, None, None)
+			.await?;
 	Ok((
 		PubSub::new_with_memory_optimization(Arc::new(driver_pub), false),
 		PubSub::new_with_memory_optimization(Arc::new(driver_sub), false),
@@ -685,7 +693,9 @@ async fn setup_pg_single() -> Result<(PubSub, PubSub)> {
 		unreachable!()
 	};
 	let url = pg.url.read().clone();
-	let driver = universalpubsub::driver::postgres::PostgresDriver::connect(url, false).await?;
+	let driver =
+		universalpubsub::driver::postgres::PostgresDriver::connect(url, false, None, None, None)
+			.await?;
 	let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), false);
 	Ok((pubsub.clone(), pubsub))
 }
@@ -701,9 +711,17 @@ async fn setup_pg_pair_mem() -> Result<(PubSub, PubSub)> {
 		unreachable!()
 	};
 	let url = pg.url.read().clone();
-	let driver_pub =
-		universalpubsub::driver::postgres::PostgresDriver::connect(url.clone(), true).await?;
-	let driver_sub = universalpubsub::driver::postgres::PostgresDriver::connect(url, true).await?;
+	let driver_pub = universalpubsub::driver::postgres::PostgresDriver::connect(
+		url.clone(),
+		true,
+		None,
+		None,
+		None,
+	)
+	.await?;
+	let driver_sub =
+		universalpubsub::driver::postgres::PostgresDriver::connect(url, true, None, None, None)
+			.await?;
 	Ok((
 		PubSub::new_with_memory_optimization(Arc::new(driver_pub), true),
 		PubSub::new_with_memory_optimization(Arc::new(driver_sub), true),
@@ -721,7 +739,9 @@ async fn setup_pg_single_mem() -> Result<(PubSub, PubSub)> {
 		unreachable!()
 	};
 	let url = pg.url.read().clone();
-	let driver = universalpubsub::driver::postgres::PostgresDriver::connect(url, true).await?;
+	let driver =
+		universalpubsub::driver::postgres::PostgresDriver::connect(url, true, None, None, None)
+			.await?;
 	let pubsub = PubSub::new_with_memory_optimization(Arc::new(driver), true);
 	Ok((pubsub.clone(), pubsub))
 }

package/packages/universalpubsub/src/chunking.rs

@@ -6,7 +6,6 @@ use rivet_ups_protocol::versioned::UpsMessage;
 use rivet_ups_protocol::{MessageBody, MessageChunk, MessageStart, PROTOCOL_VERSION};
 use vbare::OwnedVersionedData;
 
-pub const CHUNK_BUFFER_GC_INTERVAL: Duration = Duration::from_secs(60);
 pub const CHUNK_BUFFER_MAX_AGE: Duration = Duration::from_secs(300);
 
 #[derive(Debug)]
@@ -103,7 +102,7 @@ impl ChunkTracker {
 			.retain(|_, buffer| now.duration_since(buffer.last_chunk_ts) < CHUNK_BUFFER_MAX_AGE);
 		let size_after = self.chunks_in_process.len();
 
-		tracing::debug!(
+		tracing::trace!(
 			?size_before,
 			?size_after,
 			"performed chunk buffer garbage collection"
@@ -111,6 +110,17 @@ impl ChunkTracker {
 	}
 }
 
+/// Returns the number of bytes needed to encode `n` as a BARE unsigned integer (LEB128).
+fn bare_uint_len(n: usize) -> usize {
+	let mut len = 1;
+	let mut v = n >> 7;
+	while v > 0 {
+		len += 1;
+		v >>= 7;
+	}
+	len
+}
+
 /// Splits a payload into chunks that fit within message size limits.
 ///
 /// This function handles chunking by accounting for different overhead
@@ -161,9 +171,21 @@ pub fn split_payload_into_chunks(
 		.serialize_with_embedded_version(PROTOCOL_VERSION)?
 		.len();
 
-	// Calculate max payload sizes
-	let first_chunk_max_payload = max_message_size.saturating_sub(start_overhead);
-	let other_chunk_max_payload = max_message_size.saturating_sub(chunk_overhead);
+	// Calculate max payload sizes, correcting for the variable-length encoding of the
+	// data length prefix. The overhead above was computed with an empty payload
+	// (uint(0) = 1 byte). For payloads >= 128 bytes the length prefix grows (LEB128
+	// encoding), so we subtract those extra bytes to ensure every encoded chunk fits
+	// within max_message_size.
+	let first_chunk_max_payload = {
+		let raw = max_message_size.saturating_sub(start_overhead);
+		let extra = bare_uint_len(raw).saturating_sub(1);
+		raw.saturating_sub(extra)
+	};
+	let other_chunk_max_payload = {
+		let raw = max_message_size.saturating_sub(chunk_overhead);
+		let extra = bare_uint_len(raw).saturating_sub(1);
+		raw.saturating_sub(extra)
+	};
 
 	if first_chunk_max_payload == 0 || other_chunk_max_payload == 0 {
 		bail!("message overhead exceeds max message size");