From b292958d5606e2b82857099219f3ee5b9dd5e7a8 Mon Sep 17 00:00:00 2001 From: Aditya Salunkhe Date: Tue, 25 Aug 2026 23:28:59 +0530 Subject: [PATCH 1/3] Fix permanent shard loss when RevokeShards times out: drain in background, leave failed-revoke shards unassigned and reconcile the old executor --- .../src/sharding/shard_management.rs | 24 +- golem-shard-manager/tests/shard_management.rs | 257 ++++++++++++++---- golem-worker-executor/src/durable_host/mod.rs | 9 +- golem-worker-executor/src/grpc/mod.rs | 158 +++++++++-- .../src/worker/invocation_loop.rs | 15 +- golem-worker-executor/src/worker/mod.rs | 20 ++ 6 files changed, 401 insertions(+), 82 deletions(-) diff --git a/golem-shard-manager/src/sharding/shard_management.rs b/golem-shard-manager/src/sharding/shard_management.rs index 7a8eaf9b8f..95fb241551 100644 --- a/golem-shard-manager/src/sharding/shard_management.rs +++ b/golem-shard-manager/src/sharding/shard_management.rs @@ -143,7 +143,7 @@ impl ShardManagement { // - the rebalance plan is calculated, // - new and removed pods are added to the routing table and got persisted, // but the rebalance plan is NOT applied yet. The lock is then release for apply. - let (mut rebalance, full_assignment_pods) = { + let (mut rebalance, mut full_assignment_pods) = { let mut current_routing_table = routing_table.write().await; for pod in removed_pods { @@ -203,12 +203,12 @@ impl ShardManagement { "Some shards could not be assigned and will be left unassigned for retry" ); + // The executor may have applied the assignment even though the call failed + // (for example a timeout), so it always gets the authoritative assignment next. { let mut updates_guard = updates.lock().await; for (pod, _) in &rebalance_failures.failed_assignments { - if full_assignment_pods.contains(pod) { - updates_guard.retry_full_assignment(*pod); - } + updates_guard.retry_full_assignment(*pod); } } needs_retry = true; @@ -221,8 +221,16 @@ impl ShardManagement { .iter() .map(|(pod, _)| pod) .join(", "), - "Some shards could not be unassigned and rebalance will be retried" + "Some shards could not be unassigned; they are left unassigned and the pods get their authoritative assignment" ); + // A failed revoke does not mean the executor still holds the shards - it may have + // dropped them and only the response was lost. The shards are left unassigned in + // the routing table (the unassignment stays in the plan), and the pod receives its + // authoritative assignment in this pass, before the next pass hands the shards to + // another pod. If that reconciliation fails too, it is re-queued below. + for (pod, _) in &rebalance_failures.failed_unassignments { + full_assignment_pods.insert(*pod); + } needs_retry = true; } @@ -294,15 +302,15 @@ impl ShardManagement { } let failed_unassignments = revoke_shards(worker_executors.clone(), rebalance.get_unassignments()).await; - let failed_shards = failed_unassignments + let failed_shards: HashSet = failed_unassignments .iter() .flat_map(|(_, shard_ids)| shard_ids.clone()) .collect(); - rebalance.remove_shards(&failed_shards); + rebalance.remove_assignment_shards(&failed_shards); if !failed_shards.is_empty() { warn!( failed_shards = failed_shards.iter().join(", "), - "Some shards could not be unassigned and have been removed from rebalance" + "Some shards could not be unassigned and are left unassigned for retry" ); } diff --git a/golem-shard-manager/tests/shard_management.rs b/golem-shard-manager/tests/shard_management.rs index a9100d46fe..1f45968a37 100644 --- a/golem-shard-manager/tests/shard_management.rs +++ b/golem-shard-manager/tests/shard_management.rs @@ -44,6 +44,10 @@ impl TestPersistence { async fn latest(&self) -> RoutingTable { self.state.lock().await.clone() } + + async fn writes(&self) -> Vec { + self.writes.lock().await.clone() + } } #[async_trait] @@ -59,15 +63,33 @@ impl RoutingTablePersistence for TestPersistence { } } +/// One call the shard manager made to a worker executor, in order. +#[derive(Clone, Debug, PartialEq, Eq)] +enum Call { + Assign(Pod, BTreeSet), + Revoke(Pod, BTreeSet), + Set(Pod, BTreeSet), +} + #[derive(Clone, Debug, Default)] struct TestWorkerExecutors { local_assignments: Arc>>>, + calls: Arc>>, failed_assignments: Arc>>, failed_revocations: Arc>>, + applied_then_failed_revocations: Arc>>, failed_reconciliations: Arc>>, } impl TestWorkerExecutors { + async fn calls(&self) -> Vec { + self.calls.lock().await.clone() + } + + async fn record(&self, call: Call) { + self.calls.lock().await.push(call); + } + async fn set_local_assignment(&self, pod: Pod, shard_ids: &[i64]) { self.local_assignments .lock() @@ -92,6 +114,13 @@ impl TestWorkerExecutors { self.failed_revocations.lock().await.insert(pod, count); } + async fn apply_then_fail_next_revocations(&self, pod: Pod, count: usize) { + self.applied_then_failed_revocations + .lock() + .await + .insert(pod, count); + } + async fn fail_next_reconciliations(&self, pod: Pod, count: usize) { self.failed_reconciliations.lock().await.insert(pod, count); } @@ -115,6 +144,7 @@ impl WorkerExecutorService for TestWorkerExecutors { pod: &Pod, shard_ids: &BTreeSet, ) -> Result<(), ShardManagerError> { + self.record(Call::Assign(*pod, shard_ids.clone())).await; if Self::should_fail(&self.failed_assignments, *pod).await { return Err(ShardManagerError::Timeout); } @@ -137,6 +167,7 @@ impl WorkerExecutorService for TestWorkerExecutors { pod: &Pod, shard_ids: &BTreeSet, ) -> Result<(), ShardManagerError> { + self.record(Call::Revoke(*pod, shard_ids.clone())).await; if Self::should_fail(&self.failed_revocations, *pod).await { return Err(ShardManagerError::Timeout); } @@ -144,6 +175,10 @@ impl WorkerExecutorService for TestWorkerExecutors { if let Some(local_assignment) = self.local_assignments.lock().await.get_mut(pod) { local_assignment.retain(|shard_id| !shard_ids.contains(shard_id)); } + + if Self::should_fail(&self.applied_then_failed_revocations, *pod).await { + return Err(ShardManagerError::Timeout); + } Ok(()) } @@ -153,6 +188,7 @@ impl WorkerExecutorService for TestWorkerExecutors { _number_of_shards: usize, shard_ids: &BTreeSet, ) -> Result<(), ShardManagerError> { + self.record(Call::Set(*pod, shard_ids.clone())).await; if Self::should_fail(&self.failed_reconciliations, *pod).await { return Err(ShardManagerError::Timeout); } @@ -497,16 +533,111 @@ async fn failed_assignment_is_retried_from_unassigned_shards() { } #[test] -// A failed revoke must not assign the shard elsewhere, but it should be retried -// and eventually converge without another shard-manager event. -async fn failed_revoke_is_retried_without_assigning_to_new_executor_first() { +// Reconnect reconciliation failures are retried by the shard-manager worker. +async fn failed_reconnect_reconciliation_is_retried() { + let existing_pod = pod(1, 9000); + let worker_executors = Arc::new(TestWorkerExecutors::default()); + + let (shard_management, _persistence, mut join_set) = new_shard_management( + routing_table_with_pods( + 1, + vec![ + (existing_pod, "worker-executor-0", &[]), + (pod(2, 9001), "worker-executor-1", &[0]), + ], + ), + worker_executors.clone(), + ) + .await; + + worker_executors + .set_local_assignment(existing_pod, &[0]) + .await; + worker_executors + .fail_next_reconciliations(existing_pod, 1) + .await; + + shard_management + .register_pod(existing_pod, Some("worker-executor-0".to_string())) + .await; + + wait_for_local_assignment(&worker_executors, existing_pod, BTreeSet::new()).await; + + join_set.abort_all(); +} + +/// Asserts the state every failed-revoke scenario has to reach: the routing table released +/// the shards from the old pod *before* the new pod was recorded as their owner, the old pod +/// was told its authoritative assignment before the new pod was assigned, and the revoke was +/// not retried against a pod the routing table no longer credits with the shards. +async fn assert_failed_revoke_converged( + worker_executors: &TestWorkerExecutors, + persistence: &TestPersistence, + old_pod: Pod, + new_pod: Pod, +) { + wait_for_local_assignment(worker_executors, old_pod, shard_ids(&[2, 3])).await; + wait_for_local_assignment(worker_executors, new_pod, shard_ids(&[0, 1])).await; + + let routing_table = persistence.latest().await; + assert_eq!(routing_table.get_shards(old_pod), Some(shard_ids(&[2, 3]))); + assert_eq!(routing_table.get_shards(new_pod), Some(shard_ids(&[0, 1]))); + assert!(routing_table.get_unassigned_shards().is_empty()); + + let calls = worker_executors.calls().await; + let revokes = calls + .iter() + .filter(|call| matches!(call, Call::Revoke(pod, _) if *pod == old_pod)) + .count(); + assert_eq!( + revokes, 1, + "the failed revoke must not be retried against the old pod: {calls:#?}" + ); + let reconcile_idx = calls + .iter() + .position(|call| *call == Call::Set(old_pod, shard_ids(&[2, 3]))) + .unwrap_or_else(|| panic!("old pod never got its authoritative assignment: {calls:#?}")); + let assign_idx = calls + .iter() + .position(|call| *call == Call::Assign(new_pod, shard_ids(&[0, 1]))) + .unwrap_or_else(|| panic!("new pod never got the shards: {calls:#?}")); + assert!( + reconcile_idx < assign_idx, + "old pod must be reconciled before the shards are assigned elsewhere: {calls:#?}" + ); + + let writes = persistence.writes().await; + let released_idx = writes + .iter() + .position(|routing_table| { + routing_table.get_shards(old_pod) == Some(shard_ids(&[2, 3])) + && routing_table.get_unassigned_shards() == shard_ids(&[0, 1]) + }) + .expect("shards were never persisted as unassigned after the failed revoke"); + let reassigned_idx = writes + .iter() + .position(|routing_table| routing_table.get_shards(new_pod) == Some(shard_ids(&[0, 1]))) + .expect("shards were never persisted as owned by the new pod"); + assert!( + released_idx < reassigned_idx, + "shards must be persisted as unassigned before they are persisted as reassigned" + ); +} + +#[test] +// The production failure: the executor drops the shards and only the response is lost. The +// routing table must stop crediting the old pod with the shards instead of retrying the same +// revoke forever, so the new pod actually receives them. +async fn revoke_timeout_after_executor_applied_it_does_not_strand_shards() { let old_pod = pod(1, 9000); let new_pod = pod(2, 9001); let worker_executors = Arc::new(TestWorkerExecutors::default()); worker_executors .set_local_assignment(old_pod, &[0, 1, 2, 3]) .await; - worker_executors.fail_next_revocations(old_pod, 1).await; + worker_executors + .apply_then_fail_next_revocations(old_pod, usize::MAX) + .await; let (shard_management, persistence, mut join_set) = new_shard_management( routing_table_with_pods(4, vec![(old_pod, "worker-executor-0", &[0, 1, 2, 3])]), @@ -518,69 +649,97 @@ async fn failed_revoke_is_retried_without_assigning_to_new_executor_first() { .register_pod(new_pod, Some("worker-executor-1".to_string())) .await; - wait_for_local_assignment(&worker_executors, old_pod, shard_ids(&[2, 3])).await; - wait_for_local_assignment(&worker_executors, new_pod, shard_ids(&[0, 1])).await; - - assert_eq!( - worker_executors.local_assignment(old_pod).await, - shard_ids(&[2, 3]) - ); - assert_eq!( - worker_executors.local_assignment(new_pod).await, - shard_ids(&[0, 1]) - ); - - let routing_table = persistence.latest().await; - assert_eq!( - routing_table - .pod_states - .get(&old_pod) - .expect("old pod missing") - .assigned_shards, - shard_ids(&[2, 3]) - ); - assert_eq!( - routing_table - .pod_states - .get(&new_pod) - .expect("new pod missing") - .assigned_shards, - shard_ids(&[0, 1]) - ); + assert_failed_revoke_converged(&worker_executors, &persistence, old_pod, new_pod).await; join_set.abort_all(); } #[test] -// Reconnect reconciliation failures are retried by the shard-manager worker. -async fn failed_reconnect_reconciliation_is_retried() { - let existing_pod = pod(1, 9000); +// The revoke never reached the executor: the shards are still released in the routing table, +// and the old pod is brought in line by its authoritative assignment before the new pod is +// assigned - not by retrying the revoke. +async fn failed_revoke_reconciles_old_executor_before_reassigning() { + let old_pod = pod(1, 9000); + let new_pod = pod(2, 9001); let worker_executors = Arc::new(TestWorkerExecutors::default()); + worker_executors + .set_local_assignment(old_pod, &[0, 1, 2, 3]) + .await; + worker_executors.fail_next_revocations(old_pod, 1).await; - let (shard_management, _persistence, mut join_set) = new_shard_management( - routing_table_with_pods( - 1, - vec![ - (existing_pod, "worker-executor-0", &[]), - (pod(2, 9001), "worker-executor-1", &[0]), - ], - ), + let (shard_management, persistence, mut join_set) = new_shard_management( + routing_table_with_pods(4, vec![(old_pod, "worker-executor-0", &[0, 1, 2, 3])]), worker_executors.clone(), ) .await; + shard_management + .register_pod(new_pod, Some("worker-executor-1".to_string())) + .await; + + assert_failed_revoke_converged(&worker_executors, &persistence, old_pod, new_pod).await; + + join_set.abort_all(); +} + +#[test] +// The old executor is unreachable: neither the revoke nor its authoritative assignment gets +// through. The shards are still released and handed to the new executor, and the old executor +// keeps being sent its authoritative assignment (until the health check removes it) rather than +// the revoke being retried. Its local state stays stale until then - the accepted trade-off. +async fn unreachable_executor_still_releases_its_revoked_shards() { + let old_pod = pod(1, 9000); + let new_pod = pod(2, 9001); + let worker_executors = Arc::new(TestWorkerExecutors::default()); worker_executors - .set_local_assignment(existing_pod, &[0]) + .set_local_assignment(old_pod, &[0, 1, 2, 3]) .await; worker_executors - .fail_next_reconciliations(existing_pod, 1) + .fail_next_revocations(old_pod, usize::MAX) + .await; + worker_executors + .fail_next_reconciliations(old_pod, usize::MAX) .await; + let (shard_management, persistence, mut join_set) = new_shard_management( + routing_table_with_pods(4, vec![(old_pod, "worker-executor-0", &[0, 1, 2, 3])]), + worker_executors.clone(), + ) + .await; + shard_management - .register_pod(existing_pod, Some("worker-executor-0".to_string())) + .register_pod(new_pod, Some("worker-executor-1".to_string())) .await; - wait_for_local_assignment(&worker_executors, existing_pod, BTreeSet::new()).await; - + wait_for_local_assignment(&worker_executors, new_pod, shard_ids(&[0, 1])).await; + // The loop keeps retrying the reconciliation without backoff; stop it before inspecting. join_set.abort_all(); + + assert_eq!( + worker_executors.local_assignment(old_pod).await, + shard_ids(&[0, 1, 2, 3]), + "an unreachable executor cannot be told anything; its local state stays stale" + ); + + let routing_table = persistence.latest().await; + assert_eq!(routing_table.get_shards(old_pod), Some(shard_ids(&[2, 3]))); + assert_eq!(routing_table.get_shards(new_pod), Some(shard_ids(&[0, 1]))); + + let calls = worker_executors.calls().await; + let revokes = calls + .iter() + .filter(|call| matches!(call, Call::Revoke(pod, _) if *pod == old_pod)) + .count(); + assert_eq!( + revokes, 1, + "the revoke must not be retried once the routing table released the shards: {calls:#?}" + ); + let reconciles = calls + .iter() + .filter(|call| *call == &Call::Set(old_pod, shard_ids(&[2, 3]))) + .count(); + assert!( + reconciles >= 2, + "the authoritative assignment must keep being retried: {calls:#?}" + ); } diff --git a/golem-worker-executor/src/durable_host/mod.rs b/golem-worker-executor/src/durable_host/mod.rs index cdd3e30cda..56cf5cae2c 100644 --- a/golem-worker-executor/src/durable_host/mod.rs +++ b/golem-worker-executor/src/durable_host/mod.rs @@ -6011,8 +6011,15 @@ impl ExternalOperations for DurableWorkerCtx { continue; }; - // TODO: there is probably a race here between assignment changing and a suspended worker getting woken up. + // The shard set is re-checked right before starting: the assignment may have changed + // again since `get_running_workers_in_shards` read it (the shard manager moves on after + // its RPC deadline, so a revoke can land while this recovery is still running), and a + // worker must never be started on an executor that no longer owns its shard. if should_restart_after_shard_assignment_change(&latest_worker_status) + && this + .shard_service() + .check_worker(&owned_agent_id.agent_id) + .is_ok() && let Err(err) = Worker::get_or_create_running( this, &owned_agent_id, diff --git a/golem-worker-executor/src/grpc/mod.rs b/golem-worker-executor/src/grpc/mod.rs index 34b909e783..dde860cf66 100644 --- a/golem-worker-executor/src/grpc/mod.rs +++ b/golem-worker-executor/src/grpc/mod.rs @@ -32,7 +32,7 @@ use crate::services::worker_event::WorkerEventReceiver; use crate::services::{ All, HasActiveAgents, HasAll, HasComponentService, HasEvents, HasOplogService, HasPromiseService, HasRunningWorkerEnumerationService, HasShardManagerService, HasShardService, - HasWorkerEnumerationService, HasWorkerService, UsesAllDeps, + HasShutdownToken, HasWorkerEnumerationService, HasWorkerService, UsesAllDeps, }; pub use crate::worker::{ PERMISSION_CARD_INSTALL_RECIPIENT_MISMATCH, PERMISSION_CARD_TRANSFER_PAYLOAD_CONFLICT, @@ -83,7 +83,8 @@ use golem_common::model::{ AgentEvent, AgentFilter, AgentFingerprint, AgentId, AgentInvocation, AgentMetadata, AgentStatus, IdempotencyKey, OwnedAgentId, ScanCursor, ShardId, Timestamp, }; -use golem_common::{model as common_model, recorded_grpc_api_request}; +use golem_common::tracing::TraceOrigin; +use golem_common::{model as common_model, recorded_grpc_api_request, related_span}; use golem_service_base::error::worker_executor::*; use golem_service_base::grpc::{ proto_agent_id_string, proto_idempotency_key_string, proto_promise_id_string, @@ -96,13 +97,48 @@ use std::marker::PhantomData; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +use std::time::Instant; use tokio::sync::broadcast::error::RecvError; use tokio_stream::wrappers::errors::BroadcastStreamRecvError; use tonic::{Request, Response, Status}; use tracing::info_span; -use tracing::{Instrument, info}; +use tracing::{Instrument, Level, debug, error, info}; use wasmtime::Error; +/// A change of this executor's shard assignment, as received from the shard manager. +#[derive(Clone, Copy, Debug)] +enum ShardAssignmentChange { + Revoke, + Assign, + Set, +} + +impl ShardAssignmentChange { + fn trigger(self) -> &'static str { + match self { + ShardAssignmentChange::Revoke => "revoke_shards", + ShardAssignmentChange::Assign => "assign_shards", + ShardAssignmentChange::Set => "set_shard_assignment", + } + } + + /// Whether the change can take shards away, so resident agents may have to be interrupted. + fn interrupts_lost_agents(self) -> bool { + matches!( + self, + ShardAssignmentChange::Revoke | ShardAssignmentChange::Set + ) + } + + /// Whether the change can add shards, so their running agents have to be recovered. + fn recovers_gained_agents(self) -> bool { + matches!( + self, + ShardAssignmentChange::Assign | ShardAssignmentChange::Set + ) + } +} + /// This is the implementation of the Worker Executor gRPC API pub struct WorkerExecutorImpl< Ctx: WorkerCtx, @@ -113,6 +149,9 @@ pub struct WorkerExecutorImpl< /// Holds the strong Arc to the worker activator so the Weak reference /// stored in LazyWorkerActivator remains valid while the gRPC server runs. _worker_activator: Arc>, + /// Serializes the background application of shard assignment changes, so at most one + /// drain/recovery pass is in flight at a time; later changes queue behind it. + shard_assignment_change_lock: Arc>, ctx: PhantomData, } @@ -123,6 +162,7 @@ impl + UsesAllDeps + Send + Sync + Self { services: self.services.clone(), _worker_activator: self._worker_activator.clone(), + shard_assignment_change_lock: self.shard_assignment_change_lock.clone(), ctx: PhantomData, } } @@ -159,6 +199,7 @@ impl + UsesAllDeps + Send + Sync + let worker_executor = WorkerExecutorImpl { services: services.clone(), _worker_activator: worker_activator, + shard_assignment_change_lock: Arc::new(tokio::sync::Mutex::new(())), ctx: PhantomData, }; @@ -950,16 +991,7 @@ impl + UsesAllDeps + Send + Sync + let shard_ids = proto_shard_ids.into_iter().map(ShardId::from).collect(); self.shard_service().revoke_shards(&shard_ids)?; - - for (agent_id, worker_details) in self.active_agents().snapshot().await { - if self.shard_service().check_worker(&agent_id).is_err() - && let Some(mut await_interrupted) = worker_details - .set_interrupting(InterruptKind::Restart) - .await - { - await_interrupted.recv().await.unwrap(); - } - } + self.apply_shard_assignment_change_in_background(ShardAssignmentChange::Revoke); Ok(()) } @@ -973,7 +1005,7 @@ impl + UsesAllDeps + Send + Sync + let shard_ids = proto_shard_ids.into_iter().map(ShardId::from).collect(); self.shard_service().assign_shards(&shard_ids)?; - Ctx::on_shard_assignment_changed(self).await?; + self.apply_shard_assignment_change_in_background(ShardAssignmentChange::Assign); Ok(()) } @@ -990,20 +1022,100 @@ impl + UsesAllDeps + Send + Sync + self.shard_service() .set_shard_assignment(number_of_shards, &shard_ids)?; + self.apply_shard_assignment_change_in_background(ShardAssignmentChange::Set); - for (agent_id, worker_details) in self.active_agents().snapshot().await { - if self.shard_service().check_worker(&agent_id).is_err() - && let Some(mut await_interrupted) = worker_details - .set_interrupting(InterruptKind::Restart) - .await - { - await_interrupted.recv().await.unwrap(); + Ok(()) + } + + fn apply_shard_assignment_change_in_background(&self, change: ShardAssignmentChange) { + let origin = TraceOrigin::capture_current(); + let this = self.clone(); + tokio::spawn(async move { + let span = related_span!( + origin, + Level::INFO, + "shard_assignment_change", + trigger = change.trigger(), + ); + this.apply_shard_assignment_change(change) + .instrument(span) + .await + }); + } + + async fn apply_shard_assignment_change(&self, change: ShardAssignmentChange) { + let trigger = change.trigger(); + let started = Instant::now(); + let shutdown = self.shutdown_token(); + + let _guard = tokio::select! { + guard = self.shard_assignment_change_lock.lock() => guard, + _ = shutdown.cancelled() => { + info!(trigger, "Shard assignment change skipped by executor shutdown"); + return; + } + }; + + let mut signalled: usize = 0; + let mut acknowledged: usize = 0; + + if change.interrupts_lost_agents() { + for (agent_id, worker) in self.active_agents().snapshot().await { + if self.shard_service().check_worker(&agent_id).is_ok() { + continue; + } + + let interrupt = async { + match worker.set_interrupting(InterruptKind::Restart).await { + Some(mut await_interrupted) => Some(await_interrupted.recv().await), + None => None, + } + }; + + tokio::select! { + outcome = interrupt => { + signalled += 1; + match outcome { + Some(Ok(())) => acknowledged += 1, + Some(Err(err)) => debug!( + agent_id = %agent_id, + error = %err, + "Agent went away before acknowledging its interruption" + ), + None => {} + } + } + _ = shutdown.cancelled() => { + info!( + trigger, + signalled, + acknowledged, + elapsed_ms = started.elapsed().as_millis() as u64, + "Shard assignment change interrupted by executor shutdown" + ); + return; + } + } } } - Ctx::on_shard_assignment_changed(self).await?; + if change.recovers_gained_agents() + && let Err(err) = Ctx::on_shard_assignment_changed(self).await + { + error!( + trigger, + error = %err, + "Failed to recover running agents after shard assignment change" + ); + } - Ok(()) + info!( + trigger, + signalled, + acknowledged, + elapsed_ms = started.elapsed().as_millis() as u64, + "Shard assignment change applied" + ); } async fn get_agent_metadata_internal( diff --git a/golem-worker-executor/src/worker/invocation_loop.rs b/golem-worker-executor/src/worker/invocation_loop.rs index 8245f36df8..6c956ba92f 100644 --- a/golem-worker-executor/src/worker/invocation_loop.rs +++ b/golem-worker-executor/src/worker/invocation_loop.rs @@ -17,7 +17,7 @@ use crate::services::events::Event; use crate::services::golem_config::SnapshotPolicy; use crate::services::linear_memory::LinearMemoryTracker; use crate::services::oplog::{CommitLevel, EphemeralOplog, OplogOps}; -use crate::services::{HasActiveAgents, HasEvents, HasOplog, HasWorker}; +use crate::services::{HasActiveAgents, HasEvents, HasOplog, HasShardService, HasWorker}; use crate::worker::invocation::{ InvocationMode, InvokeResult, invoke_observed_and_traced, lower_invocation, }; @@ -137,6 +137,19 @@ impl InvocationLoop { let mut deferred_wakeups = VecDeque::new(); 'outer: loop { + if let Err(err @ WorkerExecutorError::InvalidShardId { .. }) = + self.parent.shard_service().check_worker(&agent_id) + { + debug!( + %agent_id, + "Invocation queue loop unloading agent whose shard is no longer owned" + ); + self.parent.acknowledge_interruption(); + self.stop_unloaded(Some(err)).await; + self.parent.remove_from_active_agents().await; + break; + } + let entity_generation = self .parent .active_agents() diff --git a/golem-worker-executor/src/worker/mod.rs b/golem-worker-executor/src/worker/mod.rs index 432fa4c184..5b48539690 100644 --- a/golem-worker-executor/src/worker/mod.rs +++ b/golem-worker-executor/src/worker/mod.rs @@ -1379,6 +1379,26 @@ impl Worker { }; } + /// Resolves a pending interruption without going through the worker's store, releasing + /// whoever is awaiting it via the receiver handed out by `set_interrupting`. Used when the + /// invocation loop unloads the worker instead of restarting it, which is the one exit that + /// does not pass through the store's `set_suspended`. + pub(crate) fn acknowledge_interruption(&self) { + let mut execution_status = self.execution_status.write().unwrap(); + if let ExecutionStatus::Interrupting { + agent_mode, + await_interruption, + .. + } = execution_status.clone() + { + *execution_status = ExecutionStatus::Suspended { + agent_mode, + timestamp: Timestamp::now_utc(), + }; + await_interruption.send(()).ok(); + } + } + pub fn get_initial_worker_metadata(&self) -> AgentMetadata { self.initial_worker_metadata.clone() } From 6f436a1ee75bba0d872ee66ad39c3505ebc78512 Mon Sep 17 00:00:00 2001 From: Aditya Salunkhe Date: Thu, 27 Aug 2026 01:28:38 +0530 Subject: [PATCH 2/3] Keep RevokeShards synchronous: signal all lost agents at once, wait until each is unloaded, raise the revoke deadline to 60s --- .../config/shard-manager.sample.env | 4 +- golem-shard-manager/config/shard-manager.toml | 4 +- golem-shard-manager/src/config.rs | 2 +- .../src/sharding/worker_executor.rs | 2 +- golem-worker-executor-test-utils/src/lib.rs | 80 ++- golem-worker-executor/src/grpc/mod.rs | 249 +++++---- golem-worker-executor/tests/lib.rs | 3 + golem-worker-executor/tests/sharding.rs | 512 ++++++++++++++++++ 8 files changed, 757 insertions(+), 99 deletions(-) create mode 100644 golem-worker-executor/tests/sharding.rs diff --git a/golem-shard-manager/config/shard-manager.sample.env b/golem-shard-manager/config/shard-manager.sample.env index 383004f127..7cadc1395c 100644 --- a/golem-shard-manager/config/shard-manager.sample.env +++ b/golem-shard-manager/config/shard-manager.sample.env @@ -79,7 +79,7 @@ GOLEM__WORKER_EXECUTORS__CONNECT_TIMEOUT="10s" GOLEM__WORKER_EXECUTORS__HEALTH_CHECK_TIMEOUT="2s" GOLEM__WORKER_EXECUTORS__MAX_MESSAGE_SIZE=33554432 #GOLEM__WORKER_EXECUTORS__REQUEST_TIMEOUT= -GOLEM__WORKER_EXECUTORS__REVOKE_SHARDS_TIMEOUT="5s" +GOLEM__WORKER_EXECUTORS__REVOKE_SHARDS_TIMEOUT="1m" GOLEM__WORKER_EXECUTORS__RETRIES__MAX_ATTEMPTS=5 GOLEM__WORKER_EXECUTORS__RETRIES__MAX_DELAY="2s" GOLEM__WORKER_EXECUTORS__RETRIES__MAX_JITTER_FACTOR=0.15 @@ -174,7 +174,7 @@ GOLEM__WORKER_EXECUTORS__CONNECT_TIMEOUT="10s" GOLEM__WORKER_EXECUTORS__HEALTH_CHECK_TIMEOUT="2s" GOLEM__WORKER_EXECUTORS__MAX_MESSAGE_SIZE=33554432 #GOLEM__WORKER_EXECUTORS__REQUEST_TIMEOUT= -GOLEM__WORKER_EXECUTORS__REVOKE_SHARDS_TIMEOUT="5s" +GOLEM__WORKER_EXECUTORS__REVOKE_SHARDS_TIMEOUT="1m" GOLEM__WORKER_EXECUTORS__RETRIES__MAX_ATTEMPTS=5 GOLEM__WORKER_EXECUTORS__RETRIES__MAX_DELAY="2s" GOLEM__WORKER_EXECUTORS__RETRIES__MAX_JITTER_FACTOR=0.15 diff --git a/golem-shard-manager/config/shard-manager.toml b/golem-shard-manager/config/shard-manager.toml index c08176c617..7823bd4ec5 100644 --- a/golem-shard-manager/config/shard-manager.toml +++ b/golem-shard-manager/config/shard-manager.toml @@ -117,7 +117,7 @@ assign_shards_timeout = "5s" connect_timeout = "10s" health_check_timeout = "2s" max_message_size = 33554432 -revoke_shards_timeout = "5s" +revoke_shards_timeout = "1m" [worker_executors.retries] max_attempts = 5 @@ -259,7 +259,7 @@ type = "Disabled" # connect_timeout = "10s" # health_check_timeout = "2s" # max_message_size = 33554432 -# revoke_shards_timeout = "5s" +# revoke_shards_timeout = "1m" # # [worker_executors.retries] # max_attempts = 5 diff --git a/golem-shard-manager/src/config.rs b/golem-shard-manager/src/config.rs index 4ab3ba8b16..23b7dcd622 100644 --- a/golem-shard-manager/src/config.rs +++ b/golem-shard-manager/src/config.rs @@ -289,7 +289,7 @@ impl Default for WorkerExecutorServiceConfig { Self { assign_shards_timeout: Duration::from_secs(5), health_check_timeout: Duration::from_secs(2), - revoke_shards_timeout: Duration::from_secs(5), + revoke_shards_timeout: Duration::from_secs(60), retries: RetryConfig::max_attempts_5(), client_config: GrpcClientConfig { connect_timeout: Duration::from_secs(10), diff --git a/golem-shard-manager/src/sharding/worker_executor.rs b/golem-shard-manager/src/sharding/worker_executor.rs index 313db8345d..ce6f660e8a 100644 --- a/golem-shard-manager/src/sharding/worker_executor.rs +++ b/golem-shard-manager/src/sharding/worker_executor.rs @@ -363,7 +363,7 @@ impl WorkerExecutorServiceDefault { }; let set_shard_assignment_response = timeout( - self.config.assign_shards_timeout, + self.config.revoke_shards_timeout, self.client.call( "set_shard_assignment", pod.uri(self.config.client_config.tls_enabled()), diff --git a/golem-worker-executor-test-utils/src/lib.rs b/golem-worker-executor-test-utils/src/lib.rs index 656c8232b0..e0339e06fd 100644 --- a/golem-worker-executor-test-utils/src/lib.rs +++ b/golem-worker-executor-test-utils/src/lib.rs @@ -812,6 +812,24 @@ impl TestWorkerExecutor { .await } + /// Arms a one-shot artificial delay on `agent_id`'s next oplog commit and records the + /// interval that commit occupied. Armed right before a shard drain, the delayed commit is the + /// agent's teardown commit, so `teardown_commit_intervals` shows whether the drain tears + /// agents down concurrently or one by one. See + /// [`AdditionalTestDeps::arm_teardown_commit_delay`]. + pub async fn arm_teardown_commit_delay(&self, agent_id: &AgentId, delay: std::time::Duration) { + self.additional_test_deps + .arm_teardown_commit_delay(agent_id.clone(), delay) + .await + } + + /// The intervals occupied by the commits delayed with `arm_teardown_commit_delay`. + pub fn teardown_commit_intervals( + &self, + ) -> Vec<(AgentId, std::time::Instant, std::time::Instant)> { + self.additional_test_deps.teardown_commit_intervals() + } + /// Returns the per-worker memory requirement that the executor uses when /// reserving from the worker memory semaphore. Lets tests sanity-check that /// they have constrained the memory budget tightly enough to force @@ -3090,7 +3108,21 @@ impl Oplog for TestOplog { async fn commit(&self, level: CommitLevel) -> BTreeMap { self.additional_test_deps .record_oplog_call(&self.owned_agent_id, "commit"); - self.oplog.commit(level).await + let delay = self + .additional_test_deps + .take_teardown_commit_delay(&self.owned_agent_id.agent_id) + .await; + let started = std::time::Instant::now(); + let result = self.oplog.commit(level).await; + if let Some(delay) = delay { + tokio::time::sleep(delay).await; + self.additional_test_deps.record_teardown_commit_interval( + &self.owned_agent_id.agent_id, + started, + std::time::Instant::now(), + ); + } + result } async fn current_oplog_index(&self) -> OplogIndex { @@ -3443,6 +3475,13 @@ pub struct AdditionalTestDeps { consume_body_scope_start_gates: Arc>>, consume_body_scope_end_gates: Arc>>, consume_body_reply_defer_gates: Arc>>, + /// One-shot artificial latency applied to an agent's next oplog commit inside the + /// [`TestOplog`] wrapper, and the intervals those delayed commits occupied. Armed right + /// before a shard drain, the delayed commit is the agent's teardown commit, so the recorded + /// intervals show whether the drain tears agents down concurrently or one by one. + teardown_commit_delays: Arc>, + teardown_commit_intervals: + Arc>>, /// Captured once on first call to [`TestWorkerCtx::create`]. Used by the /// read-only test helpers (`worker_is_loaded`, /// `worker_eviction_class`, `worker_memory_requirement`) to observe @@ -3471,10 +3510,49 @@ impl AdditionalTestDeps { consume_body_scope_start_gates: Arc::new(scc::HashMap::new()), consume_body_scope_end_gates: Arc::new(scc::HashMap::new()), consume_body_reply_defer_gates: Arc::new(scc::HashMap::new()), + teardown_commit_delays: Arc::new(scc::HashMap::new()), + teardown_commit_intervals: Arc::new(std::sync::Mutex::new(Vec::new())), active_agents: Arc::new(std::sync::OnceLock::new()), } } + /// Arms a one-shot delay on the given agent's next oplog commit; see the field documentation. + /// Re-arming replaces a delay that has not fired yet. + pub async fn arm_teardown_commit_delay(&self, agent_id: AgentId, delay: std::time::Duration) { + self.teardown_commit_delays + .entry_async(agent_id) + .await + .and_modify(|existing| *existing = delay) + .or_insert(delay); + } + + async fn take_teardown_commit_delay(&self, agent_id: &AgentId) -> Option { + self.teardown_commit_delays + .remove_async(agent_id) + .await + .map(|(_, delay)| delay) + } + + fn record_teardown_commit_interval( + &self, + agent_id: &AgentId, + started: std::time::Instant, + finished: std::time::Instant, + ) { + self.teardown_commit_intervals + .lock() + .unwrap() + .push((agent_id.clone(), started, finished)); + } + + /// The intervals occupied by the delayed commits armed with `arm_teardown_commit_delay`, in + /// completion order. + pub fn teardown_commit_intervals( + &self, + ) -> Vec<(AgentId, std::time::Instant, std::time::Instant)> { + self.teardown_commit_intervals.lock().unwrap().clone() + } + /// Arms a one-shot gate that pauses the given agent's next consume-body /// chunk `End` append after it is durable but before the wrapper's /// `Oplog::add` returns. Must be called before the invocation that diff --git a/golem-worker-executor/src/grpc/mod.rs b/golem-worker-executor/src/grpc/mod.rs index dde860cf66..3ed8dfc04a 100644 --- a/golem-worker-executor/src/grpc/mod.rs +++ b/golem-worker-executor/src/grpc/mod.rs @@ -41,6 +41,7 @@ use crate::worker::{Worker, WorkerUpdateMode}; use crate::workerctx::WorkerCtx; use futures::Stream; use futures::StreamExt; +use futures::future::join_all; use golem_api_grpc::proto::golem; use golem_api_grpc::proto::golem::worker::{Cursor, InvocationRequest, UpdateMode}; use golem_api_grpc::proto::golem::workerexecutor::v1::worker_executor_server::WorkerExecutor; @@ -97,7 +98,7 @@ use std::marker::PhantomData; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use std::time::Instant; +use std::time::{Duration, Instant}; use tokio::sync::broadcast::error::RecvError; use tokio_stream::wrappers::errors::BroadcastStreamRecvError; use tonic::{Request, Response, Status}; @@ -105,7 +106,8 @@ use tracing::info_span; use tracing::{Instrument, Level, debug, error, info}; use wasmtime::Error; -/// A change of this executor's shard assignment, as received from the shard manager. +/// The shard manager call that changed this executor's shard assignment. Only used to label the +/// drain and recovery that follow, so they can be attributed to their cause in traces and logs. #[derive(Clone, Copy, Debug)] enum ShardAssignmentChange { Revoke, @@ -121,22 +123,6 @@ impl ShardAssignmentChange { ShardAssignmentChange::Set => "set_shard_assignment", } } - - /// Whether the change can take shards away, so resident agents may have to be interrupted. - fn interrupts_lost_agents(self) -> bool { - matches!( - self, - ShardAssignmentChange::Revoke | ShardAssignmentChange::Set - ) - } - - /// Whether the change can add shards, so their running agents have to be recovered. - fn recovers_gained_agents(self) -> bool { - matches!( - self, - ShardAssignmentChange::Assign | ShardAssignmentChange::Set - ) - } } /// This is the implementation of the Worker Executor gRPC API @@ -149,9 +135,6 @@ pub struct WorkerExecutorImpl< /// Holds the strong Arc to the worker activator so the Weak reference /// stored in LazyWorkerActivator remains valid while the gRPC server runs. _worker_activator: Arc>, - /// Serializes the background application of shard assignment changes, so at most one - /// drain/recovery pass is in flight at a time; later changes queue behind it. - shard_assignment_change_lock: Arc>, ctx: PhantomData, } @@ -162,7 +145,6 @@ impl + UsesAllDeps + Send + Sync + Self { services: self.services.clone(), _worker_activator: self._worker_activator.clone(), - shard_assignment_change_lock: self.shard_assignment_change_lock.clone(), ctx: PhantomData, } } @@ -199,7 +181,6 @@ impl + UsesAllDeps + Send + Sync + let worker_executor = WorkerExecutorImpl { services: services.clone(), _worker_activator: worker_activator, - shard_assignment_change_lock: Arc::new(tokio::sync::Mutex::new(())), ctx: PhantomData, }; @@ -991,7 +972,8 @@ impl + UsesAllDeps + Send + Sync + let shard_ids = proto_shard_ids.into_iter().map(ShardId::from).collect(); self.shard_service().revoke_shards(&shard_ids)?; - self.apply_shard_assignment_change_in_background(ShardAssignmentChange::Revoke); + self.drain_lost_agents(ShardAssignmentChange::Revoke) + .await?; Ok(()) } @@ -1005,7 +987,7 @@ impl + UsesAllDeps + Send + Sync + let shard_ids = proto_shard_ids.into_iter().map(ShardId::from).collect(); self.shard_service().assign_shards(&shard_ids)?; - self.apply_shard_assignment_change_in_background(ShardAssignmentChange::Assign); + self.recover_gained_agents_in_background(ShardAssignmentChange::Assign); Ok(()) } @@ -1022,100 +1004,183 @@ impl + UsesAllDeps + Send + Sync + self.shard_service() .set_shard_assignment(number_of_shards, &shard_ids)?; - self.apply_shard_assignment_change_in_background(ShardAssignmentChange::Set); + self.drain_lost_agents(ShardAssignmentChange::Set).await?; + self.recover_gained_agents_in_background(ShardAssignmentChange::Set); Ok(()) } - fn apply_shard_assignment_change_in_background(&self, change: ShardAssignmentChange) { + async fn drain_lost_agents( + &self, + change: ShardAssignmentChange, + ) -> Result<(), WorkerExecutorError> { let origin = TraceOrigin::capture_current(); let this = self.clone(); - tokio::spawn(async move { + let handle = tokio::spawn(async move { let span = related_span!( origin, Level::INFO, - "shard_assignment_change", - trigger = change.trigger(), + "shard_drain", + trigger = change.trigger() ); - this.apply_shard_assignment_change(change) + this.drain_lost_agents_to_completion(change) .instrument(span) .await }); + handle.await.map_err(|err| { + WorkerExecutorError::runtime(format!( + "Draining the agents of the revoked shards failed: {err}" + )) + })? } - async fn apply_shard_assignment_change(&self, change: ShardAssignmentChange) { + async fn drain_lost_agents_to_completion( + &self, + change: ShardAssignmentChange, + ) -> Result<(), WorkerExecutorError> { let trigger = change.trigger(); let started = Instant::now(); let shutdown = self.shutdown_token(); - let _guard = tokio::select! { - guard = self.shard_assignment_change_lock.lock() => guard, - _ = shutdown.cancelled() => { - info!(trigger, "Shard assignment change skipped by executor shutdown"); - return; - } - }; - - let mut signalled: usize = 0; - let mut acknowledged: usize = 0; - - if change.interrupts_lost_agents() { - for (agent_id, worker) in self.active_agents().snapshot().await { - if self.shard_service().check_worker(&agent_id).is_ok() { - continue; - } - - let interrupt = async { - match worker.set_interrupting(InterruptKind::Restart).await { - Some(mut await_interrupted) => Some(await_interrupted.recv().await), - None => None, + let drain = async { + let mut lost_total: usize = 0; + let mut signalled: usize = 0; + let mut acknowledged: usize = 0; + let mut passes: usize = 0; + + loop { + passes += 1; + + // Agents whose creation is still in flight are invisible to the snapshot, which + // is why this loops until a snapshot finds nothing loaded in a lost shard. + let mut lost = Vec::new(); + for (agent_id, worker) in self.active_agents().snapshot().await { + if self.shard_service().check_worker(&agent_id).is_err() + && worker.is_loaded().await + { + lost.push((agent_id, worker)); } - }; - - tokio::select! { - outcome = interrupt => { - signalled += 1; - match outcome { - Some(Ok(())) => acknowledged += 1, - Some(Err(err)) => debug!( - agent_id = %agent_id, - error = %err, - "Agent went away before acknowledging its interruption" - ), - None => {} + } + if lost.is_empty() { + break; + } + lost_total += lost.len(); + + // Signal all of them at once. + let receivers: Vec<_> = + join_all(lost.iter().map(|(agent_id, worker)| async move { + worker + .set_interrupting(InterruptKind::Restart) + .await + .map(|receiver| (agent_id.clone(), receiver)) + })) + .await + .into_iter() + .flatten() + .collect(); + signalled += receivers.len(); + + let acks = join_all(receivers.into_iter().map( + |(agent_id, mut receiver)| async move { + match receiver.recv().await { + Ok(()) => true, + Err(err) => { + debug!( + agent_id = %agent_id, + error = %err, + "Agent went away before acknowledging its interruption" + ); + false + } } + }, + )) + .await; + acknowledged += acks.into_iter().filter(|acked| *acked).count(); + + // Wait until every lost agent has left memory. `is_loaded` only turns false after + // the worker's final oplog commit and status flush. The checks of a round are + // issued together and only for what is still loaded, with a backoff between + // rounds, so a wedged agent costs one lock touch per round rather than a busy + // wait. + let mut remaining = lost; + let mut backoff = Duration::from_millis(10); + loop { + let still_loaded = join_all(remaining.iter().map(|(agent_id, worker)| { + async move { + // A shard handed back to this executor while the drain was running + // is owned again: its agents will not unload, so stop waiting for + // them. + self.shard_service().check_worker(agent_id).is_err() + && worker.is_loaded().await + } + })) + .await; + remaining = remaining + .into_iter() + .zip(still_loaded) + .filter_map(|(entry, loaded)| loaded.then_some(entry)) + .collect(); + if remaining.is_empty() { + break; } - _ = shutdown.cancelled() => { - info!( - trigger, - signalled, - acknowledged, - elapsed_ms = started.elapsed().as_millis() as u64, - "Shard assignment change interrupted by executor shutdown" - ); - return; - } + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(Duration::from_millis(100)); } } - } - if change.recovers_gained_agents() - && let Err(err) = Ctx::on_shard_assignment_changed(self).await - { - error!( - trigger, - error = %err, - "Failed to recover running agents after shard assignment change" - ); + (lost_total, signalled, acknowledged, passes) + }; + + tokio::select! { + (lost, signalled, acknowledged, passes) = drain => { + info!( + trigger, + lost, + signalled, + acknowledged, + passes, + elapsed_ms = started.elapsed().as_millis() as u64, + "Drained agents in revoked shards" + ); + Ok(()) + } + _ = shutdown.cancelled() => { + info!( + trigger, + elapsed_ms = started.elapsed().as_millis() as u64, + "Draining agents in revoked shards interrupted by executor shutdown" + ); + Err(WorkerExecutorError::runtime("Executor is shutting down")) + } } + } - info!( - trigger, - signalled, - acknowledged, - elapsed_ms = started.elapsed().as_millis() as u64, - "Shard assignment change applied" - ); + fn recover_gained_agents_in_background(&self, change: ShardAssignmentChange) { + let trigger = change.trigger(); + let origin = TraceOrigin::capture_current(); + let this = self.clone(); + tokio::spawn(async move { + let span = related_span!(origin, Level::INFO, "shard_assignment_recovery", trigger); + async move { + let started = Instant::now(); + match Ctx::on_shard_assignment_changed(&this).await { + Ok(()) => info!( + trigger, + elapsed_ms = started.elapsed().as_millis() as u64, + "Recovered running agents of newly owned shards" + ), + Err(err) => error!( + trigger, + error = %err, + elapsed_ms = started.elapsed().as_millis() as u64, + "Failed to recover running agents of newly owned shards" + ), + } + } + .instrument(span) + .await + }); } async fn get_agent_metadata_internal( diff --git a/golem-worker-executor/tests/lib.rs b/golem-worker-executor/tests/lib.rs index 2b18b22585..257d43b6bc 100644 --- a/golem-worker-executor/tests/lib.rs +++ b/golem-worker-executor/tests/lib.rs @@ -54,6 +54,7 @@ pub mod revert; pub mod rpc; pub mod scalability; pub mod scope_cards; +pub mod sharding; pub mod storage_quota; pub mod tool_discovery; pub mod transactions; @@ -101,6 +102,7 @@ tag_suite!(instance_layer, group2); tag_suite!(transactions, group2); tag_suite!(observability, group2); tag_suite!(retry_policies, group2); +tag_suite!(sharding, group2); tag_suite!(storage_quota, storage_quota); tag_suite!(rpc, group3); @@ -124,6 +126,7 @@ tag_suite!(tool_discovery, group1); sequential_suite!(key_value_storage); sequential_suite!(namespace_routed_key_value_storage); sequential_suite!(indexed_storage); +sequential_suite!(sharding); sequential_suite!(oplog_blob_archive); sequential_suite!(resource_limits); diff --git a/golem-worker-executor/tests/sharding.rs b/golem-worker-executor/tests/sharding.rs new file mode 100644 index 0000000000..77890c0b62 --- /dev/null +++ b/golem-worker-executor/tests/sharding.rs @@ -0,0 +1,512 @@ +// Copyright 2024-2026 Golem Cloud +// +// Licensed under the Golem Source License v1.1 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Shard revocation on a running executor. `RevokeShards` and the drain leg of +//! `SetShardAssignment` must not return before every agent of the revoked shards has actually +//! left memory (the shard manager hands the shards to another executor the moment they return), +//! and they must stop those agents concurrently rather than one by one. + +use crate::Tracing; +use golem_api_grpc::proto::golem::workerexecutor::v1::{ + AssignShardsRequest, RevokeShardsRequest, SetShardAssignmentRequest, assign_shards_response, + revoke_shards_response, set_shard_assignment_response, +}; +use golem_common::model::component::ComponentDto; +use golem_common::model::oplog::OplogIndex; +use golem_common::model::{AgentId, OwnedAgentId, ShardId}; +use golem_common::{agent_id, data_value}; +use golem_test_framework::dsl::{AgentResult, TestDsl, count_agent_invocation_pair_since}; +use golem_worker_executor::worker::EvictionClass; +use golem_worker_executor_test_utils::{ + LastUniqueId, PrecompiledComponent, TestContext, TestWorkerExecutor, + WorkerExecutorTestDependencies, start, +}; +use pretty_assertions::assert_eq; +use std::time::{Duration, Instant}; +use test_r::{inherit_test_dep, test, timeout}; +use tokio::task::JoinHandle; +use tracing::Instrument; + +inherit_test_dep!(WorkerExecutorTestDependencies); +inherit_test_dep!(LastUniqueId); +inherit_test_dep!(Tracing); +inherit_test_dep!( + #[tagged_as("host_api_tests")] + PrecompiledComponent +); +inherit_test_dep!( + #[tagged_as("agent_counters")] + PrecompiledComponent +); + +/// The in-process test executor owns the single shard `0` of one, so revoking it loses every +/// agent and assigning it back restores them all. +const SHARD: i64 = 0; +const AGENTS: usize = 5; + +/// `revoke_shards` must not return until every agent of the revoked shard has left memory: if +/// it returned earlier, the shard manager would hand the shard to another executor while these +/// agents can still write to their oplogs. +#[test] +#[tracing::instrument] +#[timeout("4m")] +async fn revoke_shards_returns_only_after_lost_agents_are_unloaded( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + _tracing: &Tracing, + #[tagged_as("host_api_tests")] host_api_tests: &PrecompiledComponent, +) -> anyhow::Result<()> { + let context = TestContext::new(last_unique_id); + let executor = start(deps, &context).await?; + let component = executor + .component_dep(&context.default_environment_id, host_api_tests) + .store() + .await?; + + let agents = start_busy_agents(&executor, &context, &component, "revoke-drain").await?; + + let drain_elapsed = revoke_shards(&executor, &[SHARD]).await?; + + assert_all_unloaded( + &executor, + agents.iter().map(|agent| &agent.owned_id), + "revoke_shards", + drain_elapsed, + ) + .await; + + assign_shards(&executor, &[SHARD]).await?; + assert_invocations_completed_once(&executor, agents).await?; + + drop(executor); + Ok(()) +} + +/// The drain signals every lost agent first and waits for all of them together, so its duration +/// is the slowest agent's teardown rather than the sum of all teardowns. Each agent's teardown +/// commit is delayed artificially; the recorded commit intervals overlap under a concurrent +/// drain and cannot overlap under a sequential one (agent `k+1` is only signalled after agent +/// `k` has finished its teardown). +/// +/// The delayed commit happens after the interrupt acknowledgement, so the overlap assertion on +/// its own would also be satisfied by a drain that returns without waiting at all; the unload +/// assertion before it (the same one `revoke_shards_returns_only_after_lost_agents_are_unloaded` +/// makes) is what pins the waiting. +#[test] +#[tracing::instrument] +#[timeout("4m")] +async fn revoke_shards_drains_lost_agents_concurrently( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + _tracing: &Tracing, + #[tagged_as("host_api_tests")] host_api_tests: &PrecompiledComponent, +) -> anyhow::Result<()> { + const TEARDOWN_DELAY: Duration = Duration::from_millis(200); + + let context = TestContext::new(last_unique_id); + let executor = start(deps, &context).await?; + let component = executor + .component_dep(&context.default_environment_id, host_api_tests) + .store() + .await?; + + let agents = start_busy_agents(&executor, &context, &component, "revoke-concurrent").await?; + + // The guests are parked in `poll` and commit nothing on their own, so the next commit of each + // agent is the one its teardown performs. + for agent in &agents { + executor + .arm_teardown_commit_delay(&agent.id, TEARDOWN_DELAY) + .await; + } + + let drain_elapsed = revoke_shards(&executor, &[SHARD]).await?; + + assert_all_unloaded( + &executor, + agents.iter().map(|agent| &agent.owned_id), + "revoke_shards", + drain_elapsed, + ) + .await; + + let intervals = executor.teardown_commit_intervals(); + assert_eq!( + intervals.len(), + AGENTS, + "every lost agent must have run its teardown commit inside the RPC: {intervals:?}" + ); + let depth = max_overlap_depth(&intervals); + assert!( + depth >= 2, + "the teardowns of the lost agents did not overlap (at most {depth} at a time), so the \ + drain stops agents one by one instead of concurrently: {intervals:?}" + ); + // A sequential drain takes at least AGENTS x TEARDOWN_DELAY; a concurrent one takes about one + // TEARDOWN_DELAY plus the real teardown work and the poll granularity of the unload barrier. + let sequential_bound = TEARDOWN_DELAY * AGENTS as u32; + assert!( + drain_elapsed < sequential_bound * 3 / 4, + "revoke_shards took {drain_elapsed:?} for {AGENTS} agents with a {TEARDOWN_DELAY:?} \ + teardown each; a concurrent drain finishes well under {sequential_bound:?}" + ); + + assign_shards(&executor, &[SHARD]).await?; + assert_invocations_completed_once(&executor, agents).await?; + + drop(executor); + Ok(()) +} + +/// An agent that finished its invocation stays loaded but idle. `set_interrupting` hands out no +/// acknowledgement for it, so a drain that only waited for acknowledgements would return with the +/// agent still in memory; it must be unloaded like the others before `revoke_shards` returns. +#[test] +#[tracing::instrument] +#[timeout("2m")] +async fn revoke_shards_unloads_idle_agents( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + _tracing: &Tracing, + #[tagged_as("agent_counters")] agent_counters: &PrecompiledComponent, +) -> anyhow::Result<()> { + let context = TestContext::new(last_unique_id); + let executor = start(deps, &context).await?; + let component = executor + .component_dep(&context.default_environment_id, agent_counters) + .store() + .await?; + + let mut agents = Vec::with_capacity(AGENTS); + for i in 0..AGENTS { + let parsed_id = agent_id!("InstantiationGrowthCounter", format!("revoke-idle-{i}")); + let id = executor + .start_agent(&component.id, parsed_id.clone()) + .await?; + let owned_id = OwnedAgentId::new(context.default_environment_id, &id); + + let count: u32 = executor + .invoke_and_await_agent(&component, &parsed_id, "increment", data_value!()) + .await? + .into_typed()?; + assert_eq!(count, 1); + + wait_for_eviction_class(&executor, &owned_id, EvictionClass::LoadedIdle).await?; + assert!(executor.worker_is_loaded(&owned_id).await); + agents.push((parsed_id, owned_id)); + } + + let drain_elapsed = revoke_shards(&executor, &[SHARD]).await?; + + assert_all_unloaded( + &executor, + agents.iter().map(|(_, owned_id)| owned_id), + "revoke_shards", + drain_elapsed, + ) + .await; + + // Replay must rebuild exactly the one increment that happened before the revoke. + assign_shards(&executor, &[SHARD]).await?; + for (parsed_id, _) in &agents { + let count: u32 = executor + .invoke_and_await_agent(&component, parsed_id, "increment", data_value!()) + .await? + .into_typed()?; + assert_eq!(count, 2); + } + + drop(executor); + Ok(()) +} + +/// `set_shard_assignment` can take shards away as well, and then has to drain their agents +/// before returning exactly like `revoke_shards` does. +#[test] +#[tracing::instrument] +#[timeout("4m")] +async fn set_shard_assignment_drains_lost_agents_before_returning( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + _tracing: &Tracing, + #[tagged_as("host_api_tests")] host_api_tests: &PrecompiledComponent, +) -> anyhow::Result<()> { + let context = TestContext::new(last_unique_id); + let executor = start(deps, &context).await?; + let component = executor + .component_dep(&context.default_environment_id, host_api_tests) + .store() + .await?; + + let agents = start_busy_agents(&executor, &context, &component, "set-drain").await?; + + let drain_elapsed = set_shard_assignment(&executor, 1, &[]).await?; + + assert_all_unloaded( + &executor, + agents.iter().map(|agent| &agent.owned_id), + "set_shard_assignment", + drain_elapsed, + ) + .await; + + set_shard_assignment(&executor, 1, &[SHARD]).await?; + assert_invocations_completed_once(&executor, agents).await?; + + drop(executor); + Ok(()) +} + +/// An agent of the `Clocks` type executing its `interruption` invocation: 100 sleeps of 100ms, +/// each far below the suspend threshold, so the agent stays loaded and busy for about ten +/// seconds and is interruptible at every sleep. +struct BusyAgent { + id: AgentId, + owned_id: OwnedAgentId, + /// The `invoke_and_await` call. It stays pending while the shard is revoked and completes + /// once the agent has been recovered and finished the invocation. + invocation: JoinHandle>, + /// Finished invocations in the oplog before the invocation under test was enqueued. + finished_before: usize, +} + +async fn start_busy_agents( + executor: &TestWorkerExecutor, + context: &TestContext, + component: &ComponentDto, + name_prefix: &str, +) -> anyhow::Result> { + let mut agents = Vec::with_capacity(AGENTS); + for i in 0..AGENTS { + let parsed_id = agent_id!("Clocks", format!("{name_prefix}-{i}")); + let id = executor + .start_agent(&component.id, parsed_id.clone()) + .await?; + let owned_id = OwnedAgentId::new(context.default_environment_id, &id); + + let entries = executor.get_oplog(&id, OplogIndex::INITIAL).await?; + let (_, finished_before) = count_agent_invocation_pair_since(&entries, OplogIndex::INITIAL); + + let invocation = { + let executor = executor.clone(); + let component = component.clone(); + tokio::spawn( + async move { + executor + .invoke_and_await_agent( + &component, + &parsed_id, + "interruption", + data_value!(), + ) + .await + } + .in_current_span(), + ) + }; + + agents.push(BusyAgent { + id, + owned_id, + invocation, + finished_before, + }); + } + + for agent in &agents { + wait_until_executing(executor, &agent.owned_id, Duration::from_secs(30)).await?; + } + + Ok(agents) +} + +async fn revoke_shards( + executor: &TestWorkerExecutor, + shard_ids: &[i64], +) -> anyhow::Result { + let started = Instant::now(); + let response = executor + .client + .clone() + .revoke_shards(RevokeShardsRequest { + shard_ids: proto_shard_ids(shard_ids), + }) + .await? + .into_inner(); + match response.result { + Some(revoke_shards_response::Result::Success(_)) => Ok(started.elapsed()), + other => anyhow::bail!("revoke_shards failed: {other:?}"), + } +} + +async fn assign_shards(executor: &TestWorkerExecutor, shard_ids: &[i64]) -> anyhow::Result<()> { + let response = executor + .client + .clone() + .assign_shards(AssignShardsRequest { + shard_ids: proto_shard_ids(shard_ids), + }) + .await? + .into_inner(); + match response.result { + Some(assign_shards_response::Result::Success(_)) => Ok(()), + other => anyhow::bail!("assign_shards failed: {other:?}"), + } +} + +async fn set_shard_assignment( + executor: &TestWorkerExecutor, + number_of_shards: u32, + shard_ids: &[i64], +) -> anyhow::Result { + let started = Instant::now(); + let response = executor + .client + .clone() + .set_shard_assignment(SetShardAssignmentRequest { + number_of_shards, + shard_ids: proto_shard_ids(shard_ids), + }) + .await? + .into_inner(); + match response.result { + Some(set_shard_assignment_response::Result::Success(_)) => Ok(started.elapsed()), + other => anyhow::bail!("set_shard_assignment failed: {other:?}"), + } +} + +fn proto_shard_ids(shard_ids: &[i64]) -> Vec { + shard_ids + .iter() + .map(|shard_id| ShardId::new(*shard_id).into()) + .collect() +} + +/// Waits until the agent is loaded and actively executing. Stronger than waiting for the +/// `Running` status, which is read from the deferred status blob: the eviction class is `None` +/// exactly while the worker executes. +async fn wait_until_executing( + executor: &TestWorkerExecutor, + owned_id: &OwnedAgentId, + timeout: Duration, +) -> anyhow::Result<()> { + let deadline = Instant::now() + timeout; + loop { + if executor.worker_is_loaded(owned_id).await + && executor.worker_eviction_class(owned_id).await.is_none() + { + return Ok(()); + } + if Instant::now() > deadline { + anyhow::bail!( + "agent {owned_id} did not start executing within {timeout:?} (loaded: {}, \ + eviction class: {:?})", + executor.worker_is_loaded(owned_id).await, + executor.worker_eviction_class(owned_id).await + ); + } + tokio::time::sleep(Duration::from_millis(20)).await; + } +} + +/// Waits until `worker_eviction_class(owned_id)` is `expected`, or fails after 5s. +async fn wait_for_eviction_class( + executor: &TestWorkerExecutor, + owned_id: &OwnedAgentId, + expected: EvictionClass, +) -> anyhow::Result<()> { + let deadline = Instant::now() + Duration::from_secs(5); + loop { + if executor.worker_eviction_class(owned_id).await == Some(expected) { + return Ok(()); + } + if Instant::now() > deadline { + anyhow::bail!( + "agent {owned_id} did not reach EvictionClass::{expected:?} within 5s (current \ + class: {:?})", + executor.worker_eviction_class(owned_id).await + ); + } + tokio::time::sleep(Duration::from_millis(25)).await; + } +} + +/// Asserts, without waiting, that none of the agents is loaded any more. +async fn assert_all_unloaded( + executor: &TestWorkerExecutor, + owned_ids: impl IntoIterator, + call: &str, + elapsed: Duration, +) { + for owned_id in owned_ids { + assert!( + !executor.worker_is_loaded(owned_id).await, + "agent {owned_id} was still loaded when {call} returned (after {elapsed:?}); another \ + executor could already be recovering it" + ); + } +} + +/// Joins the invocations that were interrupted by the revoke after the shard has been given back, +/// and checks that each of them ran to completion exactly once. +async fn assert_invocations_completed_once( + executor: &TestWorkerExecutor, + agents: Vec, +) -> anyhow::Result<()> { + for BusyAgent { + id, + invocation, + finished_before, + .. + } in agents + { + let result = tokio::time::timeout(Duration::from_secs(120), invocation) + .await + .map_err(|_| { + anyhow::anyhow!( + "agent {id} did not finish its invocation after the shard was restored" + ) + })??; + let value: String = result?.into_typed()?; + assert_eq!(value, "done", "agent {id}"); + + let entries = executor.get_oplog(&id, OplogIndex::INITIAL).await?; + let (_, finished) = count_agent_invocation_pair_since(&entries, OplogIndex::INITIAL); + assert_eq!( + finished - finished_before, + 1, + "agent {id}: the interrupted invocation must be recorded as finished exactly once" + ); + } + Ok(()) +} + +/// The largest number of intervals that overlap at any single instant. +fn max_overlap_depth(intervals: &[(AgentId, Instant, Instant)]) -> usize { + let mut events: Vec<(Instant, i32)> = Vec::with_capacity(intervals.len() * 2); + for (_, started, finished) in intervals { + events.push((*started, 1)); + events.push((*finished, -1)); + } + // Ends sort before starts at the same instant, so touching intervals do not count as + // overlapping. + events.sort_by(|a, b| a.0.cmp(&b.0).then(a.1.cmp(&b.1))); + let mut current = 0i32; + let mut max = 0i32; + for (_, delta) in events { + current += delta; + max = max.max(current); + } + max as usize +} From 11a3064e9386fe8446e49e46f3ee093c5ec52044 Mon Sep 17 00:00:00 2001 From: Aditya Salunkhe Date: Tue, 1 Sep 2026 13:34:06 +0530 Subject: [PATCH 3/3] Address review findings: wait for in-flight creations in the drain, reconcile failed assignments before reassigning, single-flight the drain; survive caller-cancelled creations, never remove pending cache entries, guard the drain owner slot --- golem-common/src/cache.rs | 61 ++++++ .../src/sharding/shard_management.rs | 93 ++++++--- golem-shard-manager/tests/shard_management.rs | 151 ++++++++++++++ golem-worker-executor-test-utils/src/lib.rs | 111 ++++++++++ golem-worker-executor/src/grpc/mod.rs | 189 +++++++++++++++-- .../src/services/active_agents/mod.rs | 52 ++++- golem-worker-executor/src/worker/mod.rs | 11 + golem-worker-executor/tests/sharding.rs | 191 ++++++++++++++++++ 8 files changed, 814 insertions(+), 45 deletions(-) diff --git a/golem-common/src/cache.rs b/golem-common/src/cache.rs index d3aff356c1..cd4e34d3f9 100644 --- a/golem-common/src/cache.rs +++ b/golem-common/src/cache.rs @@ -592,6 +592,23 @@ impl< keys } + /// Keys whose entries are still pending (their producer has not resolved yet). + /// Complements `iter`, which returns only cached entries, and `keys`, which + /// returns both. + pub async fn pending_keys(&self) -> Vec { + let mut keys = vec![]; + self.state + .items + .iter_async(|key, value| { + if matches!(value, Item::Pending { .. }) { + keys.push(key.clone()); + } + true + }) + .await; + keys + } + pub async fn remove(&self, key: &K) { let removed = self.state.items.remove_async(key).await.is_some(); if removed { @@ -2121,6 +2138,50 @@ mod tests { f2_proceed.notify_one(); } + #[test] + async fn pending_keys_returns_only_pending_keys() { + let cache = test_cache("pending_keys_only"); + let f2_entered = Arc::new(tokio::sync::Notify::new()); + let f2_proceed = Arc::new(tokio::sync::Notify::new()); + + // Insert a cached value for key 1 + cache + .get_or_insert_simple(&1, || async { Ok(10u64) }) + .await + .unwrap(); + + // Start a pending insert for key 2 + let cache_clone = cache.clone(); + let entered = f2_entered.clone(); + let proceed = f2_proceed.clone(); + let producer = tokio::spawn(async move { + cache_clone + .get_or_insert_simple(&2, || async move { + entered.notify_one(); + proceed.notified().await; + Ok(20u64) + }) + .await + }); + + f2_entered.notified().await; + + assert_eq!( + cache.pending_keys().await, + vec![2], + "pending_keys() should contain exactly the pending key" + ); + + f2_proceed.notify_one(); + producer.await.unwrap().unwrap(); + + assert_eq!( + cache.pending_keys().await, + Vec::::new(), + "pending_keys() should be empty once the producer resolved" + ); + } + // ---- Remove while pending ---- #[test] diff --git a/golem-shard-manager/src/sharding/shard_management.rs b/golem-shard-manager/src/sharding/shard_management.rs index 95fb241551..f622c20d1b 100644 --- a/golem-shard-manager/src/sharding/shard_management.rs +++ b/golem-shard-manager/src/sharding/shard_management.rs @@ -62,7 +62,9 @@ impl ShardManagement { let change = Arc::new(Notify::new()); // NOTE: We consider all healthy pods as new pods to trigger full assigment, given they might be lagging: - // this can happen with interleaved shard-manager and worker restarts + // this can happen with interleaved shard-manager and worker restarts. + // The first pass reconciles them BEFORE its rebalance (the pre-rebalance stage in + // `worker`), so a lagging pod drops stale shards before they can be reassigned. let updates = Arc::new(Mutex::new(ShardManagementChanges::new( healthy_pods, unhealthy_pods, @@ -142,8 +144,9 @@ impl ShardManagement { // Getting a write lock while // - the rebalance plan is calculated, // - new and removed pods are added to the routing table and got persisted, + // - a snapshot of that persisted state is taken for the pre-rebalance reconciles, // but the rebalance plan is NOT applied yet. The lock is then release for apply. - let (mut rebalance, mut full_assignment_pods) = { + let (mut rebalance, carried_over_pods, pass_start_snapshot) = { let mut current_routing_table = routing_table.write().await; for pod in removed_pods { @@ -165,15 +168,15 @@ impl ShardManagement { } let rebalance = Rebalance::from_routing_table(¤t_routing_table, threshold); - let mut full_assignment_pods: HashSet = HashSet::new(); + let mut carried_over_pods: HashSet = HashSet::new(); for pod in send_full_assignment { - full_assignment_pods.insert(pod); + carried_over_pods.insert(pod); } for pod in retry_full_assignment_pods { if current_routing_table.has_pod(pod) { - full_assignment_pods.insert(pod); + carried_over_pods.insert(pod); } } @@ -182,9 +185,40 @@ impl ShardManagement { .await .expect("Failed to persist routing table after pod changes"); - (rebalance, full_assignment_pods) + (rebalance, carried_over_pods, current_routing_table.clone()) }; + // Pods carried over from a previous pass (failed reconciles) or (re)connecting pods + // may still hold shards the routing table no longer credits them with; they get their + // authoritative assignment BEFORE the rebalance below can hand any of those shards to + // another pod. Failures are retried once more this pass from the post-rebalance + // snapshot (the pod is still in the table mid-pass), and a failure of that retry is + // re-queued for the next pass by the handler further down. + let mut full_assignment_pods: HashSet = HashSet::new(); + if !carried_over_pods.is_empty() { + let pre_assignments = + Self::full_assignments_for(&pass_start_snapshot, &carried_over_pods); + let failed_pre_sets = if pre_assignments.is_empty() { + Vec::new() + } else { + set_shard_assignments( + worker_executors.clone(), + pass_start_snapshot.number_of_shards, + &pre_assignments, + ) + .await + }; + if !failed_pre_sets.is_empty() { + warn!( + failed_pods = failed_pre_sets.iter().map(|(pod, _)| pod).join(", "), + "Some pods could not receive their authoritative shard assignment before the rebalance; retrying after it" + ); + for (pod, _) in &failed_pre_sets { + full_assignment_pods.insert(*pod); + } + } + } + debug!(rebalance=%rebalance, "Applying rebalance plan"); let rebalance_failures = Self::execute_rebalance(worker_executors.clone(), &mut rebalance).await; @@ -200,16 +234,17 @@ impl ShardManagement { warn!( failed_shards = failed_shards.iter().join(", "), - "Some shards could not be assigned and will be left unassigned for retry" + "Some shards could not be assigned; they are left unassigned and the pods get their authoritative assignment" ); // The executor may have applied the assignment even though the call failed - // (for example a timeout), so it always gets the authoritative assignment next. - { - let mut updates_guard = updates.lock().await; - for (pod, _) in &rebalance_failures.failed_assignments { - updates_guard.retry_full_assignment(*pod); - } + // (for example a timeout), so the pod receives its authoritative assignment in + // this pass - the post-rebalance snapshot no longer credits it with the failed + // shards - before the next pass can hand them to another pod. If that + // reconciliation fails too, it is re-queued below and re-sent at the START of the + // next pass, before that pass's rebalance. + for (pod, _) in &rebalance_failures.failed_assignments { + full_assignment_pods.insert(*pod); } needs_retry = true; } @@ -227,7 +262,8 @@ impl ShardManagement { // dropped them and only the response was lost. The shards are left unassigned in // the routing table (the unassignment stays in the plan), and the pod receives its // authoritative assignment in this pass, before the next pass hands the shards to - // another pod. If that reconciliation fails too, it is re-queued below. + // another pod. If that reconciliation fails too, it is re-queued below and re-sent + // at the START of the next pass, before that pass's rebalance. for (pod, _) in &rebalance_failures.failed_unassignments { full_assignment_pods.insert(*pod); } @@ -242,16 +278,8 @@ impl ShardManagement { .await .expect("Failed to persist routing table after rebalance"); - let mut full_assignments = Assignments::new(); - for pod in &full_assignment_pods { - if let Some(mut shard_ids) = routing_table_snapshot.get_shards(*pod) { - full_assignments - .assignments - .entry(*pod) - .or_default() - .append(&mut shard_ids); - } - } + let full_assignments = + Self::full_assignments_for(&routing_table_snapshot, &full_assignment_pods); let failed_full_assignments = if full_assignments.is_empty() { Vec::new() @@ -288,6 +316,23 @@ impl ShardManagement { } } + /// The full authoritative assignment message for each of `pods`, read from `snapshot`. + /// A pod not in the snapshot gets none (it was removed meanwhile); a pod without shards + /// gets an explicit empty assignment, which tells it to drop everything it still holds. + fn full_assignments_for(snapshot: &RoutingTable, pods: &HashSet) -> Assignments { + let mut full_assignments = Assignments::new(); + for pod in pods { + if let Some(mut shard_ids) = snapshot.get_shards(*pod) { + full_assignments + .assignments + .entry(*pod) + .or_default() + .append(&mut shard_ids); + } + } + full_assignments + } + async fn execute_rebalance( worker_executors: Arc, rebalance: &mut Rebalance, diff --git a/golem-shard-manager/tests/shard_management.rs b/golem-shard-manager/tests/shard_management.rs index 1f45968a37..e66cd0dd81 100644 --- a/golem-shard-manager/tests/shard_management.rs +++ b/golem-shard-manager/tests/shard_management.rs @@ -76,6 +76,7 @@ struct TestWorkerExecutors { local_assignments: Arc>>>, calls: Arc>>, failed_assignments: Arc>>, + applied_then_failed_assignments: Arc>>, failed_revocations: Arc>>, applied_then_failed_revocations: Arc>>, failed_reconciliations: Arc>>, @@ -110,6 +111,13 @@ impl TestWorkerExecutors { self.failed_assignments.lock().await.insert(pod, count); } + async fn apply_then_fail_next_assignments(&self, pod: Pod, count: usize) { + self.applied_then_failed_assignments + .lock() + .await + .insert(pod, count); + } + async fn fail_next_revocations(&self, pod: Pod, count: usize) { self.failed_revocations.lock().await.insert(pod, count); } @@ -155,6 +163,10 @@ impl WorkerExecutorService for TestWorkerExecutors { .entry(*pod) .or_default() .extend(shard_ids.iter().copied()); + + if Self::should_fail(&self.applied_then_failed_assignments, *pod).await { + return Err(ShardManagerError::Timeout); + } Ok(()) } @@ -743,3 +755,142 @@ async fn unreachable_executor_still_releases_its_revoked_shards() { "the authoritative assignment must keep being retried: {calls:#?}" ); } + +/// Shared arrange for the failed-assignment ordering tests: three balanced pods, then pod C is +/// removed so its shards need new owners. The round-robin gives shard 4 to A and shard 5 to B; +/// B's assign is armed to fail. The failed pod must be non-empty for the ordering bug to be +/// reachable (empty pods are always re-targeted first by the rebalancer). +async fn failed_assignment_reproducer( + worker_executors: &Arc, +) -> ( + ShardManagement, + TestPersistence, + JoinSet>, +) { + let pod_a = pod(1, 9000); + let pod_b = pod(2, 9001); + let pod_c = pod(3, 9002); + worker_executors.set_local_assignment(pod_a, &[0, 1]).await; + worker_executors.set_local_assignment(pod_b, &[2, 3]).await; + worker_executors.set_local_assignment(pod_c, &[4, 5]).await; + + new_shard_management( + routing_table_with_pods( + 6, + vec![ + (pod_a, "worker-executor-0", &[0, 1]), + (pod_b, "worker-executor-1", &[2, 3]), + (pod_c, "worker-executor-2", &[4, 5]), + ], + ), + worker_executors.clone(), + ) + .await +} + +#[test] +// An assignment the executor applied while the response was lost: the pod must get its +// authoritative assignment (dropping the in-doubt shard) before the shard is assigned to +// another pod - otherwise both pods own it until the reconcile lands. +async fn assign_timeout_after_executor_applied_it_is_reconciled_before_reassignment() { + let pod_a = pod(1, 9000); + let pod_b = pod(2, 9001); + let pod_c = pod(3, 9002); + let worker_executors = Arc::new(TestWorkerExecutors::default()); + let (shard_management, persistence, mut join_set) = + failed_assignment_reproducer(&worker_executors).await; + + // The startup pass reconciles all three pods; only calls after it matter. + let baseline = worker_executors.calls().await.len(); + worker_executors + .apply_then_fail_next_assignments(pod_b, 1) + .await; + shard_management.unregister_pod(pod_c).await; + + wait_for_local_assignment(&worker_executors, pod_a, shard_ids(&[1, 4, 5])).await; + wait_for_local_assignment(&worker_executors, pod_b, shard_ids(&[0, 2, 3])).await; + join_set.abort_all(); + + let routing_table = persistence.latest().await; + assert_eq!(routing_table.get_shards(pod_a), Some(shard_ids(&[1, 4, 5]))); + assert_eq!(routing_table.get_shards(pod_b), Some(shard_ids(&[0, 2, 3]))); + assert!(routing_table.get_unassigned_shards().is_empty()); + + let calls = worker_executors.calls().await[baseline..].to_vec(); + let failed_assign_idx = calls + .iter() + .position(|call| *call == Call::Assign(pod_b, shard_ids(&[5]))) + .unwrap_or_else(|| panic!("pod B was never assigned shard 5: {calls:#?}")); + let reconcile_idx = calls + .iter() + .position(|call| *call == Call::Set(pod_b, shard_ids(&[2, 3]))) + .unwrap_or_else(|| { + panic!( + "pod B never got its authoritative assignment after the failed assign: {calls:#?}" + ) + }); + let reassign_idx = calls + .iter() + .position(|call| *call == Call::Assign(pod_a, shard_ids(&[5]))) + .unwrap_or_else(|| panic!("shard 5 was never re-assigned: {calls:#?}")); + assert!( + failed_assign_idx < reconcile_idx, + "the reconcile must come after the failed assign: {calls:#?}" + ); + assert!( + reconcile_idx < reassign_idx, + "pod B must be reconciled before shard 5 is assigned elsewhere: {calls:#?}" + ); +} + +#[test] +// If the same-pass reconcile of a failed assignment also fails, the re-queued authoritative +// assignment must go out at the START of the next pass - before that pass's rebalance can +// hand the in-doubt shard to another pod. +async fn queued_reconcile_is_sent_before_next_pass_reassigns() { + let pod_a = pod(1, 9000); + let pod_b = pod(2, 9001); + let pod_c = pod(3, 9002); + let worker_executors = Arc::new(TestWorkerExecutors::default()); + let (shard_management, persistence, mut join_set) = + failed_assignment_reproducer(&worker_executors).await; + + let baseline = worker_executors.calls().await.len(); + worker_executors + .apply_then_fail_next_assignments(pod_b, 1) + .await; + worker_executors.fail_next_reconciliations(pod_b, 1).await; + shard_management.unregister_pod(pod_c).await; + + wait_for_local_assignment(&worker_executors, pod_a, shard_ids(&[1, 4, 5])).await; + wait_for_local_assignment(&worker_executors, pod_b, shard_ids(&[0, 2, 3])).await; + join_set.abort_all(); + + let routing_table = persistence.latest().await; + assert_eq!(routing_table.get_shards(pod_a), Some(shard_ids(&[1, 4, 5]))); + assert_eq!(routing_table.get_shards(pod_b), Some(shard_ids(&[0, 2, 3]))); + + let calls = worker_executors.calls().await[baseline..].to_vec(); + let reconcile_indices: Vec = calls + .iter() + .enumerate() + .filter_map(|(idx, call)| (*call == Call::Set(pod_b, shard_ids(&[2, 3]))).then_some(idx)) + .collect(); + assert!( + reconcile_indices.len() >= 2, + "the failed reconcile must be retried (the failed attempt is recorded too): {calls:#?}" + ); + let last_reconcile_idx = *reconcile_indices.last().unwrap(); + let revoke_idx = calls + .iter() + .position(|call| *call == Call::Revoke(pod_a, shard_ids(&[0]))) + .unwrap_or_else(|| panic!("the next pass never rebalanced: {calls:#?}")); + let reassign_idx = calls + .iter() + .position(|call| *call == Call::Assign(pod_a, shard_ids(&[5]))) + .unwrap_or_else(|| panic!("shard 5 was never re-assigned: {calls:#?}")); + assert!( + last_reconcile_idx < revoke_idx && last_reconcile_idx < reassign_idx, + "the re-queued reconcile must go out before the next pass's rebalance: {calls:#?}" + ); +} diff --git a/golem-worker-executor-test-utils/src/lib.rs b/golem-worker-executor-test-utils/src/lib.rs index e0339e06fd..b9db5199ec 100644 --- a/golem-worker-executor-test-utils/src/lib.rs +++ b/golem-worker-executor-test-utils/src/lib.rs @@ -830,6 +830,17 @@ impl TestWorkerExecutor { self.additional_test_deps.teardown_commit_intervals() } + /// Arms a one-shot gate that parks `agent_id`'s next oplog commit until released. + /// Armed before the agent's first creation, the parked commit is the + /// initialization enqueue inside `Worker::new`, holding the agent's creation in + /// flight (a pending `ActiveAgents` entry). See + /// [`AdditionalTestDeps::park_next_oplog_commit`]. + pub async fn park_next_oplog_commit(&self, agent_id: &AgentId) -> CommitParkHandle { + self.additional_test_deps + .park_next_oplog_commit(agent_id.clone()) + .await + } + /// Returns the per-worker memory requirement that the executor uses when /// reserving from the worker memory semaphore. Lets tests sanity-check that /// they have constrained the memory budget tightly enough to force @@ -2911,6 +2922,31 @@ impl TestOplog { permit.forget(); } + /// Trips a one-shot [`CommitParkGate`] armed for this agent: signals the test + /// that the commit is parked, then blocks until the test releases the gate. + /// Fires at most once per gate. + async fn park_at_commit_park_gate(&self) { + let Some(gate) = self + .additional_test_deps + .commit_park_gate(&self.owned_agent_id.agent_id) + .await + else { + return; + }; + if !gate.armed.swap(false, Ordering::SeqCst) { + return; + } + if let Some(parked_tx) = gate.parked_tx.lock().unwrap().take() { + let _ = parked_tx.send(()); + } + let permit = gate + .release + .acquire() + .await + .expect("the commit park gate semaphore was closed"); + permit.forget(); + } + fn is_consume_body_scope_start(entry: &OplogEntry) -> bool { matches!(entry, OplogEntry::Start { function_name: HostFunctionName::Custom(function_name), @@ -3108,6 +3144,7 @@ impl Oplog for TestOplog { async fn commit(&self, level: CommitLevel) -> BTreeMap { self.additional_test_deps .record_oplog_call(&self.owned_agent_id, "commit"); + self.park_at_commit_park_gate().await; let delay = self .additional_test_deps .take_teardown_commit_delay(&self.owned_agent_id.agent_id) @@ -3482,6 +3519,12 @@ pub struct AdditionalTestDeps { teardown_commit_delays: Arc>, teardown_commit_intervals: Arc>>, + /// One-shot gates parking an agent's next oplog commit inside the [`TestOplog`] + /// wrapper until released. Armed before an agent's first creation, the parked + /// commit is the `AgentInitialization` enqueue inside `Worker::new`, holding the + /// agent's `ActiveAgents` entry in its `Pending` state — the seam for testing that + /// a shard drain waits for in-flight creations. + commit_park_gates: Arc>>, /// Captured once on first call to [`TestWorkerCtx::create`]. Used by the /// read-only test helpers (`worker_is_loaded`, /// `worker_eviction_class`, `worker_memory_requirement`) to observe @@ -3512,10 +3555,35 @@ impl AdditionalTestDeps { consume_body_reply_defer_gates: Arc::new(scc::HashMap::new()), teardown_commit_delays: Arc::new(scc::HashMap::new()), teardown_commit_intervals: Arc::new(std::sync::Mutex::new(Vec::new())), + commit_park_gates: Arc::new(scc::HashMap::new()), active_agents: Arc::new(std::sync::OnceLock::new()), } } + /// Arms a one-shot gate that parks the given agent's next oplog commit until + /// released; see the field documentation. Re-arming replaces a previously fired + /// gate, so a test can gate one commit per executor run. + pub async fn park_next_oplog_commit(&self, agent_id: AgentId) -> CommitParkHandle { + let (parked_tx, parked_rx) = tokio::sync::oneshot::channel(); + let gate = Arc::new(CommitParkGate { + armed: AtomicBool::new(true), + parked_tx: std::sync::Mutex::new(Some(parked_tx)), + release: tokio::sync::Semaphore::new(0), + }); + self.commit_park_gates + .entry_async(agent_id) + .await + .and_modify(|existing| *existing = gate.clone()) + .or_insert_with(|| gate.clone()); + CommitParkHandle { parked_rx, gate } + } + + async fn commit_park_gate(&self, agent_id: &AgentId) -> Option> { + self.commit_park_gates + .read_async(agent_id, |_, gate| gate.clone()) + .await + } + /// Arms a one-shot delay on the given agent's next oplog commit; see the field documentation. /// Re-arming replaces a delay that has not fired yet. pub async fn arm_teardown_commit_delay(&self, agent_id: AgentId, delay: std::time::Duration) { @@ -3913,6 +3981,49 @@ impl Drop for ConsumeBodyScopeEndGateHandle { } } +/// A one-shot pause point at an agent's next oplog commit, shared between +/// [`AdditionalTestDeps`] (which arms it) and the agent's [`TestOplog`] (which +/// trips it). The gated commit fires `parked_tx` before the underlying commit +/// runs and then blocks until a `release` permit arrives, so a test can hold a +/// `Worker::new` (whose initialization enqueue commits through the wrapper) +/// parked inside its `ActiveAgents` pending entry. +struct CommitParkGate { + armed: AtomicBool, + parked_tx: std::sync::Mutex>>, + release: tokio::sync::Semaphore, +} + +/// Test-facing side of a [`CommitParkGate`]: await [`Self::parked`] to learn +/// that the gated commit is paused, then [`Self::release`] to let it continue. +pub struct CommitParkHandle { + parked_rx: tokio::sync::oneshot::Receiver<()>, + gate: Arc, +} + +impl CommitParkHandle { + /// Resolves once the gated commit is paused inside `Oplog::commit`. + pub async fn parked(&mut self) { + (&mut self.parked_rx) + .await + .expect("the commit park gate was dropped without firing"); + } + + /// Releases the paused commit. + pub fn release(&self) { + self.gate.release.add_permits(1); + } +} + +impl Drop for CommitParkHandle { + /// A handle dropped without [`Self::release`] (e.g. by a failing test) must not + /// leave the gated commit blocked forever, so dropping releases the gate. The + /// extra permit is harmless after an explicit release because the gate fires at + /// most once. + fn drop(&mut self) { + self.gate.release.add_permits(1); + } +} + struct ConsumeBodyReplyDeferGate { armed: AtomicBool, deferred_tx: std::sync::Mutex>>, diff --git a/golem-worker-executor/src/grpc/mod.rs b/golem-worker-executor/src/grpc/mod.rs index 3ed8dfc04a..0556e50d33 100644 --- a/golem-worker-executor/src/grpc/mod.rs +++ b/golem-worker-executor/src/grpc/mod.rs @@ -39,6 +39,7 @@ pub use crate::worker::{ }; use crate::worker::{Worker, WorkerUpdateMode}; use crate::workerctx::WorkerCtx; +use futures::FutureExt; use futures::Stream; use futures::StreamExt; use futures::future::join_all; @@ -95,6 +96,7 @@ use golem_service_base::model::auth::AuthCtx; use std::cmp::min; use std::collections::HashMap; use std::marker::PhantomData; +use std::panic::AssertUnwindSafe; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -125,6 +127,77 @@ impl ShardAssignmentChange { } } +/// Single-flights the lost-agent drain across concurrent shard-assignment RPCs. +/// +/// At most one drain task exists per executor. Each RPC bumps `generation` after mutating +/// the shard assignment and then waits until a drain that *started* at a generation at +/// least as new has completed: every drain pass re-reads `check_worker`, so such a drain +/// has observed the caller's mutation. Without this, each shard-manager retry against a +/// wedged agent would park another drain task until executor shutdown. +struct DrainSingleFlight { + inner: std::sync::Mutex, + /// Publishes `(generation the finished drain started at, its result)`. Written with + /// `send_replace` so a result published before a waiter subscribes is still seen. + completed: tokio::sync::watch::Sender, +} + +type DrainCompletion = (u64, Option>); + +struct DrainSingleFlightInner { + generation: u64, + owner_running: bool, + /// The open `shard_drain` span, so joining RPCs can link their trace origins to it. + span: Option, +} + +impl DrainSingleFlight { + fn new() -> Self { + Self { + inner: std::sync::Mutex::new(DrainSingleFlightInner { + generation: 0, + owner_running: false, + span: None, + }), + completed: tokio::sync::watch::channel((0, None)).0, + } + } + + /// Poison is harmless here: every critical section only performs simple + /// assignments, so the state is consistent even after a panic inside one. + fn lock(&self) -> std::sync::MutexGuard<'_, DrainSingleFlightInner> { + self.inner + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } +} + +/// Releases the drain owner slot if the owner task dies outside its `catch_unwind` +/// window (a panic in span creation, or task abort at runtime shutdown). Without +/// this, `owner_running` would stay set with no live owner and every future drain +/// waiter would park forever. +struct DrainOwnerGuard { + drain: Arc, + defused: bool, +} + +impl Drop for DrainOwnerGuard { + fn drop(&mut self) { + if self.defused { + return; + } + let mut inner = self.drain.lock(); + inner.owner_running = false; + inner.span = None; + let generation = inner.generation; + self.drain.completed.send_replace(( + generation, + Some(Err(WorkerExecutorError::runtime( + "Drain owner terminated unexpectedly", + ))), + )); + } +} + /// This is the implementation of the Worker Executor gRPC API pub struct WorkerExecutorImpl< Ctx: WorkerCtx, @@ -135,6 +208,7 @@ pub struct WorkerExecutorImpl< /// Holds the strong Arc to the worker activator so the Weak reference /// stored in LazyWorkerActivator remains valid while the gRPC server runs. _worker_activator: Arc>, + drain: Arc, ctx: PhantomData, } @@ -145,6 +219,7 @@ impl + UsesAllDeps + Send + Sync + Self { services: self.services.clone(), _worker_activator: self._worker_activator.clone(), + drain: self.drain.clone(), ctx: PhantomData, } } @@ -181,6 +256,7 @@ impl + UsesAllDeps + Send + Sync + let worker_executor = WorkerExecutorImpl { services: services.clone(), _worker_activator: worker_activator, + drain: Arc::new(DrainSingleFlight::new()), ctx: PhantomData, }; @@ -1015,23 +1091,85 @@ impl + UsesAllDeps + Send + Sync + change: ShardAssignmentChange, ) -> Result<(), WorkerExecutorError> { let origin = TraceOrigin::capture_current(); - let this = self.clone(); - let handle = tokio::spawn(async move { + // Subscribe before the election: `wait_for` re-checks the current value, so a + // result published between the bump and the wait is never missed. + let mut rx = self.drain.completed.subscribe(); + let (my_generation, spawn_owner) = { + let mut inner = self.drain.lock(); + inner.generation += 1; + let spawn = !inner.owner_running; + if spawn { + inner.owner_running = true; + } else if let Some(span) = &inner.span { + origin.add_as_link_to(span); + } + (inner.generation, spawn) + }; + if spawn_owner { + let this = self.clone(); + // Deliberately not instrumented with the current span: the owner serves + // later callers too and must outlive this RPC. + tokio::spawn(async move { this.run_drain_owner(change, origin).await }); + } + let completion = rx + .wait_for(|(generation, _)| *generation >= my_generation) + .await + .map_err(|_| WorkerExecutorError::runtime("Drain task state dropped"))?; + completion + .1 + .clone() + .expect("every published drain generation carries a result") + } + + /// The single drain owner: repeats the drain until it has covered the newest + /// requested generation, publishing each result to the waiters, then releases the + /// owner slot. The publish and the release happen under one critical section, so a + /// caller can never observe "no owner running" while its generation is uncovered. + async fn run_drain_owner(&self, change: ShardAssignmentChange, origin: TraceOrigin) { + let mut owner_guard = DrainOwnerGuard { + drain: self.drain.clone(), + defused: false, + }; + loop { + let observed = { self.drain.lock().generation }; let span = related_span!( origin, Level::INFO, "shard_drain", trigger = change.trigger() ); - this.drain_lost_agents_to_completion(change) - .instrument(span) - .await - }); - handle.await.map_err(|err| { - WorkerExecutorError::runtime(format!( - "Draining the agents of the revoked shards failed: {err}" - )) - })? + { + self.drain.lock().span = Some(span.clone()); + } + // A panic must not leave `owner_running` set forever, wedging every future + // drain; it is published as an error instead. + let result = AssertUnwindSafe( + self.drain_lost_agents_to_completion(change) + .instrument(span), + ) + .catch_unwind() + .await + .unwrap_or_else(|_| { + Err(WorkerExecutorError::runtime( + "Draining the agents of the revoked shards panicked", + )) + }); + let exit = { + let mut inner = self.drain.lock(); + inner.span = None; + self.drain.completed.send_replace((observed, Some(result))); + if inner.generation == observed { + inner.owner_running = false; + true + } else { + false + } + }; + if exit { + break; + } + } + owner_guard.defused = true; } async fn drain_lost_agents_to_completion( @@ -1051,8 +1189,13 @@ impl + UsesAllDeps + Send + Sync + loop { passes += 1; - // Agents whose creation is still in flight are invisible to the snapshot, which - // is why this loops until a snapshot finds nothing loaded in a lost shard. + // Agents whose creation is still in flight sit in ActiveAgents as *pending* + // entries, invisible to snapshot(). All of a creation's durable writes happen + // inside the pending closure (Worker::new), so awaiting the entry's resolution + // is a write barrier; creations that begin after the revocation are rejected by + // the ownership gate at the top of Worker::new and write nothing. The drain + // therefore only finishes on a pass that finds neither a loaded agent nor a + // pending creation in a lost shard. let mut lost = Vec::new(); for (agent_id, worker) in self.active_agents().snapshot().await { if self.shard_service().check_worker(&agent_id).is_err() @@ -1061,9 +1204,27 @@ impl + UsesAllDeps + Send + Sync + lost.push((agent_id, worker)); } } - if lost.is_empty() { + + let pending_lost: Vec = self + .active_agents() + .pending_agent_ids() + .await + .into_iter() + .filter(|owned| self.shard_service().check_worker(&owned.agent_id).is_err()) + .collect(); + + if lost.is_empty() && pending_lost.is_empty() { break; } + if lost.is_empty() { + // Only in-flight creations left: park on each entry's watch until it + // resolves (Ok or Err - a failed creation has also stopped writing), then + // re-check from the top. No polling and no sleep needed. + for owned in &pending_lost { + let _ = self.active_agents().await_settled(owned).await; + } + continue; + } lost_total += lost.len(); // Signal all of them at once. diff --git a/golem-worker-executor/src/services/active_agents/mod.rs b/golem-worker-executor/src/services/active_agents/mod.rs index ffd959da96..adf370b607 100644 --- a/golem-worker-executor/src/services/active_agents/mod.rs +++ b/golem-worker-executor/src/services/active_agents/mod.rs @@ -43,7 +43,7 @@ use std::sync::{Arc, Mutex}; use std::time::Duration; use tokio_util::sync::CancellationToken; -use tracing::{Instrument, debug}; +use tracing::{Instrument, Level, debug}; use crate::services::HasAll; use crate::services::card_interest::{ @@ -73,6 +73,8 @@ use golem_common::model::environment::EnvironmentId; use golem_common::model::invocation_context::InvocationContextStack; use golem_common::model::worker::AgentConfigEntryDto; use golem_common::model::{AgentId, OwnedAgentId, Timestamp}; +use golem_common::related_span; +use golem_common::tracing::TraceOrigin; use golem_service_base::error::worker_executor::InterruptKind; use golem_service_base::error::worker_executor::WorkerExecutorError; use wasmtime::Store; @@ -526,13 +528,27 @@ impl ActiveAgents { let cache_key = owned_agent_id.clone(); let deps = deps.clone(); let invocation_context_stack = invocation_context_stack.clone(); + let card_interest_index = self.card_interest_index.clone(); + // Spawned so the creation survives cancellation of the originating call: an + // abandoned inline producer would leave the pending entry unresolved forever, + // wedging every later get-or-create for this agent and the shard-revoke drain, + // which waits on pending entries. The spawned owner outlives this caller, so + // it links back to the originating span instead of nesting inside it. + let origin = TraceOrigin::capture_current(); + let span_agent_id = owned_agent_id.agent_id.clone(); let active_agent = self .agents - .get_or_insert_simple(&cache_key, || { - Box::pin(async move { + .get_or_insert_simple_spawned(&cache_key, move || { + let span = related_span!( + origin, + Level::INFO, + "agent_creation", + agent_id = %span_agent_id, + ); + async move { let worker = Worker::new( &deps, - self.card_interest_index.clone(), + card_interest_index, owned_agent_id.clone(), worker_env, worker_agent_config, @@ -542,7 +558,6 @@ impl ActiveAgents { principal, freshness_disposition, ) - .in_current_span() .await; worker.map(|worker| { @@ -550,7 +565,8 @@ impl ActiveAgents { Worker::start_durable_stream_attachment_reconciler(&worker); Arc::new(ActiveAgent::new(worker)) }) - }) + } + .instrument(span) }) .await?; Ok(active_agent.primary()) @@ -599,7 +615,12 @@ impl ActiveAgents { .set_card_interest(worker.owned_agent_id().clone(), &[]) .await; } - self.agents.remove(owned_agent_id).await + // Only a resolved entry may be removed. The `get` above settles the entry this + // caller saw, but a new creation can insert a fresh pending entry during the + // awaits in between; deleting it would hide the in-flight `Worker::new` from the + // shard-revoke drain while it is still writing durably, and the finished + // creation would re-insert the entry anyway. + self.agents.remove_if_cached(owned_agent_id, |_| true).await; } pub async fn tracked_card_ids(&self) -> Vec { @@ -657,6 +678,23 @@ impl ActiveAgents { .collect() } + /// Owned ids of agents whose creation is still in flight (pending cache entries). + /// These are invisible to `snapshot`, which only sees resolved entries. + pub async fn pending_agent_ids(&self) -> Vec { + self.agents.pending_keys().await + } + + /// Waits until the entry for `owned_agent_id` is no longer pending. Returns the + /// worker if creation succeeded, `None` if it failed or the entry is gone; never + /// inserts. All of a creation's durable writes happen before its entry resolves, + /// so a failed creation (`None`) is as good a write barrier as a successful one. + pub async fn await_settled(&self, owned_agent_id: &OwnedAgentId) -> Option>> { + self.agents + .get(owned_agent_id) + .await + .map(|active_agent| active_agent.primary()) + } + /// Interrupts and unloads all in-memory workers whose environment matches /// `environment_id`. Called when the environment is deleted so that /// running workers stop promptly. diff --git a/golem-worker-executor/src/worker/mod.rs b/golem-worker-executor/src/worker/mod.rs index 5b48539690..3404cfb76a 100644 --- a/golem-worker-executor/src/worker/mod.rs +++ b/golem-worker-executor/src/worker/mod.rs @@ -721,6 +721,9 @@ impl Worker { principal: Principal, freshness_disposition: InvocationFreshnessDisposition, ) -> Result { + deps.shard_service() + .check_worker(&owned_agent_id.agent_id)?; + let start = std::time::Instant::now(); let GetOrCreateWorkerResult { initial_worker_metadata, @@ -1096,6 +1099,14 @@ impl Worker { Err(err.clone()) } WorkerInstance::Unloaded { .. } => { + // A revoke may have landed after the caller's ownership check (for + // example while this worker's creation was in flight and the shard + // drain waited it out as unloaded). Loading it now would run an agent + // of a lost shard after RevokeShards already returned, so re-check + // before committing to a load; the caller re-resolves the owner. + this.shard_service() + .check_worker(&this.owned_agent_id.agent_id)?; + this.mark_as_loading(); crate::metrics::workers::inc_worker_waiting_for_memory(); *instance_guard = WorkerInstance::WaitingForPermit(WaitingWorker::new( diff --git a/golem-worker-executor/tests/sharding.rs b/golem-worker-executor/tests/sharding.rs index 77890c0b62..e2e4230a1a 100644 --- a/golem-worker-executor/tests/sharding.rs +++ b/golem-worker-executor/tests/sharding.rs @@ -268,6 +268,197 @@ async fn set_shard_assignment_drains_lost_agents_before_returning( Ok(()) } +/// An agent whose creation is in flight when the revoke arrives sits in the agent cache as a +/// *pending* entry, invisible to the drain's snapshot - but all of its durable writes (oplog +/// create, cached status, the committed initialization enqueue) happen inside that window. +/// `revoke_shards` must wait for the creation to settle before returning, or another executor +/// starts recovering the agent while this one is still writing its state. The parked commit is +/// the initialization enqueue inside `Worker::new`, so the creation is held mid-flight +/// deterministically. +#[test] +#[tracing::instrument] +#[timeout("2m")] +async fn revoke_shards_waits_for_in_flight_agent_creation( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + _tracing: &Tracing, + #[tagged_as("agent_counters")] agent_counters: &PrecompiledComponent, +) -> anyhow::Result<()> { + const HOLD: Duration = Duration::from_millis(300); + + let context = TestContext::new(last_unique_id); + let executor = start(deps, &context).await?; + let component = executor + .component_dep(&context.default_environment_id, agent_counters) + .store() + .await?; + + let parsed_id = agent_id!("InstantiationGrowthCounter", "pending-creation"); + let id = AgentId { + component_id: component.id, + agent_id: parsed_id.to_string(), + }; + let owned_id = OwnedAgentId::new(context.default_environment_id, &id); + + let mut gate = executor.park_next_oplog_commit(&id).await; + + let creation = { + let executor = executor.clone(); + let component_id = component.id; + let parsed_id = parsed_id.clone(); + tokio::spawn( + async move { executor.try_start_agent(&component_id, parsed_id).await } + .in_current_span(), + ) + }; + tokio::time::timeout(Duration::from_secs(30), gate.parked()) + .await + .map_err(|_| anyhow::anyhow!("the agent creation never reached its gated commit"))?; + + let revoke = { + let executor = executor.clone(); + tokio::spawn(async move { revoke_shards(&executor, &[SHARD]).await }.in_current_span()) + }; + tokio::time::sleep(HOLD).await; + assert!( + !revoke.is_finished(), + "revoke_shards returned while an agent creation in the revoked shard was still in \ + flight; its durable writes race the next owner's recovery" + ); + + gate.release(); + let drain_elapsed = revoke.await??; + assert!( + drain_elapsed >= HOLD - Duration::from_millis(50), + "revoke_shards returned after {drain_elapsed:?}, but the in-flight creation was held \ + for {HOLD:?}" + ); + + // The creation itself must fail: by the time the created agent would be started, the shard + // is no longer owned, and loading it would run an agent of a lost shard. + let creation_result = creation.await??; + assert!( + matches!( + creation_result, + Err( + golem_service_base::error::worker_executor::WorkerExecutorError::InvalidShardId { .. } + ) + ), + "expected the racing creation to fail with InvalidShardId, got {creation_result:?}" + ); + assert!( + !executor.worker_is_loaded(&owned_id).await, + "the agent created during the revoke must not be loaded" + ); + + // The durable state written by the creation survived; once the shard is back the agent + // replays it and runs normally. + assign_shards(&executor, &[SHARD]).await?; + let count: u32 = executor + .invoke_and_await_agent(&component, &parsed_id, "increment", data_value!()) + .await? + .into_typed()?; + assert_eq!(count, 1); + + drop(executor); + Ok(()) +} + +/// Concurrent revoke RPCs (the shard manager retries after a timeout) must join one logical +/// drain instead of each spawning its own: the second call returns only once a drain covering +/// its request completes, and neither call runs the teardowns twice or serializes two whole +/// drains back to back. +#[test] +#[tracing::instrument] +#[timeout("4m")] +async fn concurrent_revoke_shards_join_one_drain( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + _tracing: &Tracing, + #[tagged_as("host_api_tests")] host_api_tests: &PrecompiledComponent, +) -> anyhow::Result<()> { + const TEARDOWN_DELAY: Duration = Duration::from_millis(200); + const STAGGER: Duration = Duration::from_millis(50); + + let context = TestContext::new(last_unique_id); + let executor = start(deps, &context).await?; + let component = executor + .component_dep(&context.default_environment_id, host_api_tests) + .store() + .await?; + + let agents = start_busy_agents(&executor, &context, &component, "revoke-join").await?; + for agent in &agents { + executor + .arm_teardown_commit_delay(&agent.id, TEARDOWN_DELAY) + .await; + } + + let first = { + let executor = executor.clone(); + tokio::spawn(async move { revoke_shards(&executor, &[SHARD]).await }.in_current_span()) + }; + tokio::time::sleep(STAGGER).await; + let second = { + let executor = executor.clone(); + tokio::spawn(async move { revoke_shards(&executor, &[SHARD]).await }.in_current_span()) + }; + + let first_elapsed = first.await??; + let second_elapsed = second.await??; + + assert_all_unloaded( + &executor, + agents.iter().map(|agent| &agent.owned_id), + "revoke_shards", + first_elapsed.max(second_elapsed), + ) + .await; + + // The second call must have joined the drain the first one started, which the armed + // teardown delays hold open well past the stagger - a joiner that returns early (for + // example satisfied by a stale completed generation) comes back in single-digit + // milliseconds instead. + assert!( + second_elapsed >= TEARDOWN_DELAY - STAGGER - Duration::from_millis(50), + "the second revoke_shards returned after {second_elapsed:?}, before the drain it must \ + join could have finished" + ); + // And joining must not mean serializing: two whole drains back to back would double the + // duration; a joined drain ends for both calls together. + let sequential_bound = TEARDOWN_DELAY * AGENTS as u32; + assert!( + first_elapsed < sequential_bound * 3 / 4, + "the first revoke_shards took {first_elapsed:?}; a single concurrent drain finishes \ + well under {sequential_bound:?}" + ); + assert!( + second_elapsed <= first_elapsed, + "the second revoke_shards ({second_elapsed:?}) outlived the first ({first_elapsed:?}), \ + which points at a second serialized drain instead of a join" + ); + + // One teardown commit per agent: the joined drain did not tear anything down twice. + let intervals = executor.teardown_commit_intervals(); + assert_eq!( + intervals.len(), + AGENTS, + "every lost agent must have exactly one teardown commit: {intervals:?}" + ); + let depth = max_overlap_depth(&intervals); + assert!( + depth >= 2, + "the teardowns of the lost agents did not overlap (at most {depth} at a time): \ + {intervals:?}" + ); + + assign_shards(&executor, &[SHARD]).await?; + assert_invocations_completed_once(&executor, agents).await?; + + drop(executor); + Ok(()) +} + /// An agent of the `Clocks` type executing its `interruption` invocation: 100 sleeps of 100ms, /// each far below the suspend threshold, so the agent stays loaded and busy for about ten /// seconds and is interruptible at every sleep.