/build/source/nativelink-worker/src/local_worker.rs

Source
// Copyright 2024 The NativeLink Authors. All rights reserved.
//
// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    See LICENSE file for details
//
// 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.

use core::hash::BuildHasher;
use core::pin::Pin;
use core::str;
use core::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
use core::time::Duration;
use std::borrow::Cow;
use std::collections::HashMap;
use std::env;
use std::process::Stdio;
use std::sync::{Arc, Weak};
use std::time::{SystemTime, UNIX_EPOCH};

use futures::future::BoxFuture;
use futures::stream::FuturesUnordered;
use futures::{Future, FutureExt, StreamExt, TryFutureExt, select};
use nativelink_config::cas_server::{EnvironmentSource, LocalWorkerConfig};
use nativelink_error::{Code, Error, ResultExt, make_err, make_input_err};
use nativelink_metric::{MetricsComponent, RootMetricsComponent};
use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::update_for_worker::Update;
use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::worker_api_client::WorkerApiClient;
use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{
    ActionResourceUsage, ExecuteAccepted, ExecuteComplete, ExecuteDeclined, ExecuteResult,
    GoingAwayRequest, KeepAliveRequest, StartExecute, UpdateForWorker, WorkerLoad,
    execute_declined, execute_result,
};
use nativelink_store::fast_slow_store::FastSlowStore;
use nativelink_util::action_messages::{ActionResult, ActionStage, OperationId};
use nativelink_util::common::fs;
use nativelink_util::digest_hasher::DigestHasherFunc;
use nativelink_util::health_utils::{HealthStatus, HealthStatusIndicator};
use nativelink_util::metrics_utils::{AsyncCounterWrapper, CounterWithTime};
use nativelink_util::shutdown_guard::ShutdownGuard;
use nativelink_util::store_trait::Store;
use nativelink_util::{background_spawn, spawn, spawn_blocking, tls_utils};
use opentelemetry::context::Context;
use tokio::sync::{broadcast, mpsc};
use tokio::{process, time};
use tokio_stream::wrappers::UnboundedReceiverStream;
use tonic::{Streaming, async_trait};
use tracing::{Level, debug, error, event, info, info_span, instrument, trace, warn};

use crate::running_actions_manager::{
    ExecutionConfiguration, MISSING_INPUT_ERROR_TIP, Metrics as RunningActionManagerMetrics,
    ResourceEnforcement, RunningAction, RunningActionsManager, RunningActionsManagerArgs,
    RunningActionsManagerImpl,
};
use crate::worker_api_client_wrapper::{WorkerApiClientTrait, WorkerApiClientWrapper};
use crate::worker_utils::make_connect_worker_request;

/// Amount of time to wait if we have actions in transit before we try to
/// 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, doubling on every failed attempt up to
/// `CONNECTION_RETRY_MAX_DELAY_S`, with jitter so a fleet that lost its
/// scheduler together does not redial together.
const CONNECTION_RETRY_DELAY_S: f32 = 0.5;
const CONNECTION_RETRY_MAX_DELAY_S: f32 = 30.0;

/// Delay before the `attempt`th consecutive reconnect (0 = first retry),
/// between half and one and a half times the exponential figure.
fn reconnect_delay(attempt: u32) -> Duration {
    let exponent = i32::try_from(attempt.min(16)).unwrap_or(16);
    let base = (CONNECTION_RETRY_DELAY_S * 2f32.powi(exponent)).min(CONNECTION_RETRY_MAX_DELAY_S);
    let nanos = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map_or(0, |d| d.subsec_nanos());
    let jitter = 0.5 + (nanos % 1_000) as f32 / 1_000.0;
    Duration::from_secs_f32(base * jitter)
}

/// Default endpoint timeout. If this value gets modified the documentation in
/// `cas_server.rs` must also be updated.
const DEFAULT_ENDPOINT_TIMEOUT_S: f32 = 5.;

/// Keepalive defaults for the scheduler connection, the same as a store's
/// gRPC endpoint (`tls_utils::endpoint`).
const DEFAULT_TCP_KEEPALIVE: Duration = Duration::from_secs(30);
const DEFAULT_HTTP2_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30);
const DEFAULT_HTTP2_KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(20);

/// A configured number of seconds, or the default when unset (0).
const fn seconds_or(configured: u64, default: Duration) -> Duration {
    if configured > 0 {
        Duration::from_secs(configured)
    } else {
        default
    }
}

/// Sends `GoingAway` within `deadline`. Logged rather than fatal: the stream
/// close that follows is what removes the worker, this only says why.
async fn send_going_away<T: WorkerApiClientTrait>(
    grpc_client: &mut T,
    drain: bool,
    deadline: Duration,
) -> Result<(), Error> {
    match time::timeout(deadline, grpc_client.going_away(GoingAwayRequest { drain })).await {
        Ok(Ok(())) => Ok(()),
        Ok(Err(e)) => {
            error!(?e, drain, "Failed to send GoingAwayRequest");
            Err(e.append("Failed to send GoingAwayRequest"))
        }
        Err(_) => {
            error!(
                drain,
                deadline_s = deadline.as_secs_f32(),
                "GoingAwayRequest did not complete; the scheduler connection is stalled"
            );
            Err(make_err!(
                Code::DeadlineExceeded,
                "GoingAwayRequest did not complete within {}s",
                deadline.as_secs_f32()
            ))
        }
    }
}

/// Default maximum amount of time a task is allowed to run for.
/// If this value gets modified the documentation in `cas_server.rs` must also be updated.
const DEFAULT_MAX_ACTION_TIMEOUT: Duration = Duration::from_mins(20);
const DEFAULT_MAX_UPLOAD_TIMEOUT: Duration = Duration::from_mins(10);
/// Default for `max_download_timeout_s`: an input fetch is quick or wedged.
const DEFAULT_MAX_DOWNLOAD_TIMEOUT: Duration = Duration::from_mins(10);
const DEFAULT_MAX_CLEANUP_WAIT: Duration = Duration::from_secs(30);
const DEFAULT_MAX_CLEANUP_BACKOFF: Duration = Duration::from_millis(500);
/// If this value gets modified the documentation in `cas_server.rs` must also be updated.
const DEFAULT_PRECONDITION_TIMEOUT: Duration = Duration::from_secs(30);
/// If this value gets modified the documentation in `cas_server.rs` must also be updated.
const DEFAULT_KILL_GRACE: Duration = Duration::from_secs(5);

struct FinishedActionResult {
    action_result: ActionResult,
    resource_usage: Option<ActionResourceUsage>,
}

struct LocalWorkerImpl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> {
    config: &'a LocalWorkerConfig,
    // According to the tonic documentation it is a cheap operation to clone this.
    grpc_client: T,
    worker_id: String,
    running_actions_manager: Arc<U>,
    // Number of actions that have been received in `Update::StartAction`, but
    // not yet processed by running_actions_manager's spawn. This number should
    // always be zero if there are no actions running and no actions being waited
    // on by the scheduler.
    actions_in_transit: Arc<AtomicU64>,
    accepted_action: Arc<AtomicBool>,
    /// The scheduler said it understands `ExecuteAccepted` and
    /// `ExecuteDeclined`. Without it the worker runs whatever it is sent,
    /// as every earlier release did.
    dispatch_ack: bool,
    memory: MemoryAdmission,
    metrics: Arc<Metrics>,
}

/// What local admission needs to refuse an action for load.
struct MemoryAdmission {
    /// The platform property the scheduler reads an action's memory
    /// reservation from, as it told us on connection; empty when it does
    /// not veto on memory. Read from the same property, a refusal for load
    /// here agrees with what the scheduler would have vetoed.
    property: String,
    /// Whether this worker admits any action while it holds nothing else,
    /// decided once at registration and sent to the scheduler then, so the
    /// two never disagree.
    admits_when_idle: bool,
}

/// Why this worker will not run an action it was just sent.
enum Refusal {
    AtCapacity { in_flight: u64, max: u64 },
    Load { needed_kb: u64, free_kb: u64 },
    ShuttingDown,
}

impl Refusal {
    fn into_declined(self, operation_id: String) -> ExecuteDeclined {
        let (reason, detail, needed_kb, free_kb) = match self {
            Self::AtCapacity { in_flight, max } => (
                execute_declined::Reason::AtCapacity,
                format!("{in_flight} of {max} in flight"),
                0,
                0,
            ),
            Self::Load { needed_kb, free_kb } => (
                execute_declined::Reason::Load,
                format!("needs {needed_kb} KiB, {free_kb} KiB free"),
                needed_kb,
                free_kb,
            ),
            Self::ShuttingDown => (
                execute_declined::Reason::ShuttingDown,
                "worker shutting down".to_string(),
                0,
                0,
            ),
        };
        ExecuteDeclined {
            operation_id,
            reason: reason as i32,
            detail,
            needed_kb,
            free_kb,
        }
    }
}

/// The memory reservation an action carries under `property`, in KiB, if
/// it declares one. An empty property name is a scheduler that does not
/// veto on memory, so nothing is ever read.
fn memory_reservation_kb(start_execute: &StartExecute, property: &str) -> Option<u64> {
    if property.is_empty() {
        return None;
    }
    start_execute
        .platform
        .as_ref()?0
        .properties
        .iter()
        .find(|p| p.name == property)
        .and_then(|p| p.value.parse::<u64>().ok())
        .filter(|kb| *kb > 0)
}

pub async fn preconditions_met<H: BuildHasher + Sync>(
    precondition_script: Option<String>,
    extra_envs: &HashMap<String, String, H>,
    timeout: Duration,
) -> Result<(), Error> {
    let Some(precondition_script3) = &precondition_script else {
        // No script means we are always ok to proceed.
        return Ok(());
    };
    // TODO: Might want to pass some information about the command to the
    //       script, but at this point it's not even been downloaded yet,
    //       so that's not currently possible.  Perhaps we'll move this in
    //       future to pass useful information through?  Or perhaps we'll
    //       have a pre-condition and a pre-execute script instead, although
    //       arguably entrypoint already gives us that.

    let maybe_split_cmd = shlex::split(precondition_script);
    let (command, args) = match &maybe_split_cmd {
        Some(split_cmd) => (&split_cmd[0], &split_cmd[1..]),
        None => {
            return Err(make_input_err!(
                "Could not parse the value of precondition_script: '{}'",
                precondition_script,
            ));
        }
    };

    let precondition_process = process::Command::new(command)
        .args(args)
        .kill_on_drop(true)
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .env_clear()
        .envs(extra_envs)
        .spawn()
        .err_tip(|| format!0("Could not execute precondition command {precondition_script:?}"))?0;
    let _owned = crate::reaper::OwnedChild::new(precondition_process.id());
    // Bounded: a script that hangs held the action forever, and
    // `kill_on_drop` ends the script when the timeout drops it.
    let output2 = match time::timeout(timeout, precondition_process.wait_with_output()).await {
        Ok(output) => output?0,
        Err(_) => {
            return Err(make_err!(
                Code::ResourceExhausted,
                "Preconditions script {precondition_script:?} did not finish within {} ms",
                timeout.as_millis()
            ));
        }
    };
    let stdout = str::from_utf8(&output.stdout).unwrap_or("");
    trace!(status = %output.status, %stdout, "Preconditions script returned");
    if output.status.code() == Some(0) {
        Ok(())
    } else {
        Err(make_err!(
            Code::ResourceExhausted,
            "Preconditions script returned status {} - {}",
            output.status,
            stdout
        ))
    }
}

impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorkerImpl<'a, T, U> {
    fn new(
        config: &'a LocalWorkerConfig,
        grpc_client: T,
        worker_id: String,
        dispatch_ack: bool,
        memory: MemoryAdmission,
        running_actions_manager: Arc<U>,
        metrics: Arc<Metrics>,
    ) -> Self {
        Self {
            config,
            grpc_client,
            worker_id,
            dispatch_ack,
            memory,
            running_actions_manager,
            // Number of actions that have been received in `Update::StartAction`, but
            // not yet processed by running_actions_manager's spawn. This number should
            // always be zero if there are no actions running and no actions being waited
            // on by the scheduler.
            actions_in_transit: Arc::new(AtomicU64::new(0)),
            accepted_action: Arc::new(AtomicBool::new(false)),
            metrics,
        }
    }

    /// Local admission: the worker's own word on whether it can take this
    /// action now. The scheduler's ledger says what it believes the worker
    /// has; this is what the worker has.
    fn admission(&self, start_execute: &StartExecute, in_flight: u64) -> Option<Refusal> {
        let max = self.config.max_inflight_tasks;
        if max > 0 && in_flight >= max5 {
            return Some(Refusal::AtCapacity { in_flight, max });
        }
        // The cgroup is read only for an action that reserves memory.
        let needed_kb2 = memory_reservation_kb(start_execute, &self.memory.property)?5;
        // With nothing else in flight under a cgroup limit, waiting frees
        // nothing: what the worker reads free now is the most it will ever
        // read, since its own process, its hot page cache and the kernel
        // always count as used. Declining here would have the scheduler
        // wait for a figure that never comes (the whole-worker escalation
        // step asks for more than any worker reports free), so the action
        // runs, and if it truly does not fit, the memory kill escalates or
        // fails it. Without a limit, free memory is the host's, which other
        // processes can give back, so waiting can help and the check stands.
        if in_flight == 0 && self.memory.admits_when_idle {
            return None;
        }
        if let Some(free_kb) = crate::capacity::free_memory_kb()
            && free_kb < needed_kb
        {
            return Some(Refusal::Load { needed_kb, free_kb });
        }
        None
    }

    /// Tells the scheduler the action is refused; only meaningful when it
    /// understands the message.
    /// The one message bounded by the endpoint timeout: `GoingAway` at
    /// shutdown, so a stalled connection cannot hold a pod past its
    /// termination grace. Everything else the worker sends rides on the
    /// `ConnectWorker` stream unbounded, and the scheduler's liveness
    /// timeout is the clock that evicts a worker it cannot hear.
    fn going_away_deadline(&self) -> Duration {
        Duration::from_secs_f32(
            self.config
                .worker_api_endpoint
                .timeout
                .unwrap_or(DEFAULT_ENDPOINT_TIMEOUT_S),
        )
    }

    /// A message to the scheduler the run loop does not wait for inline:
    /// it rides in the loop's future set, so the loop keeps reading its
    /// stream while the message waits its turn on the one-slot channel
    /// feeding the `ConnectWorker` stream. No deadline on it: a message
    /// waits behind whatever the scheduler is still processing, and a
    /// scheduler that stops hearing this worker closes the stream, which
    /// the loop now sees.
    fn scheduler_send(
        &self,
        what: &'static str,
        send: impl Future<Output = Result<(), Error>> + Send + 'static,
    ) -> BoxFuture<'static, Result<(), Error>> {
        async move { send.await.err_tip6(|| format!1("Could not send {what}")) }6.boxed()
    }

    fn decline(
        &self,
        operation_id: String,
        refusal: Refusal,
    ) -> BoxFuture<'static, Result<(), Error>> {
        self.metrics.actions_declined.inc();
        let mut grpc_client = self.grpc_client.clone();
        self.scheduler_send("ExecuteDeclined", async move {
            grpc_client
                .execute_declined(refusal.into_declined(operation_id))
                .await
        })
    }

    /// Starts a background spawn/thread that will send a message to the server every `timeout / 2`.
    async fn start_keep_alive(&self) -> Result<(), Error> {19
        // According to tonic's documentation this call should be cheap and is the same stream.
        let mut grpc_client = self.grpc_client.clone();
        let timeout = self
            .config
            .worker_api_endpoint
            .timeout
            .unwrap_or(DEFAULT_ENDPOINT_TIMEOUT_S);

        info!(timeout, "Started KeepAlive");

        // We always send 2 keep alive requests per timeout. Http2 should manage most of our
        // timeout issues, this is a secondary check to ensure we can still send data.
        let mut interval = time::interval(Duration::from_secs_f32(timeout) / 2);
        interval.set_missed_tick_behavior(time::MissedTickBehavior::Skip);

        // Skip the first interval as it happens immediately and we don't need a keep alive until timeout/2 has passed
        interval.tick().await;

        // Explicitly spawn the keep alive loop so it goes onto a different thread from the execute commands.
        // Its failure is this worker's failure: a worker whose keepalives
        // stop reaching the scheduler is evicted `worker_timeout_s` later
        // with every action it holds requeued, so ending the stream now and
        // reconnecting is the cheaper outcome. Before, the task died quietly
        // and the worker ran on without keepalives.
        spawn!("keep alives", async move {
            loop {
                interval.tick().await;
                // What the worker has to spare rides on the keepalive, so a
                // scheduler with `live_memory_veto` can skip a worker whose
                // actions declared less than they use.
                let load = crate::capacity::free_memory_kb()
                    .map(|free_memory_kb| WorkerLoad { free_memory_kb });
                if let Err(e1) = grpc_client.keep_alive(KeepAliveRequest { load }).await {
                    error!(?e, "Failed to send KeepAlive in LocalWorker");
                    return Err(e.append("KeepAlive failed; reconnecting to the scheduler"));
                }
                debug!("Sent KeepAlive");
            }
        })
        .await
        .map_err(|e| make_err!0(Code::Internal0, "KeepAlive task ended: {e:?}"))?0
    }

    async fn run(
        &self,
        update_for_worker_stream: Streaming<UpdateForWorker>,
        shutdown_rx: &mut broadcast::Receiver<ShutdownGuard>,
    ) -> Result<(), Error> {
        // This big block of logic is designed to help simplify upstream components. Upstream
        // components can write standard futures that return a `Result<(), Error>` and this block
        // will forward the error up to the client and disconnect from the scheduler.
        // It is a common use case that an item sent through update_for_worker_stream will always
        // have a response but the response will be triggered through a callback to the scheduler.
        // This can be quite tricky to manage, so what we have done here is given access to a
        // `futures` variable which because this is in a single thread as well as a channel that you
        // send a future into that makes it into the `futures` variable.
        // This means that if you want to perform an action based on the result of the future
        // you use the `.map()` method and the new action will always come to live in this spawn,
        // giving mutable access to stuff in this struct.
        // NOTE: If you ever return from this function it will disconnect from the scheduler.
        let mut futures = FuturesUnordered::new();
        futures.push(self.start_keep_alive().boxed());

        let (add_future_channel, add_future_rx) = mpsc::unbounded_channel();
        let mut add_future_rx = UnboundedReceiverStream::new(add_future_rx).fuse();

        let mut update_for_worker_stream = update_for_worker_stream.fuse();
        // A notify which is triggered every time actions_in_flight is subtracted.
        let actions_notify = Arc::new(tokio::sync::Notify::new());
        // A counter of actions that are in-flight, this is similar to actions_in_transit but
        // includes the AC upload and notification to the scheduler.
        let actions_in_flight = Arc::new(AtomicU64::new(0));
        // Set to true when shutting down, this stops any new StartAction.
        let mut shutting_down = false;

        loop {
            if self.config.single_use
                && self.accepted_action.load(Ordering::Acquire)
                && actions_in_flight.load(Ordering::Acquire) == 0
            {
                // The action, CAS/AC uploads, cleanup, and execution_response
                // acknowledgment have all completed. This container is spent.
                // Bounded: a stalled connection must not keep a spent
                // container from exiting.
                let mut grpc_client = self.grpc_client.clone();
                if let Err(err) =
                    send_going_away(&mut grpc_client, false, self.going_away_deadline()).await
                {
                    warn!(?err, "Could not unregister completed single-use worker");
                }
                return Ok(());
            }
            select! {
                maybe_update27 = update_for_worker_stream.next() => if !shutting_down27 || maybe_update0.is_some0() {
                    match maybe_update
                        .err_tip(|| "UpdateForWorker stream closed early")?6
                        .err_tip(|| "Got error in UpdateForWorker stream")?0
                        .update
                        .err_tip(|| "Expected update to exist in UpdateForWorker")?0
                    {
                        Update::ConnectionResult(_) => {
                            return Err(make_input_err!(
                                "Got ConnectionResult in LocalWorker::run which should never happen"
                            ));
                        }
                        // TODO(palfrey) We should possibly do something with this notification.
                        Update::Disconnect(()) => {
                            self.metrics.disconnects_received.inc();
                        }
                        Update::KeepAlive(()) => {
                            self.metrics.keep_alives_received.inc();
                        }
                        Update::KillOperationRequest(kill_operation_request) => {
                            let operation_id = OperationId::from(kill_operation_request.operation_id);
                            if let Err(err0) = self.running_actions_manager.kill_operation(&operation_id).await {
                                error!(
                                    %operation_id,
                                    ?err,
                                    "Failed to send kill request for operation"
                                );
                            }
                        }
                        Update::StartAction(start_execute) => {
                            // Don't accept any new requests if we're shutting down.
                            if shutting_down || (self.config.single_use
                                && self.accepted_action.load(Ordering::Acquire)) {
                                if self.dispatch_ack {
                                    futures.push(self.decline(start_execute.operation_id, Refusal::ShuttingDown));
                                } else if let Some(instance_name) = start_execute.execute_request.map(|request| request.instance_name) {
                                    let mut grpc_client = self.grpc_client.clone();
                                    let operation_id = start_execute.operation_id;
                                    futures.push(self.scheduler_send("ExecutionResponse", async move {
                                        grpc_client.execution_response(
                                            ExecuteResult{
                                                instance_name,
                                                operation_id,
                                                result: Some(execute_result::Result::InternalError(make_err!(Code::ResourceExhausted, "Worker shutting down").into())),
                                                resource_usage: None,
                                            }
                                        ).await
                                    }));
                                }
                                continue;
                            }

                            // Admission, then the acknowledgement: what the
                            // scheduler charged on the send is confirmed or
                            // handed back, and a declined action never runs.
                            // A scheduler that does not speak the
                            // acknowledgement gets the old behaviour, run
                            // whatever arrives. Neither message is awaited
                            // here: this loop is the only reader of the
                            // scheduler's stream, and a send that hangs on a
                            // stalled connection would keep it from ever
                            // seeing that stream close (the worker then sat
                            // evicted and idle until the provisioner retired
                            // it). The action starts once its acknowledgement
                            // has gone out, so the scheduler hears the
                            // acknowledgement before anything the run sends,
                            // and one that fails leaves the action unstarted.
                            let acknowledgement15 = if self.dispatch_ack {
                                if let Some(refusal3) = self.admission(
                                    &start_execute,
                                    actions_in_flight.load(Ordering::Acquire),
                                ) {
                                    futures.push(self.decline(start_execute.operation_id, refusal));
                                    continue;
                                }
                                let mut grpc_client = self.grpc_client.clone();
                                let operation_id = start_execute.operation_id.clone();
                                Some(self.scheduler_send("ExecuteAccepted", async move {
                                    grpc_client
                                        .execute_accepted(ExecuteAccepted { operation_id })
                                        .await
                                }))
                            } else {
                                None
                            };
                            // Admitted: a single-use worker is spent by this
                            // action once it is acknowledged. A decline above
                            // must not spend it, or every dispatch it turns
                            // away costs a pod.
                            let spent_mark = self.config.single_use.then(|| self.accepted_action5.clone5());

                            self.metrics.start_actions_received.inc();

                            let execute_request = start_execute.execute_request.as_ref();
                            let operation_id = start_execute.operation_id.clone();
                            let operation_id_to_log = operation_id.clone();
                            let maybe_instance_name = execute_request.map(|v| v.instance_name.clone());
                            let action_digest = execute_request.and_then(|v| v.action_digest.clone());
                            let digest_hasher = execute_request
                                .ok_or_else(|| make_input_err!0("Expected execute_request to be set"))
                                .and_then(|v| DigestHasherFunc::try_from(v.digest_function))
                                .err_tip(|| "In LocalWorkerImpl::new()")?0;

                            let start_action_fut = {
                                let precondition_script_cfg = self.config.experimental_precondition_script.clone();
                                let precondition_timeout = if self.config.precondition_timeout_ms == 0 {
                                    DEFAULT_PRECONDITION_TIMEOUT
                                } else {
                                    Duration::from_millis(self.config.precondition_timeout_ms)
                                };
                                let mut extra_envs: HashMap<String, String> = HashMap::new();
                                if let Some(ref additional_environment0) = self.config.additional_environment {
                                    for (name, source) in additional_environment {
                                        let value = match source {
                                            EnvironmentSource::Property(property) => start_execute
                                                .platform.as_ref().and_then(|p|p.properties.iter().find(|pr| &pr.name == property))
                                                .map_or_else(|| Cow::Borrowed(""), |v| Cow::Borrowed(v.value.as_str())),
                                            EnvironmentSource::Value(value) => Cow::Borrowed(value.as_str()),
                                            EnvironmentSource::FromEnvironment => Cow::Owned(env::var(name).unwrap_or_default()),
                                            other => {
                                                debug!(?other, "Worker doesn't support this type of additional environment");
                                                continue;
                                            }
                                        };
                                        extra_envs.insert(name.clone(), value.into_owned());
                                    }
                                }
                                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();
                                let complete = ExecuteComplete {
                                    operation_id: operation_id.clone(),
                                };
                                let single_use = self.config.single_use;
                                self.metrics.clone().wrap(move |metrics| async move {12
                                    metrics.preconditions.wrap(preconditions_met(precondition_script_cfg, &extra_envs, precondition_timeout))
                                    .and_then(|()| running_actions_manager11.create_and_add_action11(worker_id11, start_execute11))
                                    .map(move |r| {8
                                        // Now that we either failed or registered our action, we can
                                        // consider the action to no longer be in transit.
                                        drop(actions_in_transit_guard);
                                        r
                                    })
                                    .and_then(|action| {7
                                        debug!(
                                            operation_id = %action.get_operation_id(),
                                            "Received request to run action"
                                        );
                                        action
                                            .clone()
                                            .prepare_action()
                                            .and_then(RunningAction::execute)
                                            .and_then(|result| async move {3
                                                // Reusable workers release their slot during upload.
                                                // A single-use worker must never advertise another slot.
                                                if !single_use {
                                                    drop(grpc_client.execution_complete(complete).await);
                                                }
                                                Ok(result)
                                            })
                                            .and_then(RunningAction::upload_results)
                                            .and_then(|action| async move {3
                                                let resource_usage = action.resource_usage();
                                                let action_result = action.get_finished_result().await?0;
                                                Ok(FinishedActionResult {
                                                    action_result,
                                                    resource_usage,
                                                })
                                            })
                                            // Note: We need ensure we run cleanup even if one of the other steps fail.
                                            .then(|result| async move {5
                                                if let Err(e0) = action.cleanup().await {
                                                    return Result::<FinishedActionResult, Error>::Err(e).merge(result);
                                                }
                                                result
                                            })
                                    }7).await
                                })
                            };

                            let make_publish_future = {
                                let mut grpc_client = self.grpc_client.clone();

                                let running_actions_manager = self.running_actions_manager.clone();
                                let worker_id = self.worker_id.clone();
                                move |res: Result<FinishedActionResult, Error>| async move {
                                    let instance_name = maybe_instance_name
                                        .err_tip(|| "`instance_name` could not be resolved; this is likely an internal error in local_worker.")?0;
                                    match res {
                                        Ok(FinishedActionResult { mut action_result, resource_usage }) => {
                                            // Save in the action cache before notifying the scheduler that we've completed.
                                            if let Some(digest_info) = action_digest.clone().and_then(|action_digest| action_digest.try_into().ok()) &&
                                                let Err(err0) = running_actions_manager.cache_action_result(digest_info, &mut action_result, digest_hasher).await {
                                                    error!(
                                                        ?err,
                                                        ?action_digest,
                                                        "Error saving action in store",
                                                    );
                                                }
                                            let action_stage = ActionStage::Completed(action_result);
                                            let resource_usage = resource_usage.map(|mut resource_usage| {0
                                                resource_usage.operation_id.clone_from(&operation_id);
                                                resource_usage.worker_id.clone_from(&worker_id);
                                                resource_usage
                                            });
                                            grpc_client.execution_response(
                                                ExecuteResult{
                                                    instance_name,
                                                    operation_id,
                                                    result: Some(execute_result::Result::ExecuteResponse(action_stage.into())),
                                                    resource_usage,
                                                }
                                            )
                                            .await
                                            .err_tip(|| "Error while calling execution_response")?0;
                                        },
                                        Err(e) => {
                                            // The manager tags a NotFound on anything the client
                                            // referenced by digest; no store's wording involved.
                                            let is_cas_blob_missing = e.code == Code::NotFound
                                                && e.messages.iter()2.any2(|m| m == MISSING_INPUT_ERROR_TIP);
                                            if is_cas_blob_missing {
                                                warn!(
                                                    ?e,
                                                    "Missing CAS inputs during prepare_action, returning FAILED_PRECONDITION"
                                                );
                                                // The context names the digest; the
                                                // execute response turns it into the
                                                // PreconditionFailure detail Bazel
                                                // re-uploads and retries on.
                                                let action_result = ActionResult {
                                                    error: Some(
                                                        make_err!(
                                                            Code::FailedPrecondition,
                                                            "{}",
                                                            e.message_string()
                                                        )
                                                        .with_context(e.context.clone()),
                                                    ),
                                                    ..ActionResult::default()
                                                };
                                                let action_stage = ActionStage::Completed(action_result);
                                                grpc_client.execution_response(ExecuteResult{
                                                    instance_name,
                                                    operation_id,
                                                    result: Some(execute_result::Result::ExecuteResponse(action_stage.into())),
                                                    resource_usage: None,
                                                }).await.err_tip0(|| "Error calling execution_response with missing inputs")?0;
                                            } else {
                                                grpc_client.execution_response(ExecuteResult{
                                                    instance_name,
                                                    operation_id,
                                                    result: Some(execute_result::Result::InternalError(e.into())),
                                                    resource_usage: None,
                                                }).await.err_tip0(|| "Error calling execution_response with error")?0;
                                            }
                                        },
                                    }
                                    Ok(())
                                }
                            };

                            let add_future_channel = add_future_channel.clone();

                            // Boxed so the launch closure captures one
                            // erased type, not the whole action future, and
                            // the compiler's Send check does not overflow.
                            let start_action_fut: BoxFuture<'static, Result<FinishedActionResult, Error>> =
                                start_action_fut.boxed();
                            let launch = {
                                let span = info_span!(
                                    "worker_start_action_ctx",
                                    operation_id = operation_id_to_log,
                                    digest_function = %digest_hasher.to_string(),
                                );
                                move || {
                                    span.in_scope(|| {
                                        let _guard = Context::current_with_value(digest_hasher)
                                            .attach();
                                        spawn!("worker_start_action", start_action_fut)
                                    })
                                }
                            };

                            let actions_in_flight = actions_in_flight.clone();
                            let actions_notify = actions_notify.clone();
                            let actions_in_flight_fail = actions_in_flight.clone();
                            let actions_notify_fail = actions_notify.clone();
                            // Counted from the dispatch, so the next one is
                            // admitted against this slot while the
                            // acknowledgement is still on its way.
                            actions_in_flight.fetch_add(1, Ordering::Release);

                            futures.push(
                                async move {
                                    if let Some(acknowledgement5) = acknowledgement {
                                        acknowledgement.await?1;
                                    }
                                    if let Some(accepted_action4) = spent_mark {
                                        accepted_action.store(true, Ordering::Release);
                                    }
                                    launch().await.err_tip6(|| "Failed to launch spawn")
                                }
                                .map(move |res| {7
                                    let res6 = res?1;
                                    if let Err(err3) = &res {
                                        error!(?err, "Error executing action");
                                    }
                                    add_future_channel
                                        .send(make_publish_future(res).then(move |res| {1
                                            actions_in_flight.fetch_sub(1, Ordering::Release);
                                            actions_notify.notify_one();
                                            core::future::ready(res)
                                        }1).boxed())
                                        .map_err(|err|
                                            Error::from_std_err(Code::Internal, &err).append("LocalWorker could not send future")
                                            )?;
                                    Ok(())
                                })
                                .or_else(move |err| {1
                                    // If the make_publish_future is not run we still need to notify.
                                    actions_in_flight_fail.fetch_sub(1, Ordering::Release);
                                    actions_notify_fail.notify_one();
                                    core::future::ready(Err(err))
                                })
                                .boxed()
                            );
                        }
                    }
                },
                res6 = add_future_rx.next() => {
                    let fut = res.err_tip(|| "New future stream receives should never be closed")?0;
                    futures.push(fut);
                },
                res11 = futures.next() => res11.err_tip11(|| "Keep-alive should always pending. Likely unable to send data to scheduler")?0?2,
                complete_msg0 = shutdown_rx.recv().fuse() => {
                    warn!("Worker loop received shutdown signal. Shutting down worker...",);
                    let mut grpc_client = self.grpc_client.clone();
                    let shutdown_guard = complete_msg.map_err(|e|
                        Error::from_std_err(Code::Internal, &e).append("Failed to receive shutdown message"))?;
                    let actions_in_flight = actions_in_flight.clone();
                    let actions_notify = actions_notify.clone();
                    let drain_on_shutdown = self.config.drain_on_shutdown;
                    // The GoingAway messages below are bounded: a stalled
                    // connection must not hold the pod past its grace
                    // period. One that fails or misses is logged and the
                    // shutdown goes on, the in-flight actions still drained;
                    // the scheduler removes this worker when the stream
                    // closes either way.
                    let going_away_deadline = self.going_away_deadline();
                    let drain_deadline = if self.config.max_action_timeout_s == 0 {
                        DEFAULT_MAX_ACTION_TIMEOUT
                    } else {
                        Duration::from_secs(self.config.max_action_timeout_s as u64)
                    };
                    let shutdown_future = async move {
                        // Said first on a drain, so nothing new is dispatched
                        // here while the running actions finish; the
                        // scheduler removes this worker when the stream
                        // closes.
                        let mut going_away_result = if drain_on_shutdown {
                            send_going_away(&mut grpc_client, true, going_away_deadline).await
                        } else {
                            Ok(())
                        };
                        // Wait for in-flight operations to be fully completed,
                        // for as long as one action is allowed to run.
                        let wait = async {
                            while actions_in_flight.load(Ordering::Acquire) > 0 {
                                actions_notify.notified().await;
                            }
                        };
                        if time::timeout(drain_deadline, wait).await.is_err() {
                            error!(
                                actions_in_flight = actions_in_flight.load(Ordering::Acquire),
                                "Drain deadline passed with actions still in flight; shutting down anyway"
                            );
                        }
                        if !drain_on_shutdown {
                            // Sending this message immediately evicts all jobs from
                            // this worker, of which there should be none.
                            going_away_result = send_going_away(&mut grpc_client, false, going_away_deadline).await;
                        }
                        // Allow shutdown to occur now.
                        drop(shutdown_guard);
                        going_away_result
                    };
                    futures.push(shutdown_future.boxed());
                    shutting_down = true;
                },
            };
        }
        // Unreachable.
    }
}

type ConnectionFactory<T> = Box<dyn Fn() -> BoxFuture<'static, Result<T, Error>> + Send + Sync>;

pub struct LocalWorker<T: WorkerApiClientTrait + 'static, U: RunningActionsManager> {
    config: Arc<LocalWorkerConfig>,
    running_actions_manager: Arc<U>,
    connection_factory: ConnectionFactory<T>,
    sleep_fn: Option<Box<dyn Fn(Duration) -> BoxFuture<'static, ()> + Send + Sync>>,
    metrics: Arc<Metrics>,
    registration: Arc<WorkerRegistration>,
}

/// Whether the worker currently holds a registration with its scheduler.
/// A worker that has not registered, or lost its connection and is
/// reconnecting, cannot take work; as a health indicator this is what a
/// readiness probe reads.
#[derive(Debug)]
pub struct WorkerRegistration {
    name: String,
    registered: AtomicBool,
    ever_registered: AtomicBool,
}

impl WorkerRegistration {
    pub fn new(name: &str) -> Arc<Self> {
        Arc::new(Self {
            name: name.to_string(),
            registered: AtomicBool::new(false),
            ever_registered: AtomicBool::new(false),
        })
    }

    pub fn is_registered(&self) -> bool {
        self.registered.load(Ordering::Acquire)
    }

    fn set_registered(&self, registered: bool) {
        self.registered.store(registered, Ordering::Release);
        if registered {
            self.ever_registered.store(true, Ordering::Release);
        }9
    }
}

#[async_trait]
impl HealthStatusIndicator for WorkerRegistration {
    fn get_name(&self) -> &'static str {
        "WorkerRegistration"
    }

    async fn check_health(&self, _namespace: Cow<'static, str>) -> HealthStatus {
        if self.is_registered() {
            HealthStatus::Ok {
                struct_name: "WorkerRegistration",
                message: Cow::Owned(format!(
                    "worker '{}' registered with the scheduler",
                    self.name
                )),
            }
        } else if self.ever_registered.load(Ordering::Acquire) {
            // Lost after it was there: the worker is reconnecting, and until
            // it does it holds no work. Initializing, not Failed: Failed
            // turns the plain status check red too, and a liveness probe on
            // it would restart every worker during a scheduler roll longer
            // than its threshold, and redden a co-hosted CAS on one worker's
            // blip. This belongs to readiness alone.
            HealthStatus::Initializing {
                struct_name: "WorkerRegistration",
                message: Cow::Owned(format!(
                    "worker '{}' lost its scheduler connection, reconnecting",
                    self.name
                )),
            }
        } else {
            HealthStatus::Initializing {
                struct_name: "WorkerRegistration",
                message: Cow::Owned(format!(
                    "worker '{}' not yet registered with the scheduler",
                    self.name
                )),
            }
        }
    }
}

impl<
    T: WorkerApiClientTrait + core::fmt::Debug + 'static,
    U: RunningActionsManager + core::fmt::Debug,
> core::fmt::Debug for LocalWorker<T, U>
{
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("LocalWorker")
            .field("config", &self.config)
            .field("running_actions_manager", &self.running_actions_manager)
            .field("metrics", &self.metrics)
            .finish_non_exhaustive()
    }
}

/// Creates a new `LocalWorker`. The `cas_store` must be an instance of
/// `FastSlowStore` and will be checked at runtime.
pub async fn new_local_worker(
    config: Arc<LocalWorkerConfig>,
    cas_store: Store,
    ac_store: Option<Store>,
    historical_store: Store,
) -> Result<LocalWorker<WorkerApiClientWrapper, RunningActionsManagerImpl>, Error> {
    #[cfg(not(target_os = "linux"))]
    if config.experimental_buck2_file_capture.is_some() {
        return Err(make_input_err!(
            "Buck2 container file capture requires a Linux execution container"
        ));
    }
    let fast_slow_store = cas_store
        .downcast_ref::<FastSlowStore>(None)
        .err_tip(|| "Expected store for LocalWorker's store to be a FastSlowStore")?0
        .get_arc()
        .err_tip(|| "FastSlowStore's Arc doesn't exist")?0;

    // Log warning about CAS configuration for multi-worker setups
    event!(
        Level::INFO,
        worker_name = %config.name,
        "Starting worker '{}'. IMPORTANT: If running multiple workers, all workers \
        must share the same CAS storage path to avoid 'Object not found' errors.",
        config.name
    );

    if let Ok(path1) = fs::canonicalize(&config.work_directory).await {
        fs::remove_dir_all(&path).await.err_tip(|| {0
            format!(
                "Could not remove work_directory '{}' in LocalWorker",
                path.as_path().to_str().unwrap_or("bad path")
            )
        })?;
    }

    fs::create_dir_all(&config.work_directory)
        .await
        .err_tip(|| format!0("Could not make work_directory : {}", config.work_directory0))?0;
    let entrypoint = if config.entrypoint.is_empty() {
        None
    } else {
        Some(config.entrypoint.clone())
    };
    let max_action_timeout = if config.max_action_timeout_s == 0 {
        DEFAULT_MAX_ACTION_TIMEOUT
    } else {
        Duration::from_secs(config.max_action_timeout_s as u64)
    };
    let max_upload_timeout = if config.max_upload_timeout_s == 0 {
        DEFAULT_MAX_UPLOAD_TIMEOUT
    } else {
        Duration::from_secs(config.max_upload_timeout_s as u64)
    };
    let max_download_timeout = if config.max_download_timeout_s == 0 {
        DEFAULT_MAX_DOWNLOAD_TIMEOUT
    } else {
        Duration::from_secs(config.max_download_timeout_s as u64)
    };
    let max_cleanup_wait = if config.max_cleanup_wait_s == 0 {
        DEFAULT_MAX_CLEANUP_WAIT
    } else {
        Duration::from_secs(config.max_cleanup_wait_s as u64)
    };
    let max_cleanup_backoff = if config.max_cleanup_backoff_ms == 0 {
        DEFAULT_MAX_CLEANUP_BACKOFF
    } else {
        Duration::from_millis(config.max_cleanup_backoff_ms as u64)
    };

    // Initialize directory cache if configured
    let directory_cache = if let Some(cache_config0) = &config.directory_cache {
        use std::path::PathBuf;

        use crate::directory_cache::{
            DirectoryCache, DirectoryCacheConfig as WorkerDirCacheConfig,
        };

        let cache_root = if cache_config.cache_root.is_empty() {
            PathBuf::from(&config.work_directory).parent().map_or_else(
                || PathBuf::from("/tmp/nativelink_directory_cache"),
                |p| p.join("directory_cache"),
            )
        } else {
            PathBuf::from(&cache_config.cache_root)
        };

        let worker_cache_config = WorkerDirCacheConfig {
            max_entries: cache_config.max_entries,
            max_size_bytes: cache_config.max_size_bytes,
            cache_root,
            experimental_subtree_caching: cache_config.experimental_subtree_caching,
            max_concurrent_fetches: cache_config.max_concurrent_fetches,
            experimental_get_tree_prefetch: cache_config.experimental_get_tree_prefetch,
        };

        match DirectoryCache::new(worker_cache_config, fast_slow_store.clone()).await {
            Ok(cache) => {
                tracing::info!("Directory cache initialized successfully");
                Some(Arc::new(cache))
            }
            Err(e) => {
                tracing::warn!("Failed to initialize directory cache: {:?}", e);
                None
            }
        }
    } else {
        None
    };

    #[cfg(target_os = "linux")]
    let use_namespaces = {
        let use_mount_namespace = config.use_mount_namespace.unwrap_or_default();
        // The mount namespace only exists when `use_namespaces` is on too;
        // `isolate_tmp` is judged against that, not against the flag alone,
        // or an explicit `true` would be dropped without a word.
        let mount_namespace_on = config.use_namespaces == Some(true) && use_mount_namespace0;
        // A private /tmp is part of the mount isolation unless explicitly
        // turned off. A worker without a /tmp, such as a minimal container
        // image, has nothing for actions to collide on, so the default is
        // off there instead of failing to start.
        let isolate_tmp = config.isolate_tmp.unwrap_or_else(|| {
            let has_tmp = std::path::Path::new("/tmp").is_dir();
            if mount_namespace_on && !has_tmp0 {
                warn!("/tmp does not exist on this worker, so actions will not get a private /tmp");
            }
            mount_namespace_on && has_tmp0
        });
        if isolate_tmp && !mount_namespace_on0 {
            return Err(make_err!(
                Code::InvalidArgument,
                "isolate_tmp requires use_namespaces and use_mount_namespace to be true"
            ));
        }
        if let Some(use_namespaces0) = &config.use_namespaces {
            if *use_namespaces
                && !crate::namespace_utils::namespaces_supported(use_mount_namespace, isolate_tmp)
            {
                return Err(make_err!(Code::Unavailable, "Namespaces not supported"));
            }
            if !*use_namespaces {
                crate::running_actions_manager::UseNamespaces::No
            } else if use_mount_namespace {
                crate::running_actions_manager::UseNamespaces::YesAndMount { isolate_tmp }
            } else {
                crate::running_actions_manager::UseNamespaces::Yes
            }
        } else if use_mount_namespace {
            return Err(make_err!(
                Code::Unavailable,
                "Mount namespaces not supported"
            ));
        } else {
            crate::running_actions_manager::UseNamespaces::No
        }
    };

    #[cfg(not(target_os = "linux"))]
    if config.use_namespaces.is_some_and(core::convert::identity) {
        return Err(make_err!(
            Code::Unavailable,
            "Namespaces not supported on non-Linux OSes"
        ));
    }
    #[cfg(not(target_os = "linux"))]
    if config
        .use_mount_namespace
        .is_some_and(core::convert::identity)
    {
        return Err(make_err!(
            Code::Unavailable,
            "Mount namespaces not supported on non-Linux OSes"
        ));
    }
    #[cfg(not(target_os = "linux"))]
    if config.isolate_tmp.is_some_and(core::convert::identity) {
        return Err(make_err!(
            Code::Unavailable,
            "isolate_tmp is not supported on non-Linux OSes"
        ));
    }

    // A pooled worker process gets the same PID, user, UTS and IPC
    // namespaces a one-shot action gets, never the mount namespace.
    #[cfg(target_os = "linux")]
    let persistent_workers_namespaced = !matches!0(
        use_namespaces,
        crate::running_actions_manager::UseNamespaces::No
    );
    #[cfg(not(target_os = "linux"))]
    let persistent_workers_namespaced = false;

    let running_actions_manager =
        Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs {
            root_action_directory: config.work_directory.clone(),
            execution_configuration: ExecutionConfiguration {
                max_captured_output_bytes: config.max_captured_output_bytes,
                kill_grace: if config.kill_grace_ms == 0 {
                    DEFAULT_KILL_GRACE
                } else {
                    Duration::from_millis(config.kill_grace_ms)
                },
                set_tmpdir: config.set_tmpdir,
                persistent_workers: persistent_worker_settings(
                    config.persistent_workers.as_ref(),
                    persistent_workers_namespaced,
                ),
                buck2_file_capture: config.experimental_buck2_file_capture.clone(),
                entrypoint,
                additional_environment: config.additional_environment.clone(),
                resource_enforcement: config
                    .resource_enforcement
                    .as_ref()
                    .and_then(ResourceEnforcement::from_config),
            },
            cas_store: fast_slow_store,
            ac_store,
            historical_store,
            upload_action_result_config: &config.upload_action_result,
            max_action_timeout,
            max_upload_timeout,
            max_download_timeout,
            max_cleanup_wait,
            max_cleanup_backoff,
            timeout_handled_externally: config.timeout_handled_externally,
            directory_cache,
            active_input_leases: config.experimental_active_input_leases,
            #[cfg(target_os = "linux")]
            use_namespaces,
        })?);
    // What actions leave behind is reaped here; see `crate::reaper`.
    if let Err(err0) = crate::reaper::become_subreaper() {
        warn!(
            ?err,
            "Could not become the subreaper; orphans of actions will not be reaped here"
        );
    }
    drop(background_spawn!("orphan_reaper", async move {1
        let mut seen_last = std::collections::BTreeSet::new();
        loop {
            time::sleep(crate::reaper::REAP_INTERVAL).await;
            // The sweep reads /proc for every process; off the runtime.
            let (reaped, seen) = spawn_blocking!("orphan_reaper_sweep", move || {
                crate::reaper::reap_orphaned_zombies(&seen_last)
            })
            .await
            .unwrap_or_default();
            seen_last = seen;
            if reaped > 0 {
                info!(reaped, "Reaped zombies that actions left behind");
            }
        }
    }));
    if config.orphan_sweep_interval_s > 0 {
        let interval = Duration::from_secs(config.orphan_sweep_interval_s);
        let manager = running_actions_manager.clone();
        // Detached on purpose: the guarded spawn aborts its task when the
        // handle drops, and this one runs for the worker's whole life.
        drop(background_spawn!("orphan_sweep", async move {
            info!(interval_s = interval.as_secs(), "Orphan sweep scheduled");
            loop {
                time::sleep(interval).await;
                match manager.sweep_orphaned_directories().await {
                    Ok(0) => debug!("Orphan sweep found nothing"),
                    Ok(removed) => info!(removed, "Orphan sweep finished"),
                    Err(err) => warn!(?err, "Orphan sweep failed"),
                }
            }
        }));
    }
    let local_worker = LocalWorker::new_with_connection_factory_and_actions_manager(
        config.clone(),
        running_actions_manager,
        Box::new(move || {0
            let config = config.clone();
            Box::pin(async move {
                let timeout = config
                    .worker_api_endpoint
                    .timeout
                    .unwrap_or(DEFAULT_ENDPOINT_TIMEOUT_S);
                let timeout_duration = Duration::from_secs_f32(timeout);
                let tls_config =
                    tls_utils::load_client_config(&config.worker_api_endpoint.tls_config)
                        .err_tip(|| "Parsing local worker TLS configuration")?;
                let endpoint =
                    tls_utils::endpoint_from(&config.worker_api_endpoint.uri, tls_config)
                        .map_err(|e| {
                            Error::from_std_err(Code::InvalidArgument, &e)
                                .append("Invalid URI for worker endpoint")
                        })?
                        .connect_timeout(timeout_duration)
                        .timeout(timeout_duration)
                        // Probes under and between the messages, so a
                        // connection that has gone dead underneath fails
                        // here instead of hanging every message on it.
                        .tcp_keepalive(Some(seconds_or(
                            config.worker_api_endpoint.tcp_keepalive_s,
                            DEFAULT_TCP_KEEPALIVE,
                        )))
                        .http2_keep_alive_interval(seconds_or(
                            config.worker_api_endpoint.http2_keepalive_interval_s,
                            DEFAULT_HTTP2_KEEPALIVE_INTERVAL,
                        ))
                        .keep_alive_timeout(seconds_or(
                            config.worker_api_endpoint.http2_keepalive_timeout_s,
                            DEFAULT_HTTP2_KEEPALIVE_TIMEOUT,
                        ))
                        .keep_alive_while_idle(true);

                let transport = endpoint.connect().await.map_err(|e| {
                    Error::from_std_err(Code::Internal, &e).append(format!(
                        "Could not connect to endpoint {}",
                        config.worker_api_endpoint.uri
                    ))
                })?;
                Ok(WorkerApiClient::new(transport).into())
            })
        }),
        Box::new(move |d| Box::pin0(time::sleep0(d0))),
    );
    Ok(local_worker)
}

impl<T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorker<T, U> {
    pub fn new_with_connection_factory_and_actions_manager(
        config: Arc<LocalWorkerConfig>,
        running_actions_manager: Arc<U>,
        connection_factory: ConnectionFactory<T>,
        sleep_fn: Box<dyn Fn(Duration) -> BoxFuture<'static, ()> + Send + Sync>,
    ) -> Self {
        let metrics = Arc::new(Metrics::new(Arc::downgrade(
            running_actions_manager.metrics(),
        )));
        let registration = WorkerRegistration::new(&config.name);
        Self {
            config,
            running_actions_manager,
            connection_factory,
            sleep_fn: Some(sleep_fn),
            metrics,
            registration,
        }
    }

    /// The registration flag this worker flips, for a health registry.
    pub fn registration(&self) -> Arc<WorkerRegistration> {
        self.registration.clone()
    }

    /// Flip a flag that was registered with a health registry before the
    /// worker existed, as the server binary has to.
    #[must_use]
    pub fn with_registration(mut self, registration: Arc<WorkerRegistration>) -> Self {
        self.registration = registration;
        self
    }

    #[allow(
        clippy::missing_const_for_fn,
        reason = "False positive on stable, but not on nightly"
    )]
    pub fn name(&self) -> &String {
        &self.config.name
    }

    async fn register_worker(
        &self,
        client: &mut T,
    ) -> Result<(String, bool, MemoryAdmission, Streaming<UpdateForWorker>), Error> {
        let mut extra_envs: HashMap<String, String> = HashMap::new();
        if let Some(ref additional_environment0) = self.config.additional_environment {
            for (name, source) in additional_environment {
                let value = match source {
                    EnvironmentSource::Value(value) => Cow::Borrowed(value.as_str()),
                    EnvironmentSource::FromEnvironment => {
                        Cow::Owned(env::var(name).unwrap_or_default())
                    }
                    other => {
                        debug!(
                            ?other,
                            "Worker registration doesn't support this type of additional environment"
                        );
                        continue;
                    }
                };
                extra_envs.insert(name.clone(), value.into_owned());
            }
        }

        let mut platform_properties = self.config.platform_properties.clone();
        if let Some(capacity0) = &self.config.capacity {
            let memory_headroom_percent = crate::capacity::memory_headroom_percent(
                capacity,
                self.config.resource_enforcement.as_ref(),
            );
            crate::capacity::apply(capacity, memory_headroom_percent, &mut platform_properties)
                .err_tip(|| "Advertising capacity from the cgroup")?;
        }
        // Decided once, so what admission does and what the scheduler is
        // told always agree.
        let admits_when_idle = crate::capacity::memory_is_limited();
        let connect_worker_request = make_connect_worker_request(
            self.config.name.clone(),
            &platform_properties,
            &extra_envs,
            if self.config.single_use {
                1
            } else {
                self.config.max_inflight_tasks
            },
            admits_when_idle,
        )
        .await?0;
        let mut update_for_worker_stream21 = client
            .connect_worker(connect_worker_request)
            .await
            .err_tip(|| "Could not call connect_worker() in worker")?0
            .into_inner();

        let first_msg_update20 = update_for_worker_stream
            .next()
            .await
            .err_tip(|| "Got EOF expected UpdateForWorker")?1
            .err_tip(|| "Got error when receiving UpdateForWorker")?0
            .update;

        let (worker_id, dispatch_ack, memory_property) = match first_msg_update {
            Some(Update::ConnectionResult(connection_result)) => (
                connection_result.worker_id,
                connection_result.dispatch_ack,
                connection_result.memory_property,
            ),
            other => {
                return Err(make_input_err!(
                    "Expected first response from scheduler to be a ConnectionResult got : {:?}",
                    other
                ));
            }
        };
        Ok((
            worker_id,
            dispatch_ack,
            MemoryAdmission {
                property: memory_property,
                admits_when_idle,
            },
            update_for_worker_stream,
        ))
    }

    #[instrument(skip(self), level = Level::INFO)]
    pub async fn run(
        mut self,
        mut shutdown_rx: broadcast::Receiver<ShutdownGuard>,
    ) -> Result<(), Error> {
        // Belt-and-suspenders QoS bump: the main binary already calls
        // this before runtime creation so the tokio worker threads
        // inherit P-core preference via pthread QoS inheritance, but
        // any thread that reaches this point should also be tagged in
        // case it was spawned by a path that bypassed `on_thread_start`.
        // No-op on non-macOS.
        let _ = crate::qos::set_user_initiated();

        let sleep_fn = self
            .sleep_fn
            .take()
            .err_tip(|| "Could not unwrap sleep_fn in LocalWorker::run")?;
        let sleep_fn_pin = Pin::new(&sleep_fn);
        let attempts = AtomicU32::new(0);
        let attempts_ref = &attempts;
        let error_handler = Box::pin(move |err| async move {
            let attempt = attempts_ref.fetch_add(1, Ordering::AcqRel);
            let delay = reconnect_delay(attempt);
            error!(?err, attempt, delay_ms = delay.as_millis(), "Error");
            (sleep_fn_pin)(delay).await;
        });

        loop {
            // First connect to our endpoint.
            let mut client = match (self.connection_factory)().await {
                Ok(client) => client,
                Err(e) => {
                    (error_handler)(e).await;
                    continue; // Try to connect again.
                }
            };

            debug!("Connected to endpoint");

            // Next register our worker with the scheduler.
            let (inner, update_for_worker_stream) = match self.register_worker(&mut client).await {
                Err(e) => {
                    (error_handler)(e).await;
                    continue; // Try to connect again.
                }
                Ok((worker_id, dispatch_ack, memory, update_for_worker_stream)) => (
                    LocalWorkerImpl::new(
                        &self.config,
                        client,
                        worker_id,
                        dispatch_ack,
                        memory,
                        self.running_actions_manager.clone(),
                        self.metrics.clone(),
                    ),
                    update_for_worker_stream,
                ),
            };
            info!(
                worker_id = %inner.worker_id,
                "Worker registered with scheduler"
            );
            attempts.store(0, Ordering::Release);
            self.registration.set_registered(true);

            // Now listen for connections and run all other services.
            let run_result = inner.run(update_for_worker_stream, &mut shutdown_rx).await;
            self.registration.set_registered(false);
            if let Err(err) = run_result {
                // Give in-transit actions a chance to settle before we kill
                // them, so their results still reach the scheduler.
                const ITERATIONS: usize = 1_000;

                let sleep_duration = ACTIONS_IN_TRANSIT_TIMEOUT_S / ITERATIONS as f32;
                let mut drained = false;
                for _ in 0..ITERATIONS {
                    if inner.actions_in_transit.load(Ordering::Acquire) == 0 {
                        drained = true;
                        break;
                    }
                    (sleep_fn_pin)(Duration::from_secs_f32(sleep_duration)).await;
                }
                if !drained {
                    // Deliberately not fatal. Returning here propagates out of
                    // the worker's main loop and aborts the process, so a
                    // scheduler blip that happened to catch an action in
                    // transit took the whole worker down — and every action it
                    // held then had to run again elsewhere. At fleet scale that
                    // is a restart storm. kill_all() below discards these
                    // actions anyway, so the wait is a courtesy and overrunning
                    // it costs nothing beyond the actions we were already
                    // giving up on.
                    error!(
                        actions_in_transit = inner.actions_in_transit.load(Ordering::Acquire),
                        "Actions in transit did not reach zero before we disconnected from the scheduler"
                    );
                }

                // Kill off any existing actions because if we re-connect, we'll
                // get some more and it might resource lock us.
                self.running_actions_manager.kill_all().await;

                if self.config.single_use && inner.accepted_action.load(Ordering::Acquire) {
                    return Err(
                        err.append("Single-use worker disconnected after accepting its action")
                    );
                }
                error!(?err, "Worker disconnected from scheduler, reconnecting");
                (error_handler)(err).await; // Try to connect again.
            } else if self.config.single_use {
                return Ok(());
            }
        }
        // Unreachable.
    }
}

#[derive(Debug, MetricsComponent)]
pub struct Metrics {
    #[metric(
        help = "Total number of actions sent to this worker to process. This does not mean it started them, it just means it received a request to execute it."
    )]
    start_actions_received: CounterWithTime,
    #[metric(help = "Total number of disconnects received from the scheduler.")]
    disconnects_received: CounterWithTime,
    #[metric(
        help = "Dispatches this worker declined: at capacity, short of memory, or shutting down."
    )]
    actions_declined: CounterWithTime,
    #[metric(help = "Total number of keep-alives received from the scheduler.")]
    keep_alives_received: CounterWithTime,
    #[metric(
        help = "Stats about the calls to check if an action satisfies the config supplied script."
    )]
    preconditions: AsyncCounterWrapper,
    #[metric]
    #[allow(
        clippy::struct_field_names,
        reason = "TODO Fix this. Triggers on nightly"
    )]
    running_actions_manager_metrics: Weak<RunningActionManagerMetrics>,
}

impl RootMetricsComponent for Metrics {}

impl Metrics {
    fn new(running_actions_manager_metrics: Weak<RunningActionManagerMetrics>) -> Self {
        Self {
            start_actions_received: CounterWithTime::default(),
            disconnects_received: CounterWithTime::default(),
            actions_declined: CounterWithTime::default(),
            keep_alives_received: CounterWithTime::default(),
            preconditions: AsyncCounterWrapper::default(),
            running_actions_manager_metrics,
        }
    }
}

impl Metrics {
    async fn wrap<U, T: Future<Output = U>, F: FnOnce(Arc<Self>) -> T>(
        self: Arc<Self>,
        fut: F,
    ) -> U {12
        fut(self).await
    }
}

/// The pool settings from the worker config, defaults for what it leaves at
/// zero or unset.
fn persistent_worker_settings(
    config: Option<&nativelink_config::cas_server::PersistentWorkersConfig>,
    namespaced: bool,
) -> crate::running_actions_manager::PersistentWorkersSettings {
    use crate::persistent_worker::PoolConfig;
    let defaults = PoolConfig::default();
    let pool = config.map_or(defaults, |c| {0
        let or_default = |value: u64, default: u64| if value == 0 { default } else { value };
        PoolConfig {
            max_workers_per_key: if c.max_workers_per_key == 0 {
                defaults.max_workers_per_key
            } else {
                c.max_workers_per_key
            },
            idle_timeout: Duration::from_secs(or_default(
                c.idle_timeout_s,
                defaults.idle_timeout.as_secs(),
            )),
            max_requests_per_worker: or_default(
                c.max_requests_per_worker,
                defaults.max_requests_per_worker,
            ),
            shutdown_grace: Duration::from_millis(or_default(
                c.shutdown_grace_ms,
                u64::try_from(defaults.shutdown_grace.as_millis()).unwrap_or(5000),
            )),
            acquire_timeout: Duration::from_secs(or_default(
                c.acquire_timeout_s,
                defaults.acquire_timeout.as_secs(),
            )),
            namespaced,
        }
    });
    crate::running_actions_manager::PersistentWorkersSettings {
        enabled: config.is_none_or(|c| c.enabled),
        pool: PoolConfig { namespaced, ..pool },
    }
}

Coverage Report

Created: 2026-10-06 13:02