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

package/sdks/rust/epoxy-protocol/src/protocol.rs

@@ -0,0 +1,128 @@
+use serde::{Deserialize, Serialize};
+
+use crate::generated::v2 as raw;
+
+pub type ReplicaId = raw::ReplicaId;
+pub type ReplicaStatus = raw::ReplicaStatus;
+pub type ReplicaConfig = raw::ReplicaConfig;
+pub type ClusterConfig = raw::ClusterConfig;
+pub type Ballot = raw::Ballot;
+pub type CommittedValue = raw::CommittedValue;
+
+pub type UpdateConfigRequest = raw::UpdateConfigRequest;
+pub type UpdateConfigResponse = raw::UpdateConfigResponse;
+
+pub type PrepareRequest = raw::PrepareRequest;
+pub type PrepareResponse = raw::PrepareResponse;
+pub type PrepareResponseOk = raw::PrepareResponseOk;
+pub type PrepareResponseAlreadyCommitted = raw::PrepareResponseAlreadyCommitted;
+pub type PrepareResponseHigherBallot = raw::PrepareResponseHigherBallot;
+
+#[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone, Hash)]
+pub struct AcceptRequest {
+	pub key: Vec<u8>,
+	pub value: Vec<u8>,
+	pub ballot: Ballot,
+	pub mutable: bool,
+	pub version: u64,
+}
+
+#[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone, Hash)]
+pub struct AcceptResponseOk {
+	pub ballot: Ballot,
+}
+
+#[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone, Hash)]
+pub struct AcceptResponseAlreadyCommitted {
+	pub value: Vec<u8>,
+}
+
+#[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone, Hash)]
+pub struct AcceptResponseHigherBallot {
+	pub ballot: Ballot,
+}
+
+#[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone, Hash)]
+pub enum AcceptResponse {
+	AcceptResponseOk(AcceptResponseOk),
+	AcceptResponseAlreadyCommitted(AcceptResponseAlreadyCommitted),
+	AcceptResponseHigherBallot(AcceptResponseHigherBallot),
+}
+
+pub type CommitRequest = raw::CommitRequest;
+pub type CommitResponse = raw::CommitResponse;
+pub type CommitResponseOk = raw::CommitResponseOk;
+pub type CommitResponseAlreadyCommitted = raw::CommitResponseAlreadyCommitted;
+pub type CommitResponseStaleCommit = raw::CommitResponseStaleCommit;
+
+pub type CachingBehavior = raw::CachingBehavior;
+pub type KvGetRequest = raw::KvGetRequest;
+pub type KvGetResponse = raw::KvGetResponse;
+pub type KvPurgeCacheEntry = raw::KvPurgeCacheEntry;
+pub type KvPurgeCacheRequest = raw::KvPurgeCacheRequest;
+pub type KvPurgeCacheResponse = raw::KvPurgeCacheResponse;
+
+pub type ChangelogReadRequest = raw::ChangelogReadRequest;
+
+#[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone, Hash)]
+pub struct ChangelogEntry {
+	pub key: Vec<u8>,
+	pub value: Vec<u8>,
+	#[serde(default)]
+	pub version: u64,
+	#[serde(default)]
+	pub mutable: bool,
+}
+
+#[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone, Hash)]
+pub struct ChangelogReadResponse {
+	pub entries: Vec<ChangelogEntry>,
+	pub last_versionstamp: Vec<u8>,
+}
+
+pub type HealthCheckRequest = raw::HealthCheckRequest;
+pub type HealthCheckResponse = raw::HealthCheckResponse;
+pub type CoordinatorUpdateReplicaStatusRequest = raw::CoordinatorUpdateReplicaStatusRequest;
+pub type CoordinatorUpdateReplicaStatusResponse = raw::CoordinatorUpdateReplicaStatusResponse;
+pub type BeginLearningRequest = raw::BeginLearningRequest;
+pub type BeginLearningResponse = raw::BeginLearningResponse;
+
+#[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone, Hash)]
+pub enum RequestKind {
+	UpdateConfigRequest(UpdateConfigRequest),
+	PrepareRequest(PrepareRequest),
+	AcceptRequest(AcceptRequest),
+	CommitRequest(CommitRequest),
+	ChangelogReadRequest(ChangelogReadRequest),
+	HealthCheckRequest,
+	CoordinatorUpdateReplicaStatusRequest(CoordinatorUpdateReplicaStatusRequest),
+	BeginLearningRequest(BeginLearningRequest),
+	KvGetRequest(KvGetRequest),
+	KvPurgeCacheRequest(KvPurgeCacheRequest),
+}
+
+#[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone, Hash)]
+pub struct Request {
+	pub from_replica_id: ReplicaId,
+	pub to_replica_id: ReplicaId,
+	pub kind: RequestKind,
+}
+
+#[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone, Hash)]
+pub enum ResponseKind {
+	UpdateConfigResponse,
+	PrepareResponse(PrepareResponse),
+	AcceptResponse(AcceptResponse),
+	CommitResponse(CommitResponse),
+	ChangelogReadResponse(ChangelogReadResponse),
+	HealthCheckResponse,
+	CoordinatorUpdateReplicaStatusResponse,
+	BeginLearningResponse,
+	KvGetResponse(KvGetResponse),
+	KvPurgeCacheResponse,
+}
+
+#[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone, Hash)]
+pub struct Response {
+	pub kind: ResponseKind,
+}

package/sdks/rust/test-envoy/Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "rivet-test-envoy"
+version.workspace = true
+authors.workspace = true
+license.workspace = true
+edition.workspace = true
+description = "Rust test envoy process and harness for pegboard actor testing"
+
+[[bin]]
+name = "rivet-test-envoy"
+path = "src/main.rs"
+
+[dependencies]
+anyhow.workspace = true
+async-stream.workspace = true
+axum.workspace = true
+reqwest.workspace = true
+rivet-envoy-client.workspace = true
+rivet-envoy-protocol.workspace = true
+serde_json.workspace = true
+tokio.workspace = true
+tracing-subscriber = { workspace = true, features = ["env-filter"] }
+tracing.workspace = true

package/sdks/rust/test-envoy/Dockerfile

@@ -0,0 +1,22 @@
+FROM rust:1.88-slim AS builder
+
+WORKDIR /app
+
+COPY Cargo.toml Cargo.lock ./
+COPY engine ./engine
+
+RUN cargo build --release -p rivet-test-envoy
+
+FROM debian:bookworm-slim
+
+RUN apt-get update \
+	&& apt-get install -y --no-install-recommends ca-certificates curl \
+	&& rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+COPY --from=builder /app/target/release/rivet-test-envoy /usr/local/bin/rivet-test-envoy
+
+EXPOSE 5050
+
+CMD ["rivet-test-envoy"]

package/sdks/rust/test-envoy/src/behaviors.rs

@@ -0,0 +1,141 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
+
+use rivet_envoy_client::config::{
+	BoxFuture, EnvoyCallbacks, HttpRequest, HttpResponse, WebSocketHandler,
+};
+use rivet_envoy_client::handle::EnvoyHandle;
+use rivet_envoy_protocol as protocol;
+
+/// Default test callbacks that handle HTTP ping and WebSocket echo.
+#[derive(Default)]
+pub struct DefaultTestCallbacks {
+	pub is_shutdown: Arc<AtomicBool>,
+}
+
+impl EnvoyCallbacks for DefaultTestCallbacks {
+	fn on_actor_start(
+		&self,
+		_handle: EnvoyHandle,
+		actor_id: String,
+		generation: u32,
+		_config: protocol::ActorConfig,
+		_preloaded_kv: Option<protocol::PreloadedKv>,
+	) -> BoxFuture<anyhow::Result<()>> {
+		Box::pin(async move {
+			tracing::info!(%actor_id, generation, "actor started");
+			Ok(())
+		})
+	}
+
+	fn on_actor_stop(
+		&self,
+		_handle: EnvoyHandle,
+		actor_id: String,
+		generation: u32,
+		reason: protocol::StopActorReason,
+	) -> BoxFuture<anyhow::Result<()>> {
+		Box::pin(async move {
+			tracing::info!(%actor_id, generation, ?reason, "actor stopped");
+			Ok(())
+		})
+	}
+
+	fn on_shutdown(&self) {
+		tracing::info!("envoy shutdown");
+
+		self.is_shutdown.store(true, Ordering::SeqCst);
+	}
+
+	fn fetch(
+		&self,
+		_handle: EnvoyHandle,
+		actor_id: String,
+		_gateway_id: protocol::GatewayId,
+		_request_id: protocol::RequestId,
+		request: HttpRequest,
+	) -> BoxFuture<anyhow::Result<HttpResponse>> {
+		Box::pin(async move {
+			tracing::debug!(%actor_id, method = %request.method, path = %request.path, "handling fetch");
+
+			match request.path.as_str() {
+				"/ping" => {
+					let body = serde_json::to_vec(&serde_json::json!({
+						"actorId": actor_id,
+						"status": "ok",
+						"timestamp": std::time::SystemTime::now()
+							.duration_since(std::time::UNIX_EPOCH)
+							.unwrap_or_default()
+							.as_millis() as i64,
+					}))?;
+
+					let mut headers = HashMap::new();
+					headers.insert("content-type".to_string(), "application/json".to_string());
+					headers.insert("content-length".to_string(), body.len().to_string());
+
+					Ok(HttpResponse {
+						status: 200,
+						headers,
+						body: Some(body),
+						body_stream: None,
+					})
+				}
+				_ => {
+					let body = b"not found".to_vec();
+					let mut headers = HashMap::new();
+					headers.insert("content-length".to_string(), body.len().to_string());
+
+					Ok(HttpResponse {
+						status: 404,
+						headers,
+						body: Some(body),
+						body_stream: None,
+					})
+				}
+			}
+		})
+	}
+
+	fn websocket(
+		&self,
+		_handle: EnvoyHandle,
+		actor_id: String,
+		_gateway_id: protocol::GatewayId,
+		_request_id: protocol::RequestId,
+		_request: HttpRequest,
+		_path: String,
+		_headers: HashMap<String, String>,
+		_is_hibernatable: bool,
+		_is_restoring_hibernatable: bool,
+		_sender: rivet_envoy_client::config::WebSocketSender,
+	) -> BoxFuture<anyhow::Result<WebSocketHandler>> {
+		Box::pin(async move {
+			tracing::debug!(%actor_id, "handling websocket");
+			Ok(WebSocketHandler {
+				on_message: Box::new(move |msg| {
+					let text = format!("Echo: {}", String::from_utf8_lossy(&msg.data));
+					tracing::debug!(echo = %text, "echoing websocket message");
+					msg.sender.send_text(&text);
+					Box::pin(async {})
+				}),
+				on_close: Box::new(|code, reason| {
+					Box::pin(async move {
+						tracing::debug!(code, %reason, "websocket closed");
+					})
+				}),
+				on_open: None,
+			})
+		})
+	}
+
+	fn can_hibernate(
+		&self,
+		_actor_id: &str,
+		_gateway_id: &protocol::GatewayId,
+		_request_id: &protocol::RequestId,
+		_request: &HttpRequest,
+	) -> bool {
+		false
+	}
+}

package/sdks/rust/test-envoy/src/lib.rs

@@ -0,0 +1,11 @@
+pub mod behaviors;
+mod server;
+
+pub use rivet_envoy_client::config::{
+	BoxFuture, EnvoyCallbacks, EnvoyConfig, HttpRequest, HttpResponse, PrepopulatedActor,
+	ResponseChunk, WebSocketHandler, WebSocketMessage,
+};
+pub use rivet_envoy_client::envoy::{start_envoy, start_envoy_sync};
+pub use rivet_envoy_client::handle::EnvoyHandle;
+pub use rivet_envoy_client::protocol;
+pub use server::run_from_env;

package/sdks/rust/test-envoy/src/main.rs

@@ -0,0 +1,4 @@
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+	rivet_test_envoy::run_from_env().await
+}

package/sdks/rust/test-envoy/src/server.rs

@@ -0,0 +1,269 @@
+use std::sync::{
+	Arc,
+	atomic::{AtomicBool, Ordering},
+};
+
+use anyhow::{Context, Result};
+use async_stream::stream;
+use axum::{
+	Router,
+	body::Bytes,
+	extract::State,
+	response::{
+		IntoResponse, Json, Sse,
+		sse::{Event, KeepAlive},
+	},
+	routing::{get, post},
+};
+use rivet_envoy_protocol as protocol;
+use serde_json::json;
+use std::convert::Infallible;
+use std::time::Duration;
+use tokio::net::TcpListener;
+use tracing_subscriber::EnvFilter;
+
+use crate::behaviors::DefaultTestCallbacks;
+use rivet_envoy_client::config::EnvoyConfig;
+use rivet_envoy_client::envoy::start_envoy_sync;
+use rivet_envoy_client::handle::EnvoyHandle;
+
+#[derive(Clone)]
+struct Settings {
+	internal_server_port: u16,
+	namespace: String,
+	pool_name: String,
+	envoy_version: u32,
+	endpoint: String,
+	token: String,
+	autostart_server: bool,
+	autostart_envoy: bool,
+	autoconfigure_serverless: bool,
+}
+
+impl Settings {
+	fn from_env() -> Self {
+		Self {
+			internal_server_port: std::env::var("INTERNAL_SERVER_PORT")
+				.ok()
+				.and_then(|value| value.parse().ok())
+				.unwrap_or(5051),
+			namespace: std::env::var("RIVET_NAMESPACE").unwrap_or_else(|_| "default".to_string()),
+			pool_name: std::env::var("RIVET_POOL_NAME")
+				.unwrap_or_else(|_| "test-envoy".to_string()),
+			envoy_version: std::env::var("RIVET_ENVOY_VERSION")
+				.ok()
+				.and_then(|value| value.parse().ok())
+				.unwrap_or(1),
+			endpoint: std::env::var("RIVET_ENDPOINT")
+				.unwrap_or_else(|_| "http://127.0.0.1:6420".to_string()),
+			token: std::env::var("RIVET_TOKEN").unwrap_or_else(|_| "dev".to_string()),
+			autostart_server: read_bool_env("AUTOSTART_SERVER", true),
+			autostart_envoy: read_bool_env("AUTOSTART_ENVOY", false),
+			autoconfigure_serverless: read_bool_env("AUTOCONFIGURE_SERVERLESS", true),
+		}
+	}
+}
+
+#[derive(Clone)]
+struct AppState {
+	settings: Settings,
+	envoy_handle: Arc<tokio::sync::Mutex<Option<EnvoyHandle>>>,
+}
+
+pub async fn run_from_env() -> Result<()> {
+	init_tracing();
+
+	let settings = Settings::from_env();
+	let state = AppState {
+		settings: settings.clone(),
+		envoy_handle: Arc::new(tokio::sync::Mutex::new(None)),
+	};
+
+	if settings.autostart_envoy {
+		let (handle, _) = create_envoy(&settings);
+		*state.envoy_handle.lock().await = Some(handle);
+	} else if settings.autoconfigure_serverless {
+		auto_configure_serverless(&settings).await?;
+	}
+
+	let server = if settings.autostart_server {
+		Some(tokio::spawn(run_http_server(state.clone())))
+	} else {
+		None
+	};
+
+	install_signal_handlers();
+
+	if let Some(server) = server {
+		server.await.context("http server task failed")??;
+	} else if settings.autostart_envoy {
+		std::future::pending::<()>().await;
+	}
+
+	Ok(())
+}
+
+async fn run_http_server(state: AppState) -> Result<()> {
+	let app = Router::new()
+		.route("/health", get(health))
+		.route("/shutdown", get(shutdown))
+		.route("/api/rivet/start", post(start_serverless))
+		.route("/api/rivet/metadata", get(metadata))
+		.with_state(state.clone());
+
+	let addr = format!("0.0.0.0:{}", state.settings.internal_server_port);
+	let listener = TcpListener::bind(&addr)
+		.await
+		.with_context(|| format!("failed to bind {addr}"))?;
+
+	tracing::info!(
+		port = state.settings.internal_server_port,
+		"internal http server listening"
+	);
+
+	axum::serve(listener, app)
+		.await
+		.context("http server failed")
+}
+
+async fn health() -> &'static str {
+	"ok"
+}
+
+async fn shutdown(State(state): State<AppState>) -> &'static str {
+	if let Some(handle) = state.envoy_handle.lock().await.as_ref() {
+		handle.shutdown(false);
+	}
+	"ok"
+}
+
+async fn metadata() -> Json<serde_json::Value> {
+	Json(json!({
+		"runtime": "rivetkit",
+		"version": "1",
+		"envoyProtocolVersion": protocol::PROTOCOL_VERSION,
+	}))
+}
+
+async fn start_serverless(State(state): State<AppState>, body: Bytes) -> impl IntoResponse {
+	tracing::info!("received serverless start request");
+
+	let (handle, is_shutdown) = create_envoy(&state.settings);
+
+	// Inject the serverless start payload
+	let handle_clone = handle.clone();
+	tokio::spawn(async move {
+		if let Err(err) = handle_clone.start_serverless_actor(&body).await {
+			tracing::error!(?err, "failed to inject serverless start payload");
+		}
+	});
+
+	*state.envoy_handle.lock().await = Some(handle);
+
+	let stream = stream! {
+		let mut interval = tokio::time::interval(Duration::from_secs(1));
+		loop {
+			interval.tick().await;
+
+			if is_shutdown.load(Ordering::SeqCst) {
+				tracing::debug!("envoy shutdown, aborting SSE");
+				return;
+			}
+
+			yield Ok::<Event, Infallible>(Event::default().event("ping").data(""));
+		}
+	};
+
+	Sse::new(stream)
+		.keep_alive(KeepAlive::default())
+		.into_response()
+}
+
+fn create_envoy(settings: &Settings) -> (EnvoyHandle, Arc<AtomicBool>) {
+	let cbs = DefaultTestCallbacks::default();
+	let is_shutdown = cbs.is_shutdown.clone();
+
+	let config = EnvoyConfig {
+		version: settings.envoy_version,
+		endpoint: settings.endpoint.clone(),
+		token: Some(settings.token.clone()),
+		namespace: settings.namespace.clone(),
+		pool_name: settings.pool_name.clone(),
+		prepopulate_actor_names: std::collections::HashMap::new(),
+		metadata: None,
+		not_global: false,
+		debug_latency_ms: None,
+		callbacks: Arc::new(cbs),
+	};
+
+	(start_envoy_sync(config), is_shutdown)
+}
+
+async fn auto_configure_serverless(settings: &Settings) -> Result<()> {
+	tracing::info!("configuring serverless");
+
+	let client = reqwest::Client::new();
+	let url = format!(
+		"{}/runner-configs/{}?namespace={}",
+		settings.endpoint.trim_end_matches('/'),
+		settings.pool_name,
+		settings.namespace,
+	);
+	let body = json!({
+		"datacenters": {
+			"default": {
+				"serverless": {
+					"url": format!("http://localhost:{}/api/rivet", settings.internal_server_port),
+					"request_lifespan": 300,
+					"max_concurrent_actors": 10000,
+					"max_runners": 10000,
+					"slots_per_runner": 1
+				}
+			}
+		}
+	});
+
+	let response = client
+		.put(url)
+		.bearer_auth(&settings.token)
+		.json(&body)
+		.send()
+		.await
+		.context("failed to upsert serverless config")?;
+
+	if !response.status().is_success() {
+		let status = response.status();
+		let text = response.text().await.unwrap_or_default();
+		anyhow::bail!("serverless config request failed: {}: {}", status, text);
+	}
+
+	Ok(())
+}
+
+fn init_tracing() {
+	let filter = EnvFilter::try_from_default_env()
+		.unwrap_or_else(|_| EnvFilter::new("info,rivet_test_envoy=debug,rivet_envoy_client=debug"));
+
+	let _ = tracing_subscriber::fmt()
+		.with_env_filter(filter)
+		.with_target(false)
+		.with_ansi(true)
+		.try_init();
+}
+
+fn install_signal_handlers() {
+	tokio::spawn(async {
+		if tokio::signal::ctrl_c().await.is_ok() {
+			tracing::debug!("received stop signal, force exiting in 3s");
+			tokio::time::sleep(Duration::from_secs(3)).await;
+			std::process::exit(0);
+		}
+	});
+}
+
+fn read_bool_env(name: &str, default: bool) -> bool {
+	match std::env::var(name) {
+		Ok(value) => value == "1",
+		Err(_) => default,
+	}
+}

package/sdks/schemas/README.md

@@ -2,5 +2,4 @@
 
 These are version BARE files that are used for codegen.
 
-`sdks/rust/runner-protocol/build.rs` generates `sdks/schemas/runner-protocol/`.
-
+`packages/runner-protocol/build.rs` generates `sdks/schemas/runner-protocol/`.

package/sdks/schemas/data/namespace.runner_config.v3.bare

@@ -0,0 +1,24 @@
+type Json str
+
+type Serverless struct {
+	url: str
+	headers: map<str><str>
+	request_lifespan: u32
+	slots_per_runner: u32
+	min_runners: u32
+	max_runners: u32
+	runners_margin: u32
+}
+
+type Normal void
+
+type RunnerConfigKind union {
+	Serverless |
+	Normal
+}
+
+type RunnerConfig struct {
+	kind: RunnerConfigKind
+	metadata: optional<Json>
+	drain_on_version_upgrade: bool
+}

package/sdks/schemas/data/namespace.runner_config.v4.bare

@@ -0,0 +1,25 @@
+type Json str
+
+type Serverless struct {
+	url: str
+	headers: map<str><str>
+	request_lifespan: u32
+	slots_per_runner: u32
+	min_runners: u32
+	max_runners: u32
+	runners_margin: u32
+	metadata_poll_interval: optional<u64>
+}
+
+type Normal void
+
+type RunnerConfigKind union {
+	Serverless |
+	Normal
+}
+
+type RunnerConfig struct {
+	kind: RunnerConfigKind
+	metadata: optional<Json>
+	drain_on_version_upgrade: bool
+}

package/sdks/schemas/data/namespace.runner_config.v5.bare

@@ -0,0 +1,26 @@
+type Json str
+
+type Serverless struct {
+	url: str
+	headers: map<str><str>
+	request_lifespan: u32
+	max_concurrent_actors: u64
+	slots_per_runner: u32
+	min_runners: u32
+	max_runners: u32
+	runners_margin: u32
+	metadata_poll_interval: optional<u64>
+}
+
+type Normal void
+
+type RunnerConfigKind union {
+	Serverless |
+	Normal
+}
+
+type RunnerConfig struct {
+	kind: RunnerConfigKind
+	metadata: optional<Json>
+	drain_on_version_upgrade: bool
+}

package/sdks/schemas/data/pegboard.namespace.runner_alloc_idx.v2.bare

@@ -0,0 +1,8 @@
+type Id data
+
+type Data struct {
+	workflow_id: Id
+	remaining_slots: u32
+	total_slots: u32
+	protocol_version: u16
+}