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

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

@@ -1,169 +1,19 @@
 use anyhow::*;
-use epoxy_protocol::protocol::{ReplicaId, SlotId};
+use epoxy_protocol::protocol;
 use universaldb::prelude::*;
-use vbare::OwnedVersionedData;
-
-#[derive(Debug)]
-pub struct InstanceNumberKey;
-
-impl FormalKey for InstanceNumberKey {
-	type Value = u64;
-
-	fn deserialize(&self, raw: &[u8]) -> Result<Self::Value> {
-		Ok(u64::from_be_bytes(raw.try_into()?))
-	}
-
-	fn serialize(&self, value: Self::Value) -> Result<Vec<u8>> {
-		Ok(value.to_be_bytes().to_vec())
-	}
-}
-
-impl TuplePack for InstanceNumberKey {
-	fn pack<W: std::io::Write>(
-		&self,
-		w: &mut W,
-		tuple_depth: TupleDepth,
-	) -> std::io::Result<VersionstampOffset> {
-		let t = (INSTANCE_NUMBER,);
-		t.pack(w, tuple_depth)
-	}
-}
-
-#[derive(Debug)]
-pub struct LogEntryKey {
-	pub instance_replica_id: ReplicaId,
-	pub instance_slot_id: SlotId,
-}
-
-impl LogEntryKey {
-	pub fn new(instance_replica_id: ReplicaId, instance_slot_id: SlotId) -> Self {
-		Self {
-			instance_replica_id,
-			instance_slot_id,
-		}
-	}
-}
-
-impl FormalKey for LogEntryKey {
-	type Value = epoxy_protocol::protocol::LogEntry;
-
-	fn deserialize(&self, raw: &[u8]) -> Result<Self::Value> {
-		epoxy_protocol::versioned::LogEntry::deserialize_with_embedded_version(raw)
-	}
-
-	fn serialize(&self, value: Self::Value) -> Result<Vec<u8>> {
-		epoxy_protocol::versioned::LogEntry::wrap_latest(value)
-			.serialize_with_embedded_version(epoxy_protocol::PROTOCOL_VERSION)
-	}
-}
-
-impl TuplePack for LogEntryKey {
-	fn pack<W: std::io::Write>(
-		&self,
-		w: &mut W,
-		tuple_depth: TupleDepth,
-	) -> std::io::Result<VersionstampOffset> {
-		let t = (LOG, self.instance_replica_id, self.instance_slot_id, ENTRY);
-		t.pack(w, tuple_depth)
-	}
-}
-
-impl<'de> TupleUnpack<'de> for LogEntryKey {
-	fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> {
-		let (input, (_, instance_replica_id, instance_slot_id, _)) =
-			<(usize, ReplicaId, SlotId, usize)>::unpack(input, tuple_depth)?;
-
-		let v = LogEntryKey {
-			instance_replica_id,
-			instance_slot_id,
-		};
-
-		Result::Ok((input, v))
-	}
-}
-
-#[derive(Debug)]
-pub struct KeyInstanceKey {
-	pub key: Vec<u8>,
-	pub instance_replica_id: ReplicaId,
-	pub instance_slot_id: SlotId,
-}
-
-impl KeyInstanceKey {
-	pub fn new(key: Vec<u8>, instance_replica_id: ReplicaId, instance_slot_id: SlotId) -> Self {
-		Self {
-			key,
-			instance_replica_id,
-			instance_slot_id,
-		}
-	}
-
-	pub fn subspace(key: Vec<u8>) -> KeyInstanceSubspaceKey {
-		KeyInstanceSubspaceKey::new(key)
-	}
-
-	pub fn subspace_for_all_keys() -> KeyInstanceSubspaceForAllKeysKey {
-		KeyInstanceSubspaceForAllKeysKey
-	}
-}
-
-impl FormalKey for KeyInstanceKey {
-	type Value = ();
-
-	fn deserialize(&self, _raw: &[u8]) -> Result<Self::Value> {
-		Ok(())
-	}
-
-	fn serialize(&self, _value: Self::Value) -> Result<Vec<u8>> {
-		Ok(Vec::new())
-	}
-}
-
-impl TuplePack for KeyInstanceKey {
-	fn pack<W: std::io::Write>(
-		&self,
-		w: &mut W,
-		tuple_depth: TupleDepth,
-	) -> std::io::Result<VersionstampOffset> {
-		let t = (
-			KEY_INSTANCE,
-			&self.key,
-			INSTANCE,
-			self.instance_replica_id,
-			self.instance_slot_id,
-		);
-		t.pack(w, tuple_depth)
-	}
-}
-
-impl<'de> TupleUnpack<'de> for KeyInstanceKey {
-	fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> {
-		let (input, (_, key, _, instance_replica_id, instance_slot_id)) =
-			<(usize, Vec<u8>, usize, ReplicaId, SlotId)>::unpack(input, tuple_depth)?;
-
-		let v = KeyInstanceKey {
-			key,
-			instance_replica_id,
-			instance_slot_id,
-		};
-
-		Result::Ok((input, v))
-	}
-}
 
 #[derive(Debug)]
 pub struct ConfigKey;
 
 impl FormalKey for ConfigKey {
-	type Value = epoxy_protocol::protocol::ClusterConfig;
+	type Value = protocol::ClusterConfig;
 
 	fn deserialize(&self, raw: &[u8]) -> Result<Self::Value> {
-		epoxy_protocol::versioned::ClusterConfig::deserialize_with_embedded_version(raw)
+		serde_bare::from_slice(raw).map_err(Into::into)
 	}
 
 	fn serialize(&self, value: Self::Value) -> Result<Vec<u8>> {
-		epoxy_protocol::versioned::ClusterConfig::wrap_latest(value)
-			.serialize_with_embedded_version(epoxy_protocol::PROTOCOL_VERSION)
+		serde_bare::to_vec(&value).map_err(Into::into)
 	}
 }
 
@@ -173,111 +23,6 @@ impl TuplePack for ConfigKey {
 		w: &mut W,
 		tuple_depth: TupleDepth,
 	) -> std::io::Result<VersionstampOffset> {
-		let t = (CONFIG,);
-		t.pack(w, tuple_depth)
-	}
-}
-
-pub struct KeyInstanceSubspaceKey {
-	key: Vec<u8>,
-}
-
-impl KeyInstanceSubspaceKey {
-	pub fn new(key: Vec<u8>) -> Self {
-		KeyInstanceSubspaceKey { key }
-	}
-}
-
-impl TuplePack for KeyInstanceSubspaceKey {
-	fn pack<W: std::io::Write>(
-		&self,
-		w: &mut W,
-		tuple_depth: TupleDepth,
-	) -> std::io::Result<VersionstampOffset> {
-		let t = (KEY_INSTANCE, &self.key, INSTANCE);
-		t.pack(w, tuple_depth)
-	}
-}
-
-pub struct KeyInstanceSubspaceForAllKeysKey;
-
-impl TuplePack for KeyInstanceSubspaceForAllKeysKey {
-	fn pack<W: std::io::Write>(
-		&self,
-		w: &mut W,
-		tuple_depth: TupleDepth,
-	) -> std::io::Result<VersionstampOffset> {
-		let t = (KEY_INSTANCE,);
-		t.pack(w, tuple_depth)
-	}
-}
-
-#[derive(Debug)]
-pub struct CurrentBallotKey;
-
-impl FormalKey for CurrentBallotKey {
-	type Value = epoxy_protocol::protocol::Ballot;
-
-	fn deserialize(&self, raw: &[u8]) -> Result<Self::Value> {
-		epoxy_protocol::versioned::Ballot::deserialize_with_embedded_version(raw)
-	}
-
-	fn serialize(&self, value: Self::Value) -> Result<Vec<u8>> {
-		epoxy_protocol::versioned::Ballot::wrap_latest(value)
-			.serialize_with_embedded_version(epoxy_protocol::PROTOCOL_VERSION)
-	}
-}
-
-impl TuplePack for CurrentBallotKey {
-	fn pack<W: std::io::Write>(
-		&self,
-		w: &mut W,
-		tuple_depth: TupleDepth,
-	) -> std::io::Result<VersionstampOffset> {
-		let t = (CURRENT_BALLOT,);
-		t.pack(w, tuple_depth)
-	}
-}
-
-#[derive(Debug)]
-pub struct InstanceBallotKey {
-	instance_replica_id: ReplicaId,
-	instance_slot_id: SlotId,
-}
-
-impl InstanceBallotKey {
-	pub fn new(instance_replica_id: ReplicaId, instance_slot_id: SlotId) -> Self {
-		Self {
-			instance_replica_id,
-			instance_slot_id,
-		}
-	}
-}
-
-impl FormalKey for InstanceBallotKey {
-	type Value = epoxy_protocol::protocol::Ballot;
-
-	fn deserialize(&self, raw: &[u8]) -> Result<Self::Value> {
-		epoxy_protocol::versioned::Ballot::deserialize_with_embedded_version(raw)
-	}
-
-	fn serialize(&self, value: Self::Value) -> Result<Vec<u8>> {
-		epoxy_protocol::versioned::Ballot::wrap_latest(value)
-			.serialize_with_embedded_version(epoxy_protocol::PROTOCOL_VERSION)
-	}
-}
-
-impl TuplePack for InstanceBallotKey {
-	fn pack<W: std::io::Write>(
-		&self,
-		w: &mut W,
-		tuple_depth: TupleDepth,
-	) -> std::io::Result<VersionstampOffset> {
-		let t = (
-			INSTANCE_BALLOT,
-			self.instance_replica_id,
-			self.instance_slot_id,
-		);
-		t.pack(w, tuple_depth)
+		(CONFIG,).pack(w, tuple_depth)
 	}
 }

package/packages/epoxy/src/lib.rs

@@ -5,6 +5,7 @@ pub mod errors;
 pub mod http_client;
 pub mod http_routes;
 pub mod keys;
+pub mod metrics;
 pub mod ops;
 pub mod replica;
 pub mod types;
@@ -17,9 +18,9 @@ pub fn registry() -> WorkflowResult<Registry> {
 	use workflows::*;
 
 	let mut registry = Registry::new();
+	registry.register_workflow::<backfill::Workflow>()?;
 	registry.register_workflow::<coordinator::Workflow>()?;
 	registry.register_workflow::<replica::Workflow>()?;
-	registry.register_workflow::<purger::Workflow>()?;
 
 	Ok(registry)
 }

package/packages/epoxy/src/metrics.rs

@@ -0,0 +1,118 @@
+use epoxy_protocol::protocol;
+use rivet_metrics::{BUCKETS, REGISTRY, prometheus::*};
+
+lazy_static::lazy_static! {
+	// MARK: Proposals
+	pub static ref PROPOSAL_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"epoxy_proposals_total",
+		"Total number of per-key proposal outcomes.",
+		&["result"],
+		*REGISTRY
+	).unwrap();
+
+	pub static ref PROPOSAL_DURATION: Histogram = register_histogram_with_registry!(
+		"epoxy_proposal_duration",
+		"Duration from propose to commit in seconds.",
+		BUCKETS.to_vec(),
+		*REGISTRY
+	).unwrap();
+
+	pub static ref FAST_PATH_TOTAL: IntCounter = register_int_counter_with_registry!(
+		"epoxy_fast_path_total",
+		"Total number of fast-path proposal attempts.",
+		*REGISTRY
+	).unwrap();
+
+	pub static ref SLOW_PATH_TOTAL: IntCounter = register_int_counter_with_registry!(
+		"epoxy_slow_path_total",
+		"Total number of slow-path proposal attempts.",
+		*REGISTRY
+	).unwrap();
+
+	pub static ref PREPARE_TOTAL: IntCounter = register_int_counter_with_registry!(
+		"epoxy_prepare_total",
+		"Total number of recovery Prepare phases triggered.",
+		*REGISTRY
+	).unwrap();
+
+	pub static ref PREPARE_RETRY_TOTAL: IntCounter = register_int_counter_with_registry!(
+		"epoxy_prepare_retry_total",
+		"Total number of Prepare retries caused by contention.",
+		*REGISTRY
+	).unwrap();
+
+	pub static ref BALLOT_BUMP_TOTAL: IntCounter = register_int_counter_with_registry!(
+		"epoxy_ballot_bump_total",
+		"Total number of proposal ballot bumps caused by contention.",
+		*REGISTRY
+	).unwrap();
+
+	// MARK: Changelog
+	pub static ref CHANGELOG_SIZE: IntGauge = register_int_gauge_with_registry!(
+		"epoxy_changelog_size",
+		"Current number of changelog entries on the local replica.",
+		*REGISTRY
+	).unwrap();
+
+	// MARK: HTTP request-level
+	pub static ref REQUEST_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"epoxy_requests_total",
+		"Total number of replica HTTP requests.",
+		&["request_type", "result"],
+		*REGISTRY
+	).unwrap();
+
+	pub static ref REQUEST_DURATION: HistogramVec = register_histogram_vec_with_registry!(
+		"epoxy_request_duration",
+		"Duration of replica HTTP request handling in seconds.",
+		&["request_type"],
+		BUCKETS.to_vec(),
+		*REGISTRY
+	).unwrap();
+
+	// MARK: Cluster state
+	pub static ref REPLICAS_TOTAL: IntGaugeVec = register_int_gauge_vec_with_registry!(
+		"epoxy_replicas_total",
+		"Total number of replicas by status.",
+		&["status"],
+		*REGISTRY
+	).unwrap();
+}
+
+pub fn record_proposal_result(result: &str) {
+	PROPOSAL_TOTAL.with_label_values(&[result]).inc();
+}
+
+pub fn record_ballot_bump() {
+	BALLOT_BUMP_TOTAL.inc();
+}
+
+pub fn record_prepare_retry() {
+	PREPARE_RETRY_TOTAL.inc();
+}
+
+pub fn record_changelog_append() {
+	CHANGELOG_SIZE.inc();
+}
+
+pub fn record_request(request_type: &str, result: &str, duration: std::time::Duration) {
+	REQUEST_TOTAL
+		.with_label_values(&[request_type, result])
+		.inc();
+	REQUEST_DURATION
+		.with_label_values(&[request_type])
+		.observe(duration.as_secs_f64());
+}
+
+pub fn record_replicas(config: &protocol::ClusterConfig) {
+	REPLICAS_TOTAL.reset();
+	for replica in &config.replicas {
+		REPLICAS_TOTAL
+			.with_label_values(&[match replica.status {
+				protocol::ReplicaStatus::Active => "active",
+				protocol::ReplicaStatus::Learning => "learning",
+				protocol::ReplicaStatus::Joining => "joining",
+			}])
+			.inc();
+	}
+}

package/packages/epoxy/src/ops/kv/get_local.rs

@@ -1,9 +1,8 @@
 use anyhow::*;
 use epoxy_protocol::protocol::ReplicaId;
 use gas::prelude::*;
-use universaldb::utils::{FormalKey, IsolationLevel::*};
 
-use crate::keys;
+use super::read_value;
 
 #[derive(Debug)]
 pub struct Input {
@@ -14,31 +13,23 @@ pub struct Input {
 #[derive(Debug)]
 pub struct Output {
 	pub value: Option<Vec<u8>>,
+	pub version: Option<u64>,
+	pub mutable: bool,
 }
 
 #[operation]
 pub async fn epoxy_kv_get_local(ctx: &OperationCtx, input: &Input) -> Result<Output> {
-	// Read from local KV store only
-	let kv_key = keys::keys::KvValueKey::new(input.key.clone());
-	let subspace = keys::subspace(input.replica_id);
-	let packed_key = subspace.pack(&kv_key);
+	let committed_value =
+		read_value::read_local_value(ctx, input.replica_id, input.key.clone(), false)
+			.await?
+			.value;
 
-	let value = ctx
-		.udb()?
-		.run(|tx| {
-			let packed_key = packed_key.clone();
-			let kv_key = kv_key.clone();
-			async move {
-				let value = tx.get(&packed_key, Serializable).await?;
-				if let Some(v) = value {
-					Ok(Some(kv_key.deserialize(&v)?))
-				} else {
-					Ok(None)
-				}
-			}
-		})
-		.custom_instrument(tracing::info_span!("get_local_tx"))
-		.await?;
-
-	Ok(Output { value })
+	Ok(Output {
+		value: committed_value.as_ref().map(|value| value.value.clone()),
+		version: committed_value.as_ref().map(|value| value.version),
+		mutable: committed_value
+			.as_ref()
+			.map(|value| value.mutable)
+			.unwrap_or(false),
+	})
 }

package/packages/epoxy/src/ops/kv/get_optimistic.rs

@@ -2,14 +2,20 @@ use anyhow::*;
 use epoxy_protocol::protocol::{self, ReplicaId};
 use gas::prelude::*;
 use rivet_api_builder::ApiCtx;
-use universaldb::utils::{FormalKey, IsolationLevel::*};
 
-use crate::{http_client, keys, utils};
+use crate::{
+	http_client,
+	keys::{self, CommittedValue},
+	utils,
+};
+
+use super::read_value;
 
 #[derive(Debug)]
 pub struct Input {
 	pub replica_id: ReplicaId,
 	pub key: Vec<u8>,
+	pub caching_behavior: protocol::CachingBehavior,
 }
 
 #[derive(Debug)]
@@ -26,59 +32,41 @@ pub struct Output {
 /// WARNING: This will incorrectly return `None` in the rare case that all of the nodes that have
 /// committed the value are offline.
 ///
+/// This reads committed values from the v2 `kv/{key}/value` path, falls back to the legacy
+/// `kv/{key}/committed_value` and legacy `kv/{key}/value` paths, then checks the v2
+/// `kv/{key}/cache` path for remote values.
+///
 /// This works by:
-/// 1. Attempt to read value from optimistic cache
-/// 2. If not in cache, attempt to read value locally
-/// 3. If not locally, reach out to any datacenter, then cache & return the first datacenter that has a value
+/// 1. Attempt to read the v2 committed value locally
+/// 2. If not found, fall back to the legacy committed-value path
+/// 3. If still not found, fall back to the legacy value path
+/// 4. If still not found, check the v2 optimistic cache
+/// 5. If not in cache, reach out to any datacenter, then cache and return the first datacenter that has a value
 ///
 /// This means that if the value changes, the value will be inconsistent across all datacenters --
 /// even if it has a quorum.
 ///
-/// We cannot use quorum reads for the fanout read because of the constraints of Epaxos.
+/// We cannot use quorum reads for the fanout read because the optimistic path is intentionally a
+/// best-effort lookup.
 #[operation]
 pub async fn epoxy_kv_get_optimistic(ctx: &OperationCtx, input: &Input) -> Result<Output> {
-	// Try to read locally
-	let kv_key = keys::keys::KvValueKey::new(input.key.clone());
-	let cache_key = keys::keys::KvOptimisticCacheKey::new(input.key.clone());
-	let subspace = keys::subspace(input.replica_id);
-	let packed_key = subspace.pack(&kv_key);
-	let packed_cache_key = subspace.pack(&cache_key);
-
-	let value = ctx
-		.udb()?
-		.run(|tx| {
-			let packed_key = packed_key.clone();
-			let packed_cache_key = packed_cache_key.clone();
-			let kv_key = kv_key.clone();
-			let cache_key = cache_key.clone();
-			async move {
-				let (value, cache_value) = tokio::try_join!(
-					async {
-						let v = tx.get(&packed_key, Serializable).await?;
-						if let Some(ref bytes) = v {
-							Ok(Some(kv_key.deserialize(bytes)?))
-						} else {
-							Ok(None)
-						}
-					},
-					async {
-						let v = tx.get(&packed_cache_key, Serializable).await?;
-						if let Some(ref bytes) = v {
-							Ok(Some(cache_key.deserialize(bytes)?))
-						} else {
-							Ok(None)
-						}
-					}
-				)?;
-
-				Ok(value.or(cache_value))
-			}
-		})
-		.custom_instrument(tracing::info_span!("get_optimistic_tx"))
-		.await?;
+	let local_read = read_value::read_local_value(
+		ctx,
+		input.replica_id,
+		input.key.clone(),
+		input.caching_behavior == protocol::CachingBehavior::Optimistic,
+	)
+	.await?;
+	if local_read.value.is_some() {
+		return Ok(Output {
+			value: local_read.value.map(|value| value.value),
+		});
+	}
 
-	if value.is_some() {
-		return Ok(Output { value });
+	if let Some(value) = local_read.cache_value {
+		return Ok(Output {
+			value: Some(value.value),
+		});
 	}
 
 	// Request fanout to other datacenters, return first datacenter with any non-none value
@@ -106,7 +94,10 @@ pub async fn epoxy_kv_get_optimistic(ctx: &OperationCtx, input: &Input) -> Resul
 				let request = protocol::Request {
 					from_replica_id,
 					to_replica_id: replica_id,
-					kind: protocol::RequestKind::KvGetRequest(protocol::KvGetRequest { key }),
+					kind: protocol::RequestKind::KvGetRequest(protocol::KvGetRequest {
+						key,
+						caching_behavior: input.caching_behavior.clone(),
+					}),
 				};
 
 				// Send the message and extract the KV response
@@ -125,26 +116,73 @@ pub async fn epoxy_kv_get_optimistic(ctx: &OperationCtx, input: &Input) -> Resul
 
 	for response in responses {
 		if let Some(value) = response {
-			// Cache value
-			ctx.udb()?
-				.run(|tx| {
-					let packed_cache_key = packed_cache_key.clone();
-					let cache_key = cache_key.clone();
-					let value_to_cache = value.clone();
-
-					async move {
-						let serialized = cache_key.serialize(value_to_cache)?;
-						tx.set(&packed_cache_key, &serialized);
-						Ok(())
-					}
-				})
-				.custom_instrument(tracing::info_span!("cache_value_tx"))
-				.await?;
+			let value = CommittedValue {
+				value: value.value,
+				version: value.version,
+				mutable: value.mutable,
+			};
+
+			if input.caching_behavior == protocol::CachingBehavior::Optimistic {
+				cache_fanout_value(ctx, input.replica_id, input.key.clone(), value.clone()).await?;
+			}
 
-			return Ok(Output { value: Some(value) });
+			return Ok(Output {
+				value: Some(value.value),
+			});
 		}
 	}
 
 	// No value found in any datacenter
 	Ok(Output { value: None })
 }
+
+async fn cache_fanout_value(
+	ctx: &OperationCtx,
+	replica_id: ReplicaId,
+	key: Vec<u8>,
+	value_to_cache: CommittedValue,
+) -> Result<()> {
+	ctx.udb()?
+		.run(|tx| {
+			let value_to_cache = value_to_cache.clone();
+			let key = key.clone();
+
+			async move {
+				let tx = tx.with_subspace(keys::subspace(replica_id));
+				let committed_key = keys::KvValueKey::new(key);
+				let cache_key = keys::KvOptimisticCacheKey::new(committed_key.key().to_vec());
+
+				// Skip caching if a committed value exists with an equal or newer version.
+				// This covers the race where a commit lands between the fanout read and
+				// the cache write.
+				if let Some(committed_value) = tx
+					.read_opt(
+						&committed_key,
+						universaldb::utils::IsolationLevel::Serializable,
+					)
+					.await?
+				{
+					if committed_value.version >= value_to_cache.version {
+						return Ok(());
+					}
+				}
+
+				// Skip caching if the existing cache entry is strictly newer. This
+				// prevents a slow fanout response from overwriting a fresher cache entry
+				// written by a concurrent request.
+				if let Some(existing_cache) = tx
+					.read_opt(&cache_key, universaldb::utils::IsolationLevel::Serializable)
+					.await?
+				{
+					if existing_cache.version > value_to_cache.version {
+						return Ok(());
+					}
+				}
+
+				tx.write(&cache_key, value_to_cache)?;
+				Ok(())
+			}
+		})
+		.custom_instrument(tracing::info_span!("cache_value_tx"))
+		.await
+}

package/packages/epoxy/src/ops/kv/mod.rs

@@ -1,3 +1,4 @@
 pub mod get_local;
 pub mod get_optimistic;
 pub mod purge_local;
+pub(crate) mod read_value;