use std::sync::{
Arc,
atomic::{AtomicBool, Ordering},
};
use nautilus_common::live::dst::time::{self, Duration, Instant};
use nautilus_live::{
book::{recovery::BookRecovery as Recovery, snapshot::SnapshotGate},
task::TaskSpawner,
};
use nautilus_model::identifiers::InstrumentId;
use parking_lot::Mutex;
use tokio_util::sync::CancellationToken;
use super::{CoinApiBookError, sync::BookSyncTracker};
use crate::{
data::OhlcvCache,
websocket::{client::CoinApiWebSocketClient, messages::Subscription},
};
pub(crate) type BookRecovery = Recovery<CoinApiBookError>;
pub(crate) fn spawn_recovery_monitor(
tracker: BookSyncTracker,
ws: Arc<CoinApiWebSocketClient>,
ohlcv: Arc<Mutex<OhlcvCache>>,
snapshot_timeout: Duration,
tasks: &TaskSpawner,
) {
let cancel = tasks.cancellation_token();
let spawner = tasks.clone();
let task = async move {
tokio::select! {
biased;
() = cancel.cancelled() => {}
() = time::sleep(snapshot_timeout) => {
for signal in tracker.take_expired_snapshots(Instant::now()) {
signal.log();
start_recovery(
signal.instrument_id,
&tracker,
&ws,
&ohlcv,
snapshot_timeout,
&spawner,
);
}
}
}
};
if let Err(e) = tasks.spawn(task) {
log::warn!("Skipping book snapshot monitor after shutdown began: {e}");
}
}
pub(crate) fn start_recovery(
instrument_id: InstrumentId,
tracker: &BookSyncTracker,
ws: &Arc<CoinApiWebSocketClient>,
ohlcv: &Arc<Mutex<OhlcvCache>>,
snapshot_timeout: Duration,
tasks: &TaskSpawner,
) {
if let Some((recovery, subscription)) = tracker.claim_recovery(instrument_id) {
spawn_recovery_task(
instrument_id,
subscription,
recovery,
Arc::clone(ws),
Arc::clone(ohlcv),
snapshot_timeout,
tasks,
);
}
}
pub(crate) fn spawn_recovery_task(
instrument_id: InstrumentId,
subscription: Subscription,
recovery: Arc<BookRecovery>,
ws: Arc<CoinApiWebSocketClient>,
ohlcv: Arc<Mutex<OhlcvCache>>,
snapshot_timeout: Duration,
tasks: &TaskSpawner,
) {
let shutdown = tasks.cancellation_token();
let recovery_guard = recovery.cancellation.clone().drop_guard();
let task = async move {
let _recovery_guard = recovery_guard;
tokio::select! {
biased;
() = shutdown.cancelled() => recovery.cancellation.cancel(),
() = recovery.run(
instrument_id,
snapshot_timeout,
|attempt_cancel, gate| restart(&ws, &ohlcv, subscription, attempt_cancel, gate),
CoinApiBookError::is_retryable,
CoinApiBookError::Permanent,
|| CoinApiBookError::Retryable("book snapshot deadline expired".to_string()),
) => {}
}
};
if let Err(e) = tasks.spawn(task) {
log::warn!("Skipping book recovery for {instrument_id} after shutdown began: {e}");
}
}
async fn restart(
ws: &CoinApiWebSocketClient,
ohlcv: &Mutex<OhlcvCache>,
subscription: Subscription,
cancel: CancellationToken,
gate: SnapshotGate,
) -> Result<(), CoinApiBookError> {
let discard = || {
ohlcv
.lock()
.retain(|(symbol_id, _), _| *symbol_id != subscription.symbol_id);
};
let paused = AtomicBool::new(false);
tokio::select! {
biased;
() = cancel.cancelled() => {
Err(CoinApiBookError::Retryable("book recovery attempt cancelled".to_string()))
}
result = ws.resubscribe_book(subscription, || {
discard();
paused.store(true, Ordering::Relaxed);
}, gate) => {
let writes = result?;
if paused.load(Ordering::Relaxed) {
discard();
}
if let Some(writes) = writes { writes.commit(); }
Ok(())
}
}
}