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

package/packages/{actor-kv → pegboard}/tests/kv_operations.rs

@@ -1,8 +1,7 @@
 use anyhow::Result;
-use pegboard_actor_kv as kv;
-use rivet_runner_protocol as rp;
-use rivet_util_id::Id;
-use uuid::Uuid;
+use gas::prelude::*;
+use pegboard::actor_kv as kv;
+use rivet_envoy_protocol as ep;
 
 #[tokio::test]
 async fn test_kv_operations() -> Result<()> {
@@ -14,15 +13,20 @@ async fn test_kv_operations() -> Result<()> {
 
 	let test_id = Uuid::new_v4();
 	let dc_label = 1;
-	let datacenters = vec![rivet_config::config::topology::Datacenter {
-		name: "test-dc".to_string(),
-		datacenter_label: dc_label,
-		is_leader: true,
-		peer_url: url::Url::parse("http://127.0.0.1:8080")?,
-		public_url: url::Url::parse("http://127.0.0.1:8081")?,
-		proxy_url: None,
-		valid_hosts: None,
-	}];
+	let datacenters = [(
+		"test-dc".to_string(),
+		rivet_config::config::topology::Datacenter {
+			name: "test-dc".to_string(),
+			datacenter_label: dc_label,
+			is_leader: true,
+			peer_url: url::Url::parse("http://127.0.0.1:8080")?,
+			public_url: url::Url::parse("http://127.0.0.1:8081")?,
+			proxy_url: None,
+			valid_hosts: None,
+		},
+	)]
+	.into_iter()
+	.collect();
 
 	let api_peer_port = portpicker::pick_unused_port().expect("failed to pick api peer port");
 	let guard_port = portpicker::pick_unused_port().expect("failed to pick guard port");
@@ -38,6 +42,11 @@ async fn test_kv_operations() -> Result<()> {
 
 	let db = &test_deps.pools.udb()?;
 	let actor_id = Id::new_v1(dc_label);
+	let recipient = kv::Recipient {
+		actor_id,
+		namespace_id: Id::new_v1(dc_label),
+		name: "default".to_string(),
+	};
 
 	tracing::info!(?actor_id, "starting kv operations test");
 
@@ -58,12 +67,12 @@ async fn test_kv_operations() -> Result<()> {
 		b"other_value".to_vec(),
 	];
 
-	kv::put(db, actor_id, keys.clone(), values.clone()).await?;
+	kv::put(db, &recipient, keys.clone(), values.clone()).await?;
 	tracing::info!("successfully put {} keys", keys.len());
 
 	// Test 2: Get the keys back
 	tracing::info!("test 2: getting keys");
-	let (got_keys, got_values, got_metadata) = kv::get(db, actor_id, keys.clone()).await?;
+	let (got_keys, got_values, got_metadata) = kv::get(db, &recipient, keys.clone()).await?;
 
 	assert_eq!(got_keys.len(), 5, "should get 5 keys back");
 	assert_eq!(got_values.len(), 5, "should get 5 values back");
@@ -85,7 +94,7 @@ async fn test_kv_operations() -> Result<()> {
 			"metadata should have version"
 		);
 		assert!(
-			got_metadata[got_idx].create_ts > 0,
+			got_metadata[got_idx].update_ts > 0,
 			"metadata should have timestamp"
 		);
 	}
@@ -94,7 +103,7 @@ async fn test_kv_operations() -> Result<()> {
 	// Test 3: List all keys
 	tracing::info!("test 3: listing all keys");
 	let (list_keys, list_values, list_metadata) =
-		kv::list(db, actor_id, rp::KvListQuery::KvListAllQuery, false, None).await?;
+		kv::list(db, &recipient, ep::KvListQuery::KvListAllQuery, false, None).await?;
 
 	assert_eq!(list_keys.len(), 5, "should list 5 keys");
 	assert_eq!(list_values.len(), 5, "should list 5 values");
@@ -105,8 +114,8 @@ async fn test_kv_operations() -> Result<()> {
 	tracing::info!("test 4: listing with limit");
 	let (limited_keys, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListAllQuery,
+		&recipient,
+		ep::KvListQuery::KvListAllQuery,
 		false,
 		Some(2),
 	)
@@ -118,9 +127,9 @@ async fn test_kv_operations() -> Result<()> {
 	// Test 5: List with reverse
 	tracing::info!("test 5: listing in reverse");
 	let (forward_keys, _, _) =
-		kv::list(db, actor_id, rp::KvListQuery::KvListAllQuery, false, None).await?;
+		kv::list(db, &recipient, ep::KvListQuery::KvListAllQuery, false, None).await?;
 	let (reverse_keys, _, _) =
-		kv::list(db, actor_id, rp::KvListQuery::KvListAllQuery, true, None).await?;
+		kv::list(db, &recipient, ep::KvListQuery::KvListAllQuery, true, None).await?;
 
 	assert_eq!(forward_keys.len(), reverse_keys.len());
 	// Keys should be in opposite order
@@ -151,13 +160,13 @@ async fn test_kv_operations() -> Result<()> {
 		b"Post 2".to_vec(),
 		b"Comment 100".to_vec(),
 	];
-	kv::put(db, actor_id, prefix_keys.clone(), prefix_values.clone()).await?;
+	kv::put(db, &recipient, prefix_keys.clone(), prefix_values.clone()).await?;
 
 	// Query with "users:" prefix
 	let (users_keys, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListPrefixQuery(rp::KvListPrefixQuery {
+		&recipient,
+		ep::KvListQuery::KvListPrefixQuery(ep::KvListPrefixQuery {
 			key: b"users:".to_vec(),
 		}),
 		false,
@@ -172,8 +181,8 @@ async fn test_kv_operations() -> Result<()> {
 	// Query with "posts:" prefix
 	let (posts_keys, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListPrefixQuery(rp::KvListPrefixQuery {
+		&recipient,
+		ep::KvListQuery::KvListPrefixQuery(ep::KvListPrefixQuery {
 			key: b"posts:".to_vec(),
 		}),
 		false,
@@ -188,14 +197,14 @@ async fn test_kv_operations() -> Result<()> {
 	tracing::info!("successfully listed keys with prefix");
 
 	// Clean up the prefix test keys
-	kv::delete(db, actor_id, prefix_keys).await?;
+	kv::delete(db, &recipient, prefix_keys).await?;
 
 	// Test 7: List with range
 	tracing::info!("test 7: listing with range");
 	let (range_keys, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListRangeQuery(rp::KvListRangeQuery {
+		&recipient,
+		ep::KvListQuery::KvListRangeQuery(ep::KvListRangeQuery {
 			start: b"key1".to_vec(),
 			end: b"key2".to_vec(),
 			exclusive: false,
@@ -218,8 +227,8 @@ async fn test_kv_operations() -> Result<()> {
 	tracing::info!("test 8: listing with exclusive range");
 	let (exclusive_range_keys, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListRangeQuery(rp::KvListRangeQuery {
+		&recipient,
+		ep::KvListQuery::KvListRangeQuery(ep::KvListRangeQuery {
 			start: b"key1".to_vec(),
 			end: b"key2".to_vec(),
 			exclusive: true,
@@ -237,53 +246,73 @@ async fn test_kv_operations() -> Result<()> {
 	// Test 9: Delete specific keys
 	tracing::info!("test 9: deleting specific keys");
 	let keys_to_delete = vec![b"key1".to_vec(), b"key2".to_vec()];
-	kv::delete(db, actor_id, keys_to_delete.clone()).await?;
+	kv::delete(db, &recipient, keys_to_delete.clone()).await?;
 
 	// Verify keys are deleted
 	let (remaining_keys, _, _) =
-		kv::list(db, actor_id, rp::KvListQuery::KvListAllQuery, false, None).await?;
+		kv::list(db, &recipient, ep::KvListQuery::KvListAllQuery, false, None).await?;
 	assert_eq!(remaining_keys.len(), 3, "should have 3 keys remaining");
 	assert!(!remaining_keys.contains(&b"key1".to_vec()));
 	assert!(!remaining_keys.contains(&b"key2".to_vec()));
 	tracing::info!("successfully deleted specific keys");
 
-	// Test 10: Delete all keys
-	tracing::info!("test 10: deleting all keys");
-	kv::delete_all(db, actor_id).await?;
+	// Test 10: Delete a half-open range
+	tracing::info!("test 10: deleting a half-open range");
+	kv::delete_range(db, &recipient, b"key3".to_vec(), b"key5".to_vec()).await?;
+
+	let (post_range_keys, _, _) =
+		kv::list(db, &recipient, ep::KvListQuery::KvListAllQuery, false, None).await?;
+	assert_eq!(post_range_keys.len(), 1, "should have 1 key remaining");
+	assert_eq!(
+		post_range_keys[0],
+		b"other".to_vec(),
+		"only other should remain"
+	);
+	tracing::info!("successfully deleted keys in half-open range");
+
+	// Test 11: Delete all keys
+	tracing::info!("test 11: deleting all keys");
+	kv::delete_all(db, &recipient).await?;
 
 	// Verify all keys are deleted
 	let (all_keys, _, _) =
-		kv::list(db, actor_id, rp::KvListQuery::KvListAllQuery, false, None).await?;
+		kv::list(db, &recipient, ep::KvListQuery::KvListAllQuery, false, None).await?;
 	assert_eq!(all_keys.len(), 0, "should have no keys remaining");
 	tracing::info!("successfully deleted all keys");
 
-	// Test 11: Test storage size
-	tracing::info!("test 11: testing storage size");
-	let subspace = pegboard::keys::actor_kv_subspace().subspace(&actor_id);
-	let size = kv::get_subspace_size(db, &subspace).await?;
+	// Test 12: Test storage size
+	tracing::info!("test 12: testing storage size");
+	let size = db
+		.run(|tx| async move { kv::estimate_kv_size(&tx, actor_id).await })
+		.await
+		.unwrap();
 	assert_eq!(size, 0, "storage size should be 0 after delete_all");
 	tracing::info!("successfully verified storage size");
 
-	// Test 12: Test large value (chunking)
-	tracing::info!("test 12: testing large value chunking");
+	// Test 13: Test large value (chunking)
+	tracing::info!("test 13: testing large value chunking");
 	let large_value = vec![42u8; 50_000]; // 50 KB, will be split into chunks
 	kv::put(
 		db,
-		actor_id,
+		&recipient,
 		vec![b"large_key".to_vec()],
 		vec![large_value.clone()],
 	)
 	.await?;
 
-	let (large_keys, large_values, _) = kv::get(db, actor_id, vec![b"large_key".to_vec()]).await?;
+	let (large_keys, large_values, _) =
+		kv::get(db, &recipient, vec![b"large_key".to_vec()]).await?;
 	assert_eq!(large_keys.len(), 1);
 	assert_eq!(large_values[0], large_value, "large value should match");
 	tracing::info!("successfully stored and retrieved large value");
 
-	// Test 13: Verify storage size increased
+	// Test 14: Verify storage size increased
 	// Note: Storage size estimation may not be accurate on all backends (e.g., FileSystem)
-	tracing::info!("test 13: verifying storage size with data");
-	let size_with_data = kv::get_subspace_size(db, &subspace).await?;
+	tracing::info!("test 14: verifying storage size with data");
+	let size_with_data = db
+		.run(|tx| async move { kv::estimate_kv_size(&tx, actor_id).await })
+		.await
+		.unwrap();
 	tracing::info!(
 		?size_with_data,
 		"storage size with data (may be 0 on some backends)"

package/packages/pegboard-envoy/Cargo.toml

@@ -0,0 +1,43 @@
+[package]
+name = "pegboard-envoy"
+version.workspace = true
+authors.workspace = true
+license.workspace = true
+edition.workspace = true
+
+[dependencies]
+anyhow.workspace = true
+async-trait.workspace = true
+bytes.workspace = true
+futures-util.workspace = true
+gas.workspace = true
+http-body.workspace = true
+http-body-util.workspace = true
+# Idk how to get this working with the workspace version
+hyper = "1.6"
+hyper-tungstenite.workspace = true
+hyper-util = "0.1"
+lazy_static.workspace = true
+rand.workspace = true
+rivet-config.workspace = true
+rivet-data.workspace = true
+rivet-error.workspace = true
+rivet-guard-core.workspace = true
+rivet-metrics.workspace = true
+rivet-envoy-protocol.workspace = true
+rivet-runtime.workspace = true
+rivet-types.workspace = true
+scc.workspace = true
+serde_bare.workspace = true
+serde_json.workspace = true
+serde.workspace = true
+tokio-tungstenite.workspace = true
+tokio.workspace = true
+tracing.workspace = true
+universaldb.workspace = true
+universalpubsub.workspace = true
+url.workspace = true
+vbare.workspace = true
+
+pegboard.workspace = true
+namespace.workspace = true

package/packages/pegboard-envoy/src/actor_event_demuxer.rs

@@ -0,0 +1,165 @@
+use std::collections::HashMap;
+use std::time::{Duration, Instant};
+
+use anyhow::Result;
+use gas::prelude::*;
+use rivet_envoy_protocol as protocol;
+use tokio::sync::mpsc;
+use tokio::task::JoinHandle;
+
+use crate::metrics;
+
+struct Channel {
+	tx: mpsc::UnboundedSender<protocol::EventWrapper>,
+	handle: JoinHandle<()>,
+	last_seen: Instant,
+}
+
+pub struct ActorEventDemuxer {
+	ctx: StandaloneCtx,
+	envoy_key: String,
+	channels: HashMap<Id, Channel>,
+	last_gc: Instant,
+	gc_interval: Duration,
+	max_last_seen: Duration,
+}
+
+impl ActorEventDemuxer {
+	pub fn new(ctx: StandaloneCtx, envoy_key: String) -> Self {
+		let pegboard_config = ctx.config().pegboard();
+		let gc_interval = Duration::from_millis(pegboard_config.envoy_event_demuxer_gc_interval());
+		let max_last_seen =
+			Duration::from_millis(pegboard_config.envoy_event_demuxer_max_last_seen_threshold());
+		Self {
+			ctx,
+			envoy_key,
+			channels: HashMap::new(),
+			last_gc: Instant::now(),
+			gc_interval,
+			max_last_seen,
+		}
+	}
+
+	/// Process an event by routing it to the appropriate actor's queue
+	pub fn ingest(&mut self, actor_id: Id, event: protocol::EventWrapper) {
+		tracing::debug!(envoy_key=%self.envoy_key, ?actor_id, index=?event.checkpoint.index, "actor demuxer ingest");
+
+		if let Some(channel) = self.channels.get_mut(&actor_id) {
+			let _ = channel.tx.send(event);
+			channel.last_seen = Instant::now();
+		} else {
+			let (tx, rx) = mpsc::unbounded_channel();
+
+			let handle = tokio::spawn(channel_handler(
+				self.ctx.clone(),
+				self.envoy_key.clone(),
+				actor_id,
+				rx,
+			));
+
+			// Send initial event
+			let _ = tx.send(event);
+
+			self.channels.insert(
+				actor_id,
+				Channel {
+					tx,
+					handle,
+					last_seen: Instant::now(),
+				},
+			);
+		}
+
+		metrics::INGESTED_EVENTS_TOTAL.inc();
+
+		// Run gc periodically
+		if self.last_gc.elapsed() > self.gc_interval {
+			self.last_gc = Instant::now();
+
+			self.channels.retain(|_, channel| {
+				let keep = channel.last_seen.elapsed() < self.max_last_seen;
+
+				if !keep {
+					// TODO: Verify aborting is safe here
+					channel.handle.abort();
+				}
+
+				keep
+			});
+
+			metrics::EVENT_MULTIPLEXER_COUNT.set(self.channels.len() as i64);
+		}
+	}
+
+	/// Shutdown all tasks and wait for them to complete
+	#[tracing::instrument(skip_all)]
+	pub async fn shutdown(self) {
+		tracing::debug!(channels=?self.channels.len(), "shutting down actor demuxer");
+
+		// Drop all senders
+		let handles = self
+			.channels
+			.into_iter()
+			.map(|(_, channel)| channel.handle)
+			.collect::<Vec<_>>();
+
+		// Await remaining tasks
+		for handle in handles {
+			let _ = handle.await;
+		}
+
+		tracing::debug!("actor demuxer shut down");
+	}
+}
+
+#[tracing::instrument(name="demuxer_channel", skip_all, fields(ray_id=?ctx.ray_id(), req_id=?ctx.req_id(), %envoy_key, ?actor_id))]
+async fn channel_handler(
+	ctx: StandaloneCtx,
+	envoy_key: String,
+	actor_id: Id,
+	mut rx: mpsc::UnboundedReceiver<protocol::EventWrapper>,
+) {
+	loop {
+		let mut buffer = Vec::new();
+
+		// Batch process events
+		if rx.recv_many(&mut buffer, 1024).await == 0 {
+			break;
+		}
+
+		if let Err(err) = dispatch_events(&ctx, &envoy_key, actor_id, buffer).await {
+			tracing::error!(?err, "actor event processor failed");
+			break;
+		}
+	}
+}
+
+#[tracing::instrument(skip_all, fields(%envoy_key, ?actor_id))]
+async fn dispatch_events(
+	ctx: &StandaloneCtx,
+	envoy_key: &str,
+	actor_id: Id,
+	events: Vec<protocol::EventWrapper>,
+) -> Result<()> {
+	tracing::debug!(count=?events.len(), "actor demuxer dispatch");
+
+	let res = ctx
+		.signal(pegboard::workflows::actor2::Events {
+			envoy_key: envoy_key.to_string(),
+			events,
+		})
+		.to_workflow::<pegboard::workflows::actor2::Workflow>()
+		.tag("actor_id", actor_id)
+		.graceful_not_found()
+		.send()
+		.await
+		.with_context(|| format!("failed to forward signal to actor workflow: {}", actor_id))?;
+	if res.is_none() {
+		tracing::warn!(
+			?actor_id,
+			"failed to send signal to actor workflow, likely already stopped"
+		);
+	}
+
+	Ok(())
+}