feat(sequencer): wait for background tasks on shutdown

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<Option<Slot>>;
@@ -143,19 +148,12 @@ pub struct ZoneSdkPublisher {
     turn_rx: watch::Receiver<TurnNotification>,
     // 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<DriveTaskGuard>,
+    // 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<NodeHttpClient>,
 }
 
-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<Option<Slot>> {
         Ok(self
             .node

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<BP: BlockPublisherTrait> SequencerCore<BP> {
         &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<TaskGroup> {
+        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 {

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<TaskGroupInner>);
+
+#[derive(Default)]
+struct TaskGroupInner(Mutex<Vec<JoinHandle<()>>>);
+
+/// 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<RocksDBIO>);
+
+impl StoreRelease {
+    #[must_use]
+    pub fn new(store: &Arc<RocksDBIO>) -> 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<Vec<JoinHandle<()>>>) -> MutexGuard<'_, Vec<JoinHandle<()>>> {
+        handles.lock().unwrap_or_else(PoisonError::into_inner)
+    }
+
+    fn take(handles: &Mutex<Vec<JoinHandle<()>>>) -> Vec<JoinHandle<()>> {
+        std::mem::take(&mut *Self::handles(handles))
+    }
+}
+
+impl TaskGroup {
+    /// Takes ownership of already-spawned tasks.
+    #[must_use]
+    pub fn new(handles: Vec<JoinHandle<()>>) -> 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");
+    }
+}

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<ServerHandle>,
+    server_handle: ServerHandle,
     main_loop_handle: JoinHandle<Result<Never>>,
     /// 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<TaskGroup>,
+    /// 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<Result<Never>>,
         driver_cancellation: CancellationToken,
+        background_tasks: Vec<TaskGroup>,
+        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<Never> {
+    pub async fn failed(&mut self) -> Result<Never> {
         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<SequencerHandle> {
     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<Seq
     info!("Sequencer core set up");
 
     let driver_cancellation = sequencer_core.block_publisher().driver_cancellation();
+    // Taken while the core is still owned here: once it is behind the `Arc`
+    // below, the only owners are the RPC server and the main loop task, and
+    // neither hands it back.
+    let background_tasks = sequencer_core.background_tasks();
+    let store = sequencer_core.store_release();
     let seq_core_wrapped = Arc::new(Mutex::new(sequencer_core));
     let mempool_handle_for_server = mempool_handle.clone();
 
@@ -156,6 +247,8 @@ pub async fn run(config: SequencerConfig, listen_addr: SocketAddr) -> Result<Seq
         server_handle,
         main_loop_handle,
         driver_cancellation,
+        background_tasks,
+        store,
     ))
 }
 

lez/sequencer/service/src/main.rs

@@ -6,6 +6,7 @@ use std::{
 use anyhow::Result;
 use clap::Parser;
 use log::{error, info};
+use tokio::signal::unix::{SignalKind, signal};
 use tokio_util::sync::CancellationToken;
 
 #[derive(Debug, Parser)]
@@ -40,15 +41,13 @@ async fn main() -> 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();
     });
 

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() {