fix(epoxy): fix replica recovery to skip old ballots (#3295)

* chore(guard): add metadata to rate limit error

* chore(tracing-utils): add tracing utils for custom instrument

* feat(epoxy): add ability to get & set coordinator state

* fix(epoxy): fix replica recovery to skip old ballots

Author: NathanFlurry

engine/packages/epoxy/src/consts.rs

@@ -28,10 +28,10 @@ pub const REQUEST_TIMEOUT: Duration = Duration::from_secs(10);
 ///
 /// Our current impl is incredily slow since we do not parallelize downloads.
 ///
-/// Assume we have 10M keys with a 50k chunk size. This results in 200 round trips.
+/// Assume we have 10M keys with a 5k chunk size. This results in 2000 round trips.
 ///
-/// If we have a latency of 200 ms (worst case), this takes 40 seconds total to propagate.
-pub const DOWNLOAD_INSTANCE_COUNT: u64 = 50_000;
+/// If we have a latency of 200 ms (worst case), this takes 400 seconds total to propagate.
+pub const DOWNLOAD_INSTANCE_COUNT: u64 = 5_000;
 
 /// Number of keys to recover in a single chunk during the recovery process.
 ///

engine/packages/epoxy/src/replica/ballot.rs

@@ -65,17 +65,26 @@ pub fn compare_ballots(
 	))
 }
 
+/// Result of ballot validation with detailed context for error reporting
+#[derive(Debug)]
+pub struct BallotValidationResult {
+	pub is_valid: bool,
+	pub incoming_ballot: protocol::Ballot,
+	pub stored_ballot: protocol::Ballot,
+	pub comparison: std::cmp::Ordering,
+}
+
 /// Validate that a ballot is the highest seen for the given instance & updates the highest stored
 /// ballot if needed.
 ///
-/// Returns true if the ballot is valid (higher than previously seen).
+/// Returns detailed validation result including comparison information.
 #[tracing::instrument(skip_all)]
 pub async fn validate_and_update_ballot_for_instance(
 	tx: &Transaction,
 	replica_id: protocol::ReplicaId,
 	ballot: &protocol::Ballot,
 	instance: &protocol::Instance,
-) -> Result<bool> {
+) -> Result<BallotValidationResult> {
 	let instance_ballot_key =
 		keys::replica::InstanceBallotKey::new(instance.replica_id, instance.slot_id);
 	let subspace = keys::subspace(replica_id);
@@ -98,18 +107,24 @@ pub async fn validate_and_update_ballot_for_instance(
 	};
 
 	// Compare incoming ballot with highest seen - only accept if strictly greater
-	let is_valid = match compare_ballots(ballot, &highest_ballot) {
+	let comparison = compare_ballots(ballot, &highest_ballot);
+	let is_valid = match comparison {
 		std::cmp::Ordering::Greater => true,
 		std::cmp::Ordering::Equal | std::cmp::Ordering::Less => false,
 	};
 
 	// If the incoming ballot is higher, update our stored highest
-	if compare_ballots(ballot, &highest_ballot) == std::cmp::Ordering::Greater {
+	if comparison == std::cmp::Ordering::Greater {
 		let serialized = instance_ballot_key.serialize(ballot.clone())?;
 		tx.set(&packed_key, &serialized);
 
 		tracing::debug!(?ballot, ?instance, "updated highest ballot for instance");
 	}
 
-	Ok(is_valid)
+	Ok(BallotValidationResult {
+		is_valid,
+		incoming_ballot: ballot.clone(),
+		stored_ballot: highest_ballot,
+		comparison,
+	})
 }

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

@@ -21,10 +21,17 @@ pub async fn accept(
 
 	// Validate ballot
 	let current_ballot = ballot::get_ballot(tx, replica_id).await?;
-	let is_valid =
+	let validation =
 		ballot::validate_and_update_ballot_for_instance(tx, replica_id, &current_ballot, &instance)
 			.await?;
-	ensure!(is_valid, "ballot validation failed for pre_accept");
+	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 {

engine/packages/epoxy/src/replica/messages/pre_accept.rs

@@ -24,10 +24,17 @@ pub async fn pre_accept(
 
 	// Validate ballot
 	let current_ballot = ballot::get_ballot(tx, replica_id).await?;
-	let is_valid =
+	let validation =
 		ballot::validate_and_update_ballot_for_instance(tx, replica_id, &current_ballot, &instance)
 			.await?;
-	ensure!(is_valid, "ballot validation failed for pre_accept");
+	ensure!(
+		validation.is_valid,
+		"ballot validation failed for pre_accept: incoming ballot {:?} is not greater than stored ballot {:?} for instance {:?} (comparison: {:?})",
+		validation.incoming_ballot,
+		validation.stored_ballot,
+		instance,
+		validation.comparison
+	);
 
 	// Find interference for this key
 	let interf = utils::find_interference(tx, replica_id, &proposal.commands).await?;

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

@@ -30,10 +30,9 @@ pub async fn prepare(
 	};
 
 	// EPaxos Step 38: Validate ballot for this instance
-	let is_valid =
+	let validation =
 		ballot::validate_and_update_ballot_for_instance(tx, replica_id, &ballot, &instance).await?;
-
-	let response = if is_valid {
+	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 {

engine/packages/epoxy/src/workflows/replica/setup.rs

@@ -157,6 +157,8 @@ pub struct DownloadInstancesChunkOutput {
 }
 
 #[activity(DownloadInstancesChunk)]
+#[timeout = 300]
+#[max_retries = 100]
 pub async fn download_instances_chunk(
 	ctx: &ActivityCtx,
 	input: &DownloadInstancesChunkInput,
@@ -197,6 +199,8 @@ pub async fn download_instances_chunk(
 
 	// Apply each log entry from the downloaded instances
 	let total_entries = instances.len();
+	let mut applied_count = 0;
+	let mut skipped_count = 0;
 	for (idx, entry) in instances.iter().enumerate() {
 		tracing::debug!(
 			progress = format!("{}/{}", idx + 1, total_entries),
@@ -206,9 +210,27 @@ pub async fn download_instances_chunk(
 		);
 
 		// Apply the log entry to replay any uncommitted operations
-		apply_log_entry(ctx, &entry.log_entry, &entry.instance).await?;
+		// Skip entries that fail to apply (e.g., ballot validation failures during recovery)
+		if let Err(err) = apply_log_entry(ctx, &entry.log_entry, &entry.instance).await {
+			tracing::debug!(
+				?err,
+				?entry.instance,
+				state = ?entry.log_entry.state,
+				"skipping log entry that failed to apply during recovery"
+			);
+			skipped_count += 1;
+			continue;
+		}
+		applied_count += 1;
 	}
 
+	tracing::info!(
+		total_entries,
+		applied_count,
+		skipped_count,
+		"finished applying log entries for chunk"
+	);
+
 	// Return whether we should continue downloading chunks and the last instance
 	Ok(DownloadInstancesChunkOutput {
 		last_instance: instances.last().map(|entry| entry.instance.clone().into()),
@@ -479,7 +501,8 @@ pub async fn recover_keys_chunk(
 				// the range limit)
 				if recovered_count == 0 && scanned_count >= count {
 					bail!(
-						"single key has more than {count} instances, cannot process in one chunk",
+						"single key has more than {count} instances, cannot process in one chunk. key: {:?}",
+						current_key,
 					);
 				}