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

package/packages/pegboard/src/workflows/serverless/receiver.rs

@@ -0,0 +1,87 @@
+use gas::prelude::*;
+
+use super::conn;
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct Input {
+	pub pool_wf_id: Id,
+	pub namespace_id: Id,
+	pub runner_name: String,
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct State {
+	pub is_draining: bool,
+}
+
+impl State {
+	fn new() -> Self {
+		Self { is_draining: false }
+	}
+}
+
+/// Runs alongside the connection workflow to process signals. This is required because the connection
+/// workflow cannot listen for signals while in an activity.
+#[workflow]
+pub async fn pegboard_serverless_receiver(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> {
+	ctx.activity(InitStateInput {}).await?;
+
+	let conn_wf_id = ctx
+		.workflow(conn::Input {
+			pool_wf_id: input.pool_wf_id,
+			receiver_wf_id: ctx.workflow_id(),
+			namespace_id: input.namespace_id,
+			runner_name: input.runner_name.clone(),
+		})
+		.tag("namespace_id", input.namespace_id)
+		.tag("runner_name", input.runner_name.clone())
+		.dispatch()
+		.await?;
+
+	ctx.listen::<Drain>().await?;
+
+	ctx.activity(MarkAsDrainingInput {}).await?;
+
+	// If the connection is between retries, this will be received
+	ctx.signal(conn::Drain {})
+		.to_workflow_id(conn_wf_id)
+		.send()
+		.await?;
+
+	// if the connection is currently running an outbound req, this will be received
+	ctx.msg(conn::Drain {})
+		.topic(("workflow_id", conn_wf_id))
+		.send()
+		.await?;
+
+	// Wait for connection wf to complete so this wf's state remains readable
+	ctx.workflow::<conn::Input>(conn_wf_id).output().await?;
+
+	Ok(())
+}
+
+#[derive(Debug, Serialize, Deserialize, Hash)]
+struct InitStateInput {}
+
+#[activity(InitState)]
+async fn init_state(ctx: &ActivityCtx, input: &InitStateInput) -> Result<()> {
+	let mut state = ctx.state::<Option<State>>()?;
+
+	*state = Some(State::new());
+
+	Ok(())
+}
+
+#[derive(Debug, Serialize, Deserialize, Hash)]
+struct MarkAsDrainingInput {}
+
+#[activity(MarkAsDraining)]
+async fn mark_as_draining(ctx: &ActivityCtx, input: &MarkAsDrainingInput) -> Result<()> {
+	let mut state = ctx.state::<State>()?;
+	state.is_draining = true;
+
+	Ok(())
+}
+
+#[signal("pegboard_serverless_runner_drain")]
+pub struct Drain {}

package/packages/pegboard/tests/actor_v1_pre_migration.rs

@@ -0,0 +1,77 @@
+use std::time::Duration;
+
+use gas::prelude::*;
+use pegboard::workflows::actor::AllocationOverride;
+use test_snapshot::SnapshotTestCtx;
+use universaldb::prelude::*;
+
+#[tokio::test]
+async fn actor_v1_pre_migration() {
+	let test_ctx = SnapshotTestCtx::from_snapshot_with_coordinator("pb-actor-v1-pre-migration")
+		.await
+		.unwrap();
+	let ctx = test_ctx.get_ctx(test_ctx.leader_id);
+
+	let existing_namespace = ctx
+		.op(namespace::ops::resolve_for_name_local::Input {
+			name: "default".to_string(),
+		})
+		.await
+		.unwrap()
+		.expect("default ns should exist");
+
+	let actors_res = ctx
+		.op(pegboard::ops::actor::list_for_ns::Input {
+			namespace_id: existing_namespace.namespace_id,
+			name: "test".to_string(),
+			key: None,
+			include_destroyed: true,
+			created_before: None,
+			limit: 1,
+			fetch_error: false,
+		})
+		.await
+		.unwrap();
+	let actor = actors_res
+		.actors
+		.into_iter()
+		.next()
+		.expect("actor should exist");
+
+	ctx.signal(pegboard::workflows::actor::Wake {
+		allocation_override: AllocationOverride::default(),
+	})
+	.to_workflow::<pegboard::workflows::actor::Workflow>()
+	.tag("actor_id", actor.actor_id)
+	.send()
+	.await
+	.unwrap();
+
+	tokio::time::sleep(Duration::from_secs(3)).await;
+
+	// Get workflow id
+	let workflow_id = ctx
+		.udb()
+		.unwrap()
+		.run(|tx| async move {
+			let tx = tx.with_subspace(pegboard::keys::subspace());
+
+			tx.read(
+				&pegboard::keys::actor::WorkflowIdKey::new(actor.actor_id),
+				Serializable,
+			)
+			.await
+		})
+		.await
+		.unwrap();
+
+	let wf = ctx
+		.get_workflows(vec![workflow_id])
+		.await
+		.unwrap()
+		.into_iter()
+		.next()
+		.expect("workflow should exist");
+
+	assert!(!wf.is_dead());
+}

package/packages/{actor-kv/tests/list_edge_cases.rs → pegboard/tests/kv_list_edge_cases.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_list_edge_cases() -> Result<()> {
@@ -13,15 +12,20 @@ async fn test_list_edge_cases() -> 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");
@@ -37,18 +41,23 @@ async fn test_list_edge_cases() -> 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(),
+	};
 
 	// Test 1: List when empty
 	tracing::info!("test 1: list when empty");
 	let (empty_keys, _, _) =
-		kv::list(db, actor_id, rp::KvListQuery::KvListAllQuery, false, None).await?;
+		kv::list(db, &recipient, ep::KvListQuery::KvListAllQuery, false, None).await?;
 	assert_eq!(empty_keys.len(), 0, "should return empty list");
 
 	// Test 2: Prefix that matches nothing
 	tracing::info!("test 2: prefix that matches nothing");
 	kv::put(
 		db,
-		actor_id,
+		&recipient,
 		vec![b"foo".to_vec(), b"bar".to_vec()],
 		vec![b"1".to_vec(), b"2".to_vec()],
 	)
@@ -56,8 +65,8 @@ async fn test_list_edge_cases() -> Result<()> {
 
 	let (no_match, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListPrefixQuery(rp::KvListPrefixQuery {
+		&recipient,
+		ep::KvListQuery::KvListPrefixQuery(ep::KvListPrefixQuery {
 			key: b"xyz".to_vec(),
 		}),
 		false,
@@ -74,8 +83,8 @@ async fn test_list_edge_cases() -> Result<()> {
 	tracing::info!("test 3: range where start > end");
 	let (backwards_range, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListRangeQuery(rp::KvListRangeQuery {
+		&recipient,
+		ep::KvListQuery::KvListRangeQuery(ep::KvListRangeQuery {
 			start: b"z".to_vec(),
 			end: b"a".to_vec(),
 			exclusive: false,
@@ -94,8 +103,8 @@ async fn test_list_edge_cases() -> Result<()> {
 	tracing::info!("test 4: range where start == end");
 	let (same_inclusive, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListRangeQuery(rp::KvListRangeQuery {
+		&recipient,
+		ep::KvListQuery::KvListRangeQuery(ep::KvListRangeQuery {
 			start: b"foo".to_vec(),
 			end: b"foo".to_vec(),
 			exclusive: false,
@@ -112,8 +121,8 @@ async fn test_list_edge_cases() -> Result<()> {
 
 	let (same_exclusive, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListRangeQuery(rp::KvListRangeQuery {
+		&recipient,
+		ep::KvListQuery::KvListRangeQuery(ep::KvListRangeQuery {
 			start: b"foo".to_vec(),
 			end: b"foo".to_vec(),
 			exclusive: true,
@@ -128,28 +137,28 @@ async fn test_list_edge_cases() -> Result<()> {
 		"same key exclusive range should return 0"
 	);
 
-	kv::delete_all(db, actor_id).await?;
+	kv::delete_all(db, &recipient).await?;
 
 	// Test 5: Keys with null bytes (0x00)
 	tracing::info!("test 5: keys with null bytes");
 	let null_key = vec![b'a', 0x00, b'b'];
 	kv::put(
 		db,
-		actor_id,
+		&recipient,
 		vec![null_key.clone(), b"abc".to_vec()],
 		vec![b"null_value".to_vec(), b"normal_value".to_vec()],
 	)
 	.await?;
 
-	let (null_keys, null_values, _) = kv::get(db, actor_id, vec![null_key.clone()]).await?;
+	let (null_keys, null_values, _) = kv::get(db, &recipient, vec![null_key.clone()]).await?;
 	assert_eq!(null_keys.len(), 1, "should retrieve key with null byte");
 	assert_eq!(null_values[0], b"null_value");
 
 	// Prefix query should work with null bytes
 	let (null_prefix, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListPrefixQuery(rp::KvListPrefixQuery {
+		&recipient,
+		ep::KvListQuery::KvListPrefixQuery(ep::KvListPrefixQuery {
 			key: vec![b'a', 0x00],
 		}),
 		false,
@@ -163,29 +172,29 @@ async fn test_list_edge_cases() -> Result<()> {
 	);
 	assert_eq!(null_prefix[0], null_key);
 
-	kv::delete_all(db, actor_id).await?;
+	kv::delete_all(db, &recipient).await?;
 
 	// Test 6: Keys with 0xFF bytes
 	tracing::info!("test 6: keys with 0xFF bytes");
 	let ff_key = vec![b'a', 0xFF, b'b'];
 	kv::put(
 		db,
-		actor_id,
+		&recipient,
 		vec![ff_key.clone()],
 		vec![b"ff_value".to_vec()],
 	)
 	.await?;
 
-	let (ff_keys, _, _) = kv::get(db, actor_id, vec![ff_key.clone()]).await?;
+	let (ff_keys, _, _) = kv::get(db, &recipient, vec![ff_key.clone()]).await?;
 	assert_eq!(ff_keys.len(), 1, "should retrieve key with 0xFF byte");
 
-	kv::delete_all(db, actor_id).await?;
+	kv::delete_all(db, &recipient).await?;
 
 	// Test 7: Empty prefix (should match all keys)
 	tracing::info!("test 7: empty prefix");
 	kv::put(
 		db,
-		actor_id,
+		&recipient,
 		vec![b"a".to_vec(), b"b".to_vec(), b"c".to_vec()],
 		vec![b"1".to_vec(), b"2".to_vec(), b"3".to_vec()],
 	)
@@ -193,24 +202,24 @@ async fn test_list_edge_cases() -> Result<()> {
 
 	let (empty_prefix, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListPrefixQuery(rp::KvListPrefixQuery { key: vec![] }),
+		&recipient,
+		ep::KvListQuery::KvListPrefixQuery(ep::KvListPrefixQuery { key: vec![] }),
 		false,
 		None,
 	)
 	.await?;
 	assert_eq!(empty_prefix.len(), 3, "empty prefix should match all keys");
 
-	kv::delete_all(db, actor_id).await?;
+	kv::delete_all(db, &recipient).await?;
 
 	// Test 8: Prefix longer than any stored key
 	tracing::info!("test 8: prefix longer than stored keys");
-	kv::put(db, actor_id, vec![b"ab".to_vec()], vec![b"val".to_vec()]).await?;
+	kv::put(db, &recipient, vec![b"ab".to_vec()], vec![b"val".to_vec()]).await?;
 
 	let (long_prefix, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListPrefixQuery(rp::KvListPrefixQuery {
+		&recipient,
+		ep::KvListQuery::KvListPrefixQuery(ep::KvListPrefixQuery {
 			key: b"abcdefghijk".to_vec(),
 		}),
 		false,
@@ -223,7 +232,7 @@ async fn test_list_edge_cases() -> Result<()> {
 		"prefix longer than keys should return empty"
 	);
 
-	kv::delete_all(db, actor_id).await?;
+	kv::delete_all(db, &recipient).await?;
 
 	// Test 9: Keys that differ only in last byte
 	tracing::info!("test 9: keys differing only in last byte");
@@ -239,12 +248,12 @@ async fn test_list_edge_cases() -> Result<()> {
 		b"v2".to_vec(),
 		b"vFF".to_vec(),
 	];
-	kv::put(db, actor_id, keys.clone(), values.clone()).await?;
+	kv::put(db, &recipient, keys.clone(), values.clone()).await?;
 
 	let (prefix_match, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListPrefixQuery(rp::KvListPrefixQuery {
+		&recipient,
+		ep::KvListQuery::KvListPrefixQuery(ep::KvListPrefixQuery {
 			key: b"key".to_vec(),
 		}),
 		false,
@@ -266,8 +275,8 @@ async fn test_list_edge_cases() -> Result<()> {
 	// Range from key\x00 to key\x02 inclusive should get 3 keys
 	let (byte_range, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListRangeQuery(rp::KvListRangeQuery {
+		&recipient,
+		ep::KvListQuery::KvListRangeQuery(ep::KvListRangeQuery {
 			start: b"key\x00".to_vec(),
 			end: b"key\x02".to_vec(),
 			exclusive: false,
@@ -278,13 +287,13 @@ async fn test_list_edge_cases() -> Result<()> {
 	.await?;
 	assert_eq!(byte_range.len(), 3, "byte range should get 3 keys");
 
-	kv::delete_all(db, actor_id).await?;
+	kv::delete_all(db, &recipient).await?;
 
 	// Test 10: Limit of 0
 	tracing::info!("test 10: limit of 0");
 	kv::put(
 		db,
-		actor_id,
+		&recipient,
 		vec![b"a".to_vec(), b"b".to_vec()],
 		vec![b"1".to_vec(), b"2".to_vec()],
 	)
@@ -292,8 +301,8 @@ async fn test_list_edge_cases() -> Result<()> {
 
 	let (zero_limit, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListAllQuery,
+		&recipient,
+		ep::KvListQuery::KvListAllQuery,
 		false,
 		Some(0),
 	)
@@ -304,8 +313,8 @@ async fn test_list_edge_cases() -> Result<()> {
 	tracing::info!("test 11: limit of 1");
 	let (one_limit, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListAllQuery,
+		&recipient,
+		ep::KvListQuery::KvListAllQuery,
 		false,
 		Some(1),
 	)
@@ -316,8 +325,8 @@ async fn test_list_edge_cases() -> Result<()> {
 	tracing::info!("test 12: limit larger than total");
 	let (large_limit, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListAllQuery,
+		&recipient,
+		ep::KvListQuery::KvListAllQuery,
 		false,
 		Some(1000),
 	)
@@ -328,20 +337,26 @@ async fn test_list_edge_cases() -> Result<()> {
 		"should return all keys when limit > total"
 	);
 
-	kv::delete_all(db, actor_id).await?;
+	kv::delete_all(db, &recipient).await?;
 
 	// Test 13: Reverse with limit
 	tracing::info!("test 13: reverse with limit");
 	kv::put(
 		db,
-		actor_id,
+		&recipient,
 		vec![b"a".to_vec(), b"b".to_vec(), b"c".to_vec(), b"d".to_vec()],
 		vec![b"1".to_vec(), b"2".to_vec(), b"3".to_vec(), b"4".to_vec()],
 	)
 	.await?;
 
-	let (reverse_limited, _, _) =
-		kv::list(db, actor_id, rp::KvListQuery::KvListAllQuery, true, Some(2)).await?;
+	let (reverse_limited, _, _) = kv::list(
+		db,
+		&recipient,
+		ep::KvListQuery::KvListAllQuery,
+		true,
+		Some(2),
+	)
+	.await?;
 	assert_eq!(
 		reverse_limited.len(),
 		2,
@@ -355,8 +370,8 @@ async fn test_list_edge_cases() -> Result<()> {
 	tracing::info!("test 14: prefix with reverse");
 	let (prefix_reverse, _, _) = kv::list(
 		db,
-		actor_id,
-		rp::KvListQuery::KvListPrefixQuery(rp::KvListPrefixQuery { key: vec![] }),
+		&recipient,
+		ep::KvListQuery::KvListPrefixQuery(ep::KvListPrefixQuery { key: vec![] }),
 		true,
 		None,
 	)