From 3a7ce3f306c1ce1e5f46abeb1759fdc18c744de1 Mon Sep 17 00:00:00 2001 From: moudyellaz Date: Thu, 30 Jul 2026 22:53:17 +0200 Subject: [PATCH] feat(sequencer): wait for background tasks on shutdown --- lez/sequencer/core/src/block_publisher.rs | 36 +++--- lez/sequencer/core/src/lib.rs | 20 ++++ lez/sequencer/core/src/task_group.rs | 137 ++++++++++++++++++++++ lez/sequencer/service/src/lib.rs | 121 ++++++++++++++++--- lez/sequencer/service/src/main.rs | 42 +++++-- test_fixtures/src/lib.rs | 2 +- 6 files changed, 319 insertions(+), 39 deletions(-) create mode 100644 lez/sequencer/core/src/task_group.rs diff --git a/lez/sequencer/core/src/block_publisher.rs b/lez/sequencer/core/src/block_publisher.rs index 86d9dc0c..e2b8610e 100644 --- a/lez/sequencer/core/src/block_publisher.rs +++ b/lez/sequencer/core/src/block_publisher.rs @@ -1,4 +1,4 @@ -use std::{sync::Arc, time::Duration}; +use std::time::Duration; use anyhow::{Context as _, Result, anyhow, ensure}; use common::block::Block; @@ -34,13 +34,10 @@ use logos_blockchain_zone_sdk::{ ZoneSequencer, }, }; -use tokio::{ - sync::{mpsc, oneshot, watch}, - task::JoinHandle, -}; +use tokio::sync::{mpsc, oneshot, watch}; use tokio_util::sync::CancellationToken; -use crate::config::BedrockConfig; +use crate::{config::BedrockConfig, task_group::TaskGroup}; /// Channel capacity for the publish inbox. One publish per produced block, drained /// in microseconds by the drive task — 32 is huge headroom and just provides @@ -120,6 +117,14 @@ pub trait BlockPublisherTrait: Sized { /// are processed past that point, so the node must halt. fn driver_cancellation(&self) -> CancellationToken; + /// The publisher's background tasks, for a caller that needs to know when + /// they have actually stopped. Its sinks capture a store handle, so the + /// `RocksDB` lock outlives the sequencer until the drive task is gone. + /// Empty by default, for publishers that run no tasks. + fn background_tasks(&self) -> TaskGroup { + TaskGroup::default() + } + /// Current channel frontier slot on the connected chain, or `None` if the /// channel does not exist there. Drives the startup frontier check. async fn channel_tip_slot(&self) -> Result>; @@ -143,19 +148,12 @@ pub struct ZoneSdkPublisher { turn_rx: watch::Receiver, // Cancelled when the drive task ends for any reason, including a panic. driver_cancellation: CancellationToken, - // Aborts the drive task when the last clone is dropped. - _drive_task: Arc, + // Stops the drive task when the last clone is dropped, and lets a shutdown + // path wait until it has actually stopped. + drive_task: TaskGroup, indexer: ZoneIndexer, } -struct DriveTaskGuard(JoinHandle<()>); - -impl Drop for DriveTaskGuard { - fn drop(&mut self) { - self.0.abort(); - } -} - impl BlockPublisherTrait for ZoneSdkPublisher { async fn new( config: &BedrockConfig, @@ -318,7 +316,7 @@ impl BlockPublisherTrait for ZoneSdkPublisher { command_tx, turn_rx, driver_cancellation, - _drive_task: Arc::new(DriveTaskGuard(drive_task)), + drive_task: TaskGroup::new(vec![drive_task]), }) } @@ -359,6 +357,10 @@ impl BlockPublisherTrait for ZoneSdkPublisher { self.driver_cancellation.clone() } + fn background_tasks(&self) -> TaskGroup { + self.drive_task.clone() + } + async fn channel_tip_slot(&self) -> Result> { Ok(self .node diff --git a/lez/sequencer/core/src/lib.rs b/lez/sequencer/core/src/lib.rs index 5ed0e50c..d61f1531 100644 --- a/lez/sequencer/core/src/lib.rs +++ b/lez/sequencer/core/src/lib.rs @@ -39,6 +39,7 @@ use storage::sequencer::{ use crate::{ block_publisher::{BlockPublisherTrait, MsgId, ZoneSdkPublisher}, block_store::SequencerStore, + task_group::{StoreRelease, TaskGroup}, }; pub mod block_publisher; @@ -48,6 +49,7 @@ pub mod cross_zone_watcher; #[cfg(feature = "mock")] pub mod mock; +pub mod task_group; /// The origin of a transaction. #[derive(Clone, Copy)] @@ -828,6 +830,24 @@ impl SequencerCore { &self.block_publisher } + /// A weak reference to this sequencer's store, for a shutdown path that + /// needs to observe the database actually closing rather than infer it. + #[must_use] + pub fn store_release(&self) -> StoreRelease { + StoreRelease::new(&self.store.dbio()) + } + + /// Every background task that holds this sequencer's store handle. + /// + /// Taken before the core is shared, so a shutdown path can wait for them + /// without owning the core. Until all of them have stopped the `RocksDB` + /// lock is still held and the home directory cannot be reopened, which is + /// what a restart does. + #[must_use] + pub fn background_tasks(&self) -> Vec { + vec![self.block_publisher.background_tasks()] + } + /// Whether this sequencer is currently authorized to write to the channel. #[must_use] pub fn is_our_turn(&self) -> bool { diff --git a/lez/sequencer/core/src/task_group.rs b/lez/sequencer/core/src/task_group.rs new file mode 100644 index 00000000..8572a62f --- /dev/null +++ b/lez/sequencer/core/src/task_group.rs @@ -0,0 +1,137 @@ +//! A set of background tasks that can be stopped and waited on. + +use std::sync::{Arc, Mutex, MutexGuard, PoisonError, Weak}; + +use log::warn; +use storage::sequencer::RocksDBIO; +use tokio::task::JoinHandle; + +/// Background tasks owned by one component, stoppable on demand and stopped +/// anyway when the last handle goes away. +/// +/// `JoinHandle::abort` only *requests* cancellation, and dropping a handle +/// detaches rather than cancels, so neither on its own says when a task has +/// actually stopped. That matters because these tasks hold a store handle: +/// until they are gone the `RocksDB` lock is still held and a restarting +/// sequencer cannot reopen its home directory. [`TaskGroup::shutdown`] is the +/// answer to "have they stopped yet"; the `Drop` below stays as the best-effort +/// path for panics and tests that never call it. +/// +/// Cloneable so the owner can keep it (tying task lifetime to its own) while a +/// shutdown path elsewhere holds a clone. +#[derive(Clone, Default)] +pub struct TaskGroup(Arc); + +#[derive(Default)] +struct TaskGroupInner(Mutex>>); + +/// A weak handle to the store, for observing when it is finally closed. +/// +/// Every strong reference lives inside a task or a server that shutdown stops, +/// but the last drop runs on whichever thread owned it, not on the one awaiting +/// shutdown. Watching the count is the difference between knowing the database +/// file is closed and assuming it from another crate's drop order. +pub struct StoreRelease(Weak); + +impl StoreRelease { + #[must_use] + pub fn new(store: &Arc) -> Self { + Self(Arc::downgrade(store)) + } + + /// How many holders are left. Zero means the store is closed. + #[must_use] + pub fn holders(&self) -> usize { + self.0.strong_count() + } +} + +impl Drop for TaskGroupInner { + fn drop(&mut self) { + for task in Self::take(&self.0) { + task.abort(); + } + } +} + +impl TaskGroupInner { + /// Empties the handle list, so a second shutdown (or a drop after one) is a + /// no-op rather than a second abort. + fn handles(handles: &Mutex>>) -> MutexGuard<'_, Vec>> { + handles.lock().unwrap_or_else(PoisonError::into_inner) + } + + fn take(handles: &Mutex>>) -> Vec> { + std::mem::take(&mut *Self::handles(handles)) + } +} + +impl TaskGroup { + /// Takes ownership of already-spawned tasks. + #[must_use] + pub fn new(handles: Vec>) -> Self { + Self(Arc::new(TaskGroupInner(Mutex::new(handles)))) + } + + /// Whether any task has ended on its own. + /// + /// These tasks run for the lifetime of the sequencer, so a finished one is a + /// task that panicked, and whatever it was doing is not happening any more. + #[must_use] + pub fn any_finished(&self) -> bool { + TaskGroupInner::handles(&self.0.0) + .iter() + .any(JoinHandle::is_finished) + } + + /// Stops every task and waits for it to finish. + /// + /// Returns only once the runtime has dropped each task's future, so whatever + /// they held (a store handle, a network client) is released by the time this + /// returns. Cancellation is the expected outcome, so it is not reported; a + /// panic is, since it means the task died on its own terms earlier. + pub async fn shutdown(&self) { + let handles = TaskGroupInner::take(&self.0.0); + for handle in handles { + handle.abort(); + if let Err(err) = handle.await + && err.is_panic() + { + warn!("Background task panicked before shutdown: {err}"); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn a_task_that_ends_on_its_own_is_visible() { + let group = TaskGroup::new(vec![tokio::spawn(async {})]); + // A watcher only ends by panicking, so "finished" is the signal that a + // peer's deliveries have stopped happening. + tokio::task::yield_now().await; + assert!(group.any_finished()); + + let running = TaskGroup::new(vec![tokio::spawn(std::future::pending())]); + assert!(!running.any_finished()); + } + + #[tokio::test] + async fn shutdown_ends_a_task_that_would_never_end_on_its_own() { + let group = TaskGroup::new(vec![tokio::spawn(std::future::pending())]); + + // The watchers and the drive task are infinite loops, so awaiting one + // without cancelling it first hangs here for ever. + tokio::time::timeout(std::time::Duration::from_secs(5), group.shutdown()) + .await + .expect("shutdown must not hang on a task that never finishes by itself"); + + // Shutting down twice is a no-op rather than a second abort. + tokio::time::timeout(std::time::Duration::from_secs(5), group.shutdown()) + .await + .expect("a second shutdown must return immediately"); + } +} diff --git a/lez/sequencer/service/src/lib.rs b/lez/sequencer/service/src/lib.rs index a8aa856e..ba5b68ec 100644 --- a/lez/sequencer/service/src/lib.rs +++ b/lez/sequencer/service/src/lib.rs @@ -12,7 +12,11 @@ use sequencer_core::SequencerCore; #[cfg(feature = "standalone")] use sequencer_core::SequencerCoreWithMockClients as SequencerCore; pub use sequencer_core::config::*; -use sequencer_core::{TransactionOrigin, block_publisher::BlockPublisherTrait as _}; +use sequencer_core::{ + TransactionOrigin, + block_publisher::BlockPublisherTrait as _, + task_group::{StoreRelease, TaskGroup}, +}; use sequencer_service_rpc::RpcServer as _; use tokio::{sync::Mutex, task::JoinHandle}; use tokio_util::sync::CancellationToken; @@ -26,12 +30,19 @@ const REQUEST_BODY_MAX_SIZE: ByteSize = ByteSize::mib(10); /// Implements `Drop` to ensure all tasks are aborted and the RPC server is stopped when dropped. pub struct SequencerHandle { addr: SocketAddr, - /// Option because of `Drop` which forbids to simply move out of `self` in `stopped()`. - server_handle: Option, + server_handle: ServerHandle, main_loop_handle: JoinHandle>, /// Cancelled when the publisher's drive task terminates (e.g. a panicked /// persist sink); no channel events are processed past that point. driver_cancellation: CancellationToken, + /// The core's background tasks, taken before the core was shared. This + /// handle owns no reference to the core itself, so without these there is + /// nothing to wait on: aborting the main loop only starts the teardown. + background_tasks: Vec, + /// The store, weakly. Every strong reference lives inside something this + /// handle stops, so watching the count go to zero is how shutdown knows the + /// database file is actually closed rather than assuming it from drop order. + store: StoreRelease, } impl SequencerHandle { @@ -40,29 +51,73 @@ impl SequencerHandle { server_handle: ServerHandle, main_loop_handle: JoinHandle>, driver_cancellation: CancellationToken, + background_tasks: Vec, + store: StoreRelease, ) -> Self { Self { addr, - server_handle: Some(server_handle), + server_handle, main_loop_handle, driver_cancellation, + background_tasks, + store, } } + /// Stops the sequencer and waits for every part of it to be gone. + /// + /// `Drop` alone cannot do this: it aborts the main loop without awaiting it, + /// and the core lives behind `Arc`s held by that task and the RPC server, so + /// after a plain drop the store is still open for an unbounded stretch. That + /// is why restarting a sequencer on the same home directory used to need a + /// sleep, and why an in-process restart could fail outright with a `RocksDB` + /// lock error. + /// + /// Order matters: the main loop stops first so nothing new is produced while + /// the publisher is torn down, then the background tasks that hold the store, + /// then the server. Consuming `self` drops the last references, so the store + /// is closed by the time this returns. + pub async fn shutdown(mut self) { + self.main_loop_handle.abort(); + if let Err(err) = (&mut self.main_loop_handle).await + && err.is_panic() + { + error!("Sequencer main loop panicked before shutdown: {err}"); + } + + for tasks in &self.background_tasks { + tasks.shutdown().await; + } + + if let Err(err) = self.server_handle.stop() { + error!("An error occurred while stopping Sequencer RPC server: {err}"); + } + self.server_handle.clone().stopped().await; + + // Nothing this handle owns holds the store, so waiting here rather than + // after the drop is the same thing, and it keeps the guarantee inside + // the call the caller awaits. + wait_for_store_release(&self.store).await; + } + /// Wait for any of the sequencer tasks to fail and return the error. #[expect( clippy::integer_division_remainder_used, reason = "Generated by select! macro, can't be easily rewritten to avoid this lint" )] - pub async fn failed(mut self) -> Result { + pub async fn failed(&mut self) -> Result { let Self { addr: _, server_handle, main_loop_handle, driver_cancellation, - } = &mut self; + background_tasks: _, + store: _, + } = self; - let server_handle = server_handle.take().expect("Server handle is set"); + // Cloned rather than taken: `stopped()` consumes a handle, and taking + // this one would leave `shutdown` with no way to stop the server. + let server_handle = server_handle.clone(); tokio::select! { () = server_handle.stopped() => { Err(anyhow!("RPC Server stopped")) @@ -89,11 +144,16 @@ impl SequencerHandle { server_handle, main_loop_handle, driver_cancellation, + background_tasks, + store: _, } = self; - let stopped = server_handle.as_ref().is_none_or(ServerHandle::is_stopped) + let stopped = server_handle.is_stopped() || main_loop_handle.is_finished() - || driver_cancellation.is_cancelled(); + || driver_cancellation.is_cancelled() + // A watcher only ends by panicking, and a peer whose deliveries have + // silently stopped is exactly what this predicate exists to catch. + || background_tasks.iter().any(TaskGroup::any_finished); !stopped } @@ -110,20 +170,46 @@ impl Drop for SequencerHandle { server_handle, main_loop_handle, driver_cancellation: _, + background_tasks: _, + store: _, } = self; main_loop_handle.abort(); - let Some(handle) = server_handle else { - return; - }; - - if let Err(err) = handle.stop() { + if let Err(err) = server_handle.stop() { error!("An error occurred while stopping Sequencer RPC server: {err}"); } } } +/// Waits until nothing holds the store any more. +/// +/// Everything that holds one lives inside a task or a server this handle has +/// already stopped, but the last drop happens on whichever thread ran them, not +/// on this one. Without this the caller can reopen the database a moment too +/// early and hit a `RocksDB` lock error, which is the kind of failure that shows +/// up as an occasional flake rather than a bug. +async fn wait_for_store_release(store: &StoreRelease) { + /// Long enough for a drop that is already in flight, short enough that a + /// leak is reported rather than hung on. + const RELEASE_TIMEOUT: Duration = Duration::from_secs(10); + const POLL: Duration = Duration::from_millis(10); + + let released = tokio::time::timeout(RELEASE_TIMEOUT, async { + while store.holders() > 0 { + tokio::time::sleep(POLL).await; + } + }) + .await; + + if released.is_err() { + error!( + "Sequencer store still held by {} reference(s) after shutdown; something outlived the tasks it should have died with", + store.holders() + ); + } +} + pub async fn run(config: SequencerConfig, listen_addr: SocketAddr) -> Result { let block_timeout = config.block_create_timeout; let max_block_size = config.max_block_size; @@ -134,6 +220,11 @@ pub async fn run(config: SequencerConfig, listen_addr: SocketAddr) -> Result Result Result<()> { home, } = Args::parse(); - // TODO: handle this cancellation token more gracefully within Sequencer service - // similar to how we do in Indexer let cancellation_token = listen_for_shutdown_signal(); let mut config = sequencer_service::SequencerConfig::from_path(&config_path)?; if let Some(home) = home { config.home = home; } - let sequencer_handle = + let mut sequencer_handle = sequencer_service::run(config, SocketAddr::new(listen_address, port)).await?; tokio::select! { @@ -60,21 +59,50 @@ async fn main() -> Result<()> { } } + // Stop the watchers, the publisher's drive task, the block loop and the RPC + // server, and wait for each. Dropping the handle only asks; the store stays + // open for an unbounded stretch after that, so a restart can find its own + // home directory locked, and a watcher can be killed between recording a + // delivery and handing it over. + sequencer_handle.shutdown().await; + info!("Sequencer shutdown complete"); Ok(()) } +/// Cancelled on Ctrl-C or `SIGTERM`. +/// +/// `SIGTERM` is what a container runtime sends first, so without it every +/// orchestrated stop is the ungraceful path. +#[expect( + clippy::integer_division_remainder_used, + reason = "Generated by select! macro, can't be easily rewritten to avoid this lint" +)] fn listen_for_shutdown_signal() -> CancellationToken { let cancellation_token = CancellationToken::new(); let cancellation_token_clone = cancellation_token.clone(); tokio::spawn(async move { - if let Err(err) = tokio::signal::ctrl_c().await { - error!("Failed to listen for Ctrl-C signal: {err}"); - return; + let mut terminate = match signal(SignalKind::terminate()) { + Ok(terminate) => terminate, + Err(err) => { + error!("Failed to listen for SIGTERM: {err}"); + return; + } + }; + + tokio::select! { + result = tokio::signal::ctrl_c() => match result { + Ok(()) => info!("Received Ctrl-C signal"), + Err(err) => { + error!("Failed to listen for Ctrl-C signal: {err}"); + return; + } + }, + _ = terminate.recv() => info!("Received SIGTERM"), } - info!("Received Ctrl-C signal"); + cancellation_token_clone.cancel(); }); diff --git a/test_fixtures/src/lib.rs b/test_fixtures/src/lib.rs index 596354ab..9077feaf 100644 --- a/test_fixtures/src/lib.rs +++ b/test_fixtures/src/lib.rs @@ -217,7 +217,7 @@ impl Drop for TestContext { temp_wallet_dir: _, } = self; - let sequencer_handle = sequencer_handle + let mut sequencer_handle = sequencer_handle .take() .expect("Sequencer handle should be present in TestContext drop"); if !sequencer_handle.is_healthy() {