// Copyright 2024 The NativeLink Authors. All rights reserved.

//

// Licensed under the Apache License, Version 2.0 (the "License");

// you may not use this file except in compliance with the License.

// You may obtain a copy of the License at

//

//    http://www.apache.org/licenses/LICENSE-2.0

//

// 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 std::pin::Pin;

use std::process::Stdio;

use std::str;

use std::sync::atomic::{AtomicU64, Ordering};

use std::sync::{Arc, Weak};

use std::time::Duration;

use futures::future::BoxFuture;

use futures::stream::FuturesUnordered;

use futures::{select, Future, FutureExt, StreamExt, TryFutureExt};

use nativelink_config::cas_server::LocalWorkerConfig;

use nativelink_error::{make_err, make_input_err, Code, Error, ResultExt};

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

    execute_result, ExecuteResult, GoingAwayRequest, KeepAliveRequest, UpdateForWorker,

};

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, ACTIVE_HASHER_FUNC};

use nativelink_util::metrics_utils::{AsyncCounterWrapper, CounterWithTime};

use nativelink_util::origin_context::ActiveOriginContext;

use nativelink_util::shutdown_guard::ShutdownGuard;

use nativelink_util::store_trait::Store;

use nativelink_util::{spawn, tls_utils};

use tokio::process;

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

use tokio::time::sleep;

use tokio_stream::wrappers::UnboundedReceiverStream;

use tonic::Streaming;

use tracing::{event, info_span, instrument, Level};

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

/// 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_secs(1200); // 20 mins.

struct LocalWorkerImpl<'a, T: WorkerApiClientTrait, 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>,

}

async fn preconditions_met(precondition_script: Option<String>) -> Result<(), Error> {

    let Some(
precondition_script1
) = &precondition_script else {

  Branch (87:9): [True: 1, False: 3]

  Branch (87:9): [Folded - Ignored]

  Branch (87:9): [Folded - Ignored]

        // 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 precondition_process = process::Command::new(precondition_script)

        .kill_on_drop(true)

        .stdin(Stdio::null())

        .stdout(Stdio::piped())

        .stderr(Stdio::null())

        .env_clear()

        .spawn()

        .err_tip(|| 
format!("Could not execute precondition command {precondition_script:?}")0
)
?0
;

    let output = precondition_process.wait_with_output().await
?0
;

    if output.status.code() == Some(0) {

  Branch (106:8): [True: 0, False: 1]

  Branch (106:8): [Folded - Ignored]

  Branch (106:8): [Folded - Ignored]

        Ok(())

    } else {

        Err(make_err!(

            Code::ResourceExhausted,

            "Preconditions script returned status {} - {}",

            output.status,

            str::from_utf8(&output.stdout).unwrap_or("")

        ))

    }

}

impl<'a, T: WorkerApiClientTrait, 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();

        loop {

            let timeout = self

                .config

                .worker_api_endpoint

                .timeout

                .unwrap_or(DEFAULT_ENDPOINT_TIMEOUT_S);

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

            sleep(Duration::from_secs_f32(timeout / 2.)).
await0
;

            if let Err(e) = grpc_client

  Branch (154:20): [True: 0, False: 0]

  Branch (154:20): [Folded - Ignored]

  Branch (154:20): [Folded - Ignored]

                .keep_alive(KeepAliveRequest {

                    worker_id: self.worker_id.clone(),

                })

                .await

            {

                return Err(make_err!(

                    Code::Internal,

                    "Failed to send KeepAlive in LocalWorker : {:?}",

                    e

                ));

            }

        }

    }

    async fn run(

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

        loop {

            select! {

                
maybe_update6
 = update_for_worker_stream.next() => {

                    match maybe_update

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

                        .err_tip(|| 
"Got error in UpdateForWorker stream"0
)
?0

                        .update

                        .err_tip(|| 
"Expected update to exist in UpdateForWorker"0
)
?0

                    {

                        Update::ConnectionResult(_) => {

                            return Err(make_input_err!(

                                "Got ConnectionResult in LocalWorker::run which should never happen"

                            ));

                        }

                        // TODO(allada) 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).
await0
 {

  Branch (217:36): [True: 0, False: 1]

  Branch (217:36): [Folded - Ignored]

  Branch (217:36): [Folded - Ignored]

                                event!(

                                    Level::ERROR,

                                    ?operation_id,

                                    ?err,

                                    "Failed to send kill request for operation"

                                );

                            };

                        }

                        Update::StartAction(start_execute) => {

                            self.metrics.start_actions_received.inc();

                            let execute_request = start_execute.execute_request.as_ref();

                            let operation_id = start_execute.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(make_input_err!("Expected execute_request to be set"))

                                .and_then(|v| DigestHasherFunc::try_from(v.digest_function))

                                .err_tip(|| 
"In LocalWorkerImpl::new()"0
)
?0
;

                            let start_action_fut = {

                                let precondition_script_cfg = self.config.experimental_precondition_script.clone();

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

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

                                    metrics.preconditions.wrap(preconditions_met(precondition_script_cfg))

                                    .and_then(|_| 
running_actions_manager.create_and_add_action(worker_id, start_execute)3
)

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

                                        event!(

                                            Level::INFO,

                                            operation_id = ?action.get_operation_id(),

                                            "Received request to run action"

                                        );

                                        action

                                            .clone()

                                            .prepare_action()

                                            .and_then(RunningAction::execute)

                                            .and_then(RunningAction::upload_results)

                                            .and_then(RunningAction::get_finished_result)

                                            // Note: We need ensure we run cleanup even if one of the other steps fail.

                                            .then(|result| async move 
{2

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

  Branch (266:56): [True: 0, False: 2]

  Branch (266:56): [Folded - Ignored]

  Branch (266:56): [Folded - Ignored]

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

                                                }

                                                result

                                            })

                                    }
)4
.await

                                })

                            };

                            let make_publish_future = {

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

                                let worker_id = self.worker_id.clone();

                                let running_actions_manager = self.running_actions_manager.clone();

                                move |res: Result<ActionResult, 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
)
?0
;

                                    match res {

                                        Ok(mut action_result) => {

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

  Branch (286:52): [True: 2, False: 0]

  Branch (286:52): [Folded - Ignored]

  Branch (286:52): [Folded - Ignored]

                                                if let Err(
err0
) = running_actions_manager.cache_action_result(digest_info, &mut action_result, digest_hasher).
await0
 {

  Branch (287:56): [True: 0, False: 2]

  Branch (287:56): [Folded - Ignored]

  Branch (287:56): [Folded - Ignored]

                                                    event!(

                                                        Level::ERROR,

                                                        ?err,

                                                        ?action_digest,

                                                        "Error saving action in store",

                                                    );

                                                }

                                            }

                                            let action_stage = ActionStage::Completed(action_result);

                                            grpc_client.execution_response(

                                                ExecuteResult{

                                                    worker_id,

                                                    instance_name,

                                                    operation_id,

                                                    result: Some(execute_result::Result::ExecuteResponse(action_stage.into())),

                                                }

                                            )

                                            .await

                                            .err_tip(|| "Error while calling execution_response")?;

                                        },

                                        Err(e) => {

                                            grpc_client.execution_response(ExecuteResult{

                                                worker_id,

                                                instance_name,

                                                operation_id,

                                                result: Some(execute_result::Result::InternalError(e.into())),

                                            }).
await0
.
err_tip(0
|| 
"Error calling execution_response with error"0
)0
?0
;

                                        },

                                    }

                                    Ok(())

                                }

                            };

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

                            let futures_ref = &futures;

                            let add_future_channel = add_future_channel.clone();

                            let mut ctx = ActiveOriginContext::fork().err_tip(|| 
"Expected ActiveOriginContext to be set in local_worker::run"0
)
?0
;

                            ctx.set_value(&ACTIVE_HASHER_FUNC, Arc::new(digest_hasher));

                            ctx.run(info_span!("worker_start_action_ctx"), move || {

                                futures_ref.push(

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

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

                                        if let Err(
err1
) = &res {

  Branch (331:48): [True: 1, False: 2]

  Branch (331:48): [Folded - Ignored]

  Branch (331:48): [Folded - Ignored]

                                            event!(

                                                Level::ERROR,

                                                ?err,

                                                "Error executing action",

                                            );

                                        }

                                        add_future_channel

                                            .send(make_publish_future(res).boxed())

                                            .map_err(|_| 
make_err!(Code::Internal, "LocalWorker could not send future")0
)
?0
;

                                        Ok(())

                                    
}3
)

                                    .boxed()

                                );

                            });

                        }

                    };

                },

                
res3
 = add_future_rx.next() => {

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

                    futures.push(fut);

                },

                
res3
 = futures.next() => 
res.err_tip(3
|| 
"Keep-alive should always pending. Likely unable to send data to scheduler"0
)3
?0
?0
,

                
complete_msg0
 = shutdown_rx.recv().fuse() => {

                    event!(Level::WARN, "Worker loop reveiced shutdown signal. Shutting down worker...",);

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

                    let worker_id = self.worker_id.clone();

                    let running_actions_manager = self.running_actions_manager.clone();

                    let complete_msg_clone = complete_msg.map_err(|e| make_err!(Code::Internal, "Failed to receive shutdown message: {e:?}"))?.clone();

                    let shutdown_future = async move {

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

  Branch (361:32): [True: 0, False: 0]

  Branch (361:32): [Folded - Ignored]

  Branch (361:32): [Folded - Ignored]

                            event!(Level::ERROR, "Failed to send GoingAwayRequest: {e}",);

                            return Err(e.into());

                        }

                        running_actions_manager.complete_actions(complete_msg_clone).await;

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

                    };

                    futures.push(shutdown_future.boxed());

                },

            };

        }

        // Unreachable.

    }

}

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

pub struct LocalWorker<T: WorkerApiClientTrait, 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>,

}

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

        Arc<Metrics>,

    ),

    Error,

> {

    let fast_slow_store = cas_store

        .downcast_ref::<FastSlowStore>(None)

        .err_tip(|| 
"Expected store for LocalWorker's store to be a FastSlowStore"0
)
?0

        .get_arc()

        .err_tip(|| 
"FastSlowStore's Arc doesn't exist"0
)
?0
;

    if let Ok(
path1
) = fs::canonicalize(&config.work_directory).await {

  Branch (406:12): [True: 1, False: 1]

  Branch (406:12): [Folded - Ignored]

  Branch (406:12): [Folded - Ignored]

        fs::remove_dir_all(path)

            .await

            .err_tip(|| 
"Could not remove work_directory in LocalWorker"0
)
?0
;

    }

    fs::create_dir_all(&config.work_directory)

        .await

        .err_tip(|| 
format!("Could not make work_directory : {}", config.work_directory)0
)
?0
;

    let entrypoint = if config.entrypoint.is_empty() {

  Branch (415:25): [True: 2, False: 0]

  Branch (415:25): [Folded - Ignored]

  Branch (415:25): [Folded - Ignored]

        None

    } else {

        Some(config.entrypoint.clone())

    };

    let max_action_timeout = if config.max_action_timeout == 0 {

  Branch (420:33): [True: 2, False: 0]

  Branch (420:33): [Folded - Ignored]

  Branch (420:33): [Folded - Ignored]

        DEFAULT_MAX_ACTION_TIMEOUT

    } else {

        Duration::from_secs(config.max_action_timeout as u64)

    };

    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,

            timeout_handled_externally: config.timeout_handled_externally,

        })
?0
);

    let local_worker = LocalWorker::new_with_connection_factory_and_actions_manager(

        config.clone(),

        running_actions_manager,

        Box::new(move || {

            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| make_input_err!("Invalid URI for worker endpoint : {e:?}"))?

                        .connect_timeout(timeout_duration)

                        .timeout(timeout_duration);

                let transport = endpoint.connect().await.map_err(|e| {

                    make_err!(

                        Code::Internal,

                        "Could not connect to endpoint {}: {e:?}",

                        config.worker_api_endpoint.uri

                    )

                })?;

                Ok(WorkerApiClient::new(transport).into())

            })

        }),

        Box::new(move |d| 
Box::pin(sleep(d))0
),

    );

    let metrics = local_worker.metrics.clone();

    Ok((local_worker, metrics))

}

impl<T: WorkerApiClientTrait, 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,

        }

    }

    pub fn name(&self) -> &String {

        &self.config.name

    }

    async fn register_worker(

        &self,

        client: &mut T,

    ) -> Result<(String, Streaming<UpdateForWorker>), Error> {

        let supported_properties =

            make_supported_properties(&self.config.platform_properties).
await0
?0
;

        let 
mut update_for_worker_stream6
 = client

            .connect_worker(supported_properties)

            .await

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

            .into_inner();

        let 
first_msg_update5
 = update_for_worker_stream

            .next()

            .await

            .err_tip(|| 
"Got EOF expected UpdateForWorker"1
)
?1

            .err_tip(|| 
"Got error when receiving UpdateForWorker"0
)
?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 ConnectResult got : {:?}",

                    other

                ))

            }

        };

        Ok((worker_id, update_for_worker_stream))

    }

    #[instrument(skip(self), level = 
Level::INFO6
)
]0

    pub async fn run(

        mut self,

        mut shutdown_rx: broadcast::Receiver<ShutdownGuard>,

    ) -> Result<(), Error> {

        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 {

            event!(Level::ERROR, ?err, "Error");

            (sleep_fn_pin)(Duration::from_secs_f32(CONNECTION_RETRY_DELAY_S)).
await0
;

        });

        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.

                }

            };

            // Next register our worker with the scheduler.

            let (mut 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,

                    ),

                };

            event!(

                Level::WARN,

                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;

                    }

                    event!(Level::ERROR, ERROR_MSG);

                    return Err(err.append(ERROR_MSG));

                }

                event!(Level::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;

                continue; // Try to connect again.

            }

        }

        // Unreachable.

    }

}

#[derive(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]

    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

    }

}