fix(pb): fix various race condition edge cases (#3137)

Author: MasterPtato

packages/common/gasoline/core/src/db/kv/mod.rs

@@ -139,6 +139,7 @@ impl DatabaseKv {
 
 		let workflow_name_key = keys::workflow::NameKey::new(workflow_id);
 
+		// TODO: This does not check if the workflow is silenced
 		// Check if the workflow exists
 		let Some(workflow_name_entry) = tx
 			.get(&self.subspace.pack(&workflow_name_key), Serializable)

packages/common/util/id/src/lib.rs

@@ -289,6 +289,8 @@ fn base36_char_to_base10(c: char) -> Result<u8, IdError> {
 	match c {
 		'0'..='9' => Ok(c as u8 - b'0'),
 		'a'..='z' => Ok(c as u8 - b'a' + 10),
+		// Uppercase is treated the same as lowercase
+		'A'..='Z' => Ok(c as u8 - b'A' + 10),
 		_ => return Err(IdError::InvalidChar(c)),
 	}
 }
@@ -313,4 +315,15 @@ mod tests {
 		let parsed = Id::from_str(&s).unwrap();
 		assert_eq!(parsed, id);
 	}
+
+	#[test]
+	fn test_v1_parse_uppercase() {
+		let uuid = Uuid::new_v4();
+		let label = 0xABCD;
+		let id = Id::v1(uuid, label);
+		let s = id.to_string().to_ascii_uppercase();
+		assert_eq!(s.len(), 30);
+		let parsed = Id::from_str(&s).unwrap();
+		assert_eq!(parsed, id);
+	}
 }

packages/core/pegboard-serverless/src/lib.rs

@@ -236,7 +236,7 @@ fn spawn_connection(
 		)
 		.await
 		{
-			tracing::error!(?err, "outbound req failed");
+			tracing::warn!(?err, "outbound req failed");
 
 			// TODO: Add backoff
 			tokio::time::sleep(Duration::from_secs(1)).await;
@@ -368,15 +368,28 @@ async fn outbound_handler(
 				// If runner_id is none at this point it means we did not send the stopping signal yet, so
 				// send it now
 				if runner_id.is_none() {
-					stop_runner(ctx, Id::parse(&msg.data)?).await?;
+					let data = BASE64.decode(msg.data).context("invalid base64 message")?;
+					let payload =
+						protocol::versioned::ToServerlessServer::deserialize_with_embedded_version(
+							&data,
+						)
+						.context("invalid payload")?;
+
+					match payload {
+						protocol::ToServerlessServer::ToServerlessServerInit(init) => {
+							let runner_id =
+								Id::parse(&init.runner_id).context("invalid runner id")?;
+							stop_runner(ctx, runner_id).await?;
+						}
+					}
 				}
 			}
 			Err(sse::Error::StreamEnded) => break,
 			Err(err) => return Err(err.into()),
 		}
 	}
 
-	tracing::info!("outbound req stopped");
+	tracing::debug!("outbound req stopped");
 
 	Ok(())
 }

packages/services/pegboard/src/ops/runner/update_alloc_idx.rs

@@ -108,15 +108,20 @@ pub async fn pegboard_runner_update_alloc_idx(ctx: &OperationCtx, input: &Input)
 						continue;
 					};
 
-					// Runner is expired, updating will do nothing
+					// Runner is expired, AddIdx is invalid and UpdatePing will do nothing
 					if expired_ts_entry.is_some() {
-						notifications.push(RunnerNotification {
-							runner_id: runner.runner_id,
-							workflow_id,
-							eligibility: RunnerEligibility::Expired,
-						});
+						match runner.action {
+							Action::ClearIdx => {}
+							Action::AddIdx | Action::UpdatePing { .. } => {
+								notifications.push(RunnerNotification {
+									runner_id: runner.runner_id,
+									workflow_id,
+									eligibility: RunnerEligibility::Expired,
+								});
 
-						continue;
+								continue;
+							}
+						}
 					}
 
 					let remaining_millislots = (remaining_slots * 1000) / total_slots;

packages/services/pegboard/src/workflows/actor/mod.rs

@@ -266,6 +266,7 @@ pub async fn pegboard_actor(ctx: &mut WorkflowCtx, input: &Input) -> Result<()>
 							// Fake signal
 							Main::Lost(Lost {
 								generation: state.generation,
+								force_reschedule: false,
 							})
 						}
 					} else if let Some(alarm_ts) = state.alarm_ts {
@@ -365,7 +366,7 @@ pub async fn pegboard_actor(ctx: &mut WorkflowCtx, input: &Input) -> Result<()>
 										protocol::ActorStateStopped { code, .. },
 									) => {
 										if let Some(res) =
-											handle_stopped(ctx, &input, state, Some(code), false)
+											handle_stopped(ctx, &input, state, Some(code), false, false)
 												.await?
 										{
 											return Ok(Loop::Break(res));
@@ -386,7 +387,7 @@ pub async fn pegboard_actor(ctx: &mut WorkflowCtx, input: &Input) -> Result<()>
 									state.sleeping = false;
 									state.will_wake = false;
 
-									match runtime::reschedule_actor(ctx, &input, state).await? {
+									match runtime::reschedule_actor(ctx, &input, state, false).await? {
 										runtime::SpawnActorOutput::Allocated { .. } => {},
 										runtime::SpawnActorOutput::Sleep => {
 											state.sleeping = true;
@@ -416,6 +417,8 @@ pub async fn pegboard_actor(ctx: &mut WorkflowCtx, input: &Input) -> Result<()>
 									actor_id=?input.actor_id,
 									"cannot wake actor that is not sleeping",
 								);
+
+								state.wake_for_alarm = false;
 							}
 						}
 						Main::Lost(sig) => {
@@ -425,7 +428,7 @@ pub async fn pegboard_actor(ctx: &mut WorkflowCtx, input: &Input) -> Result<()>
 							}
 
 							if let Some(res) =
-								handle_stopped(ctx, &input, state, None, true).await?
+								handle_stopped(ctx, &input, state, None, true, sig.force_reschedule).await?
 							{
 								return Ok(Loop::Break(res));
 							}
@@ -485,8 +488,9 @@ async fn handle_stopped(
 	state: &mut runtime::LifecycleState,
 	code: Option<protocol::StopCode>,
 	lost: bool,
+	force_reschedule: bool,
 ) -> Result<Option<runtime::LifecycleRes>> {
-	tracing::debug!(?code, "actor stopped");
+	tracing::debug!(?code, %force_reschedule, "actor stopped");
 
 	// Reset retry count on successful exit
 	if let Some(protocol::StopCode::Ok) = code {
@@ -496,7 +500,7 @@ async fn handle_stopped(
 	// Clear stop gc timeout to prevent being marked as lost in the lifecycle loop
 	state.gc_timeout_ts = None;
 	state.runner_id = None;
-	state.runner_workflow_id = None;
+	let old_runner_workflow_id = state.runner_workflow_id.take();
 
 	let deallocate_res = ctx
 		.activity(runtime::DeallocateInput {
@@ -530,26 +534,44 @@ async fn handle_stopped(
 		}
 	}
 
+	// Reschedule no matter what
+	if force_reschedule {
+		match runtime::reschedule_actor(ctx, &input, state, true).await? {
+			runtime::SpawnActorOutput::Allocated { .. } => {}
+			// NOTE: This should be unreachable because force_reschedule is true
+			runtime::SpawnActorOutput::Sleep => {
+				state.sleeping = true;
+			}
+			runtime::SpawnActorOutput::Destroy => {
+				// Destroyed early
+				return Ok(Some(runtime::LifecycleRes {
+					generation: state.generation,
+					// False here because if we received the destroy signal, it is
+					// guaranteed that we did not allocate another actor.
+					kill: false,
+				}));
+			}
+		}
+	}
 	// Handle rescheduling if not marked as sleeping
-	if !state.sleeping {
+	else if !state.sleeping {
 		let failed = matches!(code, None | Some(protocol::StopCode::Error));
 
-		match (failed, input.crash_policy) {
-			(true, CrashPolicy::Restart) => {
+		match (input.crash_policy, failed) {
+			(CrashPolicy::Restart, true) => {
 				// Kill old actor immediately if lost
 				if lost {
 					destroy::kill(
 						ctx,
 						input.actor_id,
 						state.generation,
-						state
-							.runner_workflow_id
+						old_runner_workflow_id
 							.context("should have runner_workflow_id set if not sleeping")?,
 					)
 					.await?;
 				}
 
-				match runtime::reschedule_actor(ctx, &input, state).await? {
+				match runtime::reschedule_actor(ctx, &input, state, false).await? {
 					runtime::SpawnActorOutput::Allocated { .. } => {}
 					// NOTE: Its not possible for `SpawnActorOutput::Sleep` to be returned here, the crash
 					// policy is `Restart`.
@@ -564,7 +586,7 @@ async fn handle_stopped(
 					}
 				}
 			}
-			(true, CrashPolicy::Sleep) => {
+			(CrashPolicy::Sleep, true) => {
 				tracing::debug!(actor_id=?input.actor_id, "actor sleeping due to crash");
 
 				state.sleeping = true;
@@ -577,15 +599,6 @@ async fn handle_stopped(
 			_ => {
 				ctx.activity(runtime::SetCompleteInput {}).await?;
 
-				if lost {
-					ctx.msg(Failed {
-						error: errors::Actor::DestroyedWhileWaitingForReady,
-					})
-					.tag("actor_id", input.actor_id)
-					.send()
-					.await?;
-				}
-
 				return Ok(Some(runtime::LifecycleRes {
 					generation: state.generation,
 					kill: lost,
@@ -596,9 +609,8 @@ async fn handle_stopped(
 	// Rewake actor immediately after stopping if `will_wake` was set
 	else if state.will_wake {
 		state.sleeping = false;
-		state.will_wake = false;
 
-		match runtime::reschedule_actor(ctx, &input, state).await? {
+		match runtime::reschedule_actor(ctx, &input, state, false).await? {
 			runtime::SpawnActorOutput::Allocated { .. } => {}
 			runtime::SpawnActorOutput::Sleep => {
 				state.sleeping = true;
@@ -613,10 +625,11 @@ async fn handle_stopped(
 				}));
 			}
 		}
-
-		state.wake_for_alarm = false;
 	}
 
+	state.wake_for_alarm = false;
+	state.will_wake = false;
+
 	Ok(None)
 }
 
@@ -651,6 +664,7 @@ pub struct Wake {}
 #[signal("pegboard_actor_lost")]
 pub struct Lost {
 	pub generation: u32,
+	pub force_reschedule: bool,
 }
 
 #[signal("pegboard_actor_destroy")]

packages/services/pegboard/src/workflows/actor/runtime.rs

@@ -107,10 +107,11 @@ async fn update_runner(ctx: &ActivityCtx, input: &UpdateRunnerInput) -> Result<(
 struct AllocateActorInput {
 	actor_id: Id,
 	generation: u32,
-	from_alarm: bool,
+	force_allocate: bool,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
 pub enum AllocateActorOutput {
 	Allocated {
 		runner_id: Id,
@@ -315,7 +316,7 @@ async fn allocate_actor(
 
 			// At this point in the txn there is no availability
 
-			match (crash_policy, input.from_alarm, has_valid_serverless) {
+			match (crash_policy, input.force_allocate, has_valid_serverless) {
 				(CrashPolicy::Sleep, false, false) => {
 					Ok((for_serverless, AllocateActorOutput::Sleep))
 				}
@@ -373,6 +374,11 @@ async fn allocate_actor(
 		AllocateActorOutput::Pending {
 			pending_allocation_ts,
 		} => {
+			tracing::warn!(
+				actor_id=?input.actor_id,
+				"failed to allocate (no availability), waiting for allocation",
+			);
+
 			state.pending_allocation_ts = Some(*pending_allocation_ts);
 		}
 		AllocateActorOutput::Sleep => {}
@@ -473,14 +479,14 @@ pub async fn spawn_actor(
 	ctx: &mut WorkflowCtx,
 	input: &Input,
 	generation: u32,
-	from_alarm: bool,
+	force_allocate: bool,
 ) -> Result<SpawnActorOutput> {
 	// Attempt allocation
 	let allocate_res = ctx
 		.activity(AllocateActorInput {
 			actor_id: input.actor_id,
 			generation,
-			from_alarm,
+			force_allocate,
 		})
 		.await?;
 
@@ -524,11 +530,6 @@ pub async fn spawn_actor(
 		AllocateActorOutput::Pending {
 			pending_allocation_ts,
 		} => {
-			tracing::warn!(
-				actor_id=?input.actor_id,
-				"failed to allocate (no availability), waiting for allocation",
-			);
-
 			// Bump the autoscaler so it can scale up
 			ctx.msg(rivet_types::msgs::pegboard::BumpServerlessAutoscaler {})
 				.send()
@@ -611,6 +612,7 @@ pub async fn reschedule_actor(
 	ctx: &mut WorkflowCtx,
 	input: &Input,
 	state: &mut LifecycleState,
+	force_reschedule: bool,
 ) -> Result<SpawnActorOutput> {
 	tracing::debug!(actor_id=?input.actor_id, "rescheduling actor");
 
@@ -653,7 +655,13 @@ pub async fn reschedule_actor(
 	}
 
 	let next_generation = state.generation + 1;
-	let spawn_res = spawn_actor(ctx, &input, next_generation, state.wake_for_alarm).await?;
+	let spawn_res = spawn_actor(
+		ctx,
+		&input,
+		next_generation,
+		force_reschedule || state.wake_for_alarm,
+	)
+	.await?;
 
 	if let SpawnActorOutput::Allocated {
 		runner_id,