/build/source/nativelink-scheduler/src/simple_scheduler_state_manager.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::ops::Bound;
use core::time::Duration;
use std::collections::HashMap;
use std::string::ToString;
use std::sync::{Arc, Weak};
use std::time::SystemTime;

use async_lock::Mutex;
use async_trait::async_trait;
use futures::future::BoxFuture;
use futures::{StreamExt, TryStreamExt, stream};
use nativelink_error::{Code, Error, ResultExt, make_err};
use nativelink_metric::MetricsComponent;
use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::AttemptEndReason;
use nativelink_util::action_messages::{
    ActionInfo, ActionResult, ActionStage, ActionState, ActionUniqueQualifier, ExecutionMetadata,
    LOST_RESULT_ERROR_TIP, OperationId, STORE_UNREACHABLE_ERROR_TIP, WorkerId,
};
use nativelink_util::instant_wrapper::InstantWrapper;
use nativelink_util::metrics::{
    AgeBuckets, EXECUTION_METRICS, EXECUTION_RESULT, EXECUTION_STAGE, ExecutionResult,
    ExecutionStage, record_awaited_action_orphan, record_dispatch_requeue, record_queue_retired,
};
use nativelink_util::operation_state_manager::{
    ActionStateResult, ActionStateResultStream, ClientStateManager, DispatchFacts, Escalation,
    MatchingEngineStateManager, OperationFilter, OperationStageFlags, OrderDirection,
    UpdateOperationType, WorkerStateManager,
};
use nativelink_util::origin_event::OriginMetadata;
use nativelink_util::platform_properties::PlatformProperties;
use opentelemetry::KeyValue;
use tracing::{debug, info, trace, warn};

use super::awaited_action_db::{
    AwaitedAction, AwaitedActionDb, AwaitedActionSubscriber, PeerScheduler,
    SortedAwaitedActionState,
};
use crate::worker_registry::{ORPHANED_ACTION_TIMEOUT, SharedWorkerRegistry, WorkerLiveness};

/// An attempt the executing sweep ended, for the scheduler to report.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SweptAttempt {
    pub operation_id: OperationId,
    pub worker_id: WorkerId,
    pub reason: AttemptEndReason,
    pub detail: String,
}

/// What one executing sweep did.
#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct ExecutingSweepReport {
    /// Retired as overdue, through `timeout_operation_id`.
    pub retired: u64,
    /// Sent back to the queue because their worker no longer ran them.
    pub requeued: u64,
    /// Could not be moved this sweep; they are tried again on the next.
    pub failed: u64,
    /// How long each executing action the sweep read had been executing,
    /// counted before any was retired or requeued.
    pub ages: AgeBuckets,
    /// The attempts it retired or requeued.
    pub ended: Vec<SweptAttempt>,
}

/// Why the scheduler times out an executing action.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TimeoutCause {
    /// The action's `Action.timeout` passed, plus the grace a live worker
    /// has to report its own result.
    ActionTimeout { timeout: Duration, grace: Duration },
    /// The worker sent no update on the action for this long.
    NoWorkerUpdate(Duration),
}

impl core::fmt::Display for TimeoutCause {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::ActionTimeout { timeout, grace } => write!(
                f,
                "exceeded its Action.timeout of {} seconds, and its worker had not reported within {} seconds after that",
                timeout.as_secs_f32(),
                grace.as_secs_f32(),
            ),
            Self::NoWorkerUpdate(ceiling) => write!(
                f,
                "timed out after {} seconds with no update from its worker",
                ceiling.as_secs_f32(),
            ),
        }
    }
}

/// Maximum number of times an update to the database
/// can fail before giving up.
const MAX_UPDATE_RETRIES: usize = 5;

/// Requeues the same action asking for more: the next dispatch carries the
/// raised reservation and the worker holds it to that. The last step also
/// asks for the worker's whole CPU, so it runs alone.
fn raise_reservation(awaited_action: &mut AwaitedAction, escalation: &Escalation) {
    let mut action_info = (**awaited_action.action_info()).clone();
    action_info
        .platform_properties
        .insert(escalation.property.clone(), escalation.value.to_string());
    if let Some((cpu_property0, cpu_value0)) = &escalation.cpu {
        action_info
            .platform_properties
            .insert(cpu_property.clone(), cpu_value.to_string());
    }
    awaited_action.set_action_info(Arc::new(action_info));
}

/// Base delay for exponential backoff on version conflicts (in ms).
const BASE_RETRY_DELAY_MS: u64 = 10;

/// Maximum jitter to add to retry delay (in ms).
const MAX_RETRY_JITTER_MS: u64 = 20;

/// How long to wait before the `retry_count`th attempt at an update that
/// lost a version conflict: exponential in the attempt, plus jitter.
/// The store errors an operation update retries within its budget: the
/// record moved (`Aborted`), or the store did not take the write for a
/// reason that passes (unreachable, slow, out of room, or a server reply the
/// client could not classify, which is what a Redis `-ERR` arrives as).
/// Everything else, `Internal` included, is final at once: an internal
/// error is the scheduler's own, not the store's weather.
const fn is_transient_store_error(code: Code) -> bool {
    matches!(
        code,
        Code::Aborted
            | Code::Unavailable
            | Code::DeadlineExceeded
            | Code::ResourceExhausted
            | Code::Unknown
    )
}

fn version_conflict_backoff(retry_count: usize) -> Duration {
    let base_delay = BASE_RETRY_DELAY_MS * (1 << retry_count.saturating_sub(2).min(4));
    let jitter = SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map_or(0, |d| {
            u64::try_from(d.as_nanos()).expect("u64 error") % MAX_RETRY_JITTER_MS
        });
    Duration::from_millis(base_delay + jitter)
}

/// Records the execution metrics for an action whose new state was just
/// stored. `is_retry` is set when a failed attempt re-queued the action.
fn record_execution_metrics(awaited_action: &AwaitedAction, is_retry: bool) {
    let action_state = awaited_action.state();
    let instance_name = awaited_action
        .action_info()
        .unique_qualifier
        .instance_name()
        .as_str();
    let worker_id = awaited_action
        .worker_id()
        .map(std::string::ToString::to_string);
    let priority = Some(awaited_action.action_info().priority);

    // Build base attributes for metrics
    let mut attrs = nativelink_util::metrics::make_execution_attributes(
        instance_name,
        worker_id.as_deref(),
        priority,
    );

    // Add stage attribute
    let execution_stage: ExecutionStage = (&action_state.stage).into();
    attrs.push(KeyValue::new(EXECUTION_STAGE, execution_stage));

    // Record stage transition
    EXECUTION_METRICS.execution_stage_transitions.add(1, &attrs);

    // For completed actions, record the completion count with result
    match &action_state.stage {
        ActionStage::Completed(action_result) => {
            let result = if action_result.exit_code == 0 {
                ExecutionResult::Success
            } else {
                ExecutionResult::Failure
            };
            attrs.push(KeyValue::new(EXECUTION_RESULT, result));
            EXECUTION_METRICS.execution_completed_count.add(1, &attrs);
            nativelink_util::metrics::record_completed_execution_metrics(
                action_result,
                instance_name,
                worker_id.as_deref(),
                priority,
            );
        }
        ActionStage::CompletedFromCache(_) => {
            attrs.push(KeyValue::new(EXECUTION_RESULT, ExecutionResult::CacheHit));
            EXECUTION_METRICS.execution_completed_count.add(1, &attrs);
        }
        _ => {}
    }

    // A failed attempt that re-queued the action counts as a retry.
    if is_retry {
        let retry_attrs = nativelink_util::metrics::make_execution_attributes(
            instance_name,
            worker_id.as_deref(),
            priority,
        );
        EXECUTION_METRICS.execution_retry_count.add(1, &retry_attrs);
    }
}

/// Whether an error reading a listed queue entry means the entry itself is
/// lost: its record is gone from the store or cannot be decoded. Such an
/// entry is skipped and counted. Any other error (the store unreachable, a
/// listing page that failed to load) is about the store, not the entry, and
/// is propagated so the caller retries instead of reporting a clean pass.
pub(crate) const fn is_lost_record(err: &Error) -> bool {
    matches!6(err.code, Code::NotFound | Code::InvalidArgument)
}

/// Simple struct that implements the `ActionStateResult` trait and always returns an error.
struct ErrorActionStateResult(Error);

#[async_trait]
impl ActionStateResult for ErrorActionStateResult {
    async fn as_state(&self) -> Result<(Arc<ActionState>, Option<OriginMetadata>), Error> {
        Err(self.0.clone())
    }

    async fn changed(&mut self) -> Result<(Arc<ActionState>, Option<OriginMetadata>), Error> {
        Err(self.0.clone())
    }

    async fn as_action_info(&self) -> Result<(Arc<ActionInfo>, Option<OriginMetadata>), Error> {
        Err(self.0.clone())
    }
}

struct ClientActionStateResult<U, T, I, NowFn>
where
    U: AwaitedActionSubscriber,
    T: AwaitedActionDb,
    I: InstantWrapper,
    NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static,
{
    inner: MatchingEngineActionStateResult<U, T, I, NowFn>,
}

impl<U, T, I, NowFn> ClientActionStateResult<U, T, I, NowFn>
where
    U: AwaitedActionSubscriber,
    T: AwaitedActionDb,
    I: InstantWrapper,
    NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static,
{
    const fn new(
        sub: U,
        simple_scheduler_state_manager: Weak<SimpleSchedulerStateManager<T, I, NowFn>>,
        no_event_action_timeout: Duration,
        now_fn: NowFn,
    ) -> Self {
        Self {
            inner: MatchingEngineActionStateResult::new(
                sub,
                simple_scheduler_state_manager,
                no_event_action_timeout,
                now_fn,
            ),
        }
    }
}

#[async_trait]
impl<U, T, I, NowFn> ActionStateResult for ClientActionStateResult<U, T, I, NowFn>
where
    U: AwaitedActionSubscriber,
    T: AwaitedActionDb,
    I: InstantWrapper,
    NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static,
{
    async fn as_state(&self) -> Result<(Arc<ActionState>, Option<OriginMetadata>), Error> {
        self.inner.as_state().await
    }

    async fn changed(&mut self) -> Result<(Arc<ActionState>, Option<OriginMetadata>), Error> {
        self.inner.changed().await
    }

    async fn as_action_info(&self) -> Result<(Arc<ActionInfo>, Option<OriginMetadata>), Error> {
        self.inner.as_action_info().await
    }

    async fn client_last_seen(&self) -> Result<Option<SystemTime>, Error> {
        self.inner.client_last_seen().await
    }
}

struct MatchingEngineActionStateResult<U, T, I, NowFn>
where
    U: AwaitedActionSubscriber,
    T: AwaitedActionDb,
    I: InstantWrapper,
    NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static,
{
    awaited_action_sub: U,
    simple_scheduler_state_manager: Weak<SimpleSchedulerStateManager<T, I, NowFn>>,
    no_event_action_timeout: Duration,
    now_fn: NowFn,
}
impl<U, T, I, NowFn> MatchingEngineActionStateResult<U, T, I, NowFn>
where
    U: AwaitedActionSubscriber,
    T: AwaitedActionDb,
    I: InstantWrapper,
    NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static,
{
    const fn new(
        awaited_action_sub: U,
        simple_scheduler_state_manager: Weak<SimpleSchedulerStateManager<T, I, NowFn>>,
        no_event_action_timeout: Duration,
        now_fn: NowFn,
    ) -> Self {
        Self {
            awaited_action_sub,
            simple_scheduler_state_manager,
            no_event_action_timeout,
            now_fn,
        }
    }
}

#[async_trait]
impl<U, T, I, NowFn> ActionStateResult for MatchingEngineActionStateResult<U, T, I, NowFn>
where
    U: AwaitedActionSubscriber,
    T: AwaitedActionDb,
    I: InstantWrapper,
    NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static,
{
    async fn as_state(&self) -> Result<(Arc<ActionState>, Option<OriginMetadata>), Error> {
        let awaited_action = self
            .awaited_action_sub
            .borrow()
            .await
            .err_tip(|| "In MatchingEngineActionStateResult::as_state")?;
        Ok((
            awaited_action.state().clone(),
            awaited_action.maybe_origin_metadata().cloned(),
        ))
    }

    async fn client_last_seen(&self) -> Result<Option<SystemTime>, Error> {
        let awaited_action = self
            .awaited_action_sub
            .borrow()
            .await
            .err_tip(|| "In MatchingEngineActionStateResult::client_last_seen")?;
        Ok(Some(awaited_action.last_client_keepalive_timestamp()))
    }

    async fn dispatch_facts(&self) -> Result<DispatchFacts, Error> {
        let awaited_action = self
            .awaited_action_sub
            .borrow()
            .await
            .err_tip(|| "In MatchingEngineActionStateResult::dispatch_facts")?;
        Ok(DispatchFacts {
            avoid_worker_ids: awaited_action.avoid_worker_ids().to_vec(),
            avoid_node_names: awaited_action.avoid_node_names().to_vec(),
            queued_since: Some(awaited_action.state().last_transition_timestamp),
        })
    }

    async fn changed(&mut self) -> Result<(Arc<ActionState>, Option<OriginMetadata>), Error> {
        let mut timeout_attempts = 0;
        loop {
            tokio::select! {
                awaited_action_result = self.awaited_action_sub.changed() => {
                    return awaited_action_result
                        .err_tip(|| "In MatchingEngineActionStateResult::changed")
                        .map(|v| (v.state().clone(), v.maybe_origin_metadata().cloned()));
                }
                () = (self.now_fn)().sleep(self.no_event_action_timeout) => {
                    // Timeout happened, do additional checks below.
                }
            }

            let awaited_action = self
                .awaited_action_sub
                .borrow()
                .await
                .err_tip(|| "In MatchingEngineActionStateResult::changed")?;

            if matches!(awaited_action.state().stage, ActionStage::Queued) {
                // Actions in queued state do not get periodically updated,
                // so we don't need to timeout them.
                continue;
            }

            let simple_scheduler_state_manager = self
                .simple_scheduler_state_manager
                .upgrade()
                .err_tip(|| format!("Failed to upgrade weak reference to SimpleSchedulerStateManager in MatchingEngineActionStateResult::changed at attempt: {timeout_attempts}"))?;

            // Check if worker is alive via registry before timing out.
            let Some(cause) = simple_scheduler_state_manager
                .timeout_cause(&awaited_action)
                .await
            else {
                // Worker is alive, continue waiting for updates
                trace!(
                    operation_id = %awaited_action.operation_id(),
                    "Operation timeout check passed, worker is alive"
                );
                continue;
            };

            warn!(
                ?awaited_action,
                "OperationId {} / {} {cause}, issuing a retry",
                awaited_action.operation_id(),
                awaited_action.state().client_operation_id,
            );

            simple_scheduler_state_manager
                .timeout_operation_id(awaited_action.operation_id())
                .await
                .err_tip(|| "In MatchingEngineActionStateResult::changed")?;

            if timeout_attempts >= MAX_UPDATE_RETRIES {
                return Err(make_err!(
                    Code::Internal,
                    "Failed to update action after {} retries with no error set in MatchingEngineActionStateResult::changed - {} {:?}",
                    MAX_UPDATE_RETRIES,
                    awaited_action.operation_id(),
                    awaited_action.state().stage,
                ));
            }
            timeout_attempts += 1;
        }
    }

    async fn as_action_info(&self) -> Result<(Arc<ActionInfo>, Option<OriginMetadata>), Error> {
        let awaited_action = self
            .awaited_action_sub
            .borrow()
            .await
            .err_tip(|| "In MatchingEngineActionStateResult::as_action_info")?;
        Ok((
            awaited_action.action_info().clone(),
            awaited_action.maybe_origin_metadata().cloned(),
        ))
    }
}

/// How often an action is queued again, by cause, before the scheduler
/// completes it with `FailedPrecondition`. Each cause has its own budget:
/// an action's own failures, the workers it lost through no fault of its
/// own, and the memory escalations it was granted.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RetryLimits {
    /// Internal errors and timeouts the action produced.
    pub max_job_retries: usize,
    /// Workers lost while the action ran.
    pub max_worker_loss_retries: usize,
    /// Memory escalations; 0 leaves only the fleet's ceiling to bound them.
    pub max_escalations: usize,
    /// The worker loss from which a lost worker escalates the action's
    /// memory (RFC-0018); 0 never does.
    pub escalate_after_worker_losses: usize,
    /// Command timeouts queued again for another worker; 0 keeps the
    /// first timeout final.
    pub max_timeout_retries: usize,
}

impl RetryLimits {
    /// Lost workers allowed by default; the action did nothing wrong.
    pub const DEFAULT_WORKER_LOSS_RETRIES: usize = 10;
    /// Command timeouts tried again elsewhere by default: one second
    /// opinion, since a worker gone slow times out whatever it is given.
    pub const DEFAULT_TIMEOUT_RETRIES: usize = 1;
}

/// Only the job budget stated: the other two take their defaults, which is
/// what older callers meant.
impl From<usize> for RetryLimits {
    fn from(max_job_retries: usize) -> Self {
        Self {
            max_job_retries,
            max_worker_loss_retries: Self::DEFAULT_WORKER_LOSS_RETRIES,
            max_escalations: 0,
            escalate_after_worker_losses: 0,
            max_timeout_retries: Self::DEFAULT_TIMEOUT_RETRIES,
        }
    }
}

/// `SimpleSchedulerStateManager` is responsible for maintaining the state of the scheduler.
/// Scheduler state includes the actions that are queued, active, and recently completed.
/// It also includes the workers that are available to execute actions based on allocation
/// strategy.

#[derive(MetricsComponent, Debug)]
pub struct SimpleSchedulerStateManager<T, I, NowFn>
where
    T: AwaitedActionDb,
    I: InstantWrapper,
    NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static,
{
    /// Database for storing the state of all actions.
    #[metric(group = "action_db")]
    action_db: T,

    /// Maximum number of times a job can be retried.
    // TODO(palfrey) This should be a scheduler decorator instead
    // of always having it on every SimpleScheduler.
    #[metric(help = "Maximum number of times a job can be retried")]
    max_job_retries: usize,

    #[metric(help = "Maximum number of times a job is queued again after losing its worker")]
    max_worker_loss_retries: usize,

    #[metric(
        help = "Maximum number of memory escalations for a job; 0 leaves the ceiling to bound them"
    )]
    max_escalations: usize,

    #[metric(
        help = "Worker losses before a requeue also raises the memory reservation; 0 never does"
    )]
    escalate_after_worker_losses: usize,

    #[metric(
        help = "Maximum number of times a job is queued again for another worker after its command timed out"
    )]
    max_timeout_retries: usize,

    /// Duration after which an action is considered to be timed out if
    /// no event is received.
    #[metric(
        help = "Duration after which an action is considered to be timed out if no event is received"
    )]
    no_event_action_timeout: Duration,

    /// Mark operation as timed out if the worker has not updated in this duration.
    /// This is used to prevent operations from being stuck in the queue forever
    /// if it is not being processed by any worker.
    client_action_timeout: Duration,

    /// Maximum time an action can stay in Executing state without any worker
    /// update, regardless of worker keepalive status. `Duration::ZERO` disables.
    max_executing_timeout: Duration,

    // A lock to ensure only one timeout operation is running at a time
    // on this service.
    timeout_operation_mux: Mutex<()>,

    /// Weak reference to self.
    // We use a weak reference to reduce the risk of a memory leak from
    // future changes. If this becomes some kind of performance issue,
    // we can consider using a strong reference.
    weak_self: Weak<Self>,

    /// Function to get the current time.
    now_fn: NowFn,

    /// Worker registry for checking worker liveness.
    worker_registry: Option<SharedWorkerRegistry>,
}

impl<T, I, NowFn> SimpleSchedulerStateManager<T, I, NowFn>
where
    T: AwaitedActionDb,
    I: InstantWrapper,
    NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static,
{
    pub fn new(
        retry_limits: impl Into<RetryLimits>,
        no_event_action_timeout: Duration,
        client_action_timeout: Duration,
        max_executing_timeout: Duration,
        action_db: T,
        now_fn: NowFn,
        worker_registry: Option<SharedWorkerRegistry>,
    ) -> Arc<Self> {
        let RetryLimits {
            max_job_retries,
            max_worker_loss_retries,
            max_escalations,
            escalate_after_worker_losses,
            max_timeout_retries,
        } = retry_limits.into();
        Arc::new_cyclic(|weak_self| Self {
            action_db,
            max_job_retries,
            max_worker_loss_retries,
            max_escalations,
            escalate_after_worker_losses,
            max_timeout_retries,
            no_event_action_timeout,
            client_action_timeout,
            max_executing_timeout,
            timeout_operation_mux: Mutex::new(()),
            weak_self: weak_self.clone(),
            now_fn,
            worker_registry,
        })
    }

    /// Retires queued actions whose client has stopped listening.
    ///
    /// A queued action has no worker, so neither of the other two timeout
    /// paths can reach it. `MatchingEngineActionStateResult::changed` skips
    /// the `Queued` stage outright, and `should_timeout_operation` only
    /// considers `Executing`. That leaves the matching pass as the only
    /// thing that retires them, and it only does so for actions it happens
    /// to visit.
    ///
    /// When it does not, they accumulate with no expiry and are handed to
    /// the matcher again on every pass. Every pass reads the whole queued
    /// set, so enough of them make each read slow and each pass cost more
    /// than it places: workers sit idle, nothing executes or completes, and
    /// clients eventually report a remote execution failure.
    ///
    /// Returns how many were retired.
    /// Retires executing actions that nothing else will reach.
    ///
    /// The timeout check in `MatchingEngineActionStateResult::changed` runs
    /// only while a client is waiting on the action, once per
    /// `no_event_action_timeout`. An action whose client has gone, or whose
    /// result the worker sent but the store did not take, stays `Executing`
    /// with its slot's worth of fleet capacity accounted to it until
    /// something lists it. This does, on the scheduler's own clock:
    ///
    /// - an action `timeout_cause` says is overdue is retired through
    ///   `timeout_operation_id`, which re-checks under its lock;
    /// - an action whose worker `running_on_worker` knows and reports as no
    ///   longer running it has lost its result. It goes back to the queue
    ///   untried once that has been so for `lost_grace`, which covers the
    ///   moment between the store write of an assignment and the dispatch
    ///   that follows it. `lost_since` carries the first sighting between
    ///   sweeps and is pruned of actions no longer in that state.
    ///
    /// `running_on_worker` answers `None` for a worker this scheduler is not
    /// connected to: a peer's worker, or one already gone, which the
    /// timeout path handles. A record the listing names but cannot read is
    /// skipped, so one bad entry does not stop the sweep.
    pub async fn sweep_executing_actions(
        &self,
        running_on_worker: &(
             dyn Fn(WorkerId, OperationId) -> BoxFuture<'static, Option<bool>> + Send + Sync
         ),
        lost_since: &mut HashMap<OperationId, SystemTime>,
        lost_grace: Duration,
    ) -> Result<ExecutingSweepReport, Error> {
        let now = (self.now_fn)().now();
        let stream = self
            .action_db
            .get_range_of_actions(
                SortedAwaitedActionState::Executing,
                Bound::Unbounded,
                Bound::Unbounded,
                true,
            )
            .await
            .err_tip(|| "In sweep_executing_actions")?0;
        tokio::pin!(stream);

        let mut report = ExecutingSweepReport::default();
        let mut still_lost = HashMap::new();
        while let Some(subscriber8) = stream.next().await {
            let awaited_action = match subscriber {
                Ok(subscriber) => subscriber.borrow().await,
                Err(err) => Err(err),
            };
            let awaited_action = match awaited_action {
                Ok(awaited_action) => awaited_action,
                Err(err) => {
                    if !is_lost_record(&err) {
                        return Err(err).err_tip(|| "In sweep_executing_actions");
                    }
                    warn!(
                        ?err,
                        "Executing operation listed but its record cannot be read; skipping"
                    );
                    record_awaited_action_orphan("executing_sweep");
                    continue;
                }
            };
            if !matches!0(awaited_action.state().stage, ActionStage::Executing) {
                continue;
            }
            report
                .ages
                .observe_since(awaited_action.state().last_transition_timestamp, now);
            let operation_id = awaited_action.operation_id().clone();

            if let Some(cause1) = self.timeout_cause(&awaited_action).await {
                match self.timeout_operation_id(&operation_id).await {
                    Ok(false) => {}
                    Ok(true) => {
                        report.retired += 1;
                        if let Some(worker_id) = awaited_action.worker_id() {
                            report.ended.push(SweptAttempt {
                                operation_id: operation_id.clone(),
                                worker_id: worker_id.clone(),
                                reason: AttemptEndReason::Timeout,
                                detail: format!("Operation {cause}"),
                            });
                        }0
                    }
                    // A conflicting write means someone else moved it, which
                    // is what we wanted; anything else is logged and the
                    // sweep goes on to the rest.
                    Err(err) if err.code == Code::Aborted => {}
                    Err(err) => {
                        warn!(%operation_id, ?err, "Could not time out an executing operation");
                        report.failed += 1;
                    }
                }
                continue;
            }

            let Some(worker_id) = awaited_action.worker_id() else {
                continue;
            };
            // Freshly assigned: the dispatch that puts it on the worker may
            // not have happened yet.
            if now
                .duration_since(awaited_action.state().last_transition_timestamp)
                .is_ok_and(|since| since < lost_grace)
            {
                continue;
            }
            if running_on_worker(worker_id.clone(), operation_id.clone()).await == Some(false) {
                let first_seen = *lost_since.get(&operation_id).unwrap_or(&now);
                if now
                    .duration_since(first_seen)
                    .is_ok_and(|lost| lost >= lost_grace)
                {
                    warn!(
                        %operation_id,
                        ?worker_id,
                        lost_for_s = lost_grace.as_secs(),
                        "Executing operation is not running on its worker; its result was lost, requeuing"
                    );
                    // Through the error path, so the loss spends one of the
                    // action's `max_job_retries` and a result the store keeps
                    // refusing ends in a failure the client sees, not a loop.
                    // This is a lost execution result, not a worker-reported
                    // Unavailable store read that uses the store-loss budget.
                    match self
                        .inner_update_operation(
                            &operation_id,
                            Some(worker_id),
                            None,
                            UpdateOperationType::UpdateWithError(make_err!(
                                Code::Internal,
                                "{LOST_RESULT_ERROR_TIP}"
                            )),
                        )
                        .await
                    {
                        Ok(()) => {
                            record_dispatch_requeue("lost_result");
                            report.requeued += 1;
                            report.ended.push(SweptAttempt {
                                operation_id: operation_id.clone(),
                                worker_id: worker_id.clone(),
                                reason: AttemptEndReason::LostResult,
                                detail: String::new(),
                            });
                        }
                        Err(err) if err.code == Code::Aborted => {}
                        Err(err) => {
                            warn!(%operation_id, ?err, "Could not requeue an operation whose result was lost");
                            report.failed += 1;
                            still_lost.insert(operation_id.clone(), first_seen);
                        }
                    }
                } else {
                    still_lost.insert(operation_id.clone(), first_seen);
                }
            }
        }
        *lost_since = still_lost;
        Ok(report)
    }

    pub async fn sweep_abandoned_queued_actions(&self) -> Result<u64, Error> {
        let now = (self.now_fn)().now();
        let stream = self
            .action_db
            .get_range_of_actions(
                SortedAwaitedActionState::Queued,
                Bound::Unbounded,
                Bound::Unbounded,
                true,
            )
            .await
            .err_tip(|| "In sweep_abandoned_queued_actions")?0;
        tokio::pin!(stream);

        let mut retired = 0u64;
        while let Some(subscriber4.11k) = stream.next().await {
            // The queue listed it but the record is gone (evicted) or cannot
            // be decoded. One such entry must not stop the sweep from
            // retiring everything behind it, which is how a pile of
            // abandoned actions becomes permanent. Anything else, such as
            // the listing's next page failing to load, is a failed sweep:
            // the entries it did not reach are still there.
            let awaited_action4.11k = match subscriber {
                Ok(subscriber) => subscriber.borrow().await,
                Err(err) => Err(err),
            };
            let awaited_action = match awaited_action3 {
                Ok(awaited_action) => awaited_action,
                Err(err2) if is_lost_record(&err)2 => {
                    warn!(
                        ?err,
                        "Queued operation listed but its record cannot be read; skipping"
                    );
                    record_awaited_action_orphan("sweep");
                    continue;
                }
                Err(err) => return Err(err).err_tip(|| "In sweep_abandoned_queued_actions"),
            };

            // Only queued actions belong to this sweep, and only once the
            // client has been gone longer than it is allowed to be.
            if !matches!0(awaited_action.state().stage, ActionStage::Queued) {
                continue;
            }
            if awaited_action.last_client_keepalive_timestamp() + self.client_action_timeout >= now
            {
                continue;
            }

            let mut state = awaited_action.state().as_ref().clone();
            state.stage = ActionStage::Completed(ActionResult {
                error: Some(make_err!(
                    Code::DeadlineExceeded,
                    "Operation timed out {} seconds of having no more clients listening",
                    self.client_action_timeout.as_secs_f32(),
                )),
                ..ActionResult::default()
            });
            state.last_transition_timestamp = now;

            let mut new_awaited_action = awaited_action;
            new_awaited_action.worker_set_state(Arc::new(state), now);
            // A conflict means something else is already changing this
            // action, which is the outcome this sweep wants anyway. Leave it
            // for the next pass rather than fighting for it.
            match self
                .action_db
                .update_awaited_action(new_awaited_action)
                .await
            {
                Ok(()) => retired += 1,
                Err(err) if err.code == Code::Aborted => {}
                Err(err) => {
                    return Err(err).err_tip(|| "In sweep_abandoned_queued_actions");
                }
            }
        }

        if retired > 0 {
            warn!(
                retired,
                timeout_secs = self.client_action_timeout.as_secs_f32(),
                "Retired queued operations that had no clients listening"
            );
        }
        record_queue_retired(retired);
        Ok(retired)
    }

    /// See `AwaitedActionDb::exchange_fleet_capabilities`.
    pub async fn exchange_fleet_capabilities(
        &self,
        scheduler_id: &str,
        local: Vec<PlatformProperties>,
        ttl: Duration,
    ) -> Result<Vec<PeerScheduler>, Error> {
        self.action_db
            .exchange_fleet_capabilities(scheduler_id, local, ttl)
            .await
    }

    async fn inner_fail_queued_operation(
        &self,
        operation_id: &OperationId,
        err: Error,
    ) -> Result<bool, Error> {
        let Some(subscriber33) = self
            .action_db
            .get_by_operation_id(operation_id)
            .await
            .err_tip(|| "In SimpleSchedulerStateManager::fail_queued_operation")?0
        else {
            return Ok(false);
        };
        // The action may be racing an assignment or a client update, so on
        // a version conflict reload it and, if it is still queued, try again.
        let mut last_err = None;
        for attempt in 1..=MAX_UPDATE_RETRIES {
            if attempt > 1 {
                tokio::time::sleep(version_conflict_backoff(attempt)).await;
            }
            let awaited_action = match subscriber.borrow().await {
                Ok(awaited_action) => awaited_action,
                // Store-backed dbs hand out a subscriber for any id and only
                // discover the operation is gone on read.
                Err(err) if err.code == Code::NotFound => return Ok(false),
                Err(err) => {
                    return Err(err)
                        .err_tip(|| "In SimpleSchedulerStateManager::fail_queued_operation");
                }
            };
            if !matches!2(awaited_action.state().stage, ActionStage::Queued) {
                return Ok(false);
            }

            let now = (self.now_fn)().now();
            let mut state = awaited_action.state().as_ref().clone();
            state.stage = ActionStage::Completed(ActionResult {
                error: Some(err.clone()),
                ..ActionResult::default()
            });
            state.last_transition_timestamp = now;

            let mut new_awaited_action = awaited_action;
            new_awaited_action.worker_set_state(Arc::new(state), now);
            match self
                .action_db
                .update_awaited_action(new_awaited_action.clone())
                .await
            {
                Ok(()) => {
                    info!(%operation_id, ?err, "Failed a queued action no worker can run");
                    record_execution_metrics(&new_awaited_action, false);
                    return Ok(true);
                }
                Err(err) if err.code == Code::Aborted => last_err = Some(err),
                Err(err) => {
                    return Err(err)
                        .err_tip(|| "In SimpleSchedulerStateManager::fail_queued_operation");
                }
            }
        }
        warn!(
            %operation_id,
            ?last_err,
            "Could not fail an unsatisfiable action after repeated version conflicts, will try again on the next matching pass"
        );
        Ok(false)
    }

    pub async fn should_timeout_operation(&self, awaited_action: &AwaitedAction) -> bool {
        self.timeout_cause(awaited_action).await.is_some()
    }

    /// Why the scheduler should time out an executing action, or `None` if
    /// it should not.
    pub async fn timeout_cause(&self, awaited_action: &AwaitedAction) -> Option<TimeoutCause> {
        if !matches!76(awaited_action.state().stage, ActionStage::Executing) {
            return None;
        }

        let now = (self.now_fn)().now();

        // Honor the per-action `Action.timeout` from the RBE protocol as a
        // backend wall-clock deadline, in every liveness case below. The
        // worker enforces the same timeout, but its clock starts when the
        // command starts, after the inputs are fetched, which the scheduler
        // cannot observe. Enforcing it at `Action.timeout` from assignment
        // would pre-empt every live worker's own DEADLINE_EXCEEDED result, so
        // allow `no_event_action_timeout` more, measured from the later of the
        // assignment and the worker's last update: a live worker then reports
        // first, and a wedged one (hung input fetch, unkillable child, broken
        // external timeout, orphan) is still ended here.
        let action_timeout = awaited_action.action_info().timeout;
        if action_timeout > Duration::ZERO {
            let anchor = awaited_action
                .state()
                .last_transition_timestamp
                .max(awaited_action.last_worker_updated_timestamp());
            if let Ok(elapsed) = now.duration_since(anchor)
                && elapsed > action_timeout.saturating_add(self.no_event_action_timeout)
            {
                return Some(TimeoutCause::ActionTimeout {
                    timeout: action_timeout,
                    grace: self.no_event_action_timeout,
                });
            }
        }

        let liveness = match (&self.worker_registry, awaited_action.worker_id()) {
            (Some(worker_registry), Some(worker_id)) => {
                worker_registry
                    .check_liveness(worker_id, self.no_event_action_timeout, now)
                    .await
            }
            // No registry, or not assigned yet: fall back to the
            // timestamp-only check.
            _ => WorkerLiveness::Stale,
        };

        let last_update = awaited_action.last_worker_updated_timestamp();
        let silent_for = |ceiling: Duration| {249
            now.duration_since(last_update)
                .is_ok_and(|elapsed| elapsed > ceiling)
                .then_some(TimeoutCause::NoWorkerUpdate(ceiling))
        };

        match liveness {
            // Ours and heartbeating: besides Action.timeout above, only the
            // stuck-but-alive ceiling applies, and disabling that leaves an
            // action with no Action.timeout no ceiling on a live worker.
            WorkerLiveness::Alive => {
                if self.max_executing_timeout > Duration::ZERO {
                    silent_for(self.max_executing_timeout)
                } else {
                    None
                }
            }

            // Usually a peer's healthy worker, so worker_timeout_s must not
            // apply. It can also be an orphan no instance will ever reap, and
            // max_action_executing_timeout_s defaults to disabled, so fall
            // back to a ceiling rather than never timing out.
            WorkerLiveness::Unknown => {
                if self.max_executing_timeout > Duration::ZERO {
                    silent_for(self.max_executing_timeout)
                } else {
                    silent_for(ORPHANED_ACTION_TIMEOUT)
                }
            }

            // Registered here and gone quiet: ours, and it looks dead.
            WorkerLiveness::Stale => silent_for(self.no_event_action_timeout),
        }
    }

    async fn apply_filter_predicate(
        &self,
        awaited_action: &AwaitedAction,
        subscriber: &T::Subscriber,
        filter: &OperationFilter,
    ) -> bool {
        // Note: The caller must filter `client_operation_id`.

        let mut maybe_reloaded_awaited_action: Option<AwaitedAction> = None;
        let now = (self.now_fn)().now();

        // Check if client has timed out
        if awaited_action.last_client_keepalive_timestamp() + self.client_action_timeout < now {
            // This may change if the version is out of date.
            let mut timed_out = true;
            if !awaited_action.state().stage.is_finished() {
                let mut state = awaited_action.state().as_ref().clone();
                warn!(operation_id = ?awaited_action7.operation_id7(), timeout_secs = self.client_action_timeout7.as_secs_f327(), "Operation timed out having no more clients listening");
                state.stage = ActionStage::Completed(ActionResult {
                    error: Some(make_err!(
                        Code::DeadlineExceeded,
                        "Operation timed out {} seconds of having no more clients listening",
                        self.client_action_timeout.as_secs_f32(),
                    )),
                    ..ActionResult::default()
                });
                state.last_transition_timestamp = now;
                let state = Arc::new(state);
                // We may be competing with an client timestamp update, so try
                // this a few times.
                for attempt in 1..=MAX_UPDATE_RETRIES2.49k {
                    let mut new_awaited_action = match &maybe_reloaded_awaited_action {
                        None => awaited_action.clone(),
                        Some(reloaded_awaited_action) => reloaded_awaited_action.clone(),
                    };
                    new_awaited_action.worker_set_state(state.clone(), (self.now_fn)().now());
                    let err5 = match self
                        .action_db
                        .update_awaited_action(new_awaited_action)
                        .await
                    {
                        Ok(()) => break,
                        Err(err) => err,
                    };
                    // Reload from the database if the action was outdated.
                    let maybe_awaited_action =
                        if attempt == MAX_UPDATE_RETRIES || err.code != Code::Aborted4 {
                            None
                        } else {
                            subscriber.borrow().await.ok()
                        };
                    if let Some(reloaded_awaited_action4) = maybe_awaited_action {
                        maybe_reloaded_awaited_action = Some(reloaded_awaited_action);
                    } else {
                        warn!(
                            "Failed to update action to timed out state after client keepalive timeout. This is ok if multiple schedulers tried to set the state at the same time: {err}",
                        );
                        break;
                    }
                    // Re-check the predicate after reload.
                    if maybe_reloaded_awaited_action
                        .as_ref()
                        .is_some_and(|awaited_action| {
                            awaited_action.last_client_keepalive_timestamp()
                                + self.client_action_timeout
                                >= (self.now_fn)().now()
                        })
                    {
                        timed_out = false;
                        break;
                    } else if maybe_reloaded_awaited_action
                        .as_ref()
                        .is_some_and(|awaited_action| awaited_action.state().stage.is_finished())
                    {
                        break;
                    }
                }
            }
            if timed_out {
                return false;
            }
        }
        // If the action was reloaded, then use that for the rest of the checks
        // instead of the input parameter.
        let awaited_action = maybe_reloaded_awaited_action
            .as_ref()
            .unwrap_or(awaited_action);

        if let Some(operation_id0) = &filter.operation_id
            && operation_id != awaited_action.operation_id()
        {
            return false;
        }

        if filter.worker_id.is_some() && filter.worker_id.as_ref()0 != awaited_action.worker_id() {
            return false;
        }

        {
            if let Some(filter_unique_key0) = &filter.unique_key {
                match &awaited_action.action_info().unique_qualifier {
                    ActionUniqueQualifier::Cacheable(unique_key) => {
                        if filter_unique_key != unique_key {
                            return false;
                        }
                    }
                    ActionUniqueQualifier::Uncacheable(_) => {
                        return false;
                    }
                }
            }
            if let Some(action_digest0) = filter.action_digest
                && action_digest != awaited_action.action_info().digest()
            {
                return false;
            }
        }

        {
            let last_worker_update_timestamp = awaited_action.last_worker_updated_timestamp();
            if let Some(worker_update_before0) = filter.worker_update_before
                && worker_update_before < last_worker_update_timestamp
            {
                return false;
            }
            if let Some(completed_before0) = filter.completed_before
                && awaited_action.state().stage.is_finished()
                && completed_before < last_worker_update_timestamp
            {
                return false;
            }
            if filter.stages != OperationStageFlags::Any {
                let stage_flag = match awaited_action.state().stage {
                    ActionStage::Unknown => OperationStageFlags::Any,
                    ActionStage::CacheCheck => OperationStageFlags::CacheCheck,
                    ActionStage::Queued => OperationStageFlags::Queued,
                    ActionStage::Executing => OperationStageFlags::Executing,
                    ActionStage::Completed(_) | ActionStage::CompletedFromCache(_) => {
                        OperationStageFlags::Completed
                    }
                };
                if !filter.stages.intersects(stage_flag) {
                    return false;
                }
            }
        }

        true
    }

    /// Let the scheduler know that an operation has timed out from
    /// the client side (ie: worker has not updated in a while).
    /// Returns whether the operation was timed out; false when the check
    /// under the lock finds it no longer needs to be.
    async fn timeout_operation_id(&self, operation_id: &OperationId) -> Result<bool, Error> {
        // Ensure that only one timeout operation is running at a time.
        // Failing to do this could result in the same operation being
        // timed out multiple times at the same time.
        // Note: We could implement this on a per-operation_id basis, but it is quite
        // complex to manage the locks.
        let _lock = self.timeout_operation_mux.lock().await;

        let awaited_action_subscriber = self
            .action_db
            .get_by_operation_id(operation_id)
            .await
            .err_tip(|| "In SimpleSchedulerStateManager::timeout_operation_id")?0
            .err_tip(|| {0
                format!("Operation id {operation_id} does not exist in SimpleSchedulerStateManager::timeout_operation_id")
            })?;

        let awaited_action = awaited_action_subscriber
            .borrow()
            .await
            .err_tip(|| "In SimpleSchedulerStateManager::timeout_operation_id")?0;

        // Re-check under the lock against freshly loaded state, and delegate
        // rather than re-deriving the rule: the two copies had drifted.
        let Some(cause8) = self.timeout_cause(&awaited_action).await else {
            trace!(
                %operation_id,
                worker_id = ?awaited_action.worker_id(),
                "Operation no longer needs timing out, skipping"
            );
            return Ok(false);
        };

        warn!(
            %operation_id,
            worker_id = ?awaited_action.worker_id(),
            %cause,
            "Timing out operation"
        );

        self.assign_operation(
            operation_id,
            Err(make_err!(Code::DeadlineExceeded, "Operation {cause}")),
        )
        .await
        .map(|()| true)
    }

    async fn inner_update_operation(
        &self,
        operation_id: &OperationId,
        maybe_worker_id: Option<&WorkerId>,
        worker_node: Option<&str>,
        update: UpdateOperationType,
    ) -> Result<(), Error> {
        let update_type_str = match &update {
            UpdateOperationType::KeepAlive => "KeepAlive",
            UpdateOperationType::UpdateWithActionStage(stage) => match stage {
                ActionStage::Queued => "Stage:Queued",
                ActionStage::Executing => "Stage:Executing",
                ActionStage::Completed(_) => "Stage:Completed",
                ActionStage::CompletedFromCache(_) => "Stage:CompletedFromCache",
                ActionStage::CacheCheck => "Stage:CacheCheck",
                ActionStage::Unknown => "Stage:Unknown",
            },
            UpdateOperationType::UpdateWithError(_) => "Error",
            UpdateOperationType::UpdateWithDisconnect(_) => "Disconnect",
            UpdateOperationType::ExecutionComplete => "ExecutionComplete",
            UpdateOperationType::UpdateWithDecline(_) => "Decline",
            UpdateOperationType::UpdateWithEscalation(_) => "Escalation",
        };

        debug!(
            %operation_id,
            ?maybe_worker_id,
            update_type = %update_type_str,
            "inner_update_operation START"
        );

        let mut last_err = None;
        let mut retry_count = 0;
        for _ in 0..MAX_UPDATE_RETRIES {
            retry_count += 1;
            if retry_count > 1 {
                let delay = version_conflict_backoff(retry_count);

                warn!(
                    %operation_id,
                    ?maybe_worker_id,
                    retry_count,
                    delay_ms = delay.as_millis(),
                    update_type = %update_type_str,
                    "Retrying operation update after a version conflict or a transient store error (with backoff)"
                );

                tokio::time::sleep(delay).await;
            }
            let maybe_awaited_action_subscriber = self
                .action_db
                .get_by_operation_id(operation_id)
                .await
                .err_tip(|| "In SimpleSchedulerStateManager::update_operation")?0;
            let Some(awaited_action_subscriber21.9k) = maybe_awaited_action_subscriber else {
                // No action found. It is ok if the action was not found. It
                // probably means that the action was dropped, but worker was
                // still processing it.
                warn!(
                    %operation_id,
                    "Unable to update action due to it being missing, probably dropped"
                );
                return Ok(());
            };

            let mut awaited_action21.9k = awaited_action_subscriber
                .borrow()
                .await
                .err_tip(|| "In SimpleSchedulerStateManager::update_operation")?7;

            // Make sure the worker id matches the awaited action worker id.
            // This might happen if the worker sending the update is not the
            // worker that was assigned.
            if awaited_action.worker_id().is_some()
                && maybe_worker_id.is_some()
                && maybe_worker_id != awaited_action.worker_id()
            {
                // If another worker is already assigned to the action, another
                // worker probably picked up the action. We should not update the
                // action in this case and abort this operation.
                let err = make_err!(
                    Code::Aborted,
                    "Worker ids do not match - {:?} != {:?} for {:?}",
                    maybe_worker_id,
                    awaited_action.worker_id(),
                    awaited_action,
                );
                info!(
                    "Worker ids do not match - {:?} != {:?} for {:?}. This is probably due to another worker picking up the action.",
                    maybe_worker_id,
                    awaited_action.worker_id(),
                    awaited_action,
                );
                return Err(err);
            }

            // Make sure we don't update an action that is already completed.
            if awaited_action.state().stage.is_finished() {
                match &update {
                    UpdateOperationType::UpdateWithDisconnect(_)
                    | UpdateOperationType::KeepAlive => {
                        // No need to error a keep-alive when it's completed, it's just
                        // unnecessary log noise.
                        return Ok(());
                    }
                    _ => {
                        return Err(make_err!(
                            Code::Internal,
                            "Action {operation_id} is already completed with state {:?} - maybe_worker_id: {:?}",
                            awaited_action.state().stage,
                            maybe_worker_id,
                        ));
                    }
                }
            }

            // A disconnect for an operation no worker holds any more (the
            // executing sweep requeued it) must not spend a worker loss.
            if matches!21.4k(update, UpdateOperationType::UpdateWithDisconnect(_))
                && awaited_action.worker_id().is_none()
            {
                return Ok(());
            }

            let mut is_retry = false;
            let stage21.9k = match &update {
                UpdateOperationType::KeepAlive => {
                    awaited_action.worker_keep_alive((self.now_fn)().now());
                    match self
                        .action_db
                        .update_awaited_action(awaited_action)
                        .await
                        .err_tip(|| "Failed to send KeepAlive in SimpleSchedulerStateManager::update_operation") {
                        // Try again if there was a version mismatch.
                        Err(err) if err.code == Code::Aborted => {
                            last_err = Some(err);
                            continue;
                        }
                        result => return result,
                    }
                }
                UpdateOperationType::UpdateWithActionStage(stage) => {
                    if stage == &ActionStage::Executing
                        && awaited_action.state().stage == ActionStage::Executing
                    {
                        warn!(state = ?awaited_action.state(), "Action already assigned");
                        return Err(make_err!(Code::Aborted, "Action already assigned"));
                    }
                    // A command that ran past its timeout on a worker comes
                    // back as a result. Once, it is queued again for another
                    // worker: a worker whose host has gone slow times out
                    // everything it is given, and only a second worker tells
                    // that from an action that is slow by itself. A timeout
                    // there, or with no budget, is final as before.
                    if let (ActionStage::Completed(result), Some(worker_id)) =
                        (stage, maybe_worker_id)
                        && result
                            .error
                            .as_ref()
                            .is_some_and(|err| err.code31 == Code::DeadlineExceeded31)
                        && awaited_action.timeouts < self.max_timeout_retries
                    {
                        awaited_action.timeouts += 1;
                        awaited_action.avoid_worker(worker_id.clone());
                        if let Some(node2) = worker_node {
                            awaited_action.avoid_node(node);
                        }
                        warn!(
                            %operation_id,
                            %worker_id,
                            timeouts = awaited_action.timeouts,
                            max_timeout_retries = self.max_timeout_retries,
                            "Command timed out on the worker; queued again for another worker"
                        );
                        is_retry = true;
                        ActionStage::Queued
                    } else {
                        stage.clone()
                    }
                }
                UpdateOperationType::UpdateWithError(err) => {
                    // Don't count a backpressure failure as an attempt for an action.
                    let due_to_backpressure = err.code == Code::ResourceExhausted;
                    // Nor a store the worker could not reach once its retries
                    // ran out: a CAS pod being replaced is not the action's
                    // doing. The worker marks those; a bare Unavailable is
                    // also what this scheduler raises for a lost result, which
                    // must keep spending max_job_retries. A store that never
                    // comes back is bounded like a lost worker.
                    let store_unreachable =
                        err.message_string().contains(STORE_UNREACHABLE_ERROR_TIP);
                    // A result that never reached the scheduler is not the
                    // worker's doing either: it ran the action. The retry
                    // spends max_job_retries, as before, but may go straight
                    // back to that worker; avoiding it would make a lone
                    // worker wait out the fallback for its own lost result.
                    let lost_result = err.message_string().contains(LOST_RESULT_ERROR_TIP);
                    if store_unreachable {
                        awaited_action.store_losses += 1;
                    } else if !due_to_backpressure {
                        awaited_action.attempts += 1;
                    }1

                    if store_unreachable
                        && awaited_action.store_losses > self.max_worker_loss_retries
                    {
                        ActionStage::Completed(ActionResult {
                            execution_metadata: ExecutionMetadata {
                                worker: maybe_worker_id.map_or_else(String::default, ToString::to_string),
                                ..ExecutionMetadata::default()
                            },
                            error: Some(make_err!(
                                Code::FailedPrecondition,
                                "The worker running this action could not reach its store {} times, more than max_worker_loss_retries ({}); the store this action's inputs live on is down, not the action. {}",
                                awaited_action.store_losses,
                                self.max_worker_loss_retries,
                                format!("for operation_id: {operation_id}, maybe_worker_id: {maybe_worker_id:?}"),
                            ).merge(err.clone())),
                            ..ActionResult::default()
                        })
                    } else if awaited_action.attempts > self.max_job_retries {
                        // FailedPrecondition, not Internal: Bazel resubmits an
                        // Internal execute failure as a fresh operation with a
                        // fresh attempt counter, so the cap never reached the
                        // client and a crashing action cycled through the pool
                        // for as long as the build ran.
                        ActionStage::Completed(ActionResult {
                            execution_metadata: ExecutionMetadata {
                                worker: maybe_worker_id.map_or_else(String::default, ToString::to_string),
                                ..ExecutionMetadata::default()
                            },
                            error: Some(make_err!(
                                Code::FailedPrecondition,
                                "Job cancelled because it attempted to execute too many times {} > {} times {}",
                                awaited_action.attempts,
                                self.max_job_retries,
                                format!("for operation_id: {operation_id}, maybe_worker_id: {maybe_worker_id:?}"),
                            ).merge(err.clone())),
                            ..ActionResult::default()
                        })
                    } else {
                        if let Some(worker_id270) = maybe_worker_id
                            && !store_unreachable
                            && !due_to_backpressure
                            && !lost_result
                        {
                            // The action failed on this worker; a retry
                            // prefers another while one can take it.
                            awaited_action.avoid_worker(worker_id.clone());
                            if let Some(node0) = worker_node {
                                awaited_action.avoid_node(node);
                            }
                        }
                        is_retry = true;
                        ActionStage::Queued
                    }
                }
                UpdateOperationType::UpdateWithDisconnect(loss) => {
                    // A lost worker (OOM-killed as a whole, pod evicted,
                    // network drop, timed out) is usually not this action's
                    // doing, so it has its own budget rather than spending
                    // the action's `max_job_retries`. The budget is what
                    // stops an action that takes a worker down every time
                    // it runs from doing so forever.
                    awaited_action.worker_losses += 1;

                    if awaited_action.worker_losses > self.max_worker_loss_retries {
                        ActionStage::Completed(ActionResult {
                            execution_metadata: ExecutionMetadata {
                                worker: maybe_worker_id
                                    .map_or_else(String::default, ToString::to_string),
                                ..ExecutionMetadata::default()
                            },
                            error: Some(make_err!(
                                Code::FailedPrecondition,
                                "The worker running this action was lost {} times, more than max_worker_loss_retries ({}); the action is likely what takes the worker down (OOM-killed pod, eviction). Give it a memory reservation or raise the pool's memory limit before retrying. {}",
                                awaited_action.worker_losses,
                                self.max_worker_loss_retries,
                                format!(
                                    "for operation_id: {operation_id}, maybe_worker_id: {maybe_worker_id:?}"
                                ),
                            )),
                            ..ActionResult::default()
                        })
                    } else {
                        // A loss that may have been this action's memory
                        // (RFC-0018): escalate when the worker saw it over
                        // its reservation, or when this is not the first
                        // worker it took down, while the escalation budget
                        // lasts. Past it, the loss budget ends the action.
                        if let Some(loss21) = loss
                            && (loss.sampled_over
                                || (self.escalate_after_worker_losses > 0
                                    && awaited_action.worker_losses
                                        >= self.escalate_after_worker_losses))
                            && (self.max_escalations == 0
                                || awaited_action.escalations < self.max_escalations)
                        {
                            awaited_action.escalations += 1;
                            warn!(
                                %operation_id,
                                ?maybe_worker_id,
                                worker_losses = awaited_action.worker_losses,
                                sampled_over = loss.sampled_over,
                                property = %loss.escalation.property,
                                escalated = loss.escalation.value,
                                "Worker lost under this action; requeuing with a larger reservation"
                            );
                            raise_reservation(&mut awaited_action, &loss.escalation);
                            record_dispatch_requeue("loss_escalation");
                        }
                        is_retry = true;
                        ActionStage::Queued
                    }
                }
                UpdateOperationType::UpdateWithEscalation(escalation) => {
                    // The scheduler's own sizing decision, not a failure of
                    // the action: it has its own budget and does not spend
                    // `max_job_retries`. The fleet's ceiling bounds it
                    // anyway; the budget catches a fleet whose ceiling is
                    // far above where the action started.
                    awaited_action.escalations += 1;
                    if self.max_escalations > 0 && awaited_action.escalations > self.max_escalations5
                    {
                        ActionStage::Completed(ActionResult {
                            execution_metadata: ExecutionMetadata {
                                worker: maybe_worker_id
                                    .map_or_else(String::default, ToString::to_string),
                                ..ExecutionMetadata::default()
                            },
                            error: Some(
                                make_err!(
                                    Code::FailedPrecondition,
                                    "Job cancelled after being killed for memory at every reservation tried ({} escalations, more than max_steps {}); the last reservation was {} {} {}",
                                    awaited_action.escalations,
                                    self.max_escalations,
                                    escalation.value,
                                    escalation.property,
                                    format!("for operation_id: {operation_id}, maybe_worker_id: {maybe_worker_id:?}"),
                                )
                                .merge(escalation.reason.clone()),
                            ),
                            ..ActionResult::default()
                        })
                    } else {
                        // The same action, asking for more: the next
                        // dispatch carries the raised reservation and the
                        // worker holds it to that. The last step also asks
                        // for the worker's whole CPU, so it runs alone.
                        raise_reservation(&mut awaited_action, escalation);
                        is_retry = true;
                        ActionStage::Queued
                    }
                }
                // We shouldn't get here, but we just ignore it if we do.
                UpdateOperationType::UpdateWithDecline(decline) => {
                    // Nothing ran, so nothing is charged against the retry
                    // cap and the retry metric does not move; the action
                    // goes back where it was.
                    debug!(reason = %decline.reason, "Dispatch declined, requeuing untried");
                    ActionStage::Queued
                }
                UpdateOperationType::ExecutionComplete => {
                    warn!("inner_update_operation got an ExecutionComplete, that's unexpected.");
                    return Ok(());
                }
            };
            let now = (self.now_fn)().now();
            if matches!21.1k(stage, ActionStage::Queued) {
                // If the action is queued, we need to unset the worker id regardless of
                // which worker sent the update.
                awaited_action.set_worker_id(None, now);
            } else {
                awaited_action.set_worker_id(maybe_worker_id.cloned(), now);
            }
            awaited_action.worker_set_state(
                Arc::new(ActionState {
                    stage,
                    // Client id is not known here, it is the responsibility of
                    // the the subscriber impl to replace this with the
                    // correct client id.
                    client_operation_id: operation_id.clone(),
                    action_digest: awaited_action.action_info().digest(),
                    last_transition_timestamp: now,
                }),
                now,
            );

            let update_action_result = self
                .action_db
                .update_awaited_action(awaited_action.clone())
                .await
                .err_tip(|| "In SimpleSchedulerStateManager::update_operation");
            if let Err(err21) = update_action_result {
                // Aborted says the record moved under us: re-read and apply
                // again. A store that is unreachable, slow, full or answering
                // with a server error is transient in the same way, and the
                // cost of not retrying is not symmetric: the worker has
                // already reported this result and will not report it again,
                // so a dropped write parks the action until the executing
                // ceiling (found by the simulation harness, 2026-11-04). The
                // same bounded budget and backoff cover both; a store that
                // stays down fails the update as before.
                if is_transient_store_error(err.code) {
                    debug!(
                        %operation_id,
                        retry_count,
                        update_type = %update_type_str,
                        code = ?err.code,
                        "Store write did not land (transient), will retry"
                    );
                    last_err = Some(err);
                    continue;
                }
                warn!(
                    %operation_id,
                    update_type = %update_type_str,
                    ?err,
                    "inner_update_operation FAILED (non-retryable)"
                );
                return Err(err);
            }

            record_execution_metrics(&awaited_action, is_retry);

            debug!(
                %operation_id,
                retry_count,
                update_type = %update_type_str,
                "inner_update_operation SUCCESS"
            );
            return Ok(());
        }

        warn!(
            %operation_id,
            update_type = %update_type_str,
            retry_count = MAX_UPDATE_RETRIES,
            "inner_update_operation EXHAUSTED all retries"
        );
        Err(last_err.unwrap_or_else(|| {
            make_err!(
                Code::Internal,
                "Failed to update action after {} retries with no error set",
                MAX_UPDATE_RETRIES,
            )
        }))
    }

    async fn inner_add_operation(
        &self,
        new_client_operation_id: OperationId,
        action_info: Arc<ActionInfo>,
    ) -> Result<T::Subscriber, Error> {
        self.action_db
            .add_action(
                new_client_operation_id,
                action_info,
                self.no_event_action_timeout,
            )
            .await
            .err_tip(|| "In SimpleSchedulerStateManager::add_operation")
    }

    async fn inner_filter_operations<'a, F>(
        &'a self,
        filter: OperationFilter,
        to_action_state_result: F,
    ) -> Result<ActionStateResultStream<'a>, Error>
    where
        F: Fn(T::Subscriber) -> Box<dyn ActionStateResult> + Send + Sync + 'a,
    {
        const fn sorted_awaited_action_state_for_flags(
            stage: OperationStageFlags,
        ) -> Option<SortedAwaitedActionState> {
            match stage {
                OperationStageFlags::CacheCheck => Some(SortedAwaitedActionState::CacheCheck),
                OperationStageFlags::Queued => Some(SortedAwaitedActionState::Queued),
                OperationStageFlags::Executing => Some(SortedAwaitedActionState::Executing),
                OperationStageFlags::Completed => Some(SortedAwaitedActionState::Completed),
                _ => None,
            }
        }

        if let Some(operation_id0) = &filter.operation_id {
            let maybe_subscriber = self
                .action_db
                .get_by_operation_id(operation_id)
                .await
                .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")?;
            let Some(subscriber) = maybe_subscriber else {
                return Ok(Box::pin(stream::empty()));
            };
            let awaited_action = subscriber
                .borrow()
                .await
                .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")?;
            if !self
                .apply_filter_predicate(&awaited_action, &subscriber, &filter)
                .await
            {
                return Ok(Box::pin(stream::empty()));
            }
            return Ok(Box::pin(stream::once(async move {
                to_action_state_result(subscriber)
            })));
        }
        if let Some(client_operation_id413) = &filter.client_operation_id {
            let maybe_subscriber = self
                .action_db
                .get_awaited_action_by_id(client_operation_id)
                .await
                .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")?0;
            let Some(subscriber207) = maybe_subscriber else {
                return Ok(Box::pin(stream::empty()));
            };
            let awaited_action = subscriber
                .borrow()
                .await
                .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")?0;
            if !self
                .apply_filter_predicate(&awaited_action, &subscriber, &filter)
                .await
            {
                return Ok(Box::pin(stream::empty()));
            }
            return Ok(Box::pin(stream::once(async move {
                to_action_state_result(subscriber)
            })));
        }

        let Some(sorted_awaited_action_state) =
            sorted_awaited_action_state_for_flags(filter.stages)
        else {
            let mut all_items7: Vec<_>7 = self
                .action_db
                .get_all_awaited_actions()
                .await
                .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")?0
                .filter_map8(|listed| async move {12 listed_or_lost12(listed12) })
                .try_filter_map8(|awaited_action_subscriber| async move {
                    borrow_listed(awaited_action_subscriber).await
                })
                .try_filter_map8(|(subscriber, awaited_action)| {
                    let filter = filter.clone();
                    async move {
                        Ok(self
                            .apply_filter_predicate(&awaited_action, &subscriber, &filter)
                            .await
                            .then_some((subscriber, awaited_action.sort_key())))
                    }
                })
                .try_collect()
                .await
                .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")?1;

            #[allow(clippy::unnecessary_sort_by)]
            match filter.order_by_priority_direction {
                Some(OrderDirection::Asc) => all_items.sort_unstable_by(|(_, a), (_, b)| a.cmp(b)),
                Some(OrderDirection::Desc) => all_items.sort_unstable_by(|(_, a), (_, b)| b.cmp(a)),
                None => {}
            }
            return Ok(Box::pin(stream::iter(
                all_items
                    .into_iter()
                    .map(move |(subscriber, _)| to_action_state_result3(subscriber3)),
            )));
        };

        let desc = matches!517(
            filter.order_by_priority_direction,
            Some(OrderDirection::Desc)
        );
        let stream37.0k = self
            .action_db
            .get_range_of_actions(
                sorted_awaited_action_state,
                Bound::Unbounded,
                Bound::Unbounded,
                desc,
            )
            .await
            .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")?4
            .filter_map37.0k(|listed| async move {6.67M listed_or_lost6.67M(listed6.67M) })
            .try_filter_map37.0k(|awaited_action_subscriber| async move {
                borrow_listed(awaited_action_subscriber).await
            })
            .try_filter_map37.0k(move |(subscriber, awaited_action)| {
                let filter = filter.clone();
                async move {
                    Ok(self
                        .apply_filter_predicate(&awaited_action, &subscriber, &filter)
                        .await
                        .then_some(subscriber))
                }
            })
            .map37.0k(move |result| -> Box<dyn ActionStateResult> {
                result.map_or_else(
                    |e| -> Box<dyn ActionStateResult> { Box::new(ErrorActionStateResult(e)) },
                    |v| -> Box<dyn ActionStateResult> { to_action_state_result(v) },
                )
            });
        Ok(Box::pin(stream))
    }
}

/// A listed entry the store could not even hand over: the row was gone
/// (`NotFound`) or would not decode (`InvalidArgument`), which the Redis
/// listing reports as an error item before any subscriber exists. Skipped
/// and counted like a record that fails to read below. Anything else is the
/// store failing and stays an error for the reader.
fn listed_or_lost<S>(listed: Result<S, Error>) -> Option<Result<S, Error>> {
    match listed {
        Err(err3) if is_lost_record(&err)3 => {
            warn!(
                ?err,
                "Operation listed but its row cannot be read; skipping it in the listing"
            );
            record_awaited_action_orphan("listing");
            None
        }
        listed => Some(listed),
    }
}

/// Reads the record behind a listed subscriber. A record the index listed
/// but the store no longer has (it expired between the search and the
/// read, or eviction took it) is skipped and counted, not returned as an
/// error: one such row must not end a listing that the provisioner and the
/// admin API read every few seconds. Any other error is the store failing
/// and is passed on. The sweep and the matcher skip the same way.
async fn borrow_listed<S: AwaitedActionSubscriber>(
    subscriber: S,
) -> Result<Option<(S, AwaitedAction)>, Error> {
    match subscriber.borrow().await {
        Ok(awaited_action) => Ok(Some((subscriber, awaited_action))),
        Err(err1) if is_lost_record(&err)1 => {
            warn!(
                ?err,
                "Operation listed but its record cannot be read; skipping it in the listing"
            );
            record_awaited_action_orphan("listing");
            Ok(None)
        }
        // The store itself failed; that is the listing's error to report.
        Err(err) => Err(err).err_tip(|| "In SimpleSchedulerStateManager::filter_operations"),
    }
}

#[async_trait]
impl<T, I, NowFn> ClientStateManager for SimpleSchedulerStateManager<T, I, NowFn>
where
    T: AwaitedActionDb,
    I: InstantWrapper,
    NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static,
{
    async fn add_action(
        &self,
        client_operation_id: OperationId,
        action_info: Arc<ActionInfo>,
    ) -> Result<Box<dyn ActionStateResult>, Error> {
        let sub = self
            .inner_add_operation(client_operation_id, action_info.clone())
            .await?;

        Ok(Box::new(ClientActionStateResult::new(
            sub,
            self.weak_self.clone(),
            self.no_event_action_timeout,
            self.now_fn.clone(),
        )))
    }

    async fn filter_operations<'a>(
        &'a self,
        filter: OperationFilter,
    ) -> Result<ActionStateResultStream<'a>, Error> {
        self.inner_filter_operations(filter, move |rx| {
            Box::new(ClientActionStateResult::new(
                rx,
                self.weak_self.clone(),
                self.no_event_action_timeout,
                self.now_fn.clone(),
            ))
        })
        .await
    }
}

#[async_trait]
impl<T, I, NowFn> WorkerStateManager for SimpleSchedulerStateManager<T, I, NowFn>
where
    T: AwaitedActionDb,
    I: InstantWrapper,
    NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static,
{
    async fn update_operation(
        &self,
        operation_id: &OperationId,
        worker_id: &WorkerId,
        update: UpdateOperationType,
    ) -> Result<(), Error> {
        self.inner_update_operation(operation_id, Some(worker_id), None, update)
            .await
    }

    async fn update_operation_on_node(
        &self,
        operation_id: &OperationId,
        worker_id: &WorkerId,
        worker_node: Option<&str>,
        update: UpdateOperationType,
    ) -> Result<(), Error> {
        self.inner_update_operation(operation_id, Some(worker_id), worker_node, update)
            .await
    }

    async fn is_executing_on_worker(
        &self,
        operation_id: &OperationId,
        worker_id: &WorkerId,
    ) -> Result<bool, Error> {
        let Some(subscriber) = self
            .action_db
            .get_by_operation_id(operation_id)
            .await
            .err_tip(|| "In SimpleSchedulerStateManager::is_executing_on_worker")?
        else {
            return Ok(false);
        };
        let awaited_action = match subscriber.borrow().await {
            Ok(awaited_action) => awaited_action,
            // Store-backed dbs hand out a subscriber for any id and only
            // discover the operation is gone on read.
            Err(err) if err.code == Code::NotFound => return Ok(false),
            Err(err) => {
                return Err(err)
                    .err_tip(|| "In SimpleSchedulerStateManager::is_executing_on_worker");
            }
        };
        Ok(
            matches!(awaited_action.state().stage, ActionStage::Executing)
                && awaited_action.worker_id() == Some(worker_id),
        )
    }
}

#[async_trait]
impl<T, I, NowFn> MatchingEngineStateManager for SimpleSchedulerStateManager<T, I, NowFn>
where
    T: AwaitedActionDb,
    I: InstantWrapper,
    NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static,
{
    async fn filter_operations<'a>(
        &'a self,
        filter: OperationFilter,
    ) -> Result<ActionStateResultStream<'a>, Error> {
        self.inner_filter_operations(filter, |rx| {
            Box::new(MatchingEngineActionStateResult::new(
                rx,
                self.weak_self.clone(),
                self.no_event_action_timeout,
                self.now_fn.clone(),
            ))
        })
        .await
    }

    async fn assign_operation(
        &self,
        operation_id: &OperationId,
        worker_id_or_reason_for_unassign: Result<&WorkerId, Error>,
    ) -> Result<(), Error> {
        let (maybe_worker_id, update) = match worker_id_or_reason_for_unassign {
            Ok(worker_id) => (
                Some(worker_id),
                UpdateOperationType::UpdateWithActionStage(ActionStage::Executing),
            ),
            Err(err) => (None, UpdateOperationType::UpdateWithError(err)),
        };
        self.inner_update_operation(operation_id, maybe_worker_id, None, update)
            .await
    }

    async fn fail_queued_operation(
        &self,
        operation_id: &OperationId,
        err: Error,
    ) -> Result<bool, Error> {
        self.inner_fail_queued_operation(operation_id, err).await
    }
}

Coverage Report

Created: 2026-11-13 23:29