linera_core/

worker.rs

// Copyright (c) Facebook, Inc. and its affiliates.
// Copyright (c) Zefchain Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use std::{
    collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque},
    future::Future,
    pin,
    sync::{Arc, Mutex, RwLock},
    time::Duration,
};

use futures::{
    future::{self, Either, Shared, WeakShared},
    FutureExt as _,
};
use linera_base::{
    crypto::{CryptoError, CryptoHash, ValidatorPublicKey},
    data_types::{
        ApplicationDescription, ArithmeticError, Blob, BlockHeight, Epoch, Round, TimeDelta,
        Timestamp,
    },
    doc_scalar,
    identifiers::{AccountOwner, ApplicationId, BlobId, ChainId, EventId, StreamId},
};
use linera_cache::{Arc as CacheArc, UniqueValueCache, ValueCache, DEFAULT_CLEANUP_INTERVAL_SECS};
#[cfg(with_testing)]
use linera_chain::ChainExecutionContext;
use linera_chain::{
    data_types::{BlockProposal, BundleExecutionPolicy, MessageBundle, ProposedBlock},
    types::{
        Block, CertificateValue, ConfirmedBlock, ConfirmedBlockCertificate, GenericCertificate,
        LiteCertificate, Timeout, TimeoutCertificate, ValidatedBlock, ValidatedBlockCertificate,
    },
    ChainError, ChainStateView,
};
use linera_execution::{ExecutionError, ExecutionStateView, Query, QueryOutcome, ResourceTracker};
use linera_storage::{Clock as _, Storage};
use linera_views::{context::InactiveContext, ViewError};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use tokio::sync::{mpsc, oneshot, OwnedRwLockReadGuard};
use tracing::{debug, instrument, trace, warn};

/// A read guard providing access to a chain's [`ChainStateView`].
///
/// Holds a read lock on the chain worker state, preventing writes for its
/// lifetime. The `OwnedRwLockReadGuard` internally holds a strong `Arc`
/// reference to the `RwLock<ChainWorkerState>`, keeping the state alive.
/// Dereferences to `ChainStateView`.
pub struct ChainStateViewReadGuard<S: Storage>(
    OwnedRwLockReadGuard<ChainWorkerState<S>, ChainStateView<S::Context>>,
);

impl<S: Storage> std::ops::Deref for ChainStateViewReadGuard<S> {
    type Target = ChainStateView<S::Context>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

/// Re-export of [`EventSubscriptionsResult`] for use by other crate modules.
pub(crate) use crate::chain_worker::EventSubscriptionsResult;
use crate::{
    chain_worker::{
        handle,
        state::{send_result, ChainWorkerState},
        BlockOutcome, ChainWorkerConfig, CrossChainUpdateResult, DeliveryNotifier,
        ProcessConfirmedBlockMode,
    },
    client::{ChainModes, ListeningMode},
    data_types::{ChainInfoQuery, ChainInfoResponse, CrossChainRequest},
    notifier::Notifier,
};

pub const DEFAULT_BLOCK_CACHE_SIZE: usize = 5_000;
pub const DEFAULT_EXECUTION_STATE_CACHE_SIZE: usize = 10_000;

#[cfg(test)]
#[path = "unit_tests/worker_tests.rs"]
mod worker_tests;

#[cfg(all(test, feature = "rocksdb"))]
#[path = "unit_tests/worker_backup_tests.rs"]
mod worker_backup_tests;

/// Wraps a future in `SyncFuture` on non-web targets so that it satisfies `Sync` bounds.
/// On web targets the future is returned as-is.
#[cfg(not(web))]
pub(crate) fn wrap_future<F: std::future::Future>(f: F) -> sync_wrapper::SyncFuture<F> {
    sync_wrapper::SyncFuture::new(f)
}

/// Wraps a future in `SyncFuture` on non-web targets so that it satisfies `Sync` bounds.
/// On web targets the future is returned as-is.
#[cfg(web)]
pub(crate) fn wrap_future<F: std::future::Future>(f: F) -> F {
    f
}

#[cfg(with_metrics)]
mod metrics {
    use std::sync::LazyLock;

    use linera_base::prometheus_util::{
        exponential_bucket_interval, register_histogram, register_histogram_vec,
        register_int_counter, register_int_counter_vec,
    };
    use linera_chain::{data_types::MessageAction, types::ConfirmedBlockCertificate};
    use prometheus::{Histogram, HistogramVec, IntCounter, IntCounterVec};

    pub static NUM_ROUNDS_IN_CERTIFICATE: LazyLock<HistogramVec> = LazyLock::new(|| {
        register_histogram_vec(
            "num_rounds_in_certificate",
            "Number of rounds in certificate",
            &["certificate_value", "round_type"],
            exponential_bucket_interval(0.1, 50.0),
        )
    });

    pub static NUM_ROUNDS_IN_BLOCK_PROPOSAL: LazyLock<HistogramVec> = LazyLock::new(|| {
        register_histogram_vec(
            "num_rounds_in_block_proposal",
            "Number of rounds in block proposal",
            &["round_type"],
            exponential_bucket_interval(0.1, 50.0),
        )
    });

    pub static TRANSACTION_COUNT: LazyLock<IntCounterVec> =
        LazyLock::new(|| register_int_counter_vec("transaction_count", "Transaction count", &[]));

    pub static INCOMING_BUNDLE_COUNT: LazyLock<IntCounter> =
        LazyLock::new(|| register_int_counter("incoming_bundle_count", "Incoming bundle count"));

    pub static REJECTED_BUNDLE_COUNT: LazyLock<IntCounter> =
        LazyLock::new(|| register_int_counter("rejected_bundle_count", "Rejected bundle count"));

    pub static INCOMING_MESSAGE_COUNT: LazyLock<IntCounter> =
        LazyLock::new(|| register_int_counter("incoming_message_count", "Incoming message count"));

    pub static OPERATION_COUNT: LazyLock<IntCounter> =
        LazyLock::new(|| register_int_counter("operation_count", "Operation count"));

    pub static OPERATIONS_PER_BLOCK: LazyLock<Histogram> = LazyLock::new(|| {
        register_histogram(
            "operations_per_block",
            "Number of operations per block",
            exponential_bucket_interval(1.0, 10000.0),
        )
    });

    pub static INCOMING_BUNDLES_PER_BLOCK: LazyLock<Histogram> = LazyLock::new(|| {
        register_histogram(
            "incoming_bundles_per_block",
            "Number of incoming bundles per block",
            exponential_bucket_interval(1.0, 10000.0),
        )
    });

    pub static TRANSACTIONS_PER_BLOCK: LazyLock<Histogram> = LazyLock::new(|| {
        register_histogram(
            "transactions_per_block",
            "Number of transactions per block",
            exponential_bucket_interval(1.0, 10000.0),
        )
    });

    pub static NUM_BLOCKS: LazyLock<IntCounterVec> = LazyLock::new(|| {
        register_int_counter_vec("num_blocks", "Number of blocks added to chains", &[])
    });

    pub static CERTIFICATES_SIGNED: LazyLock<IntCounterVec> = LazyLock::new(|| {
        register_int_counter_vec(
            "certificates_signed",
            "Number of confirmed block certificates signed by each validator",
            &["validator_name"],
        )
    });

    pub static CHAIN_INFO_QUERIES: LazyLock<IntCounter> = LazyLock::new(|| {
        register_int_counter(
            "chain_info_queries",
            "Number of chain info queries processed",
        )
    });

    pub static CROSS_CHAIN_BATCH_SIZE: LazyLock<Histogram> = LazyLock::new(|| {
        register_histogram(
            "cross_chain_batch_size",
            "Number of cross-chain requests coalesced into a single per-chain batch",
            exponential_bucket_interval(1.0, 1000.0),
        )
    });

    /// Holds metrics data extracted from a confirmed block certificate.
    pub struct MetricsData {
        certificate_log_str: &'static str,
        round_type: &'static str,
        round_number: u32,
        confirmed_transactions: u64,
        confirmed_incoming_bundles: u64,
        confirmed_rejected_bundles: u64,
        confirmed_incoming_messages: u64,
        confirmed_operations: u64,
        validators_with_signatures: Vec<String>,
    }

    impl MetricsData {
        /// Creates a new `MetricsData` by extracting data from the certificate.
        pub fn new(certificate: &ConfirmedBlockCertificate) -> Self {
            Self {
                certificate_log_str: certificate.inner().to_log_str(),
                round_type: certificate.round.type_name(),
                round_number: certificate.round.number(),
                confirmed_transactions: certificate.block().body.transactions.len() as u64,
                confirmed_incoming_bundles: certificate.block().body.incoming_bundles().count()
                    as u64,
                confirmed_rejected_bundles: certificate
                    .block()
                    .body
                    .incoming_bundles()
                    .filter(|b| b.action == MessageAction::Reject)
                    .count() as u64,
                confirmed_incoming_messages: certificate
                    .block()
                    .body
                    .incoming_bundles()
                    .map(|b| b.messages().count())
                    .sum::<usize>() as u64,
                confirmed_operations: certificate.block().body.operations().count() as u64,
                validators_with_signatures: certificate
                    .signatures()
                    .iter()
                    .map(|(validator_name, _)| validator_name.to_string())
                    .collect(),
            }
        }

        /// Records the metrics for a processed block.
        pub fn record(self) {
            NUM_BLOCKS.with_label_values(&[]).inc();
            NUM_ROUNDS_IN_CERTIFICATE
                .with_label_values(&[self.certificate_log_str, self.round_type])
                .observe(self.round_number as f64);
            TRANSACTIONS_PER_BLOCK.observe(self.confirmed_transactions as f64);
            INCOMING_BUNDLES_PER_BLOCK.observe(self.confirmed_incoming_bundles as f64);
            OPERATIONS_PER_BLOCK.observe(self.confirmed_operations as f64);
            if self.confirmed_transactions > 0 {
                TRANSACTION_COUNT
                    .with_label_values(&[])
                    .inc_by(self.confirmed_transactions);
                if self.confirmed_incoming_bundles > 0 {
                    INCOMING_BUNDLE_COUNT.inc_by(self.confirmed_incoming_bundles);
                }
                if self.confirmed_rejected_bundles > 0 {
                    REJECTED_BUNDLE_COUNT.inc_by(self.confirmed_rejected_bundles);
                }
                if self.confirmed_incoming_messages > 0 {
                    INCOMING_MESSAGE_COUNT.inc_by(self.confirmed_incoming_messages);
                }
                if self.confirmed_operations > 0 {
                    OPERATION_COUNT.inc_by(self.confirmed_operations);
                }
            }

            for validator_name in self.validators_with_signatures {
                CERTIFICATES_SIGNED
                    .with_label_values(&[&validator_name])
                    .inc();
            }
        }
    }
}

/// Instruct the networking layer to send cross-chain requests and/or push notifications.
#[derive(Default, Debug)]
pub struct NetworkActions {
    /// The cross-chain requests
    pub cross_chain_requests: Vec<CrossChainRequest>,
    /// The push notifications.
    pub notifications: Vec<Notification>,
}

impl NetworkActions {
    pub fn extend(&mut self, other: NetworkActions) {
        self.cross_chain_requests.extend(other.cross_chain_requests);
        self.notifications.extend(other.notifications);
    }
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
/// Notification that a chain has a new certified block or a new message.
pub struct Notification {
    pub chain_id: ChainId,
    pub reason: Reason,
}

doc_scalar!(
    Notification,
    "Notify that a chain has a new certified block or a new message"
);

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
/// Reason for the notification.
pub enum Reason {
    NewBlock {
        height: BlockHeight,
        hash: CryptoHash,
    },
    NewEvents {
        height: BlockHeight,
        block_hash: CryptoHash,
        event_streams: BTreeSet<StreamId>,
    },
    NewIncomingBundle {
        origin: ChainId,
        height: BlockHeight,
    },
    NewRound {
        height: BlockHeight,
        round: Round,
    },
    BlockExecuted {
        height: BlockHeight,
        hash: CryptoHash,
    },
}

/// Error type for worker operations.
#[derive(Debug, Error, strum::IntoStaticStr)]
pub enum WorkerError {
    #[error(transparent)]
    CryptoError(#[from] CryptoError),

    #[error(transparent)]
    ArithmeticError(#[from] ArithmeticError),

    #[error(transparent)]
    ViewError(#[from] ViewError),

    #[error("Certificates referenced from chain state are missing in storage: {0:?}")]
    ReadCertificatesError(Vec<CryptoHash>),

    #[error(transparent)]
    ChainError(#[from] Box<ChainError>),

    #[error(transparent)]
    BcsError(#[from] bcs::Error),

    // Chain access control
    #[error("Block was not signed by an authorized owner")]
    InvalidOwner,

    #[error("Operations in the block are not authenticated by the proper owner: {0}")]
    InvalidSigner(AccountOwner),

    // Chaining
    #[error(
        "Chain is expecting a next block at height {expected_block_height} but the given block \
        is at height {found_block_height} instead"
    )]
    UnexpectedBlockHeight {
        expected_block_height: BlockHeight,
        found_block_height: BlockHeight,
    },
    #[error("Unexpected epoch {epoch}: chain {chain_id} is at {chain_epoch}")]
    InvalidEpoch {
        chain_id: ChainId,
        chain_epoch: Epoch,
        epoch: Epoch,
    },

    #[error("Events not found: {0:?}")]
    EventsNotFound(Vec<EventId>),

    // Other server-side errors
    #[error("Invalid cross-chain request")]
    InvalidCrossChainRequest,
    #[error("The block does not contain the hash that we expected for the previous block")]
    InvalidBlockChaining,
    #[error(
        "Block timestamp ({block_timestamp}) is further in the future from local time \
        ({local_time}) than block time grace period ({block_time_grace_period:?})"
    )]
    InvalidTimestamp {
        block_timestamp: Timestamp,
        local_time: Timestamp,
        block_time_grace_period: Duration,
    },
    #[error("We don't have the value for the certificate.")]
    MissingCertificateValue,
    #[error("The hash certificate doesn't match its value.")]
    InvalidLiteCertificate,
    #[error("Fast blocks cannot query oracles")]
    FastBlockUsingOracles,
    #[error("Blobs not found: {0:?}")]
    BlobsNotFound(Vec<BlobId>),
    /// Variant raised when the chain references these block hashes via a
    /// verified-checkpoint trust mark (`pre_checkpoint_block_trust`) but the
    /// actual content isn't in storage yet. The caller is expected to upload
    /// each missing block via `handle_confirmed_certificate`; the trust-mark
    /// accept path verifies the cert against its own (possibly revoked)
    /// epoch's committee and writes it through.
    #[error("Blocks not found: {0:?}")]
    BlocksNotFound(Vec<CryptoHash>),
    #[error("Block hash at height {height} for chain {chain_id} not found")]
    BlockHashNotFound {
        height: BlockHeight,
        chain_id: ChainId,
    },
    #[error("Block at height {height} on chain {chain_id} not found in local storage")]
    LocalBlockNotFound {
        height: BlockHeight,
        chain_id: ChainId,
    },
    #[error("The block proposal is invalid: {0}")]
    InvalidBlockProposal(String),
    #[error("Blob was not required by any pending block")]
    UnexpectedBlob,
    #[error("Number of published blobs per block must not exceed {0}")]
    TooManyPublishedBlobs(u64),
    #[error("Missing network description")]
    MissingNetworkDescription,
    #[error("thread error: {0}")]
    Thread(#[from] web_thread_pool::Error),
    #[error("Chain worker was poisoned by a journal resolution failure")]
    PoisonedWorker,
    #[error("Cross-chain batch was rolled back due to an error in another request")]
    BatchRolledBack,
}

impl WorkerError {
    /// Returns whether this error is caused by an issue in the local node.
    ///
    /// Returns `false` whenever the error could be caused by a bad message from a peer.
    pub fn is_local(&self) -> bool {
        match self {
            WorkerError::CryptoError(_)
            | WorkerError::ArithmeticError(_)
            | WorkerError::InvalidOwner
            | WorkerError::InvalidSigner(_)
            | WorkerError::UnexpectedBlockHeight { .. }
            | WorkerError::InvalidEpoch { .. }
            | WorkerError::EventsNotFound(_)
            | WorkerError::InvalidBlockChaining
            | WorkerError::InvalidTimestamp { .. }
            | WorkerError::MissingCertificateValue
            | WorkerError::InvalidLiteCertificate
            | WorkerError::FastBlockUsingOracles
            | WorkerError::BlobsNotFound(_)
            | WorkerError::BlocksNotFound(_)
            | WorkerError::InvalidBlockProposal(_)
            | WorkerError::UnexpectedBlob
            | WorkerError::TooManyPublishedBlobs(_)
            | WorkerError::ViewError(ViewError::NotFound(_)) => false,
            WorkerError::BcsError(_)
            | WorkerError::InvalidCrossChainRequest
            | WorkerError::ViewError(_)
            | WorkerError::BlockHashNotFound { .. }
            | WorkerError::LocalBlockNotFound { .. }
            | WorkerError::MissingNetworkDescription
            | WorkerError::Thread(_)
            | WorkerError::ReadCertificatesError(_)
            | WorkerError::PoisonedWorker
            | WorkerError::BatchRolledBack => true,
            WorkerError::ChainError(chain_error) => chain_error.is_local(),
        }
    }

    /// Returns the qualified error variant name for the `error_type` metric label,
    /// e.g. `"WorkerError::UnexpectedBlockHeight"`.
    ///
    /// For `ChainError` variants, delegates to `ChainError::error_type()` to
    /// surface the underlying error name rather than just `"ChainError"`.
    pub fn error_type(&self) -> String {
        match self {
            WorkerError::ChainError(chain_error) => chain_error.error_type(),
            other => {
                let variant: &'static str = other.into();
                format!("WorkerError::{variant}")
            }
        }
    }

    /// Returns `true` if this error indicates that the chain worker's in-memory
    /// state may be inconsistent and must be evicted from the cache.
    fn must_reload_view(&self) -> bool {
        matches!(
            self,
            WorkerError::PoisonedWorker
                | WorkerError::ViewError(ViewError::StoreError {
                    must_reload_view: true,
                    ..
                })
        )
    }

    /// Returns `true` if this error indicates that the chain's persisted state is
    /// internally inconsistent, so the worker should consider resetting and
    /// re-executing it from storage.
    fn indicates_corrupted_chain_state(&self) -> bool {
        matches!(
            self,
            WorkerError::ChainError(chain_error)
                if matches!(chain_error.as_ref(), ChainError::CorruptedChainState(_))
        )
    }
}

impl From<ChainError> for WorkerError {
    #[instrument(level = "trace", skip(chain_error))]
    fn from(chain_error: ChainError) -> Self {
        match chain_error {
            ChainError::ExecutionError(execution_error, context) => match *execution_error {
                ExecutionError::BlobsNotFound(blob_ids) => Self::BlobsNotFound(blob_ids),
                ExecutionError::EventsNotFound(event_ids) => Self::EventsNotFound(event_ids),
                _ => Self::ChainError(Box::new(ChainError::ExecutionError(
                    execution_error,
                    context,
                ))),
            },
            error => Self::ChainError(Box::new(error)),
        }
    }
}

#[cfg(with_testing)]
impl WorkerError {
    /// Returns the inner [`ExecutionError`] in this error.
    ///
    /// # Panics
    ///
    /// If this is not caused by an [`ExecutionError`].
    pub fn expect_execution_error(self, expected_context: ChainExecutionContext) -> ExecutionError {
        let WorkerError::ChainError(chain_error) = self else {
            panic!("Expected an `ExecutionError`. Got: {self:#?}");
        };

        let ChainError::ExecutionError(execution_error, context) = *chain_error else {
            panic!("Expected an `ExecutionError`. Got: {chain_error:#?}");
        };

        assert_eq!(context, expected_context);

        *execution_error
    }
}

type ChainWorkerArc<S> = Arc<tokio::sync::RwLock<ChainWorkerState<S>>>;
type ChainWorkerWeak<S> = std::sync::Weak<tokio::sync::RwLock<ChainWorkerState<S>>>;
type ChainWorkerFuture<S> = Shared<oneshot::Receiver<ChainWorkerWeak<S>>>;

/// Each map entry is a `Shared<oneshot::Receiver<Weak<...>>>`:
///
/// - `peek()` returns `None` while a task is loading the worker from storage.
/// - `peek()` returns `Some(Ok(weak))` once the worker is loaded.
/// - `peek()` returns `Some(Err(_))` if loading failed (sender dropped).
///
/// Callers that find a pending entry clone the `Shared` future and await it.
type ChainWorkerMap<S> = Arc<papaya::HashMap<ChainId, ChainWorkerFuture<S>>>;

/// A cross-chain request waiting to be processed in a batch.
pub(crate) enum BatchRequest {
    Update {
        origin: ChainId,
        bundles: Vec<(Epoch, MessageBundle)>,
        previous_height: Option<BlockHeight>,
        result_sender: oneshot::Sender<Result<CrossChainUpdateResult, WorkerError>>,
    },
    Confirm {
        recipient: ChainId,
        latest_height: BlockHeight,
        result_sender: oneshot::Sender<Result<NetworkActions, WorkerError>>,
    },
}

/// The inner future type for cross-chain batch processing.
///
/// Wrapped in `Shared<BatchFuture>` so that all tasks waiting for
/// cross-chain operations on the same chain can cooperatively poll a
/// single driver. The driver loops: wait for an item from the request channel, acquire the
/// write lock, drain the channel, process all requests in one batch, repeat.
#[cfg(not(web))]
type BatchFuture = pin::Pin<Box<dyn Future<Output = ()> + Send>>;
#[cfg(web)]
type BatchFuture = pin::Pin<Box<dyn Future<Output = ()>>>;

#[derive(Clone)]
struct ChainBatchRequestProcessor {
    /// Sender half of the channel whose receiver lives inside the driver future.
    sender: mpsc::UnboundedSender<BatchRequest>,
    /// Weak handle to the shared driver future. Upgraded to `Shared<BatchFuture>`
    /// by callers who need to poll it.
    future: WeakShared<BatchFuture>,
}

impl ChainBatchRequestProcessor {
    fn create<StorageClient>(
        state: ChainWorkerArc<StorageClient>,
        batch_size_limit: usize,
    ) -> (ChainBatchRequestProcessor, Shared<BatchFuture>)
    where
        StorageClient: Storage + Clone + 'static,
    {
        let (sender, mut receiver) = mpsc::unbounded_channel();
        let future: BatchFuture = Box::pin(async move {
            while let Some(first) = receiver.recv().await {
                let mut requests = vec![first];
                match handle::write_lock(&state).await {
                    Ok(mut guard) => {
                        while requests.len() < batch_size_limit {
                            match receiver.try_recv() {
                                Ok(request) => requests.push(request),
                                Err(_) => break,
                            }
                        }
                        #[cfg(with_metrics)]
                        metrics::CROSS_CHAIN_BATCH_SIZE.observe(requests.len() as f64);
                        guard.process_batch(requests).await
                    }
                    Err(error) => {
                        tracing::error!(%error, "failed to obtain write lock");
                        for request in requests {
                            match request {
                                BatchRequest::Update { result_sender, .. } => {
                                    send_result(result_sender, Err(WorkerError::PoisonedWorker));
                                }
                                BatchRequest::Confirm { result_sender, .. } => {
                                    send_result(result_sender, Err(WorkerError::PoisonedWorker));
                                }
                            }
                        }
                    }
                }
            }
        });
        let shared = future.shared();
        let weak = shared.downgrade().expect("future has not been polled yet");
        let batch_processor = ChainBatchRequestProcessor {
            sender,
            future: weak,
        };
        (batch_processor, shared)
    }
}

type ChainBatchMap = Arc<papaya::HashMap<ChainId, ChainBatchRequestProcessor>>;

/// Starts a background task that periodically removes dead weak references
/// from the chain handle map. The actual lifetime management is handled by
/// each handle's keep-alive task.
fn start_sweep<S: Storage + Clone + 'static>(
    chain_workers: &ChainWorkerMap<S>,
    config: &ChainWorkerConfig,
) {
    // Sweep at the smaller of the two TTLs. If both are None, workers
    // live forever so there's nothing to sweep.
    let interval = match (config.ttl, config.sender_chain_ttl) {
        (None, None) => return,
        (Some(d), None) | (None, Some(d)) => d,
        (Some(a), Some(b)) => a.min(b),
    };
    let weak_map = Arc::downgrade(chain_workers);
    linera_base::Task::spawn(async move {
        loop {
            linera_base::time::timer::sleep(interval).await;
            let Some(map) = weak_map.upgrade() else {
                break;
            };
            map.pin_owned().retain(|_, shared| match shared.peek() {
                Some(Ok(weak)) => weak.strong_count() > 0,
                Some(Err(_)) => false, // Loading failed; clean up.
                None => true,          // Still loading; keep.
            });
        }
    })
    .forget();
}

/// State of a worker in a validator or a local node.
pub struct WorkerState<StorageClient: Storage> {
    /// Access to local persistent storage.
    storage: StorageClient,
    /// Configuration options for chain workers.
    chain_worker_config: ChainWorkerConfig,
    block_cache: Arc<ValueCache<CryptoHash, ConfirmedBlock>>,
    execution_state_cache:
        Option<Arc<UniqueValueCache<CryptoHash, ExecutionStateView<InactiveContext>>>>,
    /// Chains tracked by a worker, along with their listening modes.
    pub(crate) chain_modes: Option<Arc<RwLock<ChainModes>>>,
    /// One-shot channels to notify callers when messages of a particular chain have been
    /// delivered.
    delivery_notifiers: Arc<Mutex<DeliveryNotifiers>>,
    /// The cache of loaded chain workers. Stores weak references; each worker
    /// manages its own lifetime via a keep-alive task. A background sweep
    /// periodically removes dead entries.
    chain_workers: ChainWorkerMap<StorageClient>,
    /// Per-chain batch processing state for cross-chain requests.
    chain_batches: ChainBatchMap,
    /// Shard-routing dispatcher for outbound cross-chain requests. Used when we need
    /// to send cross-chain requests outside of the normal `NetworkActions` return
    /// path — in particular, the `RevertConfirm`s emitted after resetting a
    /// corrupted chain. The RPC server layer installs this; without it, we fall
    /// back to dispatching locally through `handle_cross_chain_request`.
    outbound_cross_chain_sender: Option<OutboundCrossChainSender>,
}

/// Dispatcher for outbound cross-chain requests that handles the source-shard-to-
/// target-shard routing that the worker itself is not aware of.
pub type OutboundCrossChainSender = Arc<dyn Fn(CrossChainRequest) + Send + Sync>;

impl<StorageClient> Clone for WorkerState<StorageClient>
where
    StorageClient: Storage + Clone,
{
    fn clone(&self) -> Self {
        WorkerState {
            storage: self.storage.clone(),
            chain_worker_config: self.chain_worker_config.clone(),
            block_cache: self.block_cache.clone(),
            execution_state_cache: self.execution_state_cache.clone(),
            chain_modes: self.chain_modes.clone(),
            delivery_notifiers: self.delivery_notifiers.clone(),
            chain_workers: self.chain_workers.clone(),
            chain_batches: self.chain_batches.clone(),
            outbound_cross_chain_sender: self.outbound_cross_chain_sender.clone(),
        }
    }
}

pub(crate) type DeliveryNotifiers = HashMap<ChainId, DeliveryNotifier>;

impl<StorageClient> WorkerState<StorageClient>
where
    StorageClient: Storage,
{
    /// Returns an instance with the specified cross-chain message chunk limit.
    #[cfg(with_testing)]
    #[instrument(level = "trace", skip(self))]
    pub fn with_cross_chain_message_chunk_limit(mut self, limit: usize) -> Self {
        self.chain_worker_config.cross_chain_message_chunk_limit = limit;
        self
    }

    /// Sets the cross-chain message chunk limit.
    #[cfg(with_testing)]
    pub fn set_cross_chain_message_chunk_limit(&mut self, limit: usize) {
        self.chain_worker_config.cross_chain_message_chunk_limit = limit;
    }

    #[cfg(with_testing)]
    #[instrument(level = "trace", skip(self, value))]
    pub fn with_allow_revert_confirm(mut self, value: bool) -> Self {
        self.chain_worker_config.allow_revert_confirm = value;
        self
    }

    #[instrument(level = "trace", skip(self))]
    pub fn nickname(&self) -> &str {
        &self.chain_worker_config.nickname
    }

    /// Returns the storage client so that it can be manipulated or queried.
    #[instrument(level = "trace", skip(self))]
    #[cfg(not(feature = "test"))]
    pub(crate) fn storage_client(&self) -> &StorageClient {
        &self.storage
    }

    /// Returns the storage client so that it can be manipulated or queried by tests in other
    /// crates.
    #[instrument(level = "trace", skip(self))]
    #[cfg(feature = "test")]
    pub fn storage_client(&self) -> &StorageClient {
        &self.storage
    }

    #[instrument(level = "trace", skip(self, certificate))]
    pub(crate) async fn full_certificate(
        &self,
        certificate: LiteCertificate<'_>,
    ) -> Result<Either<ConfirmedBlockCertificate, ValidatedBlockCertificate>, WorkerError> {
        let block = self
            .block_cache
            .get(&certificate.value.value_hash)
            .ok_or(WorkerError::MissingCertificateValue)?;
        let block = CacheArc::unwrap_or_clone(block);

        match certificate.value.kind {
            linera_chain::types::CertificateKind::Confirmed => Ok(Either::Left(
                certificate
                    .with_value(block)
                    .ok_or(WorkerError::InvalidLiteCertificate)?,
            )),
            linera_chain::types::CertificateKind::Validated => {
                let value = ValidatedBlock::from_hashed(block.into_inner());
                Ok(Either::Right(
                    certificate
                        .with_value(value)
                        .ok_or(WorkerError::InvalidLiteCertificate)?,
                ))
            }
            _ => Err(WorkerError::InvalidLiteCertificate),
        }
    }
}

#[allow(async_fn_in_trait)]
#[cfg_attr(not(web), trait_variant::make(Send))]
pub trait ProcessableCertificate: CertificateValue + Sized + 'static {
    async fn process_certificate<S: Storage + Clone + 'static>(
        worker: &WorkerState<S>,
        certificate: GenericCertificate<Self>,
    ) -> Result<(ChainInfoResponse, NetworkActions), WorkerError>;
}

impl ProcessableCertificate for ConfirmedBlock {
    async fn process_certificate<S: Storage + Clone + 'static>(
        worker: &WorkerState<S>,
        certificate: ConfirmedBlockCertificate,
    ) -> Result<(ChainInfoResponse, NetworkActions), WorkerError> {
        Box::pin(worker.handle_confirmed_certificate(
            certificate,
            ProcessConfirmedBlockMode::Auto,
            None,
        ))
        .await
    }
}

impl ProcessableCertificate for ValidatedBlock {
    async fn process_certificate<S: Storage + Clone + 'static>(
        worker: &WorkerState<S>,
        certificate: ValidatedBlockCertificate,
    ) -> Result<(ChainInfoResponse, NetworkActions), WorkerError> {
        Box::pin(worker.handle_validated_certificate(certificate)).await
    }
}

impl ProcessableCertificate for Timeout {
    async fn process_certificate<S: Storage + Clone + 'static>(
        worker: &WorkerState<S>,
        certificate: TimeoutCertificate,
    ) -> Result<(ChainInfoResponse, NetworkActions), WorkerError> {
        worker.handle_timeout_certificate(certificate).await
    }
}

impl<StorageClient> WorkerState<StorageClient>
where
    StorageClient: Storage + Clone + 'static,
{
    /// Creates a new `WorkerState`.
    ///
    /// The `chain_worker_config` must be fully configured before calling this, because the
    /// TTL sweep task is started immediately based on the config's TTL values.
    #[instrument(level = "trace", skip(storage, chain_worker_config))]
    pub fn new(
        storage: StorageClient,
        chain_worker_config: ChainWorkerConfig,
        chain_modes: Option<Arc<RwLock<ChainModes>>>,
    ) -> Self {
        let chain_workers = Arc::new(papaya::HashMap::new());
        start_sweep(&chain_workers, &chain_worker_config);
        let block_cache_size = chain_worker_config.block_cache_size;
        let execution_state_cache_size = chain_worker_config.execution_state_cache_size;
        WorkerState {
            storage,
            chain_worker_config,
            block_cache: Arc::new(ValueCache::new(
                "worker_block",
                block_cache_size,
                DEFAULT_CLEANUP_INTERVAL_SECS,
            )),
            execution_state_cache: (execution_state_cache_size > 0)
                .then(|| Arc::new(UniqueValueCache::new(execution_state_cache_size))),
            chain_modes,
            delivery_notifiers: Arc::default(),
            chain_workers,
            // On wasm, `ChainBatchRequestProcessor` is not `Send`/`Sync` (the inner future
            // isn't `Send`), but `Arc` is still correct: `WorkerState` clones
            // share this map and wasm is single-threaded.
            #[cfg_attr(web, expect(clippy::arc_with_non_send_sync))]
            chain_batches: Arc::new(papaya::HashMap::new()),
            outbound_cross_chain_sender: None,
        }
    }

    /// Installs a shard-routing dispatcher used for outbound cross-chain requests
    /// generated outside the normal response path (specifically, after resetting a
    /// corrupted chain). Without it, such requests are dispatched locally in a
    /// loop via `handle_cross_chain_request`.
    pub fn with_outbound_cross_chain_sender(mut self, sender: OutboundCrossChainSender) -> Self {
        self.outbound_cross_chain_sender = Some(sender);
        self
    }

    #[instrument(level = "trace", skip(self, certificate, notifier))]
    #[inline]
    pub async fn fully_handle_certificate_with_notifications<T>(
        &self,
        certificate: GenericCertificate<T>,
        notifier: &impl Notifier,
    ) -> Result<ChainInfoResponse, WorkerError>
    where
        T: ProcessableCertificate,
    {
        let notifications = (*notifier).clone();
        let this = self.clone();
        linera_base::Task::spawn(async move {
            let (response, actions) =
                ProcessableCertificate::process_certificate(&this, certificate).await?;
            notifications.notify(&actions.notifications);
            let mut requests = VecDeque::from(actions.cross_chain_requests);
            while let Some(request) = requests.pop_front() {
                let actions = this.handle_cross_chain_request(request).await?;
                requests.extend(actions.cross_chain_requests);
                notifications.notify(&actions.notifications);
            }
            Ok(response)
        })
        .await
    }

    /// Same as [`Self::fully_handle_certificate_with_notifications`] but for a
    /// confirmed block certificate and with an explicit [`ProcessConfirmedBlockMode`].
    /// The generic variant always uses [`ProcessConfirmedBlockMode::Auto`].
    #[instrument(level = "trace", skip(self, certificate, notifier))]
    #[inline]
    pub async fn fully_handle_confirmed_certificate_with_notifications(
        &self,
        certificate: ConfirmedBlockCertificate,
        mode: ProcessConfirmedBlockMode,
        notifier: &impl Notifier,
    ) -> Result<ChainInfoResponse, WorkerError> {
        let notifications = (*notifier).clone();
        let this = self.clone();
        linera_base::Task::spawn(async move {
            let (response, actions) =
                Box::pin(this.handle_confirmed_certificate(certificate, mode, None)).await?;
            notifications.notify(&actions.notifications);
            let mut requests = VecDeque::from(actions.cross_chain_requests);
            while let Some(request) = requests.pop_front() {
                let actions = this.handle_cross_chain_request(request).await?;
                requests.extend(actions.cross_chain_requests);
                notifications.notify(&actions.notifications);
            }
            Ok(response)
        })
        .await
    }

    /// Acquires a read lock on the chain worker and executes the given closure.
    ///
    /// The future is boxed to keep deeply nested types off the stack. On non-web
    /// targets it is also wrapped in `SyncFuture` to satisfy `Sync` bounds.
    async fn chain_read<R, F, Fut>(&self, chain_id: ChainId, f: F) -> Result<R, WorkerError>
    where
        F: FnOnce(OwnedRwLockReadGuard<ChainWorkerState<StorageClient>>) -> Fut,
        Fut: std::future::Future<Output = Result<R, WorkerError>>,
    {
        let state = self.get_or_create_chain_worker(chain_id).await?;
        let state_ref = &state;
        let result = Box::pin(wrap_future(async move {
            let guard = handle::read_lock(state_ref).await?;
            f(guard).await
        }))
        .await;
        if let Err(error) = &result {
            if error.must_reload_view() {
                self.evict_poisoned_worker(chain_id, &state);
            }
        }
        result
    }

    /// Acquires a write lock on the chain worker and executes the given closure.
    ///
    /// The write work runs on a detached task (via [`linera_base::task::run_detached`])
    /// so that caller cancellation does not unwind the task mid-save. The
    /// [`RollbackGuard`] lives inside the detached task, so the write lock is held
    /// until the DB round-trip and `post_save` have fully completed — subsequent
    /// readers, including a freshly-loaded replacement worker, only see the
    /// committed state.
    async fn chain_write<R, F, Fut>(&self, chain_id: ChainId, f: F) -> Result<R, WorkerError>
    where
        F: FnOnce(handle::RollbackGuard<StorageClient>) -> Fut
            + linera_base::task::MaybeSend
            + 'static,
        Fut: std::future::Future<Output = Result<R, WorkerError>> + linera_base::task::MaybeSend,
        R: linera_base::task::MaybeSend + 'static,
    {
        let state = self.get_or_create_chain_worker(chain_id).await?;
        let state_for_task = state.clone();
        let result = Box::pin(wrap_future(linera_base::task::run_detached(async move {
            let guard = handle::write_lock(&state_for_task).await?;
            f(guard).await
        })))
        .await;
        if let Err(error) = &result {
            if error.must_reload_view() {
                self.evict_poisoned_worker(chain_id, &state);
            } else if error.indicates_corrupted_chain_state() {
                self.spawn_reset_corrupted_chain_state(chain_id, state);
            }
        }
        result
    }

    /// Spawns a detached task that re-acquires the write lock and recovers the
    /// chain from a detected state corruption. Running the recovery in a separate
    /// task ensures it survives cancellation of the originating request: if the
    /// caller's future is dropped mid-way through re-execution, the chain would
    /// otherwise be left at a partial tip and our safety snapshot would be lost.
    /// Generated `RevertConfirm` requests are dispatched via the installed
    /// shard-routing sender when present (sharded validators), or locally
    /// through `handle_cross_chain_request` otherwise (client nodes and tests).
    /// Errors are logged; the caller already has the original error to
    /// propagate.
    fn spawn_reset_corrupted_chain_state(
        &self,
        chain_id: ChainId,
        state: ChainWorkerArc<StorageClient>,
    ) where
        StorageClient: Clone,
    {
        let this = self.clone();
        linera_base::Task::spawn(async move {
            let requests = {
                let mut guard = match handle::write_lock(&state).await {
                    Ok(guard) => guard,
                    Err(error) => {
                        tracing::error!(
                            %chain_id, %error,
                            "Failed to acquire write lock to reset corrupted chain state"
                        );
                        return;
                    }
                };
                match guard.maybe_reset_corrupted_chain_state().await {
                    Ok(Some(requests)) => requests,
                    Ok(None) => return,
                    Err(error) => {
                        tracing::error!(
                            %chain_id, %error, "Failed to reset corrupted chain state"
                        );
                        return;
                    }
                }
            };
            if let Some(sender) = &this.outbound_cross_chain_sender {
                // Sharded validator path: let the RPC layer route each request to
                // the shard that owns the target chain.
                for request in requests {
                    sender(request);
                }
            } else {
                // No routing dispatcher is installed (client node or test), so all
                // involved chains are co-located on this worker. Dispatch locally
                // in a loop, following any cascading cross-chain requests.
                let mut queue = VecDeque::from(requests);
                while let Some(request) = queue.pop_front() {
                    match this.handle_cross_chain_request(request).await {
                        Ok(actions) => queue.extend(actions.cross_chain_requests),
                        Err(error) => {
                            warn!(
                                %chain_id, %error,
                                "Failed to dispatch cross-chain request after \
                                resetting corrupted chain state"
                            );
                        }
                    }
                }
            }
        })
        .forget();
    }

    /// Evicts a poisoned chain worker from the cache, but only if the entry still
    /// points to the same instance. This avoids removing a fresh replacement that
    /// another task may have already loaded.
    fn evict_poisoned_worker(&self, chain_id: ChainId, poisoned: &ChainWorkerArc<StorageClient>) {
        tracing::warn!(%chain_id, "Evicting poisoned chain worker from cache");
        let pin = self.chain_workers.pin();
        let weak_poisoned = Arc::downgrade(poisoned);
        let removed = pin.remove_if(&chain_id, |_key, future| {
            future
                .peek()
                .and_then(|r| r.clone().ok())
                .is_some_and(|weak| weak.ptr_eq(&weak_poisoned))
        });
        if removed.is_err() {
            tracing::trace!(%chain_id, "Poisoned worker entry already replaced; skipping eviction");
        }
    }

    /// Returns or creates the per-chain batch state.
    async fn get_or_create_chain_batch(
        &self,
        chain_id: ChainId,
    ) -> Result<(mpsc::UnboundedSender<BatchRequest>, Shared<BatchFuture>), WorkerError> {
        // Fast path: reuse an existing live driver. The pin guard is !Send,
        // so it must be dropped before any .await point.
        if let Some(batch_processor) = self.chain_batches.pin().get(&chain_id) {
            if let Some(future) = batch_processor.future.upgrade() {
                return Ok((batch_processor.sender.clone(), future));
            }
        }
        let state = self.get_or_create_chain_worker(chain_id).await?;
        let (new_request_processor, new_future) = ChainBatchRequestProcessor::create(
            state,
            self.chain_worker_config.cross_chain_batch_size_limit,
        );
        match self
            .chain_batches
            .pin()
            .compute(chain_id, |existing| match existing {
                Some((_, batch_processor)) => {
                    if let Some(future) = batch_processor.future.upgrade() {
                        papaya::Operation::Abort((batch_processor.sender.clone(), future))
                    } else {
                        papaya::Operation::Insert(new_request_processor.clone())
                    }
                }
                None => papaya::Operation::Insert(new_request_processor.clone()),
            }) {
            papaya::Compute::Aborted((sender, future)) => Ok((sender, future)),
            papaya::Compute::Inserted(_, batch_processor)
            | papaya::Compute::Updated {
                new: (_, batch_processor),
                ..
            } => Ok((batch_processor.sender.clone(), new_future)),
            papaya::Compute::Removed { .. } => unreachable!(),
        }
    }

    /// Gets or creates a chain worker for the given chain.
    ///
    /// The oneshot channel is created outside the `compute` closure to keep
    /// the closure pure (papaya may call it more than once on CAS retry and
    /// may memoize the output). If the fast path hits, the unused channel is
    /// dropped harmlessly.
    ///
    /// Returns a type-erased future to keep `!Sync` intermediate types (e.g.
    /// `std::sync::mpsc::Receiver` from `handle::ServiceRuntimeActor::spawn`) out of
    /// the caller's future type.
    fn get_or_create_chain_worker(
        &self,
        chain_id: ChainId,
    ) -> std::pin::Pin<
        Box<
            impl std::future::Future<Output = Result<ChainWorkerArc<StorageClient>, WorkerError>> + '_,
        >,
    > {
        Box::pin(wrap_future(async move {
            loop {
                // Create the channel outside the closure so that the
                // sender/receiver always match regardless of CAS retries.
                let (sender, receiver) = oneshot::channel();
                let shared_receiver = receiver.shared();

                // The papaya guard is !Send, so it must be dropped before
                // any .await point.
                let wait_or_sender = {
                    let pin = self.chain_workers.pin();
                    match pin.compute(chain_id, |existing| match existing {
                        Some((_, entry)) => match entry.peek() {
                            Some(Ok(weak)) => match weak.upgrade() {
                                Some(arc) => papaya::Operation::Abort(Ok(arc)),
                                None => papaya::Operation::Insert(shared_receiver.clone()),
                            },
                            Some(Err(_)) => papaya::Operation::Insert(shared_receiver.clone()),
                            None => papaya::Operation::Abort(Err(entry.clone())),
                        },
                        None => papaya::Operation::Insert(shared_receiver.clone()),
                    }) {
                        papaya::Compute::Aborted(Ok(arc), ..) => return Ok(arc),
                        papaya::Compute::Aborted(Err(wait), ..) => Either::Left(wait),
                        papaya::Compute::Inserted { .. } | papaya::Compute::Updated { .. } => {
                            Either::Right(sender)
                        }
                        papaya::Compute::Removed { .. } => unreachable!(),
                    }
                };

                match wait_or_sender {
                    Either::Left(wait) => {
                        // Another task is loading. Await the shared future.
                        if let Ok(weak) = wait.await {
                            if let Some(arc) = weak.upgrade() {
                                return Ok(arc);
                            }
                        }
                        // Loading failed or worker already dead; retry.
                    }
                    Either::Right(sender) => {
                        // We claimed the loading slot. Load from storage.
                        // On success, send the Weak through the channel.
                        // On error, dropping sender wakes waiters so they can retry.
                        let worker = self.load_chain_worker(chain_id).await?;
                        if sender.send(Arc::downgrade(&worker)).is_err() {
                            tracing::error!(%chain_id, "Receiver dropped while loading worker state.");
                            continue;
                        }
                        return Ok(worker);
                    }
                }
            }
        }))
    }

    /// Loads a chain worker state from storage and wraps it in an Arc.
    async fn load_chain_worker(
        &self,
        chain_id: ChainId,
    ) -> Result<ChainWorkerArc<StorageClient>, WorkerError> {
        let delivery_notifier = self
            .delivery_notifiers
            .lock()
            .unwrap()
            .entry(chain_id)
            .or_default()
            .clone();

        // `chain_modes=None` means "no tracked/sender distinction" (validators) — treat
        // every chain as tracked so it routes through `config.ttl`, not the unset
        // `config.sender_chain_ttl` which would skip `spawn_keep_alive` entirely.
        let is_tracked = self.chain_modes.as_ref().is_none_or(|chain_modes| {
            chain_modes
                .read()
                .unwrap()
                .get(&chain_id)
                .is_some_and(ListeningMode::is_full)
        });

        let (service_runtime_endpoint, service_runtime_task) =
            if self.chain_worker_config.long_lived_services {
                let actor =
                    handle::ServiceRuntimeActor::spawn(chain_id, self.storage.thread_pool()).await;
                (Some(actor.endpoint), Some(actor.task))
            } else {
                (None, None)
            };

        let state = crate::chain_worker::state::ChainWorkerState::load(
            self.chain_worker_config.clone(),
            self.storage.clone(),
            self.block_cache.clone(),
            self.execution_state_cache.clone(),
            self.chain_modes.clone(),
            delivery_notifier,
            chain_id,
            service_runtime_endpoint,
            service_runtime_task,
        )
        .await?;

        Ok(handle::create_chain_worker(
            state,
            is_tracked,
            &self.chain_worker_config,
        ))
    }

    /// Tries to execute a block proposal with a policy for handling bundle failures.
    ///
    /// Returns the modified block (bundles may be rejected/removed), the executed block,
    /// chain info response, and resource tracker.
    #[instrument(level = "trace", skip(self, block))]
    pub async fn stage_block_execution(
        &self,
        block: ProposedBlock,
        round: Option<u32>,
        published_blobs: Vec<Blob>,
        policy: BundleExecutionPolicy,
    ) -> Result<
        (
            ProposedBlock,
            Block,
            ChainInfoResponse,
            ResourceTracker,
            HashSet<ChainId>,
        ),
        WorkerError,
    > {
        let chain_id = block.chain_id;
        self.chain_write(chain_id, move |mut guard| async move {
            guard
                .stage_block_execution(block, round, &published_blobs, policy)
                .await
        })
        .await
    }

    /// Executes a [`Query`] for an application's state on a specific chain.
    ///
    /// If `block_hash` is specified, system will query the application's state
    /// at that block. If it doesn't exist, it uses latest state.
    #[instrument(level = "trace", skip(self, chain_id, query))]
    pub async fn query_application(
        &self,
        chain_id: ChainId,
        query: Query,
        block_hash: Option<CryptoHash>,
    ) -> Result<(QueryOutcome, BlockHeight), WorkerError> {
        self.chain_write(chain_id, move |mut guard| async move {
            guard.query_application(query, block_hash).await
        })
        .await
    }

    #[instrument(level = "trace", skip(self, chain_id, application_id), fields(
        nickname = %self.nickname(),
        chain_id = %chain_id,
        application_id = %application_id
    ))]
    pub async fn describe_application(
        &self,
        chain_id: ChainId,
        application_id: ApplicationId,
    ) -> Result<ApplicationDescription, WorkerError> {
        let state = self.get_or_create_chain_worker(chain_id).await?;
        let guard = handle::read_lock_initialized(&state).await?;
        guard.describe_application_readonly(application_id).await
    }

    /// Processes a confirmed block (aka a commit).
    #[instrument(
        level = "trace",
        skip(self, certificate, notify_when_messages_are_delivered),
        fields(
            nickname = %self.nickname(),
            chain_id = %certificate.block().header.chain_id,
            block_height = %certificate.block().header.height
        )
    )]
    async fn process_confirmed_block(
        &self,
        certificate: ConfirmedBlockCertificate,
        mode: ProcessConfirmedBlockMode,
        notify_when_messages_are_delivered: Option<oneshot::Sender<()>>,
    ) -> Result<(ChainInfoResponse, NetworkActions, BlockOutcome), WorkerError> {
        let chain_id = certificate.block().header.chain_id;
        self.chain_write(chain_id, move |mut guard| async move {
            guard
                .process_confirmed_block(certificate, mode, notify_when_messages_are_delivered)
                .await
        })
        .await
    }

    /// Processes a validated block issued from a multi-owner chain.
    #[instrument(level = "trace", skip(self, certificate), fields(
        nickname = %self.nickname(),
        chain_id = %certificate.block().header.chain_id,
        block_height = %certificate.block().header.height
    ))]
    async fn process_validated_block(
        &self,
        certificate: ValidatedBlockCertificate,
    ) -> Result<(ChainInfoResponse, NetworkActions, BlockOutcome), WorkerError> {
        let chain_id = certificate.block().header.chain_id;
        self.chain_write(chain_id, move |mut guard| async move {
            guard.process_validated_block(certificate).await
        })
        .await
    }

    /// Processes a leader timeout issued from a multi-owner chain.
    #[instrument(level = "trace", skip(self, certificate), fields(
        nickname = %self.nickname(),
        chain_id = %certificate.value().chain_id(),
        height = %certificate.value().height()
    ))]
    async fn process_timeout(
        &self,
        certificate: TimeoutCertificate,
    ) -> Result<(ChainInfoResponse, NetworkActions), WorkerError> {
        let chain_id = certificate.value().chain_id();
        self.chain_write(chain_id, move |mut guard| async move {
            guard.process_timeout(certificate).await
        })
        .await
    }

    /// Enqueues a cross-chain update request and cooperatively drives the
    /// batch-processing future until the result is ready.
    #[instrument(level = "trace", skip(self, origin, recipient, bundles), fields(
        nickname = %self.nickname(),
        origin = %origin,
        recipient = %recipient,
        num_bundles = %bundles.len()
    ))]
    async fn process_cross_chain_update(
        &self,
        origin: ChainId,
        recipient: ChainId,
        bundles: Vec<(Epoch, MessageBundle)>,
        previous_height: Option<BlockHeight>,
    ) -> Result<CrossChainUpdateResult, WorkerError> {
        let (result_sender, receiver) = oneshot::channel();
        let request = BatchRequest::Update {
            origin,
            bundles,
            previous_height,
            result_sender,
        };
        self.enqueue_and_drive(recipient, request, receiver).await
    }

    /// Enqueues a confirmation request and cooperatively drives the
    /// batch-processing future until the result is ready.
    async fn confirm_updated_recipient(
        &self,
        sender: ChainId,
        recipient: ChainId,
        latest_height: BlockHeight,
    ) -> Result<NetworkActions, WorkerError> {
        let (result_sender, receiver) = oneshot::channel();
        let request = BatchRequest::Confirm {
            recipient,
            latest_height,
            result_sender,
        };
        self.enqueue_and_drive(sender, request, receiver).await
    }

    /// Sends a [`BatchRequest`] to the per-chain driver and cooperatively
    /// polls the shared processing future until our result arrives.
    async fn enqueue_and_drive<R>(
        &self,
        chain_id: ChainId,
        request: BatchRequest,
        mut receiver: oneshot::Receiver<Result<R, WorkerError>>,
    ) -> Result<R, WorkerError> {
        let mut pending = Some(request);
        loop {
            let (sender, future) = self.get_or_create_chain_batch(chain_id).await?;
            if let Some(request) = pending.take() {
                if let Err(mpsc::error::SendError(request)) = sender.send(request) {
                    pending = Some(request);
                    continue; // Driver died; retry will create a new one.
                }
            }
            // Poll the receiver first (biased): if the result is already available,
            // return it immediately without driving the batch future further.
            match future::select(pin::pin!(&mut receiver), future).await {
                Either::Left((result, _)) => {
                    return result.expect("batch result sender dropped");
                }
                Either::Right(((), _)) => match receiver.try_recv() {
                    Ok(result) => return result,
                    Err(oneshot::error::TryRecvError::Empty) => {}
                    Err(oneshot::error::TryRecvError::Closed) => {
                        return Err(WorkerError::ChainError(Box::new(
                            ChainError::InternalError("batch driver stopped".into()),
                        )));
                    }
                },
            }
        }
    }

    /// Returns a stored [`ConfirmedBlockCertificate`] for a chain's block.
    #[instrument(level = "trace", skip(self, chain_id, height), fields(
        nickname = %self.nickname(),
        chain_id = %chain_id,
        height = %height
    ))]
    #[cfg(with_testing)]
    pub async fn read_certificate(
        &self,
        chain_id: ChainId,
        height: BlockHeight,
    ) -> Result<Option<CacheArc<ConfirmedBlockCertificate>>, WorkerError> {
        let state = self.get_or_create_chain_worker(chain_id).await?;
        let guard = handle::read_lock_initialized(&state).await?;
        guard.read_certificate(height).await
    }

    /// Test helper that runs `ChainWorkerState::select_message_bundles` for the given
    /// recipient chain.
    #[cfg(with_testing)]
    pub async fn select_message_bundles(
        &self,
        recipient: ChainId,
        origin: &ChainId,
        next_height_to_receive: BlockHeight,
        last_anticipated_block_height: Option<BlockHeight>,
        bundles: Vec<(Epoch, MessageBundle)>,
    ) -> Result<Vec<MessageBundle>, WorkerError> {
        let state = self.get_or_create_chain_worker(recipient).await?;
        let guard = handle::read_lock(&state).await?;
        guard
            .select_message_bundles(
                origin,
                next_height_to_receive,
                last_anticipated_block_height,
                bundles,
            )
            .await
    }

    /// Returns a read-only view of the [`ChainStateView`] of a chain referenced by its
    /// [`ChainId`].
    ///
    /// The returned guard holds a read lock on the chain state, preventing writes for
    /// its lifetime. Multiple concurrent readers are allowed.
    #[instrument(level = "trace", skip(self), fields(
        nickname = %self.nickname(),
        chain_id = %chain_id
    ))]
    pub async fn chain_state_view(
        &self,
        chain_id: ChainId,
    ) -> Result<ChainStateViewReadGuard<StorageClient>, WorkerError> {
        let state = self.get_or_create_chain_worker(chain_id).await?;
        let guard = handle::read_lock(&state).await?;
        Ok(ChainStateViewReadGuard(OwnedRwLockReadGuard::map(
            guard,
            |s| s.chain(),
        )))
    }

    #[instrument(skip_all, fields(
        nick = self.nickname(),
        chain_id = format!("{:.8}", proposal.content.block.chain_id),
        height = %proposal.content.block.height,
    ))]
    pub async fn handle_block_proposal(
        &self,
        proposal: BlockProposal,
    ) -> (Result<ChainInfoResponse, WorkerError>, NetworkActions) {
        trace!("{} <-- {:?}", self.nickname(), proposal);
        #[cfg(with_metrics)]
        let round = proposal.content.round;

        let chain_id = proposal.content.block.chain_id;
        // Delay if block timestamp is in the future but within grace period.
        let now = self.storage.clock().current_time();
        let block_timestamp = proposal.content.block.timestamp;
        let delta = block_timestamp.delta_since(now);
        let grace_period = TimeDelta::from_micros(
            u64::try_from(self.chain_worker_config.block_time_grace_period.as_micros())
                .unwrap_or(u64::MAX),
        );
        if delta > TimeDelta::ZERO && delta <= grace_period {
            self.storage.clock().sleep_until(block_timestamp).await;
        }

        let outcome = self
            .chain_write(chain_id, move |mut guard| async move {
                Ok::<_, WorkerError>(guard.handle_block_proposal(proposal).await)
            })
            .await;
        let (result, actions) = match outcome {
            Ok((result, actions)) => (result, actions),
            Err(err) => (Err(err), NetworkActions::default()),
        };
        #[cfg(with_metrics)]
        if result.is_ok() {
            metrics::NUM_ROUNDS_IN_BLOCK_PROPOSAL
                .with_label_values(&[round.type_name()])
                .observe(round.number() as f64);
        }
        (result, actions)
    }

    /// Processes a certificate, e.g. to extend a chain with a confirmed block.
    // Other fields will be included in the caller's span.
    #[instrument(skip_all, fields(
        chain_id = %certificate.value.chain_id,
        hash = %certificate.value.value_hash,
    ))]
    pub async fn handle_lite_certificate(
        &self,
        certificate: LiteCertificate<'_>,
        notify_when_messages_are_delivered: Option<oneshot::Sender<()>>,
    ) -> Result<(ChainInfoResponse, NetworkActions), WorkerError> {
        match self.full_certificate(certificate).await? {
            Either::Left(confirmed) => {
                Box::pin(self.handle_confirmed_certificate(
                    confirmed,
                    ProcessConfirmedBlockMode::Auto,
                    notify_when_messages_are_delivered,
                ))
                .await
            }
            Either::Right(validated) => {
                if let Some(notifier) = notify_when_messages_are_delivered {
                    // Nothing to wait for.
                    if let Err(()) = notifier.send(()) {
                        debug!("Failed to notify message delivery to caller (validation cert)");
                    }
                }
                Box::pin(self.handle_validated_certificate(validated)).await
            }
        }
    }

    /// Processes a confirmed block certificate.
    #[instrument(skip_all, fields(
        nick = self.nickname(),
        chain_id = format!("{:.8}", certificate.block().header.chain_id),
        height = %certificate.block().header.height,
    ))]
    pub async fn handle_confirmed_certificate(
        &self,
        certificate: ConfirmedBlockCertificate,
        mode: ProcessConfirmedBlockMode,
        notify_when_messages_are_delivered: Option<oneshot::Sender<()>>,
    ) -> Result<(ChainInfoResponse, NetworkActions), WorkerError> {
        trace!("{} <-- {:?}", self.nickname(), certificate);
        #[cfg(with_metrics)]
        let metrics_data = metrics::MetricsData::new(&certificate);

        #[allow(unused_variables)]
        let (info, actions, outcome) = Box::pin(self.process_confirmed_block(
            certificate,
            mode,
            notify_when_messages_are_delivered,
        ))
        .await?;

        #[cfg(with_metrics)]
        if matches!(outcome, BlockOutcome::Processed) {
            metrics_data.record();
        }
        Ok((info, actions))
    }

    /// Processes a validated block certificate.
    #[instrument(skip_all, fields(
        nick = self.nickname(),
        chain_id = format!("{:.8}", certificate.block().header.chain_id),
        height = %certificate.block().header.height,
    ))]
    pub async fn handle_validated_certificate(
        &self,
        certificate: ValidatedBlockCertificate,
    ) -> Result<(ChainInfoResponse, NetworkActions), WorkerError> {
        trace!("{} <-- {:?}", self.nickname(), certificate);

        #[cfg(with_metrics)]
        let round = certificate.round;
        #[cfg(with_metrics)]
        let cert_str = certificate.inner().to_log_str();

        #[allow(unused_variables)]
        let (info, actions, outcome) = Box::pin(self.process_validated_block(certificate)).await?;
        #[cfg(with_metrics)]
        {
            if matches!(outcome, BlockOutcome::Processed) {
                metrics::NUM_ROUNDS_IN_CERTIFICATE
                    .with_label_values(&[cert_str, round.type_name()])
                    .observe(round.number() as f64);
            }
        }
        Ok((info, actions))
    }

    /// Processes a timeout certificate
    #[instrument(skip_all, fields(
        nick = self.nickname(),
        chain_id = format!("{:.8}", certificate.inner().chain_id()),
        height = %certificate.inner().height(),
    ))]
    pub async fn handle_timeout_certificate(
        &self,
        certificate: TimeoutCertificate,
    ) -> Result<(ChainInfoResponse, NetworkActions), WorkerError> {
        trace!("{} <-- {:?}", self.nickname(), certificate);
        self.process_timeout(certificate).await
    }

    #[instrument(skip_all, fields(
        nick = self.nickname(),
        chain_id = format!("{:.8}", query.chain_id)
    ))]
    pub async fn handle_chain_info_query(
        &self,
        query: ChainInfoQuery,
    ) -> Result<ChainInfoResponse, WorkerError> {
        trace!("{} <-- {:?}", self.nickname(), query);
        #[cfg(with_metrics)]
        metrics::CHAIN_INFO_QUERIES.inc();
        let chain_id = query.chain_id;
        let result = self
            .chain_write(chain_id, move |mut guard| async move {
                guard.handle_chain_info_query(query).await
            })
            .await;
        trace!("{} --> {:?}", self.nickname(), result);
        result
    }

    #[instrument(skip_all, fields(
        nick = self.nickname(),
        chain_id = format!("{:.8}", chain_id)
    ))]
    pub async fn download_pending_blob(
        &self,
        chain_id: ChainId,
        blob_id: BlobId,
    ) -> Result<CacheArc<Blob>, WorkerError> {
        trace!("{} <-- download_pending_blob({blob_id:8})", self.nickname());
        let result = self
            .chain_read(chain_id, |guard| async move {
                guard.download_pending_blob(blob_id).await
            })
            .await;
        trace!(
            "{} --> {:?}",
            self.nickname(),
            result.as_ref().map(|_| blob_id)
        );
        result
    }

    #[instrument(skip_all, fields(
        nick = self.nickname(),
        chain_id = format!("{:.8}", chain_id)
    ))]
    pub async fn handle_pending_blob(
        &self,
        chain_id: ChainId,
        blob: Blob,
    ) -> Result<ChainInfoResponse, WorkerError> {
        let blob_id = blob.id();
        trace!("{} <-- handle_pending_blob({blob_id:8})", self.nickname());
        let result = self
            .chain_write(chain_id, move |mut guard| async move {
                guard.handle_pending_blob(blob).await
            })
            .await;
        trace!(
            "{} --> {:?}",
            self.nickname(),
            result.as_ref().map(|_| blob_id)
        );
        result
    }

    #[instrument(skip_all, fields(
        nick = self.nickname(),
        chain_id = format!("{:.8}", request.target_chain_id())
    ))]
    pub async fn handle_cross_chain_request(
        &self,
        request: CrossChainRequest,
    ) -> Result<NetworkActions, WorkerError> {
        trace!("{} <-- {:?}", self.nickname(), request);
        match request {
            CrossChainRequest::UpdateRecipient {
                sender,
                recipient,
                bundles,
                previous_height,
            } => {
                let mut actions = NetworkActions::default();
                let origin = sender;
                match self
                    .process_cross_chain_update(origin, recipient, bundles, previous_height)
                    .await?
                {
                    CrossChainUpdateResult::NothingToDo => {}
                    CrossChainUpdateResult::Updated(height) => {
                        actions.notifications.push(Notification {
                            chain_id: recipient,
                            reason: Reason::NewIncomingBundle { origin, height },
                        });
                        actions.cross_chain_requests.push(
                            CrossChainRequest::ConfirmUpdatedRecipient {
                                sender,
                                recipient,
                                latest_height: height,
                            },
                        );
                    }
                    CrossChainUpdateResult::GapDetected {
                        origin,
                        retransmit_from,
                    } => {
                        actions
                            .cross_chain_requests
                            .push(CrossChainRequest::RevertConfirm {
                                sender: origin,
                                recipient,
                                retransmit_from,
                            });
                    }
                }
                Ok(actions)
            }
            CrossChainRequest::ConfirmUpdatedRecipient {
                sender,
                recipient,
                latest_height,
            } => {
                let actions = self
                    .confirm_updated_recipient(sender, recipient, latest_height)
                    .await?;
                Ok(actions)
            }
            CrossChainRequest::RevertConfirm {
                sender,
                recipient,
                retransmit_from,
            } => {
                self.chain_write(sender, move |mut guard| async move {
                    guard
                        .handle_revert_confirm(recipient, retransmit_from)
                        .await
                })
                .await
            }
        }
    }

    /// Updates the received certificate trackers to at least the given values.
    #[instrument(skip_all, fields(
        nickname = %self.nickname(),
        chain_id = %chain_id,
        num_trackers = %new_trackers.len()
    ))]
    pub async fn update_received_certificate_trackers(
        &self,
        chain_id: ChainId,
        new_trackers: BTreeMap<ValidatorPublicKey, u64>,
    ) -> Result<(), WorkerError> {
        self.chain_write(chain_id, move |mut guard| async move {
            guard
                .update_received_certificate_trackers(new_trackers)
                .await
        })
        .await
    }

    /// Gets preprocessed block hashes in a given height range.
    #[instrument(skip_all, fields(
        nickname = %self.nickname(),
        chain_id = %chain_id,
        start = %start,
        end = %end
    ))]
    pub async fn get_preprocessed_block_hashes(
        &self,
        chain_id: ChainId,
        start: BlockHeight,
        end: BlockHeight,
    ) -> Result<Vec<CryptoHash>, WorkerError> {
        self.chain_read(chain_id, |guard| async move {
            guard.get_preprocessed_block_hashes(start, end).await
        })
        .await
    }

    /// Gets the next block height to receive from an inbox.
    #[instrument(skip_all, fields(
        nickname = %self.nickname(),
        chain_id = %chain_id,
        origin = %origin
    ))]
    pub async fn get_inbox_next_height(
        &self,
        chain_id: ChainId,
        origin: ChainId,
    ) -> Result<BlockHeight, WorkerError> {
        self.chain_read(chain_id, |guard| async move {
            guard.get_inbox_next_height(origin).await
        })
        .await
    }

    /// Gets locking blobs for specific blob IDs.
    /// Returns `Ok(None)` if any of the blobs is not found.
    #[instrument(skip_all, fields(
        nickname = %self.nickname(),
        chain_id = %chain_id,
        num_blob_ids = %blob_ids.len()
    ))]
    pub async fn get_locking_blobs(
        &self,
        chain_id: ChainId,
        blob_ids: Vec<BlobId>,
    ) -> Result<Option<Vec<Blob>>, WorkerError> {
        self.chain_read(chain_id, |guard| async move {
            guard.get_locking_blobs(blob_ids).await
        })
        .await
    }

    /// Gets block hashes for the given heights.
    pub async fn get_block_hashes(
        &self,
        chain_id: ChainId,
        heights: Vec<BlockHeight>,
    ) -> Result<Vec<CryptoHash>, WorkerError> {
        self.chain_read(chain_id, |guard| async move {
            guard.get_block_hashes(heights).await
        })
        .await
    }

    /// Gets proposed blobs from the manager for specified blob IDs.
    pub async fn get_proposed_blobs(
        &self,
        chain_id: ChainId,
        blob_ids: Vec<BlobId>,
    ) -> Result<Vec<Blob>, WorkerError> {
        self.chain_read(chain_id, |guard| async move {
            guard.get_proposed_blobs(blob_ids).await
        })
        .await
    }

    /// Gets event subscriptions from the chain.
    pub async fn get_event_subscriptions(
        &self,
        chain_id: ChainId,
    ) -> Result<EventSubscriptionsResult, WorkerError> {
        self.chain_read(chain_id, |guard| async move {
            guard.get_event_subscriptions().await
        })
        .await
    }

    /// Gets the next expected event index for a stream.
    pub async fn get_next_expected_event(
        &self,
        chain_id: ChainId,
        stream_id: StreamId,
    ) -> Result<Option<u32>, WorkerError> {
        self.chain_read(chain_id, |guard| async move {
            guard.get_next_expected_event(stream_id).await
        })
        .await
    }

    /// Gets the `next_expected_events` indices for the given streams.
    pub async fn next_expected_events(
        &self,
        chain_id: ChainId,
        stream_ids: Vec<StreamId>,
    ) -> Result<BTreeMap<StreamId, u32>, WorkerError> {
        self.chain_read(chain_id, |guard| async move {
            guard.get_next_expected_events(stream_ids).await
        })
        .await
    }

    /// Gets received certificate trackers.
    pub async fn get_received_certificate_trackers(
        &self,
        chain_id: ChainId,
    ) -> Result<HashMap<ValidatorPublicKey, u64>, WorkerError> {
        self.chain_read(chain_id, |guard| async move {
            guard.get_received_certificate_trackers().await
        })
        .await
    }

    /// Returns the pending cross-chain network actions for this chain without
    /// initializing its execution state. Safe to call on chains whose
    /// `ChainDescription` blob is not available locally.
    pub async fn cross_chain_network_actions(
        &self,
        chain_id: ChainId,
    ) -> Result<NetworkActions, WorkerError> {
        // Fast path: when the outbox index is already reconciled to the current tracked set,
        // the network actions are built from a read-only view, so a shared lock with no save
        // suffices — avoiding the write lock, task spawn and `save()` of the slow path.
        if let Some(actions) = self
            .chain_read(chain_id, |guard| async move {
                guard.cross_chain_network_actions_if_reconciled().await
            })
            .await?
        {
            return Ok(actions);
        }
        // Slow path (first load after migration, or the tracked set changed): reconcile and
        // persist the index under an exclusive lock before building.
        self.chain_write(chain_id, |mut guard| async move {
            guard.reconcile_and_cross_chain_network_actions().await
        })
        .await
    }

    /// Gets tip state and outbox info for next_outbox_heights calculation.
    pub async fn get_tip_state_and_outbox_info(
        &self,
        chain_id: ChainId,
        receiver_id: ChainId,
    ) -> Result<(BlockHeight, Option<BlockHeight>), WorkerError> {
        self.chain_read(chain_id, |guard| async move {
            guard.get_tip_state_and_outbox_info(receiver_id).await
        })
        .await
    }

    /// Gets the next height to preprocess.
    pub async fn get_next_height_to_preprocess(
        &self,
        chain_id: ChainId,
    ) -> Result<BlockHeight, WorkerError> {
        self.chain_read(chain_id, |guard| async move {
            Ok(guard.get_next_height_to_preprocess())
        })
        .await
    }
}

#[cfg(with_testing)]
impl<StorageClient> WorkerState<StorageClient>
where
    StorageClient: Storage + Clone + 'static,
{
    /// Gets a reference to the validator's [`ValidatorPublicKey`].
    ///
    /// # Panics
    ///
    /// If the validator doesn't have a key pair assigned to it.
    #[instrument(level = "trace", skip(self))]
    pub fn public_key(&self) -> ValidatorPublicKey {
        self.chain_worker_config
            .key_pair()
            .expect(
                "Test validator should have a key pair assigned to it \
                in order to obtain its public key",
            )
            .public()
    }
}