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/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/shard_management.rs b/golem-shard-manager/src/sharding/shard_management.rs index 7a8eaf9b8f..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, 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" ); - { - 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); - } - } + // The executor may have applied the assignment even though the call failed + // (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; } @@ -221,8 +256,17 @@ 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 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); + } needs_retry = true; } @@ -234,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() @@ -280,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, @@ -294,15 +347,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/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-shard-manager/tests/shard_management.rs b/golem-shard-manager/tests/shard_management.rs index a9100d46fe..e66cd0dd81 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,34 @@ 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>>, + applied_then_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() @@ -88,10 +111,24 @@ 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); } + 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 +152,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); } @@ -125,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(()) } @@ -137,6 +179,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 +187,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 +200,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,9 +545,132 @@ 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 + .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])]), + 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 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()); @@ -518,69 +689,208 @@ 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; + 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(old_pod, &[0, 1, 2, 3]) + .await; + worker_executors + .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(new_pod, Some("worker-executor-1".to_string())) + .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(&[2, 3]) - ); - assert_eq!( - worker_executors.local_assignment(new_pod).await, - shard_ids(&[0, 1]) + 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!( - routing_table - .pod_states - .get(&old_pod) - .expect("old pod missing") - .assigned_shards, - shard_ids(&[2, 3]) + revokes, 1, + "the revoke must not be retried once the routing table released the shards: {calls:#?}" ); - assert_eq!( - routing_table - .pod_states - .get(&new_pod) - .expect("new pod missing") - .assigned_shards, - shard_ids(&[0, 1]) + 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:#?}" ); - - 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); - let worker_executors = Arc::new(TestWorkerExecutors::default()); - - let (shard_management, _persistence, mut join_set) = new_shard_management( +/// 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( - 1, + 6, vec![ - (existing_pod, "worker-executor-0", &[]), - (pod(2, 9001), "worker-executor-1", &[0]), + (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; + .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 - .set_local_assignment(existing_pod, &[0]) - .await; - worker_executors - .fail_next_reconciliations(existing_pod, 1) + .apply_then_fail_next_assignments(pod_b, 1) .await; + shard_management.unregister_pod(pod_c).await; - shard_management - .register_pod(existing_pod, Some("worker-executor-0".to_string())) - .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(); - wait_for_local_assignment(&worker_executors, existing_pod, BTreeSet::new()).await; + 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 656c8232b0..b9db5199ec 100644 --- a/golem-worker-executor-test-utils/src/lib.rs +++ b/golem-worker-executor-test-utils/src/lib.rs @@ -812,6 +812,35 @@ 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() + } + + /// 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 @@ -2893,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), @@ -3090,7 +3144,22 @@ 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 + self.park_at_commit_park_gate().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 +3512,19 @@ 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>>, + /// 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 @@ -3471,10 +3553,74 @@ 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())), + 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) { + 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 @@ -3835,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/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..0556e50d33 100644 --- a/golem-worker-executor/src/grpc/mod.rs +++ b/golem-worker-executor/src/grpc/mod.rs @@ -32,15 +32,17 @@ 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, }; use crate::worker::{Worker, WorkerUpdateMode}; use crate::workerctx::WorkerCtx; +use futures::FutureExt; 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; @@ -83,7 +85,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, @@ -93,16 +96,108 @@ 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}; +use std::time::{Duration, 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; +/// 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, + Assign, + Set, +} + +impl ShardAssignmentChange { + fn trigger(self) -> &'static str { + match self { + ShardAssignmentChange::Revoke => "revoke_shards", + ShardAssignmentChange::Assign => "assign_shards", + ShardAssignmentChange::Set => "set_shard_assignment", + } + } +} + +/// 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, @@ -113,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, } @@ -123,6 +219,7 @@ impl + UsesAllDeps + Send + Sync + Self { services: self.services.clone(), _worker_activator: self._worker_activator.clone(), + drain: self.drain.clone(), ctx: PhantomData, } } @@ -159,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, }; @@ -950,16 +1048,8 @@ 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.drain_lost_agents(ShardAssignmentChange::Revoke) + .await?; Ok(()) } @@ -973,7 +1063,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.recover_gained_agents_in_background(ShardAssignmentChange::Assign); Ok(()) } @@ -990,20 +1080,268 @@ impl + UsesAllDeps + Send + Sync + self.shard_service() .set_shard_assignment(number_of_shards, &shard_ids)?; + self.drain_lost_agents(ShardAssignmentChange::Set).await?; + self.recover_gained_agents_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 + Ok(()) + } + + async fn drain_lost_agents( + &self, + change: ShardAssignmentChange, + ) -> Result<(), WorkerExecutorError> { + let origin = TraceOrigin::capture_current(); + // 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() + ); { - await_interrupted.recv().await.unwrap(); + 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( + &self, + change: ShardAssignmentChange, + ) -> Result<(), WorkerExecutorError> { + let trigger = change.trigger(); + let started = Instant::now(); + let shutdown = self.shutdown_token(); + + 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 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() + && worker.is_loaded().await + { + lost.push((agent_id, worker)); + } + } - Ctx::on_shard_assignment_changed(self).await?; + 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(); - Ok(()) + 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. + 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; + } + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(Duration::from_millis(100)); + } + } + + (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")) + } + } + } + + 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/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/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..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( @@ -1379,6 +1390,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() } 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..e2e4230a1a --- /dev/null +++ b/golem-worker-executor/tests/sharding.rs @@ -0,0 +1,703 @@ +// 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 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. +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 +}