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

package/packages/epoxy/src/replica/message_request.rs

@@ -1,5 +1,5 @@
 use anyhow::*;
-use epoxy_protocol::protocol::{self};
+use epoxy_protocol::protocol;
 use gas::prelude::*;
 use rivet_api_builder::prelude::*;
 
@@ -9,6 +9,14 @@ use crate::{ops, replica};
 pub async fn message_request(
 	ctx: &ApiCtx,
 	request: protocol::Request,
+) -> Result<protocol::Response> {
+	message_request_inner(ctx, request).await
+}
+
+#[tracing::instrument(skip_all)]
+async fn message_request_inner(
+	ctx: &ApiCtx,
+	request: protocol::Request,
 ) -> Result<protocol::Response> {
 	let current_replica_id = ctx.config().epoxy_replica_id();
 
@@ -20,7 +28,6 @@ pub async fn message_request(
 				"received configuration update request"
 			);
 
-			// Store the configuration
 			ctx.udb()?
 				.run(|tx| {
 					let req = req.clone();
@@ -31,16 +38,16 @@ pub async fn message_request(
 
 			protocol::ResponseKind::UpdateConfigResponse
 		}
-		protocol::RequestKind::PreAcceptRequest(req) => {
+		protocol::RequestKind::PrepareRequest(req) => {
 			let response = ctx
 				.udb()?
 				.run(|tx| {
 					let req = req.clone();
-					async move { replica::messages::pre_accept(&*tx, current_replica_id, req).await }
+					async move { replica::messages::prepare(&*tx, current_replica_id, req).await }
 				})
-				.custom_instrument(tracing::info_span!("pre_accept_tx"))
+				.custom_instrument(tracing::info_span!("prepare_tx"))
 				.await?;
-			protocol::ResponseKind::PreAcceptResponse(response)
+			protocol::ResponseKind::PrepareResponse(response)
 		}
 		protocol::RequestKind::AcceptRequest(req) => {
 			let response = ctx
@@ -54,59 +61,36 @@ pub async fn message_request(
 			protocol::ResponseKind::AcceptResponse(response)
 		}
 		protocol::RequestKind::CommitRequest(req) => {
-			// Commit and update KV store
-			ctx.udb()?
-				.run(|tx| {
-					let req = req.clone();
-					async move {
-						replica::messages::commit(&*tx, current_replica_id, req, true).await?;
-						Result::Ok(())
-					}
-				})
-				.custom_instrument(tracing::info_span!("commit_tx"))
-				.await?;
-
-			protocol::ResponseKind::CommitResponse
-		}
-		protocol::RequestKind::PrepareRequest(req) => {
 			let response = ctx
 				.udb()?
 				.run(|tx| {
 					let req = req.clone();
-					async move { replica::messages::prepare(&*tx, current_replica_id, req).await }
+					async move { replica::messages::commit(&*tx, current_replica_id, req).await }
 				})
-				.custom_instrument(tracing::info_span!("prepare_tx"))
+				.custom_instrument(tracing::info_span!("commit_tx"))
 				.await?;
-			protocol::ResponseKind::PrepareResponse(response)
+			protocol::ResponseKind::CommitResponse(response)
 		}
-		protocol::RequestKind::DownloadInstancesRequest(req) => {
-			// Handle download instances request - read from UDB and return instances
-			let instances = ctx
+		protocol::RequestKind::ChangelogReadRequest(req) => {
+			let response = ctx
 				.udb()?
 				.run(|tx| {
 					let req = req.clone();
-					async move {
-						replica::messages::download_instances(&*tx, current_replica_id, req).await
-					}
+					async move { replica::changelog::read(&*tx, current_replica_id, req).await }
 				})
-				.custom_instrument(tracing::info_span!("download_instances_tx"))
+				.custom_instrument(tracing::info_span!("changelog_read_tx"))
 				.await?;
-
-			protocol::ResponseKind::DownloadInstancesResponse(protocol::DownloadInstancesResponse {
-				instances,
-			})
+			protocol::ResponseKind::ChangelogReadResponse(response)
 		}
 		protocol::RequestKind::HealthCheckRequest => {
-			// Simple health check - just return success
 			tracing::debug!("received health check request");
 			protocol::ResponseKind::HealthCheckResponse
 		}
 		protocol::RequestKind::CoordinatorUpdateReplicaStatusRequest(req) => {
-			// Send signal to coordinator workflow
 			tracing::debug!(
 				?current_replica_id,
-				update_replica_id=?req.replica_id,
-				update_status=?req.status,
+				update_replica_id = ?req.replica_id,
+				update_status = ?req.status,
 				"received coordinator update replica status request"
 			);
 
@@ -123,11 +107,10 @@ pub async fn message_request(
 			protocol::ResponseKind::CoordinatorUpdateReplicaStatusResponse
 		}
 		protocol::RequestKind::BeginLearningRequest(req) => {
-			// Send signal to replica workflow
 			tracing::debug!(?current_replica_id, "received begin learning request");
 
 			ctx.signal(crate::workflows::replica::BeginLearning {
-				config: req.config.clone().into(),
+				config: req.config.into(),
 			})
 			.bypass_signal_from_workflow_I_KNOW_WHAT_IM_DOING()
 			.to_workflow::<crate::workflows::replica::Workflow>()
@@ -138,27 +121,29 @@ pub async fn message_request(
 			protocol::ResponseKind::BeginLearningResponse
 		}
 		protocol::RequestKind::KvGetRequest(req) => {
-			// Handle KV get request
 			let result = ctx
 				.op(ops::kv::get_local::Input {
 					replica_id: current_replica_id,
-					key: req.key.clone(),
+					key: req.key,
 				})
 				.await?;
 
 			protocol::ResponseKind::KvGetResponse(protocol::KvGetResponse {
-				value: result.value,
+				value: result.value.map(|value| protocol::CommittedValue {
+					value,
+					version: result.version.unwrap_or(0),
+					mutable: result.mutable,
+				}),
 			})
 		}
-		protocol::RequestKind::KvPurgeRequest(req) => {
-			// Handle KV purge request
+		protocol::RequestKind::KvPurgeCacheRequest(req) => {
 			ctx.op(ops::kv::purge_local::Input {
 				replica_id: current_replica_id,
-				keys: req.keys.clone(),
+				entries: req.entries,
 			})
 			.await?;
 
-			protocol::ResponseKind::KvPurgeResponse
+			protocol::ResponseKind::KvPurgeCacheResponse
 		}
 	};
 

package/packages/epoxy/src/replica/messages/accept.rs

@@ -1,50 +1,91 @@
-use anyhow::{Result, ensure};
+use anyhow::Result;
 use epoxy_protocol::protocol;
-use universaldb::Transaction;
+use universaldb::{Transaction, utils::IsolationLevel::Serializable};
 
-use crate::replica::ballot;
+use crate::{
+	keys::{self, KvAcceptedKey, KvAcceptedValue, KvBallotKey, KvValueKey},
+	replica::ballot::Ballot,
+};
 
-#[tracing::instrument(skip_all)]
+#[tracing::instrument(skip_all, fields(%replica_id, key = ?accept_req.key))]
 pub async fn accept(
 	tx: &Transaction,
 	replica_id: protocol::ReplicaId,
 	accept_req: protocol::AcceptRequest,
 ) -> Result<protocol::AcceptResponse> {
-	let protocol::Payload {
-		proposal,
-		seq,
-		deps,
-		instance,
-	} = accept_req.payload;
-
-	tracing::debug!(?replica_id, ?instance, "handling accept message");
-
-	// Validate ballot
-	let current_ballot = ballot::get_ballot(tx, replica_id).await?;
-	let validation =
-		ballot::validate_and_update_ballot_for_instance(tx, replica_id, &current_ballot, &instance)
-			.await?;
-	ensure!(
-		validation.is_valid,
-		"ballot validation failed for accept: incoming ballot {:?} is not greater than stored ballot {:?} for instance {:?} (comparison: {:?})",
-		validation.incoming_ballot,
-		validation.stored_ballot,
-		instance,
-		validation.comparison
-	);
-
-	// EPaxos Step 18
-	let log_entry = protocol::LogEntry {
-		commands: proposal.commands.clone(),
-		seq,
-		deps,
-		state: protocol::State::Accepted,
-		ballot: current_ballot,
-	};
-	crate::replica::update_log(tx, replica_id, log_entry, &instance).await?;
-
-	// EPaxos Step 19
-	Ok(protocol::AcceptResponse {
-		payload: protocol::AcceptOKPayload { proposal, instance },
-	})
+	let tx = tx.with_subspace(keys::subspace(replica_id));
+	let protocol::AcceptRequest {
+		key,
+		value,
+		ballot,
+		mutable,
+		version,
+	} = accept_req;
+
+	let value_key = KvValueKey::new(key.clone());
+	let ballot_key = KvBallotKey::new(key.clone());
+	let accepted_key = KvAcceptedKey::new(key);
+	let request_ballot = Ballot::from(ballot.clone());
+
+	let (committed_value, current_ballot, accepted_value) = tokio::try_join!(
+		tx.read_opt(&value_key, Serializable),
+		tx.read_opt(&ballot_key, Serializable),
+		tx.read_opt(&accepted_key, Serializable),
+	)?;
+
+	if let Some(committed_value) = committed_value {
+		if !committed_value.mutable || !mutable || version <= committed_value.version {
+			return Ok(protocol::AcceptResponse::AcceptResponseAlreadyCommitted(
+				protocol::AcceptResponseAlreadyCommitted {
+					value: committed_value.value,
+				},
+			));
+		}
+	}
+
+	if let Some(current_ballot) = current_ballot.map(Ballot::from) {
+		if current_ballot > request_ballot {
+			return Ok(protocol::AcceptResponse::AcceptResponseHigherBallot(
+				protocol::AcceptResponseHigherBallot {
+					ballot: current_ballot.into(),
+				},
+			));
+		}
+	}
+
+	if let Some(existing_accepted) = accepted_value {
+		let existing_ballot = Ballot::from(existing_accepted.ballot.clone());
+		if existing_ballot == request_ballot {
+			let same_value = existing_accepted.value == value;
+			let same_version = existing_accepted.version == version;
+			let same_mutability = existing_accepted.mutable == mutable;
+			if same_value && same_version && same_mutability {
+				tx.write(&ballot_key, ballot.clone())?;
+				return Ok(protocol::AcceptResponse::AcceptResponseOk(
+					protocol::AcceptResponseOk { ballot },
+				));
+			}
+
+			return Ok(protocol::AcceptResponse::AcceptResponseHigherBallot(
+				protocol::AcceptResponseHigherBallot {
+					ballot: existing_ballot.into(),
+				},
+			));
+		}
+	}
+
+	tx.write(&ballot_key, ballot.clone())?;
+	tx.write(
+		&accepted_key,
+		KvAcceptedValue {
+			value,
+			ballot: ballot.clone(),
+			version,
+			mutable,
+		},
+	)?;
+
+	Ok(protocol::AcceptResponse::AcceptResponseOk(
+		protocol::AcceptResponseOk { ballot },
+	))
 }

package/packages/epoxy/src/replica/messages/commit.rs

@@ -2,45 +2,33 @@ use anyhow::Result;
 use epoxy_protocol::protocol;
 use universaldb::Transaction;
 
-use crate::replica::ballot;
+use crate::replica::commit_kv::{self, CommitKvOutcome};
 
-// EPaxos Step 24
-#[tracing::instrument(skip_all)]
+#[tracing::instrument(skip_all, fields(%replica_id, key = ?commit_req.key))]
 pub async fn commit(
 	tx: &Transaction,
 	replica_id: protocol::ReplicaId,
 	commit_req: protocol::CommitRequest,
-	commit_to_kv: bool,
-) -> Result<()> {
-	let protocol::Payload {
-		proposal,
-		seq,
-		deps,
-		instance,
-	} = commit_req.payload;
+) -> Result<protocol::CommitResponse> {
+	let protocol::CommitRequest {
+		key,
+		value,
+		ballot,
+		mutable,
+		version,
+	} = commit_req;
 
-	tracing::debug!(?replica_id, ?instance, "handling commit message");
-
-	// EPaxos Step 24
-	let current_ballot = ballot::get_ballot(tx, replica_id).await?;
-	let log_entry = protocol::LogEntry {
-		commands: proposal.commands.clone(),
-		seq,
-		deps,
-		state: protocol::State::Committed,
-		ballot: current_ballot,
-	};
-	crate::replica::update_log(tx, replica_id, log_entry, &instance).await?;
-
-	// Commit commands if requested
-	let cmd_err = if commit_to_kv {
-		crate::replica::commit_kv::commit_kv(&*tx, replica_id, &proposal.commands).await?
-	} else {
-		tracing::debug!(?replica_id, ?instance, "skipping kv commit");
-		None
+	let response = match commit_kv::commit_kv(tx, replica_id, key, value, ballot, mutable, version)
+		.await?
+	{
+		CommitKvOutcome::Committed => protocol::CommitResponse::CommitResponseOk,
+		CommitKvOutcome::AlreadyCommitted { value, .. } => {
+			protocol::CommitResponse::CommitResponseAlreadyCommitted(
+				protocol::CommitResponseAlreadyCommitted { value },
+			)
+		}
+		CommitKvOutcome::StaleBallot { .. } => protocol::CommitResponse::CommitResponseStaleCommit,
 	};
 
-	tracing::debug!(?replica_id, ?instance, ?cmd_err, "committed");
-
-	Ok(())
+	Ok(response)
 }

package/packages/epoxy/src/replica/messages/mod.rs

@@ -1,15 +1,7 @@
 pub mod accept;
-pub mod accepted;
 pub mod commit;
-pub mod committed;
-pub mod download_instances;
-pub mod pre_accept;
 pub mod prepare;
 
 pub use accept::accept;
-pub use accepted::accepted;
 pub use commit::commit;
-pub use committed::committed;
-pub use download_instances::download_instances;
-pub use pre_accept::pre_accept;
 pub use prepare::prepare;

package/packages/epoxy/src/replica/messages/prepare.rs

@@ -1,89 +1,88 @@
 use anyhow::Result;
 use epoxy_protocol::protocol;
-use universaldb::Transaction;
-use universaldb::utils::{FormalKey, IsolationLevel::*};
+use universaldb::{Transaction, utils::IsolationLevel::Serializable};
 
-use crate::{keys, replica::ballot};
+use crate::{
+	keys::{self, KvAcceptedKey, KvBallotKey, KvValueKey},
+	replica::ballot::Ballot,
+};
 
-#[tracing::instrument(skip_all)]
+#[tracing::instrument(skip_all, fields(%replica_id, key = ?prepare_req.key))]
 pub async fn prepare(
 	tx: &Transaction,
 	replica_id: protocol::ReplicaId,
 	prepare_req: protocol::PrepareRequest,
 ) -> Result<protocol::PrepareResponse> {
-	tracing::debug!(?replica_id, "handling prepare message");
+	let tx = tx.with_subspace(keys::subspace(replica_id));
+	let protocol::PrepareRequest {
+		key,
+		ballot,
+		mutable,
+		version,
+	} = prepare_req;
 
-	let protocol::PrepareRequest { ballot, instance } = prepare_req;
+	let value_key = KvValueKey::new(key.clone());
+	let ballot_key = KvBallotKey::new(key.clone());
+	let accepted_key = KvAcceptedKey::new(key);
+	let request_ballot = Ballot::from(ballot.clone());
 
-	tracing::debug!(?ballot, ?instance, "processing prepare request");
+	let (committed_value, current_ballot, accepted_value) = tokio::try_join!(
+		tx.read_opt(&value_key, Serializable),
+		tx.read_opt(&ballot_key, Serializable),
+		tx.read_opt(&accepted_key, Serializable),
+	)?;
 
-	// Look up the current log entry for this instance
-	let subspace = keys::subspace(replica_id);
-	let log_key = keys::replica::LogEntryKey::new(instance.replica_id, instance.slot_id);
-
-	let current_entry = match tx.get(&subspace.pack(&log_key), Serializable).await? {
-		Some(bytes) => {
-			// Deserialize the existing log entry
-			Some(log_key.deserialize(&bytes)?)
+	if let Some(committed_value) = committed_value {
+		if !committed_value.mutable || !mutable || version <= committed_value.version {
+			return Ok(protocol::PrepareResponse::PrepareResponseAlreadyCommitted(
+				protocol::PrepareResponseAlreadyCommitted {
+					value: committed_value.value,
+				},
+			));
 		}
-		None => None,
-	};
-
-	// EPaxos Step 38: Validate ballot for this instance
-	let validation =
-		ballot::validate_and_update_ballot_for_instance(tx, replica_id, &ballot, &instance).await?;
-	let response = if validation.is_valid {
-		// EPaxos Step 39: Reply PrepareOK with current log entry data
-		match current_entry {
-			Some(entry) => protocol::PrepareResponse::PrepareOk(protocol::PrepareOk {
-				data: Some(protocol::PrepareResponseData {
-					commands: entry.commands,
-					seq: entry.seq,
-					deps: entry.deps,
-					state: entry.state,
-					ballot: entry.ballot.clone(),
-				}),
-				previous_ballot: entry.ballot,
-				instance,
-			}),
-			None => {
-				// No existing entry
-				let default_ballot = protocol::Ballot {
-					epoch: 0,
-					ballot: 0,
-					replica_id: instance.replica_id,
-				};
+	}
 
-				protocol::PrepareResponse::PrepareOk(protocol::PrepareOk {
-					data: None,
-					previous_ballot: default_ballot,
-					instance,
-				})
-			}
+	if let Some(current_ballot) = current_ballot.map(Ballot::from) {
+		// Equal-ballot Prepare is idempotent. This allows a proposer that has already
+		// reserved a ballot locally to include itself in the prepare quorum.
+		if request_ballot < current_ballot {
+			return Ok(protocol::PrepareResponse::PrepareResponseHigherBallot(
+				protocol::PrepareResponseHigherBallot {
+					ballot: current_ballot.into(),
+				},
+			));
 		}
-	} else {
-		// EPaxos Step 40: Ballot validation failed
-		// EPaxos Step 41: Reply NACK with highest ballot stored for this instance
-		let instance_ballot_key =
-			keys::replica::InstanceBallotKey::new(instance.replica_id, instance.slot_id);
-		let subspace = keys::subspace(replica_id);
-		let packed_key = subspace.pack(&instance_ballot_key);
+	}
 
-		let highest_ballot = match tx.get(&packed_key, Serializable).await? {
-			Some(bytes) => instance_ballot_key.deserialize(&bytes)?,
-			None => {
-				// Default ballot for the original replica
-				protocol::Ballot {
-					epoch: 0,
-					ballot: 0,
-					replica_id: instance.replica_id,
-				}
-			}
-		};
+	// Promise not to accept lower ballots by persisting the request ballot.
+	// After this write, highest_ballot == request.ballot by construction.
+	tx.write(&ballot_key, ballot.clone())?;
 
-		protocol::PrepareResponse::PrepareNack(protocol::PrepareNack { highest_ballot })
+	let (accepted_value, accepted_ballot) = match accepted_value {
+		Some(accepted_value) => {
+			let crate::keys::KvAcceptedValue {
+				value,
+				ballot,
+				version,
+				mutable,
+			} = accepted_value;
+			(
+				Some(protocol::CommittedValue {
+					value,
+					version,
+					mutable,
+				}),
+				Some(ballot),
+			)
+		}
+		None => (None, None),
 	};
 
-	tracing::debug!(?response, "prepare response generated");
-	Ok(response)
+	Ok(protocol::PrepareResponse::PrepareResponseOk(
+		protocol::PrepareResponseOk {
+			highest_ballot: ballot,
+			accepted_value,
+			accepted_ballot,
+		},
+	))
 }

package/packages/epoxy/src/replica/mod.rs

@@ -1,11 +1,6 @@
 pub mod ballot;
+pub mod changelog;
 pub mod commit_kv;
-pub mod decide_path;
-pub mod lead_consensus;
-pub mod log;
 pub mod message_request;
 pub mod messages;
 pub mod update_config;
-pub mod utils;
-
-pub use log::update_log;

package/packages/epoxy/src/replica/update_config.rs

@@ -3,7 +3,7 @@ use epoxy_protocol::protocol::{self, ReplicaId};
 use universaldb::Transaction;
 use universaldb::utils::FormalKey;
 
-use crate::keys;
+use crate::{keys, metrics};
 
 pub fn update_config(
 	tx: &Transaction,
@@ -12,6 +12,8 @@ pub fn update_config(
 ) -> Result<()> {
 	tracing::debug!("updating config");
 
+	metrics::record_replicas(&update_config_req.config);
+
 	// Store config in UDB
 	let config_key = keys::replica::ConfigKey;
 	let subspace = keys::subspace(replica_id);