/build/source/nativelink-service/src/worker_api_server.rs

Source
// Copyright 2024 Trace Machina, Inc. 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::convert::Into;
use core::pin::Pin;
use core::time::Duration;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};

use futures::stream::unfold;
use futures::{Stream, StreamExt};
use nativelink_config::cas_server::WorkerApiConfig;
use nativelink_error::{make_err, Code, Error, ResultExt};
use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::update_for_scheduler::Update;
use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::worker_api_server::{
    WorkerApi, WorkerApiServer as Server,
};
use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{
    execute_declined, execute_result, ExecuteAccepted, ExecuteComplete, ExecuteDeclined, ExecuteResult, GoingAwayRequest, KeepAliveRequest, UpdateForScheduler, UpdateForWorker
};
use nativelink_scheduler::worker::Worker;
use nativelink_scheduler::worker_scheduler::WorkerScheduler;
use nativelink_util::background_spawn;
use nativelink_util::action_messages::{OperationId, WorkerId};
use nativelink_util::operation_state_manager::UpdateOperationType;
use nativelink_util::platform_properties::PlatformProperties;
use rand::RngCore;
use tokio::sync::mpsc;
use tokio::time::interval;
use tonic::{Response, Status};
use tracing::{debug, error, warn, instrument, Level};
use uuid::Uuid;

pub type ConnectWorkerStream =
    Pin<Box<dyn Stream<Item = Result<UpdateForWorker, Status>> + Send + Sync + 'static>>;

pub type NowFn = Box<dyn Fn() -> Result<Duration, Error> + Send + Sync>;

/// Makes the unique part of a new worker's id.
pub type WorkerIdFn = Box<dyn Fn() -> String + Send + Sync>;

/// How often workers are told to kill operations the scheduler no longer
/// has executing on them, unless overridden by
/// `kill_revoked_operations_interval_s` in the worker API config.
const DEFAULT_KILL_REVOKED_OPERATIONS_INTERVAL_S: u64 = 5;

pub struct WorkerApiServer {
    scheduler: Arc<dyn WorkerScheduler>,
    now_fn: Arc<NowFn>,
    node_id: [u8; 6],
    worker_id_fn: Option<WorkerIdFn>,
}

impl core::fmt::Debug for WorkerApiServer {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("WorkerApiServer")
            .field("node_id", &self.node_id)
            .finish_non_exhaustive()
    }
}

impl WorkerApiServer {
    pub fn new(
        config: &WorkerApiConfig,
        schedulers: &HashMap<String, Arc<dyn WorkerScheduler>>,
    ) -> Result<Self, Error> {
        let node_id = {
            let mut out = [0; 6];
            rand::rng().fill_bytes(&mut out);
            out
        };
        let now_fn: Arc<NowFn> = Arc::new(Box::new(move || {
            SystemTime::now().duration_since(UNIX_EPOCH).map_err(|err| {
                Error::from_std_err(Code::Internal, &err)
                    .append("System time is now behind unix epoch")
            })
        }));
        Self::spawn_timers(config, schedulers, &now_fn);
        Self::with_shared_now_fn(config, schedulers, now_fn, node_id)
    }

    /// Starts the worker timeout and revoked-operation timers for every
    /// scheduler, reading the time from `now_fn`. `new()` calls it; a server
    /// built with `new_with_now_fn()` has no timers until this is called.
    pub fn spawn_timers(
        config: &WorkerApiConfig,
        schedulers: &HashMap<String, Arc<dyn WorkerScheduler>>,
        now_fn: &Arc<NowFn>,
    ) {
        let kill_revoked_enabled = !config.disable_kill_revoked_operations;
        let kill_revoked_interval_s = if config.kill_revoked_operations_interval_s == 0 {
            DEFAULT_KILL_REVOKED_OPERATIONS_INTERVAL_S
        } else {
            config.kill_revoked_operations_interval_s
        };
        for scheduler in schedulers.values() {
            // This will protect us from holding a reference to the scheduler forever in the
            // event our ExecutionServer dies. Our scheduler is a weak ref, so the spawn will
            // eventually see the Arc went away and return.
            let weak_scheduler = Arc::downgrade(scheduler);
            let now_fn = now_fn.clone();
            background_spawn!("worker_api_server", async move {
                let mut timeout_ticker = interval(Duration::from_secs(1));
                let mut kill_ticker = interval(Duration::from_secs(kill_revoked_interval_s));
                loop {
                    tokio::select! {
                        _ = timeout_ticker.tick() => {
                            let timestamp = match now_fn() {
                                Ok(timestamp) => timestamp,
                                Err(err) => {
                                    error!(?err, "Failed to read the time for remove_timedout_workers");
                                    continue;
                                }
                            };
                            match weak_scheduler.upgrade() {
                                Some(scheduler) => {
                                    if let Err(err) =
                                        scheduler.remove_timedout_workers(timestamp.as_secs()).await
                                    {
                                        error!(?err, "Failed to remove_timedout_workers",);
                                    }
                                }
                                // If we fail to upgrade, our service is probably destroyed, so return.
                                None => return,
                            }
                        }
                        _ = kill_ticker.tick(), if kill_revoked_enabled => {
                            match weak_scheduler.upgrade() {
                                Some(scheduler) => {
                                    if let Err(err0) = scheduler.kill_revoked_operations().await {
                                        error!(?err, "Failed to kill_revoked_operations");
                                    }
                                }
                                None => return,
                            }
                        }
                    }
                }
            });
        }
    }

    /// Same as `new()`, but you can pass a custom `now_fn`, that returns a Duration since `UNIX_EPOCH`
    /// representing the current time. Used mostly in  unit tests.
    pub fn new_with_now_fn(
        config: &WorkerApiConfig,
        schedulers: &HashMap<String, Arc<dyn WorkerScheduler>>,
        now_fn: NowFn,
        node_id: [u8; 6],
    ) -> Result<Self, Error> {
        Self::with_shared_now_fn(config, schedulers, Arc::new(now_fn), node_id)
    }

    fn with_shared_now_fn(
        config: &WorkerApiConfig,
        schedulers: &HashMap<String, Arc<dyn WorkerScheduler>>,
        now_fn: Arc<NowFn>,
        node_id: [u8; 6],
    ) -> Result<Self, Error> {
        let scheduler = schedulers
            .get(&config.scheduler)
            .err_tip(|| {0
                format!(
                    "Scheduler needs config for '{}' because it exists in worker_api",
                    config.scheduler
                )
            })?
            .clone();
        Ok(Self {
            scheduler,
            now_fn,
            node_id,
            worker_id_fn: None,
        })
    }

    /// Takes the unique part of each worker id from `worker_id_fn` instead
    /// of a time-based UUID, for tests that need the same ids every run.
    #[must_use]
    pub fn with_worker_id_fn(mut self, worker_id_fn: WorkerIdFn) -> Self {
        self.worker_id_fn = Some(worker_id_fn);
        self
    }

    pub fn into_service(self) -> Server<Self> {
        Server::new(self)
    }

    async fn inner_connect_worker(
        &self,
        mut update_stream: impl Stream<Item = Result<UpdateForScheduler, Status>>
        + Unpin
        + Send
        + 'static,
    ) -> Result<Response<ConnectWorkerStream>, Error> {
        let first_message = update_stream
            .next()
            .await
            .err_tip(|| "Missing first message for connect_worker")?0
            .err_tip(|| "Error reading first message for connect_worker")?0;
        let Some(Update::ConnectWorkerRequest(connect_worker_request)) = first_message.update
        else {
            return Err(make_err!(
                Code::Internal,
                "First message was not a ConnectWorkerRequest"
            ));
        };

        let (tx, rx) = mpsc::channel(nativelink_scheduler::worker::channel_capacity(
            connect_worker_request.max_inflight_tasks,
        ));

        // First convert our proto platform properties into one our scheduler understands.
        let platform_properties = {
            let mut platform_properties = PlatformProperties::default();
            for property7 in connect_worker_request.properties {
                let platform_property_value = self
                    .scheduler
                    .get_platform_property_manager()
                    .make_prop_value(&property.name, &property.value)
                    .err_tip(|| "Bad Property during connect_worker()")?0;
                platform_properties
                    .properties
                    .insert(property.name.clone(), platform_property_value);
            }
            platform_properties
        };

        // Now register the worker with the scheduler.
        let worker_id = {
            let unique = self.worker_id_fn.as_ref().map_or_else(
                || Uuid::now_v6(&self.node_id).hyphenated().to_string(),
                |worker_id_fn| worker_id_fn(),
            );
            let worker_id = WorkerId(format!(
                "{}{unique}",
                connect_worker_request.worker_id_prefix,
            ));
            let mut worker = Worker::new(
                worker_id.clone(),
                platform_properties,
                tx,
                (self.now_fn)()?0.as_secs(),
                connect_worker_request.max_inflight_tasks,
            );
            worker.admits_when_idle = connect_worker_request.admits_when_idle;
            self.scheduler
                .add_worker(worker)
                .await
                .err_tip(|| "Failed to add worker in inner_connect_worker()")?0;
            worker_id
        };

        WorkerConnection::start(
            self.scheduler.clone(),
            self.now_fn.clone(),
            worker_id.clone(),
            update_stream,
        );

        Ok(Response::new(Box::pin(unfold(
            (rx, worker_id),
            move |state| async move {
                let (mut rx, worker_id) = state;
                if let Some(update_for_worker13.0k) = rx.recv().await {
                    return Some((Ok(update_for_worker), (rx, worker_id)));
                }
                warn!(
                    ?worker_id,
                    "UpdateForWorker channel was closed, thus closing connection to worker node",
                );

                None
            },
        ))))
    }

    pub async fn inner_connect_worker_for_testing(
        &self,
        update_stream: impl Stream<Item = Result<UpdateForScheduler, Status>> + Unpin + Send + 'static,
    ) -> Result<Response<ConnectWorkerStream>, Error> {
        self.inner_connect_worker(update_stream).await
    }
}

#[tonic::async_trait]
impl WorkerApi for WorkerApiServer {
    type ConnectWorkerStream = ConnectWorkerStream;

    #[instrument(
        err,
        level = Level::ERROR,
        skip_all,
        fields(request = ?grpc_request.get_ref())
    )]
    async fn connect_worker(
        &self,
        grpc_request: tonic::Request<tonic::Streaming<UpdateForScheduler>>,
    ) -> Result<Response<Self::ConnectWorkerStream>, Status> {
        let resp = self
            .inner_connect_worker(grpc_request.into_inner())
            .await
            .map_err(Into::into);
        if resp.is_ok() {
            debug!(return = "Ok(<stream>)");
        }
        resp
    }
}

struct WorkerConnection {
    scheduler: Arc<dyn WorkerScheduler>,
    now_fn: Arc<NowFn>,
    worker_id: WorkerId,
}

impl WorkerConnection {
    fn start(
        scheduler: Arc<dyn WorkerScheduler>,
        now_fn: Arc<NowFn>,
        worker_id: WorkerId,
        mut connection: impl Stream<Item = Result<UpdateForScheduler, Status>> + Unpin + Send + 'static,
    ) {
        let instance = Self {
            scheduler,
            now_fn,
            worker_id,
        };

        background_spawn!("worker_api", async move {2.00k
            // The stream is read by a task of its own that only hands each
            // message on; they are processed here, in order. Processing one
            // inline in the reader (a result's store write, say) stopped the
            // reads for its duration, and a worker that filled the stream's
            // flow-control window in the meantime had nothing it was allowed
            // to put on the wire: its keepalives waited behind the stall
            // until this scheduler evicted it. The queue is unbounded on
            // purpose: a worker's messages are bounded by its in-flight
            // actions and its keepalive interval, and backpressure here is
            // exactly the stall.
            let (tx, mut rx) = mpsc::unbounded_channel();
            let reader_worker_id = instance.worker_id.clone();
            background_spawn!("worker_api_reader", async move {
                while let Some(maybe_update99.8k) = connection.next().await {
                    let update = match maybe_update.map(|u| u.update) {
                        Ok(Some(update)) => update,
                        Ok(None) => {
                            tracing::warn!(worker_id=?reader_worker_id, "Empty update");
                            continue;
                        }
                        Err(err) => {
                            tracing::warn!(worker_id=?reader_worker_id, ?err, "Error from worker");
                            break;
                        }
                    };
                    if tx.send(update).is_err() {
                        break;
                    }
                }
            });
            let mut had_going_away = false;
            while let Some(update97.9k) = rx.recv().await {
                let result97.8k = match update {
                    Update::ConnectWorkerRequest(_connect_worker_request) => Err(make_err!(
                        Code::Internal,
                        "Got ConnectWorkerRequest after initial message for {}",
                        instance.worker_id
                    )),
                    Update::KeepAliveRequest(keep_alive_request) => {
                        instance.inner_keep_alive(keep_alive_request).await
                    }
                    Update::GoingAwayRequest(going_away_request) => {
                        // A drain keeps the worker until its stream closes,
                        // so that close still has to remove it.
                        had_going_away = !going_away_request.drain;
                        instance.inner_going_away(going_away_request).await
                    }
                    Update::ExecuteResult(execute_result) => {
                        instance.inner_execution_response(execute_result).await
                    }
                    Update::ExecuteComplete(execute_complete) => {
                        instance.execution_complete(execute_complete).await
                    }
                    Update::ExecuteAccepted(execute_accepted) => {
                        instance.dispatch_accepted(execute_accepted).await
                    }
                    Update::ExecuteDeclined(execute_declined) => {
                        instance.dispatch_declined(execute_declined).await
                    }
                };
                if let Err(err2.54k) = result {
                    tracing::warn!(worker_id=?instance.worker_id, ?err, "Error processing worker message");
                }
            }
            tracing::debug!(worker_id=?instance.worker_id, "Update for scheduler dropped");
            if !had_going_away {
                drop(
                    instance
                        .scheduler
                        .worker_disconnected(&instance.worker_id)
                        .await,
                );
            }
        });
    }

    async fn inner_keep_alive(&self, keep_alive_request: KeepAliveRequest) -> Result<(), Error> {
        self.scheduler
            .worker_keep_alive_received(
                &self.worker_id,
                (self.now_fn)()?0.as_secs(),
                keep_alive_request.load,
            )
            .await
            .err_tip(|| "Could not process keep_alive from worker in inner_keep_alive()")?2.53k;
        if let Some(running64.2k) = keep_alive_request.running_operations {
            // The memory first: the running report can release an operation
            // the worker left out, and a peak for it is then dropped anyway.
            let peaks: Vec<(OperationId, u64)> = running
                .memory
                .into_iter()
                .map(|memory| {1
                    (
                        OperationId::from(memory.operation_id),
                        memory.peak_memory_kb,
                    )
                })
                .collect();
            if !peaks.is_empty() {
                self.scheduler
                    .worker_operation_memory_reported(&self.worker_id, peaks)
                    .await
                    .err_tip(|| "Could not process the memory report in inner_keep_alive()")?0;
            }
            self.scheduler
                .worker_running_operations_reported(
                    &self.worker_id,
                    running
                        .operation_ids
                        .into_iter()
                        .map(OperationId::from)
                        .collect(),
                )
                .await
                .err_tip(
                    || "Could not process the running operations report in inner_keep_alive()",
                )?;
        }
        Ok(())
    }

    async fn inner_going_away(&self, going_away_request: GoingAwayRequest) -> Result<(), Error> {
        if going_away_request.drain {
            self.scheduler
                .set_drain_worker(&self.worker_id, true)
                .await
                .err_tip(|| "While draining worker in WorkerApiServer::inner_going_away")?0;
            return Ok(());
        }
        self.scheduler
            .remove_worker(&self.worker_id)
            .await
            .err_tip(|| "While calling WorkerApiServer::inner_going_away")?;
        Ok(())
    }

    /// Any message from the worker proves it is alive, not only keepalives.
    /// Only a keepalive carries a load report and lifts a pause; this
    /// refreshes the timestamp alone, so a decline cannot undo the pause
    /// it just took.
    ///
    /// The refresh is best-effort: it is a side effect of the message,
    /// never the point of it. A result or completion the worker already
    /// produced must reach `update_action` even if the refresh fails, or
    /// the operation is stranded in `Executing` with nothing left to
    /// finish it.
    async fn touch_liveness_best_effort(&self) {
        let result = async {
            self.scheduler
                .worker_liveness_refreshed(&self.worker_id, (self.now_fn)()?0.as_secs())
                .await
        }
        .await;
        if let Err(err2) = result {
            warn!(
                worker_id = %self.worker_id,
                ?err,
                "Could not refresh worker liveness; continuing to process the message",
            );
        }
    }

    async fn inner_execution_response(&self, execute_result: ExecuteResult) -> Result<(), Error> {
        self.touch_liveness_best_effort().await;
        let operation_id = OperationId::from(execute_result.operation_id.clone());

        if let Some(resource_usage0) = execute_result.resource_usage {
            self.scheduler
                .record_action_resource_usage(&self.worker_id, &operation_id, resource_usage)
                .await
                .err_tip(|| {
                    format!("Failed to record resource usage for operation {operation_id}")
                })?;
        }

        match execute_result
            .result
            .err_tip(|| "Expected result to exist in ExecuteResult")?0
        {
            execute_result::Result::ExecuteResponse(finished_result) => {
                let action_stage = finished_result
                    .try_into()
                    .err_tip(|| "Failed to convert ExecuteResponse into an ActionStage")?0;
                self.scheduler
                    .update_action(
                        &self.worker_id,
                        &operation_id,
                        UpdateOperationType::UpdateWithActionStage(action_stage),
                    )
                    .await
                    .err_tip(|| format!9("Failed to operation {operation_id}"))?9;
            }
            execute_result::Result::InternalError(e) => {
                self.scheduler
                    .update_action(
                        &self.worker_id,
                        &operation_id,
                        UpdateOperationType::UpdateWithError(e.into()),
                    )
                    .await
                    .err_tip(|| format!1("Failed to operation {operation_id}"))?1;
            }
        }
        Ok(())
    }

    /// The answer is recorded before the liveness refresh: the refresh runs
    /// the unacknowledged sweep, which would otherwise requeue an
    /// acknowledgement that arrives right at the timeout.
    async fn dispatch_accepted(&self, execute_accepted: ExecuteAccepted) -> Result<(), Error> {
        let operation_id = OperationId::from(execute_accepted.operation_id);
        self.scheduler
            .worker_dispatch_accepted(&self.worker_id, &operation_id)
            .await
            .err_tip(|| format!0("Failed to record acceptance of operation {operation_id}"))?0;
        self.touch_liveness_best_effort().await;
        Ok(())
    }

    async fn dispatch_declined(&self, execute_declined: ExecuteDeclined) -> Result<(), Error> {
        let operation_id = OperationId::from(execute_declined.operation_id);
        let reason = execute_declined::Reason::try_from(execute_declined.reason)
            .unwrap_or(execute_declined::Reason::Unspecified);
        let needs_kb = (reason == execute_declined::Reason::Load && execute_declined.needed_kb > 0)
            .then_some(execute_declined.needed_kb);
        let mut why = reason.as_str_name().to_ascii_lowercase();
        if !execute_declined.detail.is_empty() {
            why.push_str(": ");
            why.push_str(&execute_declined.detail);
        }
        self.scheduler
            .worker_dispatch_declined(&self.worker_id, &operation_id, why, needs_kb)
            .await
            .err_tip(|| format!0("Failed to record decline of operation {operation_id}"))?0;
        self.touch_liveness_best_effort().await;
        Ok(())
    }

    async fn execution_complete(&self, execute_complete: ExecuteComplete) -> Result<(), Error> {
        self.touch_liveness_best_effort().await;
        let operation_id = OperationId::from(execute_complete.operation_id);
        self.scheduler
            .update_action(
                &self.worker_id,
                &operation_id,
                UpdateOperationType::ExecutionComplete,
            )
            .await
            .err_tip(|| format!0("Failed to operation {operation_id}"))?0;
        Ok(())
    }
}

Coverage Report

Created: 2026-11-13 23:29