Fix disconnect_with_action_in_transit_reconnects_test

Author: cormacrelf

nativelink-worker/src/local_worker.rs

@@ -61,6 +61,28 @@ use crate::worker_utils::make_connect_worker_request;
 /// consider an error to have occurred.
 const ACTIONS_IN_TRANSIT_TIMEOUT_S: f32 = 10.;
 
+/// Increments `actions_in_transit` on creation and decrements it on drop, so
+/// the count stays accurate even when the owning action future is aborted by
+/// a disconnect. A stranded count makes the disconnect handler conclude the
+/// in-transit actions never drained, turning every disconnect-with-work into
+/// a fatal error instead of a reconnect.
+struct ActionsInTransitGuard {
+    actions_in_transit: Arc<AtomicU64>,
+}
+
+impl ActionsInTransitGuard {
+    fn new(actions_in_transit: Arc<AtomicU64>) -> Self {
+        actions_in_transit.fetch_add(1, Ordering::Release);
+        Self { actions_in_transit }
+    }
+}
+
+impl Drop for ActionsInTransitGuard {
+    fn drop(&mut self) {
+        self.actions_in_transit.fetch_sub(1, Ordering::Release);
+    }
+}
+
 /// If we lose connection to the worker api server we will wait this many seconds
 /// before trying to connect.
 const CONNECTION_RETRY_DELAY_S: f32 = 0.5;
@@ -315,7 +337,8 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke
                                         extra_envs.insert(name.clone(), value.into_owned());
                                     }
                                 }
-                                let actions_in_transit = self.actions_in_transit.clone();
+                                let actions_in_transit_guard =
+                                    ActionsInTransitGuard::new(self.actions_in_transit.clone());
                                 let worker_id = self.worker_id.clone();
                                 let running_actions_manager = self.running_actions_manager.clone();
                                 let mut grpc_client = self.grpc_client.clone();
@@ -328,7 +351,7 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke
                                     .map(move |r| {
                                         // Now that we either failed or registered our action, we can
                                         // consider the action to no longer be in transit.
-                                        actions_in_transit.fetch_sub(1, Ordering::Release);
+                                        drop(actions_in_transit_guard);
                                         r
                                     })
                                     .and_then(|action| {
@@ -438,8 +461,6 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke
                                 }
                             };
 
-                            self.actions_in_transit.fetch_add(1, Ordering::Release);
-
                             let add_future_channel = add_future_channel.clone();
 
                             info_span!(