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/ops/actor/get_for_key.rs

@@ -1,7 +1,12 @@
 use anyhow::Result;
+use futures_util::TryStreamExt;
 use gas::prelude::*;
 use rivet_api_util::{Method, request_remote_datacenter};
 use rivet_types::actors::Actor;
+use universaldb::options::StreamingMode;
+use universaldb::utils::IsolationLevel::Serializable;
+
+use crate::keys;
 
 #[derive(Debug)]
 pub struct Input {
@@ -34,21 +39,67 @@ pub async fn pegboard_actor_get_for_key(ctx: &OperationCtx, input: &Input) -> Re
 
 	// Check if the actor is in the current datacenter
 	if reservation_id.label() == ctx.config().dc_label() {
-		// Local datacenter - get the actor directly
-		let actors_res = ctx
-			.op(crate::ops::actor::list_for_ns::Input {
-				namespace_id: input.namespace_id,
-				name: input.name.clone(),
-				key: Some(input.key.clone()),
-				include_destroyed: false,
-				created_before: None,
-				limit: 1,
-				fetch_error: input.fetch_error,
+		// Local datacenter - resolve with a serializable read so actors that are
+		// already tearing down do not leak through as a valid get/getOrCreate
+		// target.
+		let actors_with_wf_ids = ctx
+			.udb()?
+			.run(|tx| async move {
+				let tx = tx.with_subspace(keys::subspace());
+
+				let actor_subspace = keys::subspace().subspace(&keys::ns::ActorByKeyKey::subspace(
+					input.namespace_id,
+					input.name.clone(),
+					input.key.clone(),
+				));
+				let (start, end) = actor_subspace.range();
+
+				let mut stream = tx.get_ranges_keyvalues(
+					universaldb::RangeOption {
+						mode: StreamingMode::Iterator,
+						reverse: true,
+						..(start, end).into()
+					},
+					Serializable,
+				);
+
+				let mut results = Vec::new();
+
+				while let Some(entry) = stream.try_next().await? {
+					let (idx_key, data) = tx.read_entry::<keys::ns::ActorByKeyKey>(&entry)?;
+
+					if !data.is_destroyed {
+						results.push((idx_key.actor_id, data.workflow_id));
+						break;
+					}
+				}
+
+				Ok(results)
 			})
 			.await?;
 
+		let wfs = ctx
+			.get_workflows(
+				actors_with_wf_ids
+					.iter()
+					.map(|(_, workflow_id)| *workflow_id)
+					.collect(),
+			)
+			.await?;
+
+		let dc_name = ctx.config().dc_name()?.to_string();
+
+		let actors = super::util::build_actors_from_workflows(
+			ctx,
+			actors_with_wf_ids,
+			wfs,
+			&dc_name,
+			input.fetch_error,
+		)
+		.await?;
+
 		Ok(Output {
-			actor: actors_res.actors.into_iter().next(),
+			actor: actors.into_iter().next(),
 		})
 	} else {
 		// Remote datacenter - request from the correct datacenter

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 {