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

package/packages/test-snapshot-gen/src/lib.rs

@@ -0,0 +1,328 @@
+use std::collections::HashMap;
+use std::path::{Path, PathBuf};
+
+use anyhow::{Context, Result};
+use epoxy_protocol::protocol::ReplicaId;
+use gas::prelude::{TestCtx as WorkflowTestCtx, *};
+use url::Url;
+
+/// Base directory where snapshots are stored.
+fn snapshots_dir() -> PathBuf {
+	Path::new(env!("CARGO_MANIFEST_DIR")).join("snapshots")
+}
+
+/// Return the snapshot directory for a scenario.
+pub fn snapshot_dir(scenario: &str) -> PathBuf {
+	snapshots_dir().join(scenario)
+}
+
+/// Copy a snapshot's RocksDB directories into temp locations and return
+/// per-replica paths suitable for constructing test infrastructure.
+///
+/// Returns a map of `replica_id -> temp_path` where each temp_path is a
+/// copy of that replica's RocksDB data directory.
+pub fn load_snapshot(scenario: &str, test_id: uuid::Uuid) -> Result<HashMap<ReplicaId, PathBuf>> {
+	let dir = snapshot_dir(scenario);
+	load_snapshot_from(&dir, test_id)
+}
+
+/// Like `load_snapshot` but from an explicit directory.
+pub fn load_snapshot_from(
+	snapshot_dir: &Path,
+	test_id: uuid::Uuid,
+) -> Result<HashMap<ReplicaId, PathBuf>> {
+	let mut replicas = HashMap::new();
+
+	for entry in std::fs::read_dir(snapshot_dir)
+		.with_context(|| format!("failed to read snapshot dir {}", snapshot_dir.display()))?
+	{
+		let entry = entry?;
+		let name = entry.file_name().to_string_lossy().to_string();
+		if !name.starts_with("replica-") {
+			continue;
+		}
+
+		let replica_id: ReplicaId = name
+			.strip_prefix("replica-")
+			.unwrap()
+			.parse()
+			.context("invalid replica id in snapshot dir name")?;
+
+		// Copy to the standard test path format so the test infrastructure picks it up.
+		let dest = std::env::temp_dir().join(format!("rivet-test-{}-{}", test_id, replica_id));
+		if dest.exists() {
+			std::fs::remove_dir_all(&dest)?;
+		}
+		copy_dir_recursive(&entry.path(), &dest)
+			.with_context(|| format!("failed to copy snapshot for replica {replica_id}"))?;
+
+		replicas.insert(replica_id, dest);
+	}
+
+	Ok(replicas)
+}
+
+fn copy_dir_recursive(src: &Path, dst: &Path) -> Result<()> {
+	std::fs::create_dir_all(dst)?;
+	for entry in std::fs::read_dir(src)? {
+		let entry = entry?;
+		let ty = entry.file_type()?;
+		let dest_path = dst.join(entry.file_name());
+		if ty.is_dir() {
+			copy_dir_recursive(&entry.path(), &dest_path)?;
+		} else {
+			std::fs::copy(entry.path(), &dest_path)?;
+		}
+	}
+	Ok(())
+}
+
+/// Helper to build a full test cluster from a pre-existing snapshot.
+///
+/// This copies the snapshot RocksDB directories into the standard temp
+/// locations so that when `rivet_test_deps::setup_single_datacenter` runs
+/// it finds the pre-populated database.
+pub struct SnapshotTestCtx {
+	pub test_id: uuid::Uuid,
+	pub leader_id: ReplicaId,
+	pub coordinator_workflow_id: rivet_util::Id,
+	replica_metadata: HashMap<ReplicaId, ReplicaMetadata>,
+	replica_contexts: HashMap<ReplicaId, ReplicaContext>,
+}
+
+struct ReplicaMetadata {
+	api_peer_port: u16,
+	guard_port: u16,
+}
+
+struct ReplicaContext {
+	wf_ctx: WorkflowTestCtx,
+	api_server_handle: tokio::task::JoinHandle<()>,
+}
+
+impl SnapshotTestCtx {
+	/// Create a test cluster from a snapshot scenario.
+	///
+	/// Loads the latest snapshot for `scenario`, copies the RocksDB dirs
+	/// into temp, then boots the full cluster pointing at that data.
+	pub async fn from_snapshot(scenario: &str) -> Result<Self> {
+		let test_id = uuid::Uuid::new_v4();
+		let snapshot_replicas = load_snapshot(scenario, test_id)?;
+
+		let mut replica_ids: Vec<_> = snapshot_replicas.keys().copied().collect();
+		replica_ids.sort_unstable();
+
+		let leader_id = replica_ids[0];
+		let mut ctx = SnapshotTestCtx {
+			test_id,
+			leader_id,
+			coordinator_workflow_id: rivet_util::Id::new_v1(leader_id as u16),
+			replica_metadata: HashMap::new(),
+			replica_contexts: HashMap::new(),
+		};
+
+		// Assign ports.
+		for &replica_id in &replica_ids {
+			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")?;
+			ctx.replica_metadata.insert(
+				replica_id,
+				ReplicaMetadata {
+					api_peer_port,
+					guard_port,
+				},
+			);
+		}
+
+		// Start replicas pointing at the snapshot data.
+		for &replica_id in &replica_ids {
+			ctx.start_replica(replica_id).await?;
+		}
+
+		Ok(ctx)
+	}
+
+	/// Create a test cluster from a snapshot and also start the coordinator + wait for Active.
+	pub async fn from_snapshot_with_coordinator(scenario: &str) -> Result<Self> {
+		let mut ctx = Self::from_snapshot(scenario).await?;
+
+		let mut config_sub = ctx
+			.get_ctx(ctx.leader_id)
+			.subscribe::<epoxy::workflows::coordinator::ConfigChangeMessage>((
+				"replica",
+				ctx.leader_id,
+			))
+			.await?;
+
+		let coordinator_workflow_id =
+			setup_coordinator(&ctx.replica_contexts, ctx.leader_id).await?;
+		ctx.coordinator_workflow_id = coordinator_workflow_id;
+
+		tracing::info!("waiting for replicas to become ready");
+		loop {
+			let config_msg = config_sub.next().await?;
+			let all_active = config_msg
+				.config
+				.replicas
+				.iter()
+				.all(|r| r.status == epoxy::types::ReplicaStatus::Active);
+			if all_active {
+				break;
+			}
+		}
+
+		Ok(ctx)
+	}
+
+	pub fn get_ctx(&self, replica_id: ReplicaId) -> &WorkflowTestCtx {
+		&self
+			.replica_contexts
+			.get(&replica_id)
+			.expect("replica not started")
+			.wf_ctx
+	}
+
+	pub fn replica_ids(&self) -> Vec<ReplicaId> {
+		let mut ids = self.replica_metadata.keys().copied().collect::<Vec<_>>();
+		ids.sort_unstable();
+		ids
+	}
+
+	pub async fn shutdown(&mut self) -> Result<()> {
+		tokio::time::sleep(std::time::Duration::from_millis(200)).await;
+
+		let mut ids = self.replica_contexts.keys().copied().collect::<Vec<_>>();
+		ids.sort_unstable();
+
+		for replica_id in ids {
+			let mut rc = self.replica_contexts.remove(&replica_id).unwrap();
+			rc.wf_ctx.shutdown().await?;
+			rc.api_server_handle.abort();
+			let _ = (&mut rc.api_server_handle).await;
+		}
+
+		Ok(())
+	}
+
+	async fn start_replica(&mut self, replica_id: ReplicaId) -> Result<()> {
+		tracing::info!(%replica_id, "starting replica from snapshot");
+
+		let metadata = self.replica_metadata.get(&replica_id).unwrap();
+		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?;
+
+		let reg = epoxy::registry()?
+			.merge(namespace::registry()?)?
+			.merge(pegboard::registry()?)?;
+		let test_ctx = WorkflowTestCtx::new_with_deps(reg, test_deps).await?;
+
+		let api_handle = setup_api_server(
+			test_ctx.config().clone(),
+			test_ctx.pools().clone(),
+			metadata.api_peer_port,
+		)
+		.await?;
+
+		// Start the replica workflow.
+		let workflow_id = test_ctx
+			.workflow(epoxy::workflows::replica::Input {})
+			.tag("replica", replica_id)
+			.dispatch()
+			.await?;
+		tracing::info!(%workflow_id, %replica_id, "created epoxy replica");
+
+		self.replica_contexts.insert(
+			replica_id,
+			ReplicaContext {
+				wf_ctx: test_ctx,
+				api_server_handle: api_handle,
+			},
+		);
+
+		Ok(())
+	}
+
+	fn build_datacenters(
+		&self,
+	) -> Result<HashMap<String, rivet_config::config::topology::Datacenter>> {
+		let mut datacenters = HashMap::new();
+		let mut ids: Vec<_> = self.replica_metadata.keys().copied().collect();
+		ids.sort_unstable();
+
+		for &id in &ids {
+			let metadata = &self.replica_metadata[&id];
+			let name = format!("dc-{}", id);
+			datacenters.insert(
+				name.clone(),
+				rivet_config::config::topology::Datacenter {
+					name: format!("dc-{}", id),
+					datacenter_label: id as u16,
+					is_leader: 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)
+	}
+}
+
+async fn setup_coordinator(
+	replica_contexts: &HashMap<ReplicaId, ReplicaContext>,
+	leader_id: ReplicaId,
+) -> Result<rivet_util::Id> {
+	let leader_ctx = &replica_contexts
+		.get(&leader_id)
+		.expect("leader not in replica contexts")
+		.wf_ctx;
+
+	let workflow_id = leader_ctx
+		.workflow(epoxy::workflows::coordinator::Input {})
+		.tag("replica", leader_id)
+		.dispatch()
+		.await?;
+
+	let mut sub = leader_ctx
+		.subscribe::<epoxy::workflows::coordinator::ConfigChangeMessage>(("replica", leader_id))
+		.await?;
+	leader_ctx
+		.signal(epoxy::workflows::coordinator::Reconfigure {})
+		.to_workflow_id(workflow_id)
+		.send()
+		.await?;
+	sub.next().await?;
+
+	Ok(workflow_id)
+}
+
+async fn setup_api_server(
+	config: rivet_config::Config,
+	pools: rivet_pools::Pools,
+	port: u16,
+) -> Result<tokio::task::JoinHandle<()>> {
+	let addr = std::net::SocketAddr::from(([127, 0, 0, 1], port));
+
+	let router = rivet_api_builder::create_router("api-peer-test", config, pools, |router| {
+		epoxy::http_routes::mount_routes(router)
+	})
+	.await?;
+
+	let listener = tokio::net::TcpListener::bind(addr).await?;
+	let handle = tokio::spawn(async move {
+		axum::serve(listener, router).await.unwrap();
+	});
+
+	Ok(handle)
+}

package/packages/test-snapshot-gen/src/main.rs

@@ -0,0 +1,145 @@
+use std::path::{Path, PathBuf};
+
+use anyhow::{Context, Result};
+use clap::{Parser, Subcommand};
+use serde::Serialize;
+
+mod scenarios;
+mod test_cluster;
+
+/// Generate UDB snapshots for integration and migration tests.
+///
+/// Each scenario sets up a multi-replica cluster, writes state through
+/// normal APIs, then checkpoints each replica's RocksDB into a snapshot
+/// directory that can be loaded by tests.
+#[derive(Parser)]
+#[command(name = "test-snapshot-gen")]
+struct Cli {
+	#[command(subcommand)]
+	command: Command,
+}
+
+#[derive(Subcommand)]
+enum Command {
+	/// Build a single scenario's snapshot.
+	Build {
+		/// Scenario name (e.g. "epoxy-v1").
+		scenario: String,
+	},
+	/// List available scenarios.
+	List,
+}
+
+#[derive(Serialize)]
+struct SnapshotMetadata {
+	commit: String,
+	branch: String,
+	generated_at: String,
+}
+
+fn snapshot_base_dir() -> PathBuf {
+	Path::new(env!("CARGO_MANIFEST_DIR")).join("snapshots")
+}
+
+fn get_commit_hash() -> Result<String> {
+	let output = std::process::Command::new("git")
+		.args(["rev-parse", "--short", "HEAD"])
+		.output()
+		.context("failed to run git")?;
+	Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
+}
+
+fn get_branch_name() -> Result<String> {
+	let output = std::process::Command::new("git")
+		.args(["rev-parse", "--abbrev-ref", "HEAD"])
+		.output()
+		.context("failed to run git")?;
+	Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
+}
+
+fn get_iso_timestamp() -> String {
+	let secs = std::time::SystemTime::now()
+		.duration_since(std::time::UNIX_EPOCH)
+		.unwrap()
+		.as_secs();
+	// Simple ISO-ish format from epoch. Not worth adding chrono for this.
+	format!("{secs}")
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+	gas::ctx::test::setup_logging();
+
+	let cli = Cli::parse();
+	let all_scenarios = scenarios::all();
+
+	match cli.command {
+		Command::List => {
+			for s in &all_scenarios {
+				println!(
+					"{} ({} replica{})",
+					s.name(),
+					s.replica_count(),
+					if s.replica_count() == 1 { "" } else { "s" }
+				);
+			}
+		}
+		Command::Build { scenario } => {
+			let s = all_scenarios
+				.iter()
+				.find(|s| s.name() == scenario)
+				.with_context(|| format!("unknown scenario: {scenario}"))?;
+			run_scenario(s.as_ref()).await?;
+		}
+	}
+
+	Ok(())
+}
+
+async fn run_scenario(scenario: &dyn scenarios::Scenario) -> Result<()> {
+	let name = scenario.name();
+	tracing::info!(%name, "running scenario");
+
+	let scenario_dir = snapshot_base_dir().join(name);
+
+	// Clean previous snapshot if it exists.
+	if scenario_dir.exists() {
+		std::fs::remove_dir_all(&scenario_dir).context("failed to remove old snapshot")?;
+	}
+	std::fs::create_dir_all(&scenario_dir).context("failed to create snapshot directory")?;
+
+	// Build the cluster.
+	let replica_ids: Vec<u64> = (1..=scenario.replica_count() as u64).collect();
+	let mut cluster = test_cluster::TestCluster::new(&replica_ids).await?;
+
+	// Run the scenario to populate state.
+	scenario.populate(&cluster).await?;
+
+	// Checkpoint each replica.
+	for &replica_id in &replica_ids {
+		let checkpoint_path = scenario_dir.join(format!("replica-{replica_id}"));
+		tracing::info!(%replica_id, path = %checkpoint_path.display(), "checkpointing replica");
+
+		let ctx = cluster.get_ctx(replica_id);
+		ctx.udb()?
+			.checkpoint(&checkpoint_path)
+			.context("failed to checkpoint")?;
+	}
+
+	cluster.shutdown().await?;
+
+	// Write metadata.
+	let metadata = SnapshotMetadata {
+		commit: get_commit_hash()?,
+		branch: get_branch_name()?,
+		generated_at: get_iso_timestamp(),
+	};
+	let metadata_path = scenario_dir.join("metadata.json");
+	std::fs::write(&metadata_path, serde_json::to_string_pretty(&metadata)?)
+		.context("failed to write metadata")?;
+
+	tracing::info!(path = %scenario_dir.display(), "snapshot saved");
+	println!("snapshot: {}", scenario_dir.display());
+
+	Ok(())
+}

package/packages/test-snapshot-gen/src/scenarios/epoxy_keys.rs

@@ -0,0 +1,60 @@
+use anyhow::{Result, ensure};
+use async_trait::async_trait;
+use epoxy::ops::propose::{self, Command, CommandKind, Proposal, ProposalResult, SetCommand};
+
+use crate::test_cluster::TestCluster;
+
+use super::Scenario;
+
+/// Scenario that writes epoxy keys through the normal consensus path.
+/// When generated from a v1 branch, these keys live in the v1 format.
+/// Tests can load this snapshot to verify migration/backfill to v2.
+pub struct EpoxyKeys;
+
+#[async_trait(?Send)]
+impl Scenario for EpoxyKeys {
+	fn name(&self) -> &'static str {
+		"epoxy-v1"
+	}
+
+	fn replica_count(&self) -> usize {
+		2
+	}
+
+	async fn populate(&self, cluster: &TestCluster) -> Result<()> {
+		let leader_ctx = cluster.get_ctx(cluster.leader_id());
+
+		let test_keys: &[(&[u8], &[u8])] = &[
+			(b"actor:abc123", b"running"),
+			(b"actor:def456", b"stopped"),
+			(b"config:version", b"42"),
+		];
+
+		for &(key, value) in test_keys {
+			let result = leader_ctx
+				.op(propose::Input {
+					proposal: Proposal {
+						commands: vec![Command {
+							kind: CommandKind::SetCommand(SetCommand {
+								key: key.to_vec(),
+								value: Some(value.to_vec()),
+							}),
+						}],
+					},
+					mutable: false,
+					purge_cache: false,
+					target_replicas: None,
+				})
+				.await?;
+
+			ensure!(
+				matches!(result, ProposalResult::Committed),
+				"proposal failed: {result:?}"
+			);
+		}
+
+		tracing::info!(count = test_keys.len(), "wrote keys via consensus");
+
+		Ok(())
+	}
+}

package/packages/test-snapshot-gen/src/scenarios/mod.rs

@@ -0,0 +1,27 @@
+use anyhow::Result;
+use async_trait::async_trait;
+
+use crate::test_cluster::TestCluster;
+
+mod epoxy_keys;
+mod pb_actor_v1_pre_migration;
+
+#[async_trait(?Send)]
+pub trait Scenario {
+	/// Unique name for this scenario (used as directory name).
+	fn name(&self) -> &'static str;
+
+	/// Number of replicas in the cluster.
+	fn replica_count(&self) -> usize;
+
+	/// Populate state in the cluster. Called after the cluster is fully
+	/// initialized with all replicas active.
+	async fn populate(&self, cluster: &TestCluster) -> Result<()>;
+}
+
+pub fn all() -> Vec<Box<dyn Scenario>> {
+	vec![
+		Box::new(epoxy_keys::EpoxyKeys),
+		Box::new(pb_actor_v1_pre_migration::PbActorV1PreMigration),
+	]
+}

package/packages/test-snapshot-gen/src/scenarios/pb_actor_v1_pre_migration.rs

@@ -0,0 +1,56 @@
+use std::time::Duration;
+
+use anyhow::{Context, Result};
+use async_trait::async_trait;
+use gas::prelude::*;
+use rivet_types::actors::CrashPolicy;
+
+use crate::test_cluster::TestCluster;
+
+use super::Scenario;
+
+/// Scenario that creates a sleeping actor v1 before envoys were introduced. When loaded after envoys
+/// introduced the wf should not die when awoken.
+pub struct PbActorV1PreMigration;
+
+#[async_trait(?Send)]
+impl Scenario for PbActorV1PreMigration {
+	fn name(&self) -> &'static str {
+		"pb-actor-v1-pre-migration"
+	}
+
+	fn replica_count(&self) -> usize {
+		2
+	}
+
+	async fn populate(&self, cluster: &TestCluster) -> Result<()> {
+		let ctx = cluster.get_ctx(cluster.leader_id());
+
+		let existing_namespace = ctx
+			.op(namespace::ops::resolve_for_name_local::Input {
+				name: "default".to_string(),
+			})
+			.await?
+			.context("default ns should exist")?;
+
+		let actor_id = Id::new_v1(ctx.config().dc_label());
+
+		ctx.op(pegboard::ops::actor::create::Input {
+			actor_id,
+			namespace_id: existing_namespace.namespace_id,
+			name: "test".to_string(),
+			key: None,
+			runner_name_selector: "default".to_string(),
+			input: None,
+			crash_policy: CrashPolicy::Sleep,
+			forward_request: false,
+			datacenter_name: None,
+		})
+		.await?;
+
+		// Wait for wf to sleep
+		tokio::time::sleep(Duration::from_secs(5)).await;
+
+		Ok(())
+	}
+}