fix: handle edge cases with gateway + actor sleep (#3248)

Author: MasterPtato

engine/packages/engine/src/util/wf/signal.rs

@@ -54,10 +54,6 @@ pub async fn print_signals(signals: Vec<SignalData>, pretty: bool) -> Result<()>
 				);
 			}
 
-			if let Some(workflow_id) = &signal.workflow_id {
-				println!("  {} {}", style("to workflow id").bold(), workflow_id,);
-			}
-
 			if let Some(workflow_id) = signal.workflow_id {
 				println!("  {} {}", style("to workflow id").bold(), workflow_id);
 			}

engine/packages/guard/src/routing/pegboard_gateway.rs

@@ -101,6 +101,9 @@ pub async fn route_request(
 	let mut ready_sub = ctx
 		.subscribe::<pegboard::workflows::actor::Ready>(("actor_id", actor_id))
 		.await?;
+	let mut stopped_sub = ctx
+		.subscribe::<pegboard::workflows::actor::Stopped>(("actor_id", actor_id))
+		.await?;
 	let mut fail_sub = ctx
 		.subscribe::<pegboard::workflows::actor::Failed>(("actor_id", actor_id))
 		.await?;
@@ -122,6 +125,8 @@ pub async fn route_request(
 
 	// Wake actor if sleeping
 	if actor.sleeping {
+		tracing::debug!(?actor_id, "actor sleeping, waking");
+
 		ctx.signal(pegboard::workflows::actor::Wake {})
 			.to_workflow_id(actor.workflow_id)
 			.send()
@@ -133,20 +138,51 @@ pub async fn route_request(
 	} else {
 		tracing::debug!(?actor_id, "waiting for actor to become ready");
 
+		let mut wake_retries = 0;
+
 		// Wait for ready, fail, or destroy
-		tokio::select! {
-			res = ready_sub.next() => { res?.runner_id },
-			res = fail_sub.next() => {
-				let msg = res?;
-				return Err(msg.error.clone().build());
-			}
-			res = destroy_sub.next() => {
-				res?;
-				return Err(pegboard::errors::Actor::DestroyedWhileWaitingForReady.build());
-			}
-			// Ready timeout
-			_ = tokio::time::sleep(ACTOR_READY_TIMEOUT) => {
-				return Err(errors::ActorReadyTimeout { actor_id }.build());
+		loop {
+			tokio::select! {
+				res = ready_sub.next() => break res?.runner_id,
+				res = stopped_sub.next() => {
+					res?;
+
+					// Attempt to rewake once
+					if wake_retries < 3 {
+						tracing::debug!(?actor_id, ?wake_retries, "actor stopped while we were waiting for it to beocme ready, attempting rewake");
+						wake_retries += 1;
+
+						let res = ctx.signal(pegboard::workflows::actor::Wake {})
+							.to_workflow_id(actor.workflow_id)
+							.send()
+							.await;
+
+						if let Some(WorkflowError::WorkflowNotFound) = res
+							.as_ref()
+							.err()
+							.and_then(|x| x.chain().find_map(|x| x.downcast_ref::<WorkflowError>()))
+						{
+							tracing::warn!(
+								?actor_id,
+								"actor workflow not found for rewake"
+							);
+						} else {
+							res?;
+						}
+					}
+				}
+				res = fail_sub.next() => {
+					let msg = res?;
+					return Err(msg.error.clone().build());
+				}
+				res = destroy_sub.next() => {
+					res?;
+					return Err(pegboard::errors::Actor::DestroyedWhileWaitingForReady.build());
+				}
+				// Ready timeout
+				_ = tokio::time::sleep(ACTOR_READY_TIMEOUT) => {
+					return Err(errors::ActorReadyTimeout { actor_id }.build());
+				}
 			}
 		}
 	};

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

@@ -26,7 +26,7 @@ const X_RIVET_ENDPOINT: HeaderName = HeaderName::from_static("x-rivet-endpoint")
 const X_RIVET_TOKEN: HeaderName = HeaderName::from_static("x-rivet-token");
 const X_RIVET_TOTAL_SLOTS: HeaderName = HeaderName::from_static("x-rivet-total-slots");
 const X_RIVET_RUNNER_NAME: HeaderName = HeaderName::from_static("x-rivet-runner-name");
-const X_RIVET_NAMESPACE_ID: HeaderName = HeaderName::from_static("x-rivet-namespace-id");
+const X_RIVET_NAMESPACE_NAME: HeaderName = HeaderName::from_static("x-rivet-namespace-name");
 
 const DRAIN_GRACE_PERIOD: Duration = Duration::from_secs(5);
 
@@ -109,7 +109,9 @@ async fn tick(
 
 	// Process each runner config with error handling
 	for (ns_id, runner_name, desired_slots) in &serverless_data {
-		let runner_config = runner_configs.iter().find(|rc| rc.namespace_id == *ns_id);
+		let runner_config = runner_configs
+			.iter()
+			.find(|rc| rc.namespace_id == *ns_id && &rc.name == runner_name);
 
 		let Some(runner_config) = runner_config else {
 			tracing::debug!(
@@ -225,7 +227,7 @@ async fn tick_runner_config(
 
 		for conn in draining_connections {
 			if conn.shutdown_tx.send(()).is_err() {
-				tracing::warn!(
+				tracing::debug!(
 					"serverless connection shutdown channel dropped, likely already stopped"
 				);
 			}
@@ -307,8 +309,6 @@ async fn outbound_handler(
 	shutdown_rx: oneshot::Receiver<()>,
 	draining: Arc<AtomicBool>,
 ) -> Result<()> {
-	tracing::debug!(%url, "sending outbound req");
-
 	let current_dc = ctx.config().topology().current_dc()?;
 
 	let client = rivet_pools::reqwest::client_no_timeout().await?;
@@ -341,12 +341,21 @@ async fn outbound_handler(
 				HeaderValue::try_from(slots_per_runner)?,
 			),
 			(X_RIVET_RUNNER_NAME, HeaderValue::try_from(runner_name)?),
-			(X_RIVET_NAMESPACE_ID, HeaderValue::try_from(namespace_name)?),
+			(
+				X_RIVET_NAMESPACE_NAME,
+				HeaderValue::try_from(namespace_name.clone())?,
+			),
+			// Deprecated
+			(
+				HeaderName::from_static("x-rivet-namespace-id"),
+				HeaderValue::try_from(namespace_name)?,
+			),
 		])
 		.chain(token)
 		.collect();
 
 	let endpoint_url = format!("{}/start", url.trim_end_matches('/'));
+	tracing::debug!(%endpoint_url, "sending outbound req");
 	let req = client.get(endpoint_url).headers(headers);
 
 	let mut source = sse::EventSource::new(req).context("failed creating event source")?;

engine/packages/pegboard/src/workflows/actor/destroy.rs

@@ -35,7 +35,15 @@ pub(crate) async fn pegboard_actor_destroy(ctx: &mut WorkflowCtx, input: &Input)
 
 	// Destroy actor
 	if let (Some(runner_workflow_id), true) = (res.runner_workflow_id, &input.kill) {
-		kill(ctx, input.actor_id, input.generation, runner_workflow_id).await?;
+		ctx.signal(crate::workflows::runner::Command {
+			inner: protocol::Command::CommandStopActor(protocol::CommandStopActor {
+				actor_id: input.actor_id.to_string(),
+				generation: input.generation,
+			}),
+		})
+		.to_workflow_id(runner_workflow_id)
+		.send()
+		.await?;
 	}
 
 	// If a slot was allocated at the time of actor destruction then bump the serverless autoscaler so it can scale down
@@ -274,22 +282,3 @@ pub(crate) async fn clear_slot(
 
 	Ok(())
 }
-
-pub(crate) async fn kill(
-	ctx: &mut WorkflowCtx,
-	actor_id: Id,
-	generation: u32,
-	runner_workflow_id: Id,
-) -> Result<()> {
-	ctx.signal(crate::workflows::runner::Command {
-		inner: protocol::Command::CommandStopActor(protocol::CommandStopActor {
-			actor_id: actor_id.to_string(),
-			generation,
-		}),
-	})
-	.to_workflow_id(runner_workflow_id)
-	.send()
-	.await?;
-
-	Ok(())
-}

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

@@ -295,13 +295,18 @@ pub async fn pegboard_actor(ctx: &mut WorkflowCtx, input: &Input) -> Result<()>
 								return Ok(Loop::Continue);
 							}
 
+							let (Some(runner_id), Some(runner_workflow_id)) = (state.runner_id, state.runner_workflow_id) else {
+								tracing::warn!("actor not allocated, ignoring event");
+								return Ok(Loop::Continue);
+							};
+
 							match sig.inner {
 								protocol::Event::EventActorIntent(protocol::EventActorIntent {
 									intent,
 									..
 								}) => match intent {
 									protocol::ActorIntent::ActorIntentSleep => {
-										if let Some(runner_workflow_id) = state.runner_workflow_id {
+										if !state.sleeping {
 											state.gc_timeout_ts =
 												Some(util::timestamp::now() + ACTOR_STOP_THRESHOLD_MS);
 											state.sleeping = true;
@@ -312,37 +317,35 @@ pub async fn pegboard_actor(ctx: &mut WorkflowCtx, input: &Input) -> Result<()>
 											.await?;
 
 											// Send signal to kill actor now that we know it will be sleeping
-											destroy::kill(
-												ctx,
-												input.actor_id,
-												state.generation,
-												runner_workflow_id,
-											)
+											ctx.signal(crate::workflows::runner::Command {
+												inner: protocol::Command::CommandStopActor(protocol::CommandStopActor {
+													actor_id: input.actor_id.to_string(),
+													generation: state.generation,
+												}),
+											})
+											.to_workflow_id(runner_workflow_id)
+											.send()
 											.await?;
-										} else {
-											tracing::warn!("actor not allocated, ignoring sleep intent");
 										}
 									}
 									protocol::ActorIntent::ActorIntentStop => {
-										if let Some(runner_workflow_id) = state.runner_workflow_id {
-											state.gc_timeout_ts =
-												Some(util::timestamp::now() + ACTOR_STOP_THRESHOLD_MS);
+										state.gc_timeout_ts =
+											Some(util::timestamp::now() + ACTOR_STOP_THRESHOLD_MS);
 
-											ctx.activity(runtime::SetNotConnectableInput {
-												actor_id: input.actor_id,
-											})
-											.await?;
+										ctx.activity(runtime::SetNotConnectableInput {
+											actor_id: input.actor_id,
+										})
+										.await?;
 
-											destroy::kill(
-												ctx,
-												input.actor_id,
-												state.generation,
-												runner_workflow_id,
-											)
-											.await?;
-										} else {
-											tracing::warn!("actor not allocated, ignoring stop intent");
-										}
+										ctx.signal(crate::workflows::runner::Command {
+											inner: protocol::Command::CommandStopActor(protocol::CommandStopActor {
+												actor_id: input.actor_id.to_string(),
+												generation: state.generation,
+											}),
+										})
+										.to_workflow_id(runner_workflow_id)
+										.send()
+										.await?;
 									}
 								},
 								protocol::Event::EventActorStateUpdate(
@@ -351,23 +354,19 @@ pub async fn pegboard_actor(ctx: &mut WorkflowCtx, input: &Input) -> Result<()>
 									},
 								) => match actor_state {
 									protocol::ActorState::ActorStateRunning => {
-										if let Some(runner_id) = state.runner_id {
-											state.gc_timeout_ts = None;
-
-											ctx.activity(runtime::SetStartedInput {
-												actor_id: input.actor_id,
-											})
-											.await?;
-
-											ctx.msg(Ready {
-												runner_id,
-											})
-											.tag("actor_id", input.actor_id)
-											.send()
-											.await?;
-										} else {
-											tracing::warn!("actor not allocated, ignoring running event");
-										}
+										state.gc_timeout_ts = None;
+
+										ctx.activity(runtime::SetStartedInput {
+											actor_id: input.actor_id,
+										})
+										.await?;
+
+										ctx.msg(Ready {
+											runner_id,
+										})
+										.tag("actor_id", input.actor_id)
+										.send()
+										.await?;
 									}
 									protocol::ActorState::ActorStateStopped(
 										protocol::ActorStateStopped { code, .. },
@@ -541,6 +540,19 @@ async fn handle_stopped(
 		}
 	}
 
+	// Kill old actor if lost (just in case it ended up allocating)
+	if let (true, Some(old_runner_workflow_id)) = (lost, old_runner_workflow_id) {
+		ctx.signal(crate::workflows::runner::Command {
+			inner: protocol::Command::CommandStopActor(protocol::CommandStopActor {
+				actor_id: input.actor_id.to_string(),
+				generation: state.generation,
+			}),
+		})
+		.to_workflow_id(old_runner_workflow_id)
+		.send()
+		.await?;
+	}
+
 	// Reschedule no matter what
 	if force_reschedule {
 		match runtime::reschedule_actor(ctx, &input, state, true).await? {
@@ -566,18 +578,6 @@ async fn handle_stopped(
 
 		match (input.crash_policy, failed) {
 			(CrashPolicy::Restart, true) => {
-				// Kill old actor immediately if lost
-				if lost {
-					destroy::kill(
-						ctx,
-						input.actor_id,
-						state.generation,
-						old_runner_workflow_id
-							.context("should have runner_workflow_id set if not sleeping")?,
-					)
-					.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
@@ -637,6 +637,11 @@ async fn handle_stopped(
 	state.wake_for_alarm = false;
 	state.will_wake = false;
 
+	ctx.msg(Stopped {})
+		.tag("actor_id", input.actor_id)
+		.send()
+		.await?;
+
 	Ok(None)
 }
 
@@ -653,6 +658,9 @@ pub struct Ready {
 	pub runner_id: Id,
 }
 
+#[message("pegboard_actor_stopped")]
+pub struct Stopped {}
+
 #[signal("pegboard_actor_allocate")]
 #[derive(Debug)]
 pub struct Allocate {