Worker: a stalled scheduler connection reconnects instead of leaving the worker evicted and idle

nativelink-config/src/cas_server.rs

@@ -877,10 +877,35 @@ pub struct EndpointConfig {
     #[serde(deserialize_with = "convert_string_with_shellexpand")]
     pub uri: String,
 
-    /// Timeout in seconds that a request should take.
+    /// Timeout in seconds for connecting to the endpoint and for a request
+    /// to be answered. The worker's messages to the scheduler (keepalives,
+    /// acknowledgements, results) ride on the one `ConnectWorker` stream
+    /// opened at registration, which this does not bound; the `GoingAway`
+    /// sent at shutdown is bounded by it, so a stalled connection cannot
+    /// hold a pod past its grace period. A keepalive goes out every half of
+    /// it.
     /// Default: 5 seconds
     pub timeout: Option<f32>,
 
+    /// TCP keepalive interval (seconds), probing the socket under the
+    /// scheduler connection so a dead peer is noticed at the OS level.
+    /// If not set or 0, defaults to 30 seconds.
+    #[serde(default, deserialize_with = "convert_duration_with_shellexpand")]
+    pub tcp_keepalive_s: u64,
+
+    /// HTTP/2 keepalive interval (seconds): PING frames on the scheduler
+    /// connection, so one that has gone dead underneath fails instead of
+    /// hanging every message on it.
+    /// If not set or 0, defaults to 30 seconds.
+    #[serde(default, deserialize_with = "convert_duration_with_shellexpand")]
+    pub http2_keepalive_interval_s: u64,
+
+    /// HTTP/2 keepalive timeout (seconds): a PING unanswered for this long
+    /// ends the connection, and the worker reconnects.
+    /// If not set or 0, defaults to 20 seconds.
+    #[serde(default, deserialize_with = "convert_duration_with_shellexpand")]
+    pub http2_keepalive_timeout_s: u64,
+
     /// The TLS configuration to use to connect to the endpoint.
     pub tls_config: Option<ClientTlsConfig>,
 }

nativelink-service/src/worker_api_server.rs

@@ -303,19 +303,38 @@ impl WorkerConnection {
         };
 
         background_spawn!("worker_api", async move {
-            let mut had_going_away = false;
-            while let Some(maybe_update) = connection.next().await {
-                let update = match maybe_update.map(|u| u.update) {
-                    Ok(Some(update)) => update,
-                    Ok(None) => {
-                        tracing::warn!(worker_id=?instance.worker_id, "Empty update");
-                        continue;
-                    }
-                    Err(err) => {
-                        tracing::warn!(worker_id=?instance.worker_id, ?err, "Error from worker");
+            // The stream is read by a task of its own that only hands each
+            // message on; they are processed here, in order. Processing one
+            // inline in the reader (a result's store write, say) stopped the
+            // reads for its duration, and a worker that filled the stream's
+            // flow-control window in the meantime had nothing it was allowed
+            // to put on the wire: its keepalives waited behind the stall
+            // until this scheduler evicted it. The queue is unbounded on
+            // purpose: a worker's messages are bounded by its in-flight
+            // actions and its keepalive interval, and backpressure here is
+            // exactly the stall.
+            let (tx, mut rx) = mpsc::unbounded_channel();
+            let reader_worker_id = instance.worker_id.clone();
+            background_spawn!("worker_api_reader", async move {
+                while let Some(maybe_update) = connection.next().await {
+                    let update = match maybe_update.map(|u| u.update) {
+                        Ok(Some(update)) => update,
+                        Ok(None) => {
+                            tracing::warn!(worker_id=?reader_worker_id, "Empty update");
+                            continue;
+                        }
+                        Err(err) => {
+                            tracing::warn!(worker_id=?reader_worker_id, ?err, "Error from worker");
+                            break;
+                        }
+                    };
+                    if tx.send(update).is_err() {
                         break;
                     }
-                };
+                }
+            });
+            let mut had_going_away = false;
+            while let Some(update) = rx.recv().await {
                 let result = match update {
                     Update::ConnectWorkerRequest(_connect_worker_request) => Err(make_err!(
                         Code::Internal,

nativelink-service/tests/worker_api_server_test.rs

@@ -1176,3 +1176,70 @@ pub async fn a_full_channel_defers_the_kill_instead_of_evicting_test()
     );
     Ok(())
 }
+
+/// A message the scheduler is slow to process must not stop it reading the
+/// worker's stream: the keepalives behind it still come off the wire (here,
+/// a one-slot channel the test could not fill otherwise), and are taken in,
+/// in order, once the slow one is done. Reading inline, the stream's
+/// flow-control window filled behind a slow result and the worker's
+/// keepalives never left it, so this scheduler evicted a worker that was
+/// only waiting for it.
+#[nativelink_test]
+pub async fn a_slow_result_does_not_stop_the_stream_reads_test()
+-> Result<(), Box<dyn core::error::Error>> {
+    let mut test_context = setup_api_server(BASE_WORKER_TIMEOUT_S, Box::new(static_now_fn)).await?;
+    let operation_id = dispatch(&test_context, 1).await?;
+    let update_for_worker = test_context
+        .connection_worker_stream
+        .next()
+        .await
+        .expect("Worker stream ended early")?
+        .update
+        .expect("Expected update field to be populated");
+    let update_for_worker::Update::StartAction(start_execute) = update_for_worker else {
+        panic!("Expected StartAction message");
+    };
+    assert_eq!(operation_id.to_string(), start_execute.operation_id);
+
+    // The result's store write is not answered yet: processing it blocks.
+    test_context
+        .worker_stream
+        .send(Update::ExecuteResult(ExecuteResult {
+            instance_name: "instance_name".to_string(),
+            operation_id: operation_id.to_string(),
+            result: Some(execute_result::Result::InternalError(
+                make_err!(tonic::Code::Internal, "the action failed").into(),
+            )),
+            resource_usage: None,
+        }))
+        .await?;
+
+    // The keepalives behind it are still read off the stream.
+    for free_memory_kb in 1..=8u64 {
+        tokio::time::timeout(
+            Duration::from_secs(5),
+            test_context
+                .worker_stream
+                .send(Update::KeepAliveRequest(KeepAliveRequest {
+                    load: Some(WorkerLoad { free_memory_kb }),
+                })),
+        )
+        .await
+        .map_err(|_| {
+            make_err!(
+                tonic::Code::DeadlineExceeded,
+                "the stream read stalled behind the result"
+            )
+        })??;
+    }
+
+    // Once the result is through, the keepalives are taken in after it.
+    let (operation_id_seen, worker_id, _) = test_context
+        .state_manager
+        .expect_update_operation(Ok(()))
+        .await;
+    assert_eq!(operation_id_seen, operation_id);
+    assert_eq!(worker_id, test_context.worker_id);
+    summary_after_report(&test_context, 8).await;
+    Ok(())
+}

nativelink-store/src/filesystem_store.rs

@@ -1518,7 +1518,7 @@ impl<Fe: FileEntry> FilesystemStore<Fe> {
 
     fn get_and_update_generation(&self) -> Result<Generation, Error> {
         self.next_generation
-            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |generation| {
+            .try_update(Ordering::Relaxed, Ordering::Relaxed, |generation| {
                 generation.checked_add(1)
             })
             .map(Generation::new)