wip: fix whatever is going on with envoys

Author: NathanFlurry

engine/packages/pegboard-envoy/src/conn.rs

@@ -147,7 +147,7 @@ pub async fn handle_init(
 				// Read existing data
 				let (create_ts_entry, old_last_ping_ts_entry, version_entry) = tokio::try_join!(
 					tx.read_opt(&create_ts_key, Serializable),
-					tx.read_opt(&create_ts_key, Serializable),
+					tx.read_opt(&last_ping_ts_key, Serializable),
 					tx.read_opt(&version_key, Serializable),
 				)?;
 
@@ -228,6 +228,7 @@ pub async fn handle_init(
 					);
 
 					tx.add_conflict_key(&old_lb_key, ConflictRangeType::Read)?;
+					tx.delete(&old_lb_key);
 				}
 
 				// Insert into LB

engine/packages/pegboard-kv-channel/src/lib.rs

@@ -434,9 +434,8 @@ async fn actor_request_task(
 	open_actors: Arc<Mutex<HashSet<String>>>,
 	mut rx: mpsc::Receiver<protocol::ToServerRequest>,
 ) {
-	// Cached actor resolution. Populated on first KV request, reused for all
-	// subsequent requests. Actor name is immutable so this never goes stale.
-	let mut cached_actor: Option<(Id, String)> = None;
+	// Cache keyed by actor id since a single connection multiplexes many actors.
+	let mut cached_actors: HashMap<String, (Id, String)> = HashMap::new();
 
 	while let Some(req) = rx.recv().await {
 		let is_close = matches!(req.data, protocol::RequestData::ActorCloseRequest);
@@ -464,11 +463,10 @@ async fn actor_request_task(
 						)
 					}
 				} else {
-					// Lazy-resolve and cache.
-					if cached_actor.is_none() {
+					if !cached_actors.contains_key(&req.actor_id) {
 						match resolve_actor(&ctx, &req.actor_id, namespace_id).await {
 							Ok(v) => {
-								cached_actor = Some(v);
+								cached_actors.insert(req.actor_id.clone(), v);
 							}
 							Err(resp) => {
 								// Don't cache failures. Next request will retry.
@@ -480,7 +478,8 @@ async fn actor_request_task(
 							}
 						}
 					}
-					let (parsed_id, actor_name) = cached_actor.as_ref().unwrap();
+					let (parsed_id, actor_name) =
+						cached_actors.get(&req.actor_id).unwrap();
 
 					let recipient = actor_kv::Recipient {
 						actor_id: *parsed_id,
@@ -795,10 +794,7 @@ async fn handle_kv_delete_range(
 
 /// Look up an actor by ID and return the parsed ID and actor name.
 ///
-/// Defense-in-depth: verifies the actor belongs to the authenticated namespace.
-/// The admin_token is a global credential, so this is not strictly necessary
-/// today, but prevents cross-namespace access if a less-privileged auth
-/// mechanism is introduced in the future.
+/// Verifies the actor belongs to the authenticated namespace.
 async fn resolve_actor(
 	ctx: &StandaloneCtx,
 	actor_id: &str,
@@ -812,11 +808,15 @@ async fn resolve_actor(
 	})?;
 
 	let actor = ctx
-		.op(pegboard::ops::actor::get_for_runner::Input {
-			actor_id: parsed_id,
+		.op(pegboard::ops::actor::get::Input {
+			actor_ids: vec![parsed_id],
+			fetch_error: false,
 		})
 		.await
-		.map_err(|err| internal_error(&err))?;
+		.map_err(|err| internal_error(&err))?
+		.actors
+		.into_iter()
+		.next();
 
 	match actor {
 		Some(actor) => {

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

@@ -672,22 +672,30 @@ async fn process_signal(
 							} = &mut state.transition
 							{
 								if *destroy_after_start {
-									// Transition to destroying
-									state.transition = Transition::Destroying {
-										envoy: runtime::EnvoyState::new(sig.envoy_key.clone()),
-										lost_timeout_ts: now
-											+ ctx.config().pegboard().actor_stop_threshold(),
+									state.transition = Transition::Running {
+										envoy: runtime::EnvoyState::new(
+											sig.envoy_key.clone(),
+										),
+										last_liveness_check_ts: now,
 									};
 
-									ctx.activity(runtime::InsertAndSendCommandsInput {
+									ctx.activity(runtime::SetConnectableInput {
+										envoy_key: sig.envoy_key.clone(),
 										generation: state.generation,
+									})
+									.await?;
+
+									ctx.msg(Ready {
 										envoy_key: sig.envoy_key.clone(),
-										commands: vec![protocol::Command::CommandStopActor(
-											protocol::CommandStopActor {
-												reason: protocol::StopActorReason::Destroy,
-											},
-										)],
 									})
+									.topic(("actor_id", input.actor_id))
+									.send()
+									.await?;
+
+									let workflow_id = ctx.workflow_id();
+									ctx.signal(Destroy {})
+										.to_workflow_id(workflow_id)
+										.send()
 									.await?;
 								} else {
 									// Transition to starting
@@ -804,10 +812,15 @@ async fn process_signal(
 					runtime::reschedule_actor(ctx, &input, state, metrics_workflow_id).await?;
 				}
 				Transition::SleepIntent {
-					rewake_after_stop, ..
+					lost_timeout_ts,
+					rewake_after_stop,
+					..
 				} => {
 					if !*rewake_after_stop {
 						*rewake_after_stop = true;
+						let now = ctx.activity(runtime::GetTsInput {}).await?;
+						*lost_timeout_ts =
+							now + ctx.config().pegboard().actor_stop_threshold();
 
 						tracing::debug!(
 							actor_id=?input.actor_id,
@@ -983,6 +996,8 @@ async fn process_signal(
 						lost_timeout_ts: now + ctx.config().pegboard().actor_stop_threshold(),
 					};
 
+					ctx.activity(runtime::SetNotConnectableInput {}).await?;
+
 					ctx.activity(runtime::InsertAndSendCommandsInput {
 						generation: state.generation,
 						envoy_key,
@@ -1004,6 +1019,8 @@ async fn process_signal(
 						lost_timeout_ts: now + ctx.config().pegboard().actor_stop_threshold(),
 					};
 
+					ctx.activity(runtime::SetNotConnectableInput {}).await?;
+
 					// Stop command was already sent
 				}
 				Transition::Allocating {

engine/packages/pegboard/src/workflows/actor2/runtime.rs

@@ -565,6 +565,40 @@ pub async fn handle_stopped(
 
 		StoppedResult::Continue
 	} else if try_reallocate {
+		// A wake can arrive after the actor has already committed to sleeping
+		// but before the stop completes. In that case, mirror the v1 workflow
+		// and immediately reschedule after the stop instead of treating a
+		// graceful sleep exit as terminal destruction.
+		if let Transition::SleepIntent {
+			rewake_after_stop: true,
+			..
+		} = state.transition
+		{
+			let allocate_res = ctx.activity(AllocateInput {}).await?;
+
+			if let Some(allocation) = allocate_res.allocation {
+				state.generation += 1;
+
+				ctx.activity(SendOutboundInput {
+					generation: state.generation,
+					input: input.input.clone(),
+					allocation,
+				})
+				.await?;
+
+				state.transition = Transition::Allocating {
+					destroy_after_start: false,
+					lost_timeout_ts: allocate_res.now
+						+ ctx.config().pegboard().actor_allocation_threshold(),
+				};
+			} else {
+				state.transition = Transition::Reallocating {
+					since_ts: allocate_res.now,
+				};
+			}
+
+			StoppedResult::Continue
+		} else {
 		// An actor stopping with `StopCode::Ok` indicates a graceful exit
 		let graceful_exit = matches!(
 			variant,
@@ -574,63 +608,64 @@ pub async fn handle_stopped(
 			}
 		);
 
-		match (input.crash_policy, graceful_exit) {
-			(CrashPolicy::Restart, false) => {
-				let allocate_res = ctx.activity(AllocateInput {}).await?;
-
-				if let Some(allocation) = allocate_res.allocation {
-					state.generation += 1;
+			match (input.crash_policy, graceful_exit) {
+				(CrashPolicy::Restart, false) => {
+					let allocate_res = ctx.activity(AllocateInput {}).await?;
+
+					if let Some(allocation) = allocate_res.allocation {
+						state.generation += 1;
+
+						ctx.activity(SendOutboundInput {
+							generation: state.generation,
+							input: input.input.clone(),
+							allocation,
+						})
+						.await?;
+
+						// Transition to allocating
+						state.transition = Transition::Allocating {
+							destroy_after_start: false,
+							lost_timeout_ts: allocate_res.now
+								+ ctx.config().pegboard().actor_allocation_threshold(),
+						};
+					} else {
+						// Transition to retry backoff
+						state.transition = Transition::Reallocating {
+							since_ts: allocate_res.now,
+						};
+					}
 
-					ctx.activity(SendOutboundInput {
-						generation: state.generation,
-						input: input.input.clone(),
-						allocation,
-					})
-					.await?;
-
-					// Transition to allocating
-					state.transition = Transition::Allocating {
-						destroy_after_start: false,
-						lost_timeout_ts: allocate_res.now
-							+ ctx.config().pegboard().actor_allocation_threshold(),
-					};
-				} else {
-					// Transition to retry backoff
-					state.transition = Transition::Reallocating {
-						since_ts: allocate_res.now,
-					};
+					StoppedResult::Continue
 				}
+				(CrashPolicy::Sleep, false) => {
+					tracing::debug!(actor_id=?input.actor_id, "actor sleeping due to ungraceful exit");
 
-				StoppedResult::Continue
-			}
-			(CrashPolicy::Sleep, false) => {
-				tracing::debug!(actor_id=?input.actor_id, "actor sleeping due to ungraceful exit");
-
-				// Clear alarm
-				if let Some(alarm_ts) = state.alarm_ts {
-					let now = ctx.activity(GetTsInput {}).await?;
+					// Clear alarm
+					if let Some(alarm_ts) = state.alarm_ts {
+						let now = ctx.activity(GetTsInput {}).await?;
 
-					if now >= alarm_ts {
-						state.alarm_ts = None;
+						if now >= alarm_ts {
+							state.alarm_ts = None;
+						}
 					}
-				}
 
-				// Transition to sleeping
-				state.transition = Transition::Sleeping;
+					// Transition to sleeping
+					state.transition = Transition::Sleeping;
 
-				StoppedResult::Continue
-			}
-			_ => {
-				let now = ctx.activity(GetTsInput {}).await?;
+					StoppedResult::Continue
+				}
+				_ => {
+					let now = ctx.activity(GetTsInput {}).await?;
 
-				// Don't destroy on failed allocation, retry instead
-				if let StoppedVariant::FailedAllocation { .. } = &variant {
-					// Transition to retry backoff
-					state.transition = Transition::Reallocating { since_ts: now };
+					// Don't destroy on failed allocation, retry instead
+					if let StoppedVariant::FailedAllocation { .. } = &variant {
+						// Transition to retry backoff
+						state.transition = Transition::Reallocating { since_ts: now };
 
-					StoppedResult::Continue
-				} else {
-					StoppedResult::Destroy
+						StoppedResult::Continue
+					} else {
+						StoppedResult::Destroy
+					}
 				}
 			}
 		}