newton-chainio 0.5.2

newton prover chainio
//! Owned, signer-wide nonce allocator.
//!
//! # Why this exists
//!
//! Alloy's default [`CachedNonceManager`] allocates a nonce *inside* the filler
//! stack's `prepare` step — concurrently (via `try_join!`) with the gas filler's
//! `eth_estimateGas`. The increment is committed to its internal cache before the
//! transaction is ever broadcast, and is **never rolled back** if a sibling filler
//! (gas-estimate revert) or the subsequent `send` fails. A failed broadcast
//! therefore silently burns a cached nonce: the slot is consumed in the cache but
//! nothing lands on-chain and no receipt tracker pursues it — a permanent
//! "blocked-from-below" gap that wedges every higher nonce until the process
//! restarts (only the startup sweep heals it).
//!
//! This allocator moves nonce assignment *out* of the filler stack. Callers
//! reserve a nonce explicitly and set it on the request via `.nonce(N)` before
//! `fill()`; alloy's [`NonceFiller::status`] then reports `Finished` (it only
//! fills when the nonce is absent), so the [`CachedNonceManager`] never allocates.
//!
//! # Invariant
//!
//! **Every nonce handed out is eventually consumed on-chain (mined or cancelled)
//! and never handed out twice concurrently.** The allocator owns only the "handed
//! out / not handed out twice" half; the caller drives a stuck nonce to on-chain
//! resolution (confirm-or-cancel) for the "eventually consumed" half.
//!
//! The counter is **strictly monotonic** — `reserve` only ever increments. The
//! single exception is [`release`](Self::release): a *LIFO un-bump* that reclaims
//! the most-recent reservation **iff nothing newer was reserved after it**. This
//! is the only nonce that is provably safe to reuse — there is no committed nonce
//! above it, so the chain cannot have advanced past it. There is no free-list of
//! arbitrary holes (which is what made stale-reuse possible); a release that is
//! *not* the latest reservation returns `false`, and the caller resolves that
//! nonce on-chain instead. This removes the entire class of "is this cached hole
//! still valid?" races by construction.
//!
//! # Single source of truth
//!
//! The signer is shared between the batch submitter and the state-root committer
//! (`commit_state_root`). A per-path allocator would race the other path's
//! implicit filler allocation, so this allocator is injected as a shared
//! `Arc<NonceAllocator>` and used by **every** broadcast on that signer. It is the
//! one coordinator for the address's nonce sequence.
//!
//! [`CachedNonceManager`]: alloy::providers::fillers::CachedNonceManager
//! [`NonceFiller::status`]: alloy::providers::fillers::NonceFiller

use crate::error::ChainIoError;
use alloy::{
    primitives::Address,
    providers::{Provider, WalletProvider},
};
use eigensdk::common::SdkSigner;
use tokio::sync::Mutex;
use tracing::{info, warn};

/// Hands out a strictly-monotonic per-signer nonce sequence, seeded lazily from
/// the chain's pending transaction count.
///
/// The only mutable state is `next` (the next nonce to hand out), behind a single
/// async `Mutex` so concurrent pipeline slots and the state-root committer can
/// never be handed the same nonce. `None` until lazily seeded on the first
/// [`reserve`](Self::reserve) (so construction performs no I/O).
#[derive(Debug)]
pub struct NonceAllocator {
    provider: SdkSigner,
    address: Address,
    next: Mutex<Option<u64>>,
}

impl NonceAllocator {
    /// Create an allocator for the provider's default signer address. The counter
    /// is seeded lazily on the first [`reserve`](Self::reserve).
    pub fn new(provider: SdkSigner) -> Self {
        let address = provider.default_signer_address();
        Self {
            provider,
            address,
            next: Mutex::new(None),
        }
    }

    /// The signer address this allocator manages.
    pub fn address(&self) -> Address {
        self.address
    }

    /// Reserve the next nonce, seeding from the chain's *pending* count on first
    /// use. Strictly monotonic: each call hands out a distinct, increasing value,
    /// so two concurrent callers can never collide.
    ///
    /// Seeding from `pending` (not `latest`) starts us *after* any txs the signer
    /// already has in flight (a prior process, or txs sent before this allocator
    /// was installed), avoiding an immediate "nonce too low".
    pub async fn reserve(&self) -> Result<u64, ChainIoError> {
        let mut guard = self.next.lock().await;
        let nonce = match *guard {
            Some(n) => n,
            None => {
                let seed = self.pending_count().await?;
                info!(address = %self.address, seed, "nonce allocator seeded from pending count");
                seed
            }
        };
        *guard = Some(nonce + 1);
        Ok(nonce)
    }

    /// LIFO un-bump: reclaim `nonce` **iff it was the most recent reservation**
    /// (`next == nonce + 1`), returning `true`. Otherwise returns `false` — a
    /// newer nonce has already been handed out, so `nonce` is now a genuine gap
    /// below committed work and the caller MUST resolve it on-chain (cancel)
    /// rather than reuse it.
    ///
    /// Call this only when the broadcast at `nonce` **never reached the chain**
    /// (the `fill` gas-estimate reverted — nothing was signed or sent). Reclaiming
    /// the latest reservation is provably safe: with nothing committed above it,
    /// the chain cannot have advanced past `nonce`, so the next `reserve` re-hands
    /// the same value with no gap and no on-chain cost. A non-latest release can
    /// never be safely reused (the holes between it and `next` are live), which is
    /// exactly why this returns `false` instead of caching it.
    pub async fn release(&self, nonce: u64) -> bool {
        let mut guard = self.next.lock().await;
        if *guard == Some(nonce + 1) {
            *guard = Some(nonce);
            true
        } else {
            warn!(
                address = %self.address,
                nonce,
                next = ?*guard,
                "nonce allocator: release is not the latest reservation; caller must resolve nonce on-chain"
            );
            false
        }
    }

    /// Re-seed the counter from the chain's pending count, never moving it
    /// backwards past an in-flight reservation. Call this when a slot is detected
    /// stuck so a transient desync — or a nonce burned before this allocator owned
    /// the sequence — self-heals.
    ///
    /// `max(prev, pending)`: an in-flight broadcast may sit at or above `pending`
    /// (its tx is in the mempool, counted by `pending`, while a concurrent
    /// `reserve` advanced our counter further), so taking the max guarantees a
    /// reserved-but-unmined nonce is never re-handed-out.
    pub async fn resync(&self) -> Result<u64, ChainIoError> {
        let pending = self.pending_count().await?;
        let mut guard = self.next.lock().await;
        let resynced = match *guard {
            Some(prev) => prev.max(pending),
            None => pending,
        };
        if *guard != Some(resynced) {
            warn!(address = %self.address, prev = ?*guard, resynced, "nonce allocator resynced to chain");
        }
        *guard = Some(resynced);
        Ok(resynced)
    }

    /// Current pending transaction count for the signer (mined + mempool).
    async fn pending_count(&self) -> Result<u64, ChainIoError> {
        self.provider
            .get_transaction_count(self.address)
            .pending()
            .await
            .map_err(ChainIoError::RpcError)
    }
}

#[cfg(test)]
impl NonceAllocator {
    /// Test-only constructor that pre-seeds the counter so the REAL
    /// `reserve`/`release`/`resync` logic can be exercised without a live provider
    /// (the only thing a provider is needed for is the lazy seed + resync RPC).
    /// `resync` is covered by integration tests against an anvil node.
    pub(crate) fn seeded_for_test(provider: SdkSigner, next: u64) -> Self {
        let address = provider.default_signer_address();
        Self {
            provider,
            address,
            next: Mutex::new(Some(next)),
        }
    }

    /// Current counter value (test inspection).
    pub(crate) async fn peek_next(&self) -> Option<u64> {
        *self.next.lock().await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use eigensdk::common::get_signer;
    use std::sync::Arc;

    // A throwaway signer/provider; these tests never hit the network (the counter
    // is pre-seeded via `seeded_for_test`, so no RPC is performed). The key is a
    // well-known anvil dev key.
    fn test_alloc(next: u64) -> NonceAllocator {
        let signer = get_signer(
            "ac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80",
            "http://127.0.0.1:8545",
        );
        NonceAllocator::seeded_for_test(signer, next)
    }

    #[tokio::test]
    async fn reservations_are_monotonic() {
        let a = test_alloc(100);
        let mut handed = Vec::new();
        for _ in 0..5 {
            handed.push(a.reserve().await.unwrap());
        }
        assert_eq!(handed, vec![100, 101, 102, 103, 104]);
        assert_eq!(a.peek_next().await, Some(105));
    }

    #[tokio::test]
    async fn release_of_latest_reservation_reclaims() {
        let a = test_alloc(100);
        let n = a.reserve().await.unwrap(); // 100, next=101
        assert_eq!(n, 100);
        // Broadcast at 100 failed during fill, and it's the latest → reclaim.
        assert!(a.release(n).await, "latest reservation must be reclaimable");
        assert_eq!(a.peek_next().await, Some(100));
        // Next reservation re-hands the same nonce — gap-free, no on-chain cost.
        assert_eq!(a.reserve().await.unwrap(), 100);
    }

    #[tokio::test]
    async fn release_of_non_latest_is_rejected() {
        let a = test_alloc(100);
        let lower = a.reserve().await.unwrap(); // 100
        let _higher = a.reserve().await.unwrap(); // 101, next=102
                                                  // 100 is NOT the latest (101 was reserved after) → cannot reclaim; the
                                                  // caller must resolve 100 on-chain. Counter is untouched.
        assert!(!a.release(lower).await, "non-latest release must return false");
        assert_eq!(a.peek_next().await, Some(102));
        // The next fresh reservation is still 102 — 100 is not re-handed-out.
        assert_eq!(a.reserve().await.unwrap(), 102);
    }

    #[tokio::test]
    async fn repeated_latest_release_reclaims_each_time() {
        // Reserve→release→reserve→release of the tail keeps reusing the same slot.
        let a = test_alloc(50);
        for _ in 0..3 {
            let n = a.reserve().await.unwrap();
            assert_eq!(n, 50);
            assert!(a.release(n).await);
            assert_eq!(a.peek_next().await, Some(50));
        }
    }

    #[tokio::test]
    async fn concurrent_reservations_never_collide() {
        let a = Arc::new(test_alloc(0));
        let mut tasks = Vec::new();
        for _ in 0..64 {
            let a = a.clone();
            tasks.push(tokio::spawn(async move { a.reserve().await.unwrap() }));
        }
        let mut got = Vec::new();
        for t in tasks {
            got.push(t.await.unwrap());
        }
        got.sort_unstable();
        assert_eq!(
            got,
            (0..64).collect::<Vec<u64>>(),
            "every nonce handed out exactly once"
        );
    }
}