@rivetkit/engine 1.0.0

package/packages/epoxy/tests/common/api.rs

@@ -0,0 +1,25 @@
+use anyhow::*;
+use rivet_api_builder::create_router;
+use std::net::SocketAddr;
+
+pub async fn setup_api_server(
+	config: rivet_config::Config,
+	pools: rivet_pools::Pools,
+	port: u16,
+) -> Result<tokio::task::JoinHandle<()>> {
+	let addr = SocketAddr::from(([127, 0, 0, 1], port));
+
+	let router = create_router("api-peer-test", config, pools, |router| {
+		epoxy::http_routes::mount_routes(router)
+	})
+	.await?;
+
+	let listener = tokio::net::TcpListener::bind(addr).await?;
+	tracing::info!(?port, "test api-peer server listening");
+
+	let handle = tokio::spawn(async move {
+		axum::serve(listener, router).await.unwrap();
+	});
+
+	Ok(handle)
+}

package/packages/epoxy/tests/common/mod.rs

@@ -0,0 +1,301 @@
+use anyhow::{self, Context};
+use epoxy::types;
+use epoxy_protocol::protocol::ReplicaId;
+use gas::prelude::{TestCtx as WorkflowTestCtx, *};
+// Note: workflows module no longer exposed in tests
+use rivet_util::Id;
+use serde_json::json;
+use std::collections::HashMap;
+use url::Url;
+
+pub mod api;
+pub mod utils;
+
+#[allow(dead_code)]
+pub const THREE_REPLICAS: &[ReplicaId] = &[1, 2, 3];
+#[allow(dead_code)]
+pub const DEFAULT_REPLICA_IDS: &[ReplicaId] = &[1, 2];
+
+#[derive(Clone)]
+pub struct ReplicaMetadata {
+	pub api_peer_port: u16,
+	pub guard_port: u16,
+}
+
+pub struct TestCtx {
+	test_id: Uuid,
+
+	pub leader_id: ReplicaId,
+	pub coordinator_workflow_id: Id,
+
+	/// Metadata about each replica. This is used to build the config.
+	///
+	/// Created with add_replica.
+	replica_metadata: HashMap<ReplicaId, ReplicaMetadata>,
+
+	/// Actual context for each replica.
+	///
+	/// Must first create metadata with add_replica.
+	///
+	/// Created with start_replica & stopped with stop_replica.
+	replica_contexts: HashMap<ReplicaId, ReplicaContext>,
+}
+
+#[allow(dead_code)]
+struct ReplicaContext {
+	wf_ctx: WorkflowTestCtx,
+	api_server_handle: tokio::task::JoinHandle<()>,
+}
+
+impl TestCtx {
+	#[allow(dead_code)]
+	pub async fn new() -> anyhow::Result<Self> {
+		Self::new_with(DEFAULT_REPLICA_IDS).await
+	}
+
+	pub async fn new_with(replica_ids: &[ReplicaId]) -> anyhow::Result<Self> {
+		gas::ctx::test::setup_logging();
+
+		let leader_id = replica_ids[0];
+
+		let mut test_ctx = TestCtx {
+			test_id: Uuid::new_v4(),
+			leader_id,
+			coordinator_workflow_id: Id::new_v1(leader_id as u16),
+			replica_contexts: HashMap::new(),
+			replica_metadata: HashMap::new(),
+		};
+
+		// Add replicas. Must do this before starting since start_replica will generate a config
+		// from the replica metadata.
+		tracing::info!("adding replicas");
+		for &replica_id in replica_ids {
+			test_ctx.add_replica(replica_id).await?;
+		}
+
+		// Start replicas.
+		tracing::info!("starting replicas");
+		for &replica_id in replica_ids {
+			test_ctx.start_replica(replica_id).await?;
+		}
+
+		// Start coordinator workflow on leader
+		let mut config_sub = test_ctx
+			.get_ctx(leader_id)
+			.subscribe::<epoxy::workflows::coordinator::ConfigChangeMessage>(
+				json!({ "replica": leader_id }),
+			)
+			.await
+			.unwrap();
+
+		let coordinator_workflow_id =
+			setup_epoxy_coordinator_wf(&test_ctx.replica_contexts, leader_id).await?;
+		test_ctx.coordinator_workflow_id = coordinator_workflow_id;
+
+		tracing::info!("waiting for replicas to become ready");
+		loop {
+			let config_msg = config_sub.next().await?;
+			tracing::info!(?config_msg.config, "received config");
+			let all_active = config_msg
+				.config
+				.replicas
+				.iter()
+				.all(|r| r.status == types::ReplicaStatus::Active);
+			if all_active {
+				break;
+			}
+		}
+
+		Ok(test_ctx)
+	}
+
+	pub fn get_ctx(&self, replica_id: ReplicaId) -> &WorkflowTestCtx {
+		&self
+			.replica_contexts
+			.get(&replica_id)
+			.expect("replica not started")
+			.wf_ctx
+	}
+
+	pub async fn add_replica(&mut self, replica_id: ReplicaId) -> anyhow::Result<()> {
+		tracing::info!(?replica_id, "adding replica");
+
+		assert!(
+			!self.replica_contexts.contains_key(&replica_id),
+			"replica already exists"
+		);
+
+		let api_peer_port =
+			portpicker::pick_unused_port().context("failed to pick API peer port")?;
+		let guard_port = portpicker::pick_unused_port().context("failed to pick guard port")?;
+
+		self.replica_metadata.insert(
+			replica_id,
+			ReplicaMetadata {
+				api_peer_port,
+				guard_port,
+			},
+		);
+
+		Ok(())
+	}
+
+	pub async fn start_replica(&mut self, replica_id: ReplicaId) -> anyhow::Result<()> {
+		assert!(
+			!self.replica_contexts.contains_key(&replica_id),
+			"replica already running"
+		);
+
+		tracing::info!(?replica_id, "starting replica");
+
+		let metadata = self.replica_metadata.get(&replica_id).unwrap();
+
+		// Create test deps
+		let datacenters = self.build_datacenters()?;
+		let test_deps = rivet_test_deps::setup_single_datacenter(
+			self.test_id,
+			replica_id as u16,
+			datacenters,
+			metadata.api_peer_port,
+			metadata.guard_port,
+		)
+		.await?;
+
+		// Create test context
+		let reg = epoxy::registry()?;
+		let test_ctx = WorkflowTestCtx::new_with_deps(reg, test_deps).await?;
+
+		// Start the API server
+		let api_handle = api::setup_api_server(
+			test_ctx.config().clone(),
+			test_ctx.pools().clone(),
+			metadata.api_peer_port,
+		)
+		.await?;
+
+		// Start replica set
+		setup_replica_wf(replica_id, &test_ctx).await?;
+
+		// Create context
+		self.replica_contexts.insert(
+			replica_id,
+			ReplicaContext {
+				wf_ctx: test_ctx,
+				api_server_handle: api_handle,
+			},
+		);
+
+		Ok(())
+	}
+
+	#[allow(dead_code)]
+	pub async fn stop_replica(
+		&mut self,
+		replica_id: ReplicaId,
+		preserve_data: bool,
+	) -> anyhow::Result<()> {
+		tracing::info!(?replica_id, "stopping replica");
+
+		let mut replica_context = self.replica_contexts.remove(&replica_id).unwrap();
+
+		// Retain containers so the databases are still there when we restart this replica
+		if preserve_data {
+			replica_context
+				.wf_ctx
+				.test_deps
+				.dont_stop_docker_containers_on_drop();
+		}
+
+		// Stop the workflow worker
+		tracing::info!(replica_id, "stopping replica workflow worker");
+		replica_context.wf_ctx.shutdown().await?;
+
+		// Stop the API server
+		tracing::info!(replica_id, "stopping replica API server");
+		replica_context.api_server_handle.abort();
+		let _ = (&mut replica_context.api_server_handle).await; // Ignore the result since we aborted it
+
+		tracing::info!(replica_id, "replica stopped");
+		Ok(())
+	}
+
+	fn build_datacenters(&self) -> Result<Vec<rivet_config::config::topology::Datacenter>> {
+		// Build datacenters using our predetermined metadata
+		let mut datacenters = Vec::new();
+		let replica_ids: Vec<_> = self.replica_metadata.keys().copied().collect();
+
+		for &other_replica_id in &replica_ids {
+			let metadata = &self.replica_metadata[&other_replica_id];
+			datacenters.push(rivet_config::config::topology::Datacenter {
+				name: format!("dc-{}", other_replica_id),
+				datacenter_label: other_replica_id as u16,
+				is_leader: other_replica_id == self.leader_id,
+				peer_url: Url::parse(&format!("http://127.0.0.1:{}", metadata.api_peer_port))?,
+				public_url: Url::parse(&format!("http://127.0.0.1:{}", metadata.guard_port))?,
+				proxy_url: None,
+				valid_hosts: None,
+			});
+		}
+
+		Ok(datacenters)
+	}
+
+	#[allow(dead_code)]
+	pub fn replica_ids(&self) -> Vec<ReplicaId> {
+		self.replica_metadata.keys().cloned().collect()
+	}
+}
+
+async fn setup_replica_wf(replica_id: ReplicaId, ctx: &WorkflowTestCtx) -> anyhow::Result<Id> {
+	tracing::info!(?replica_id, "setting up epoxy replica");
+
+	// Create coordinator if does not exist
+	let workflow_id = ctx
+		.workflow(epoxy::workflows::replica::Input {})
+		.tag("replica", replica_id)
+		.dispatch()
+		.await?;
+	tracing::info!(%workflow_id, ?replica_id, "created epoxy replica");
+
+	Ok(workflow_id)
+}
+
+async fn setup_epoxy_coordinator_wf(
+	replica_contexts: &HashMap<ReplicaId, ReplicaContext>,
+	leader_id: ReplicaId,
+) -> anyhow::Result<Id> {
+	let leader_ctx = &replica_contexts
+		.get(&leader_id)
+		.expect("leader not in replica contexts")
+		.wf_ctx;
+
+	tracing::info!(
+		replica_id = leader_id,
+		"setting up epoxy coordinator for leader replica"
+	);
+
+	// Create coordinator if does not exist
+	let workflow_id = leader_ctx
+		.workflow(epoxy::workflows::coordinator::Input {})
+		.tag("replica", leader_id)
+		.dispatch()
+		.await?;
+	tracing::info!(%workflow_id, replica_id = leader_id, "created epoxy coordinator");
+
+	// Trigger reconfiguration
+	let mut sub = leader_ctx
+		.subscribe::<epoxy::workflows::coordinator::ConfigChangeMessage>(
+			json!({ "replica": leader_id }),
+		)
+		.await?;
+	leader_ctx
+		.signal(epoxy::workflows::coordinator::Reconfigure {})
+		.to_workflow_id(workflow_id)
+		.send()
+		.await?;
+	tracing::info!(%workflow_id, replica_id = leader_id, "sent reconfigure signal to epoxy coordinator");
+	sub.next().await?;
+	tracing::info!(%workflow_id, replica_id = leader_id, "reconfigure complete");
+
+	Ok(workflow_id)
+}

package/packages/epoxy/tests/common/utils.rs

@@ -0,0 +1,23 @@
+use anyhow::*;
+use epoxy::ops::propose::ProposalResult;
+use epoxy_protocol::protocol;
+
+/// Helper function to execute a command through the ops API
+#[allow(dead_code)]
+pub async fn execute_command(
+	ctx: &gas::prelude::TestCtx,
+	command: protocol::CommandKind,
+	_wait_for_propagation: bool,
+) -> Result<ProposalResult> {
+	// Call the propose operation directly
+	let result = ctx
+		.op(epoxy::ops::propose::Input {
+			proposal: protocol::Proposal {
+				commands: vec![protocol::Command { kind: command }],
+			},
+			purge_cache: false,
+		})
+		.await?;
+
+	Ok(result)
+}

package/packages/epoxy/tests/kv.rs

@@ -0,0 +1,187 @@
+mod common;
+
+use epoxy::ops::propose::{CommandError, ProposalResult};
+use epoxy_protocol::protocol;
+use gas::prelude::*;
+
+use common::{THREE_REPLICAS, TestCtx, utils::execute_command};
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_set_operations() {
+	let test_ctx = TestCtx::new_with(THREE_REPLICAS).await.unwrap();
+	let replica_id = THREE_REPLICAS[0];
+	let ctx = test_ctx.get_ctx(replica_id);
+
+	let key = b"test_key";
+
+	// First set
+	let result = execute_command(
+		ctx,
+		protocol::CommandKind::SetCommand(protocol::SetCommand {
+			key: key.to_vec(),
+			value: Some(b"value1".to_vec()),
+		}),
+		true,
+	)
+	.await
+	.unwrap();
+	assert!(matches!(result, ProposalResult::Committed));
+
+	// Note: Direct workflow state checking is no longer available through ops API
+	// Verify through KV get operation instead
+
+	// Second set
+	let result = execute_command(
+		ctx,
+		protocol::CommandKind::SetCommand(protocol::SetCommand {
+			key: key.to_vec(),
+			value: Some(b"value2".to_vec()),
+		}),
+		true,
+	)
+	.await
+	.unwrap();
+	assert!(matches!(result, ProposalResult::Committed));
+
+	// Note: Direct workflow state checking is no longer available through ops API
+	// Verify through KV get operation instead
+
+	// Third set with None (delete)
+	let result = execute_command(
+		ctx,
+		protocol::CommandKind::SetCommand(protocol::SetCommand {
+			key: key.to_vec(),
+			value: None,
+		}),
+		true,
+	)
+	.await
+	.unwrap();
+	assert!(matches!(result, ProposalResult::Committed));
+
+	// Note: Direct workflow state checking is no longer available through ops API
+	// Verify through KV get operation instead
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_check_and_set_operations() {
+	let test_ctx = TestCtx::new_with(THREE_REPLICAS).await.unwrap();
+	let replica_id = THREE_REPLICAS[0];
+	let ctx = test_ctx.get_ctx(replica_id);
+
+	let key = b"cas_test_key";
+
+	// Set initial value
+	let result = execute_command(
+		ctx,
+		protocol::CommandKind::SetCommand(protocol::SetCommand {
+			key: key.to_vec(),
+			value: Some(b"initial".to_vec()),
+		}),
+		false,
+	)
+	.await
+	.unwrap();
+	assert!(matches!(result, ProposalResult::Committed));
+
+	// Successful check-and-set
+	let result = execute_command(
+		ctx,
+		protocol::CommandKind::CheckAndSetCommand(protocol::CheckAndSetCommand {
+			key: key.to_vec(),
+			expect_one_of: vec![Some(b"initial".to_vec())],
+			new_value: Some(b"updated".to_vec()),
+		}),
+		false,
+	)
+	.await
+	.unwrap();
+	assert!(matches!(result, ProposalResult::Committed));
+
+	// Note: Direct workflow state checking is no longer available through ops API
+	// Verify through KV get operation instead
+
+	// Failed check-and-set (wrong expected value)
+	let result = execute_command(
+		ctx,
+		protocol::CommandKind::CheckAndSetCommand(protocol::CheckAndSetCommand {
+			key: key.to_vec(),
+			expect_one_of: vec![Some(b"wrong".to_vec())],
+			new_value: Some(b"should_not_apply".to_vec()),
+		}),
+		false,
+	)
+	.await
+	.unwrap();
+
+	// Verify it returns CommandError with ExpectedValueDoesNotMatch
+	match result {
+		ProposalResult::CommandError(CommandError::ExpectedValueDoesNotMatch {
+			current_value: _,
+		}) => {
+			// Success - got the expected error variant
+		}
+		_ => panic!(
+			"Expected CommandError::ExpectedValueDoesNotMatch, got {:?}",
+			result
+		),
+	}
+
+	// Note: Direct workflow state checking is no longer available through ops API
+	// Verify through KV get operation instead
+
+	// Check-and-set with None expected (should fail since value exists)
+	let result = execute_command(
+		ctx,
+		protocol::CommandKind::CheckAndSetCommand(protocol::CheckAndSetCommand {
+			key: key.to_vec(),
+			expect_one_of: vec![None],
+			new_value: Some(b"new_value".to_vec()),
+		}),
+		false,
+	)
+	.await
+	.unwrap();
+
+	match result {
+		ProposalResult::CommandError(CommandError::ExpectedValueDoesNotMatch {
+			current_value: _,
+		}) => {
+			// Success - got the expected error variant
+		}
+		_ => panic!(
+			"Expected CommandError::ExpectedValueDoesNotMatch, got {:?}",
+			result
+		),
+	}
+
+	// Delete the value
+	let result = execute_command(
+		ctx,
+		protocol::CommandKind::SetCommand(protocol::SetCommand {
+			key: key.to_vec(),
+			value: None,
+		}),
+		false,
+	)
+	.await
+	.unwrap();
+	assert!(matches!(result, ProposalResult::Committed));
+
+	// Check-and-set with None expected (should succeed now)
+	let result = execute_command(
+		ctx,
+		protocol::CommandKind::CheckAndSetCommand(protocol::CheckAndSetCommand {
+			key: key.to_vec(),
+			expect_one_of: vec![None],
+			new_value: Some(b"created".to_vec()),
+		}),
+		false,
+	)
+	.await
+	.unwrap();
+	assert!(matches!(result, ProposalResult::Committed));
+
+	// Note: Direct workflow state checking is no longer available through ops API
+	// Verify through KV get operation instead
+}

package/packages/epoxy/tests/kv_get_optimistic.rs

@@ -0,0 +1,179 @@
+mod common;
+
+use epoxy::ops::propose::ProposalResult;
+use epoxy_protocol::protocol;
+use gas::prelude::*;
+
+use common::{DEFAULT_REPLICA_IDS, THREE_REPLICAS, TestCtx, utils::execute_command};
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_kv_get_optimistic_local() {
+	let test_ctx = TestCtx::new_with(THREE_REPLICAS).await.unwrap();
+	let replica_id = test_ctx.leader_id;
+	let ctx = test_ctx.get_ctx(replica_id);
+
+	let key = b"test_optimistic_local";
+	let value = b"local_value";
+
+	// Set value locally
+	let result = execute_command(
+		ctx,
+		protocol::CommandKind::SetCommand(protocol::SetCommand {
+			key: key.to_vec(),
+			value: Some(value.to_vec()),
+		}),
+		true, // Wait for propagation to ensure value is committed
+	)
+	.await
+	.unwrap();
+	assert!(matches!(result, ProposalResult::Committed));
+
+	let output = ctx
+		.op(epoxy::ops::kv::get_optimistic::Input {
+			replica_id,
+			key: key.to_vec(),
+		})
+		.await
+		.unwrap();
+	assert_eq!(output.value, Some(value.to_vec()));
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_kv_get_optimistic_remote() {
+	// Create test context with multiple replicas
+	let test_ctx = TestCtx::new_with(THREE_REPLICAS).await.unwrap();
+
+	// Use replica 1 to write data
+	let writer_replica_id = THREE_REPLICAS[0];
+	let writer_ctx = test_ctx.get_ctx(writer_replica_id);
+
+	// Use replica 2 to read (simulating remote DC fetch)
+	let reader_replica_id = THREE_REPLICAS[1];
+	let reader_ctx = test_ctx.get_ctx(reader_replica_id);
+
+	let key = b"test_optimistic_remote";
+	let value = b"remote_value";
+
+	// Set value on replica 1
+	let result = execute_command(
+		writer_ctx,
+		protocol::CommandKind::SetCommand(protocol::SetCommand {
+			key: key.to_vec(),
+			value: Some(value.to_vec()),
+		}),
+		true, // Wait for propagation
+	)
+	.await
+	.unwrap();
+	assert!(matches!(result, ProposalResult::Committed));
+
+	let output = reader_ctx
+		.op(epoxy::ops::kv::get_optimistic::Input {
+			replica_id: reader_replica_id,
+			key: key.to_vec(),
+		})
+		.await
+		.unwrap();
+	assert_eq!(output.value, Some(value.to_vec()));
+
+	// Second read from replica 2 - should now be cached
+	let output2 = reader_ctx
+		.op(epoxy::ops::kv::get_optimistic::Input {
+			replica_id: reader_replica_id,
+			key: key.to_vec(),
+		})
+		.await
+		.unwrap();
+	assert_eq!(output2.value, Some(value.to_vec()));
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_kv_get_optimistic_not_found() {
+	let test_ctx = TestCtx::new().await.unwrap();
+	let replica_id = DEFAULT_REPLICA_IDS[0];
+	let ctx = test_ctx.get_ctx(replica_id);
+
+	let key = b"nonexistent_key";
+
+	// Read a key that doesn't exist
+	let output = ctx
+		.op(epoxy::ops::kv::get_optimistic::Input {
+			replica_id,
+			key: key.to_vec(),
+		})
+		.await
+		.unwrap();
+	assert_eq!(output.value, None);
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_kv_get_optimistic_cache_persistence() {
+	// Create test context with multiple replicas
+	let test_ctx = TestCtx::new().await.unwrap();
+
+	let writer_replica_id = DEFAULT_REPLICA_IDS[0];
+	let writer_ctx = test_ctx.get_ctx(writer_replica_id);
+
+	let reader_replica_id = DEFAULT_REPLICA_IDS[1];
+	let reader_ctx = test_ctx.get_ctx(reader_replica_id);
+
+	let key = b"test_cache_persistence";
+	let initial_value = b"initial";
+	let updated_value = b"updated";
+
+	// Set initial value on replica 1
+	let result = execute_command(
+		writer_ctx,
+		protocol::CommandKind::SetCommand(protocol::SetCommand {
+			key: key.to_vec(),
+			value: Some(initial_value.to_vec()),
+		}),
+		true,
+	)
+	.await
+	.unwrap();
+	assert!(matches!(result, ProposalResult::Committed));
+
+	// Wait for propagation
+	tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+
+	// Read from replica 2 - this will cache the initial value
+	let output = reader_ctx
+		.op(epoxy::ops::kv::get_optimistic::Input {
+			replica_id: reader_replica_id,
+			key: key.to_vec(),
+		})
+		.await
+		.unwrap();
+	assert_eq!(output.value, Some(initial_value.to_vec()));
+
+	// Update value on replica 1
+	let result = execute_command(
+		writer_ctx,
+		protocol::CommandKind::SetCommand(protocol::SetCommand {
+			key: key.to_vec(),
+			value: Some(updated_value.to_vec()),
+		}),
+		true,
+	)
+	.await
+	.unwrap();
+	assert!(matches!(result, ProposalResult::Committed));
+
+	// Wait for propagation
+	tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+
+	// Read from replica 2 again
+	// Since the value has been propagated to replica 2, it will now find it locally
+	// and return the updated value (not from cache)
+	let output2 = reader_ctx
+		.op(epoxy::ops::kv::get_optimistic::Input {
+			replica_id: reader_replica_id,
+			key: key.to_vec(),
+		})
+		.await
+		.unwrap();
+
+	// The value should be updated since it's now available locally on replica 2
+	assert_eq!(output2.value, Some(updated_value.to_vec()));
+}