use crate::error::ChainIoError;
use alloy::{
primitives::Address,
providers::{Provider, WalletProvider},
};
use eigensdk::common::SdkSigner;
use tokio::sync::Mutex;
use tracing::{info, warn};
#[derive(Debug)]
pub struct NonceAllocator {
provider: SdkSigner,
address: Address,
next: Mutex<Option<u64>>,
}
impl NonceAllocator {
pub fn new(provider: SdkSigner) -> Self {
let address = provider.default_signer_address();
Self {
provider,
address,
next: Mutex::new(None),
}
}
pub fn address(&self) -> Address {
self.address
}
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)
}
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
}
}
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)
}
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 {
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)),
}
}
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;
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(); assert_eq!(n, 100);
assert!(a.release(n).await, "latest reservation must be reclaimable");
assert_eq!(a.peek_next().await, Some(100));
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(); let _higher = a.reserve().await.unwrap(); assert!(!a.release(lower).await, "non-latest release must return false");
assert_eq!(a.peek_next().await, Some(102));
assert_eq!(a.reserve().await.unwrap(), 102);
}
#[tokio::test]
async fn repeated_latest_release_reclaims_each_time() {
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"
);
}
}