// 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::{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 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, ExecuteComplete, ExecuteResult, GoingAwayRequest, KeepAliveRequest,

    UpdateForWorker, 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::metrics_utils::{AsyncCounterWrapper, CounterWithTime};

use nativelink_util::shutdown_guard::ShutdownGuard;

use nativelink_util::store_trait::Store;

use nativelink_util::{spawn, tls_utils};

use opentelemetry::context::Context;

use tokio::sync::{broadcast, mpsc};

use tokio::{process, time};

use tokio_stream::wrappers::UnboundedReceiverStream;

use tonic::Streaming;

use tracing::{Level, debug, error, event, info, info_span, instrument, trace, warn};

use crate::running_actions_manager::{

    ExecutionConfiguration, Metrics as RunningActionManagerMetrics, 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.;

/// 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;

/// 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.;

/// 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);

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>,

    metrics: Arc<Metrics>,

}

pub async fn preconditions_met<H: BuildHasher + Sync>(

    precondition_script: Option<String>,

    extra_envs: &HashMap<String, String, H>,

) -> Result<(), Error> {

    let Some(
precondition_script2
) = &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 output = precondition_process.wait_with_output().await
?0
;

    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,

        running_actions_manager: Arc<U>,

        metrics: Arc<Metrics>,

    ) -> Self {

        Self {

            config,

            grpc_client,

            worker_id,

            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)),

            metrics,

        }

    }

    /// Starts a background spawn/thread that will send a message to the server every `timeout / 2`.

    async fn start_keep_alive(&self) -> Result<(), Error> {

        // 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

        drop(

            spawn!("keep alives", async move {

                loop {

                    interval.tick().await;

                    if let Err(
e1
) = grpc_client.keep_alive(KeepAliveRequest {}).await {

                        error!(?e, "Failed to send KeepAlive in LocalWorker");

                        return;

                    }

                    debug!("Sent KeepAlive");

                }

            })

            .await,

        );

        Ok(())

    }

    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 {

            select! {

                
maybe_update8
 = update_for_worker_stream.next() => if 
!shutting_down8
 || 
maybe_update0
.
is_some0
() {

                    match maybe_update

                        .err_tip(|| "UpdateForWorker stream closed early")
?1

                        .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 {

                                if let Some(instance_name) = start_execute.execute_request.map(|request| request.instance_name) {

                                    self.grpc_client.clone().execution_response(

                                        ExecuteResult{

                                            instance_name,

                                            operation_id: start_execute.operation_id,

                                            result: Some(execute_result::Result::InternalError(make_err!(Code::ResourceExhausted, "Worker shutting down").into())),

                                            resource_usage: None,

                                        }

                                    ).await?;

                                }

                                continue;

                            }

                            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 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 = 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(),

                                };

                                self.metrics.clone().wrap(move |metrics| async move {

                                    metrics.preconditions.wrap(preconditions_met(precondition_script_cfg, &extra_envs))

                                    .and_then(|()| 
running_actions_manager5
.
create_and_add_action5
(
worker_id5
, 
start_execute5
))

                                    .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);

                                        r

                                    })

                                    .and_then(|action| 
{5

                                        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 
{2

                                                // Notify that execution has completed so it can schedule a new action.

                                                drop(grpc_client.execution_complete(complete).await);

                                                Ok(result)

                                            })

                                            .and_then(RunningAction::upload_results)

                                            .and_then(|action| async move 
{2

                                                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 
{4

                                                if let Err(
e0
) = action.cleanup().await {

                                                    return Result::<FinishedActionResult, Error>::Err(e).merge(result);

                                                }

                                                result

                                            })

                                    
}5
).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")?;

                                        },

                                        Err(e) => {

                                            let is_cas_blob_missing = e.code == Code::NotFound

                                                && e.message_string().contains("not found in either fast or slow store");

                                            if is_cas_blob_missing {

                                                warn!(

                                                    ?e,

                                                    "Missing CAS inputs during prepare_action, returning FAILED_PRECONDITION"

                                                );

                                                let action_result = ActionResult {

                                                    error: Some(make_err!(

                                                        Code::FailedPrecondition,

                                                        "{}",

                                                        e.message_string()

                                                    )),

                                                    ..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(())

                                }

                            };

                            self.actions_in_transit.fetch_add(1, Ordering::Release);

                            let add_future_channel = add_future_channel.clone();

                            info_span!(

                                "worker_start_action_ctx",

                                operation_id = operation_id_to_log,

                                digest_function = %digest_hasher.to_string(),

                            ).in_scope(|| {

                                let _guard = Context::current_with_value(digest_hasher)

                                    .attach();

                                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();

                                actions_in_flight.fetch_add(1, Ordering::Release);

                                futures.push(

                                    spawn!("worker_start_action", start_action_fut).map(move |res| 
{5

                                        let res = res.err_tip(|| "Failed to launch spawn")
?0
;

                                        if let Err(
err3
) = &res {

                                            error!(?err, "Error executing action");

                                        }

                                        add_future_channel

                                            .send(make_publish_future(res).then(move |res| 
{0

                                                actions_in_flight.fetch_sub(1, Ordering::Release);

                                                actions_notify.notify_one();

                                                core::future::ready(res)

                                            
}0
).boxed())

                                            .map_err(|err|

                                                Error::from_std_err(Code::Internal, &err).append("LocalWorker could not send future")

                                                )?;

                                        Ok(())

                                    })

                                    .or_else(move |err| 
{0

                                        // 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()

                                );

                            });

                        }

                    }

                },

                
res5
 = add_future_rx.next() => {

                    let fut = res.err_tip(|| "New future stream receives should never be closed")
?0
;

                    futures.push(fut);

                },

                
res7
 = futures.next() => 
res7
.
err_tip7
(|| "Keep-alive should always pending. Likely unable to send data to scheduler")
?1
?0
,

                
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 shutdown_future = async move {

                        // Wait for in-flight operations to be fully completed.

                        while actions_in_flight.load(Ordering::Acquire) > 0 {

                            actions_notify.notified().await;

                        }

                        // Sending this message immediately evicts all jobs from

                        // this worker, of which there should be none.

                        if let Err(e) = grpc_client.going_away(GoingAwayRequest {}).await {

                            error!("Failed to send GoingAwayRequest: {e}",);

                            return Err(e);

                        }

                        // Allow shutdown to occur now.

                        drop(shutdown_guard);

                        Ok::<(), Error>(())

                    };

                    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>,

}

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> {

    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 == 0 {

        DEFAULT_MAX_ACTION_TIMEOUT

    } else {

        Duration::from_secs(config.max_action_timeout as u64)

    };

    let max_upload_timeout = if config.max_upload_timeout == 0 {

        DEFAULT_MAX_UPLOAD_TIMEOUT

    } else {

        Duration::from_secs(config.max_upload_timeout 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,

        };

        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 = if let Some(
use_namespaces0
) = &config.use_namespaces {

        if *use_namespaces

            && !crate::namespace_utils::namespaces_supported(

                config.use_mount_namespace.unwrap_or_default(),

            )

        {

            return Err(make_err!(Code::Unavailable, "Namespaces not supported"));

        }

        if !*use_namespaces {

            crate::running_actions_manager::UseNamespaces::No

        } else if config.use_mount_namespace.unwrap_or_default() {

            crate::running_actions_manager::UseNamespaces::YesAndMount

        } else {

            crate::running_actions_manager::UseNamespaces::Yes

        }

    } else if config

        .use_mount_namespace

        .is_some_and(core::convert::identity)

    {

        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"

        ));

    }

    let running_actions_manager =

        Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs {

            root_action_directory: config.work_directory.clone(),

            execution_configuration: ExecutionConfiguration {

                entrypoint,

                additional_environment: config.additional_environment.clone(),

            },

            cas_store: fast_slow_store,

            ac_store,

            historical_store,

            upload_action_result_config: &config.upload_action_result,

            max_action_timeout,

            max_upload_timeout,

            timeout_handled_externally: config.timeout_handled_externally,

            directory_cache,

            #[cfg(target_os = "linux")]

            use_namespaces,

        })
?0
);

    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);

                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(),

        )));

        Self {

            config,

            running_actions_manager,

            connection_factory,

            sleep_fn: Some(sleep_fn),

            metrics,

        }

    }

    #[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, 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 connect_worker_request = make_connect_worker_request(

            self.config.name.clone(),

            &self.config.platform_properties,

            &extra_envs,

            self.config.max_inflight_tasks,

        )

        .await
?0
;

        let 
mut update_for_worker_stream9
 = client

            .connect_worker(connect_worker_request)

            .await

            .err_tip(|| "Could not call connect_worker() in worker")
?0

            .into_inner();

        let 
first_msg_update8
 = 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 = match first_msg_update {

            Some(Update::ConnectionResult(connection_result)) => connection_result.worker_id,

            other => {

                return Err(make_input_err!(

                    "Expected first response from scheduler to be a ConnectionResult got : {:?}",

                    other

                ));

            }

        };

        Ok((worker_id, 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 error_handler = Box::pin(move |err| async move {

            error!(?err, "Error");

            (sleep_fn_pin)(Duration::from_secs_f32(CONNECTION_RETRY_DELAY_S)).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, update_for_worker_stream)) => (

                    LocalWorkerImpl::new(

                        &self.config,

                        client,

                        worker_id,

                        self.running_actions_manager.clone(),

                        self.metrics.clone(),

                    ),

                    update_for_worker_stream,

                ),

            };

            info!(

                worker_id = %inner.worker_id,

                "Worker registered with scheduler"

            );

            // Now listen for connections and run all other services.

            if let Err(err) = inner.run(update_for_worker_stream, &mut shutdown_rx).await {

                'no_more_actions: {

                    // Ensure there are no actions in transit before we try to kill

                    // all our actions.

                    const ITERATIONS: usize = 1_000;

                    const ERROR_MSG: &str = "Actions in transit did not reach zero before we disconnected from the scheduler";

                    let sleep_duration = ACTIONS_IN_TRANSIT_TIMEOUT_S / ITERATIONS as f32;

                    for _ in 0..ITERATIONS {

                        if inner.actions_in_transit.load(Ordering::Acquire) == 0 {

                            break 'no_more_actions;

                        }

                        (sleep_fn_pin)(Duration::from_secs_f32(sleep_duration)).await;

                    }

                    error!(ERROR_MSG);

                    return Err(err.append(ERROR_MSG));

                }

                error!(?err, "Worker disconnected from 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;

                (error_handler)(err).await; // Try to connect again.

            }

        }

        // 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 = "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(),

            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 {

        fut(self).await

    }

}