diff --git a/Cargo.lock b/Cargo.lock index 0bd194bb..1c21dc68 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9760,6 +9760,7 @@ dependencies = [ "tempfile", "testcontainers", "tokio", + "tokio-util", "url", "wallet", ] diff --git a/lez/indexer/service/src/lib.rs b/lez/indexer/service/src/lib.rs index 355d7801..945c75be 100644 --- a/lez/indexer/service/src/lib.rs +++ b/lez/indexer/service/src/lib.rs @@ -5,6 +5,7 @@ pub use indexer_core::config::*; use indexer_service_rpc::RpcServer as _; use jsonrpsee::server::{Server, ServerHandle}; use log::{error, info}; +use tokio_util::sync::CancellationToken; pub mod service; @@ -69,9 +70,10 @@ pub async fn run_server( config: IndexerConfig, storage_dir: &Path, port: u16, + shutdown: CancellationToken, ) -> Result { #[cfg(feature = "mock-responses")] - let _ = (config, storage_dir); + let _ = (config, storage_dir, shutdown); let server = Server::builder() .build(SocketAddr::from(([0, 0, 0, 0], port))) @@ -86,7 +88,7 @@ pub async fn run_server( #[cfg(not(feature = "mock-responses"))] let handle = { - let service = service::IndexerService::new(config, storage_dir) + let service = service::IndexerService::new(config, storage_dir, shutdown) .await .context("Failed to initialize indexer service")?; server.start(service.into_rpc()) diff --git a/lez/indexer/service/src/main.rs b/lez/indexer/service/src/main.rs index 3e36967d..52f195e9 100644 --- a/lez/indexer/service/src/main.rs +++ b/lez/indexer/service/src/main.rs @@ -34,7 +34,9 @@ async fn main() -> Result<()> { let cancellation_token = listen_for_shutdown_signal(); let config = indexer_service::IndexerConfig::from_path(&config_path)?; - let indexer_handle = indexer_service::run_server(config, data_dir.as_path(), port).await?; + let indexer_handle = + indexer_service::run_server(config, data_dir.as_path(), port, cancellation_token.clone()) + .await?; tokio::select! { () = cancellation_token.cancelled() => { diff --git a/lez/indexer/service/src/service.rs b/lez/indexer/service/src/service.rs index 767560b1..fedef895 100644 --- a/lez/indexer/service/src/service.rs +++ b/lez/indexer/service/src/service.rs @@ -2,7 +2,7 @@ use std::{path::Path, pin::pin, sync::Arc}; use anyhow::{Context as _, Result, bail}; use arc_swap::ArcSwap; -use futures::{StreamExt as _, never::Never}; +use futures::StreamExt as _; use indexer_core::{IndexerCore, config::IndexerConfig}; use indexer_service_protocol::{ Account, AccountId, Block, BlockId, HashType, IndexerStatus, Transaction, @@ -14,6 +14,7 @@ use jsonrpsee::{ }; use log::{debug, error, info, warn}; use tokio::sync::mpsc::UnboundedSender; +use tokio_util::sync::CancellationToken; pub struct IndexerService { subscription_service: SubscriptionService, @@ -21,9 +22,13 @@ pub struct IndexerService { } impl IndexerService { - pub async fn new(config: IndexerConfig, storage_dir: &Path) -> Result { + pub async fn new( + config: IndexerConfig, + storage_dir: &Path, + shutdown: CancellationToken, + ) -> Result { let indexer = IndexerCore::new(config, storage_dir).await?; - let subscription_service = SubscriptionService::spawn_new(indexer.clone()); + let subscription_service = SubscriptionService::spawn_new(indexer.clone(), shutdown); Ok(Self { subscription_service, @@ -170,15 +175,17 @@ impl indexer_service_rpc::RpcServer for IndexerService { struct SubscriptionService { parts: ArcSwap, indexer: IndexerCore, + shutdown: CancellationToken, } impl SubscriptionService { - pub fn spawn_new(indexer: IndexerCore) -> Self { - let parts = Self::spawn_respond_subscribers_loop(indexer.clone()); + pub fn spawn_new(indexer: IndexerCore, shutdown: CancellationToken) -> Self { + let parts = Self::spawn_respond_subscribers_loop(indexer.clone(), shutdown.clone()); Self { parts: ArcSwap::new(Arc::new(parts)), indexer, + shutdown, } } @@ -192,12 +199,16 @@ impl SubscriptionService { // Respawn the subscription service loop if it has finished (either with error or panic) if guard.handle.is_finished() { drop(guard); - let new_parts = Self::spawn_respond_subscribers_loop(self.indexer.clone()); + let new_parts = Self::spawn_respond_subscribers_loop( + self.indexer.clone(), + self.shutdown.clone(), + ); let old_handle_and_sender = self.parts.swap(Arc::new(new_parts)); let old_parts = Arc::into_inner(old_handle_and_sender) .expect("There should be no other references to the old handle and sender"); match old_parts.handle.await { + Ok(Ok(())) => {} Ok(Err(err)) => { error!( "Subscription service loop has unexpectedly finished with error: {err:#}" @@ -215,7 +226,10 @@ impl SubscriptionService { Ok(()) } - fn spawn_respond_subscribers_loop(indexer: IndexerCore) -> SubscriptionLoopParts { + fn spawn_respond_subscribers_loop( + indexer: IndexerCore, + shutdown: CancellationToken, + ) -> SubscriptionLoopParts { let (new_subscription_sender, mut sub_receiver) = tokio::sync::mpsc::unbounded_channel::>(); @@ -231,6 +245,10 @@ impl SubscriptionService { )] loop { tokio::select! { + () = shutdown.cancelled() => { + info!("Shutdown requested; stopping block ingestion"); + return Ok(()); + } sub = sub_receiver.recv() => { let Some(subscription) = sub else { bail!("Subscription receiver closed unexpectedly"); @@ -259,10 +277,11 @@ impl SubscriptionService { } } }; - let res: anyhow::Result = run_loop.await; - let Err(err) = res; - error!("Subscription service loop has unexpectedly finished with error: {err:#?}"); - Err(err) + let res: anyhow::Result<()> = run_loop.await; + if let Err(err) = &res { + error!("Subscription service loop has unexpectedly finished with error: {err:#?}"); + } + res }); SubscriptionLoopParts { handle, @@ -278,7 +297,7 @@ impl Drop for SubscriptionService { } struct SubscriptionLoopParts { - handle: tokio::task::JoinHandle>, + handle: tokio::task::JoinHandle>, new_subscription_sender: UnboundedSender>, } diff --git a/test_fixtures/Cargo.toml b/test_fixtures/Cargo.toml index 71f887b6..c83bdff9 100644 --- a/test_fixtures/Cargo.toml +++ b/test_fixtures/Cargo.toml @@ -31,4 +31,5 @@ serde_json.workspace = true tempfile.workspace = true testcontainers = { version = "0.27.3", features = ["docker-compose"] } tokio = { workspace = true, features = ["rt-multi-thread", "macros"] } +tokio-util.workspace = true url.workspace = true diff --git a/test_fixtures/src/setup.rs b/test_fixtures/src/setup.rs index 6559192f..af12a381 100644 --- a/test_fixtures/src/setup.rs +++ b/test_fixtures/src/setup.rs @@ -101,10 +101,15 @@ pub async fn setup_indexer(bedrock_addr: SocketAddr) -> Result<(IndexerHandle, T let indexer_config = config::indexer_config(bedrock_addr).context("Failed to create Indexer config")?; - indexer_service::run_server(indexer_config, temp_indexer_dir.path(), 0) - .await - .context("Failed to run Indexer Service") - .map(|handle| (handle, temp_indexer_dir)) + indexer_service::run_server( + indexer_config, + temp_indexer_dir.path(), + 0, + tokio_util::sync::CancellationToken::new(), + ) + .await + .context("Failed to run Indexer Service") + .map(|handle| (handle, temp_indexer_dir)) } pub async fn setup_sequencer(