diff --git a/Cargo.lock b/Cargo.lock index a13ea66c6..551f2e52f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3664,6 +3664,7 @@ dependencies = [ "magicblock-program", "solana-hash 3.1.0", "solana-pubkey 3.0.0", + "solana-signature 3.4.0", "solana-transaction", "solana-transaction-error 3.2.0", "thiserror 2.0.18", diff --git a/magicblock-accounts/Cargo.toml b/magicblock-accounts/Cargo.toml index aaf0879d4..80e857e14 100644 --- a/magicblock-accounts/Cargo.toml +++ b/magicblock-accounts/Cargo.toml @@ -28,7 +28,5 @@ thiserror = { workspace = true } url = { workspace = true } [dev-dependencies] -magicblock-committor-service = { workspace = true, features = [ - "dev-context-only-utils", -] } +solana-signature = { workspace = true, features = ["rand"] } tokio-util = { workspace = true } diff --git a/magicblock-accounts/src/scheduled_commits_processor.rs b/magicblock-accounts/src/scheduled_commits_processor.rs index 2df5e7f16..74e5ccfc7 100644 --- a/magicblock-accounts/src/scheduled_commits_processor.rs +++ b/magicblock-accounts/src/scheduled_commits_processor.rs @@ -1,6 +1,7 @@ use std::{ collections::{HashMap, HashSet}, sync::{Arc, Mutex}, + time::Duration, }; use async_trait::async_trait; @@ -260,6 +261,25 @@ impl ScheduledCommitsProcessorImpl { } } + fn pending_intent_results_count( + intents_meta_map: &Mutex>, + ) -> usize { + intents_meta_map.lock().expect(POISONED_MUTEX_MSG).len() + } + + async fn wait_until_pending_results_drained(&self) { + while Self::pending_intent_results_count(&self.intents_meta_map) > 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + } + + pub async fn shutdown(&self) { + if !self.cancellation_token.is_cancelled() { + self.cancellation_token.cancel(); + } + self.wait_until_pending_results_drained().await; + } + #[instrument(skip( result_subscriber, cancellation_token, @@ -281,17 +301,27 @@ impl ScheduledCommitsProcessorImpl { let mut result_receiver = result_subscriber.await.expect(SUBSCRIPTION_ERR_MSG); + let mut draining = false; loop { let execution_result = tokio::select! { biased; - _ = cancellation_token.cancelled() => { - info!("Shutting down"); - return; + _ = cancellation_token.cancelled(), if !draining => { + draining = true; + info!("ScheduledCommitsProcessor draining in-flight results"); + continue; } execution_result = result_receiver.recv() => { match execution_result { Ok(result) => result, Err(broadcast::error::RecvError::Closed) => { + if Self::pending_intent_results_count(&intents_meta_map) > 0 { + warn!( + pending_count = Self::pending_intent_results_count( + &intents_meta_map, + ), + "Intent execution service shut down with pending commit results", + ); + } info!("Intent execution service shut down"); break; } @@ -330,6 +360,8 @@ impl ScheduledCommitsProcessorImpl { ) .await; } + + info!("ScheduledCommitsProcessor result processor shutdown"); } #[instrument( @@ -446,6 +478,10 @@ impl ScheduledCommitsProcessorImpl { impl ScheduledCommitsProcessor for ScheduledCommitsProcessorImpl { #[instrument(skip(self))] async fn process(&self) -> ScheduledCommitsProcessorResult<()> { + if self.cancellation_token.is_cancelled() { + return Ok(()); + } + let intent_bundles = self.transaction_scheduler.take_scheduled_intent_bundles(); @@ -471,6 +507,10 @@ impl ScheduledCommitsProcessor for ScheduledCommitsProcessorImpl { fn stop(&self) { self.cancellation_token.cancel(); } + + async fn shutdown(&self) { + ScheduledCommitsProcessorImpl::shutdown(self).await; + } } struct ScheduledBaseIntentMeta { @@ -502,3 +542,182 @@ impl ScheduledBaseIntentMeta { } } } + +#[cfg(test)] +mod shutdown_tests { + use std::{sync::Arc, time::Duration}; + + use magicblock_program::magic_scheduled_base_intent::ScheduledIntentBundle; + use solana_hash::Hash; + use solana_pubkey::Pubkey; + use tokio::sync::broadcast; + + use super::*; + + fn test_intent(id: u64) -> ScheduledIntentBundle { + ScheduledIntentBundle { + id, + slot: 0, + blockhash: Hash::default(), + sent_transaction: Transaction::default(), + payer: Pubkey::default(), + intent_bundle: Default::default(), + } + } + + /// Mirrors the drain-mode branch of `result_processor` without the full + /// commit-signaling pipeline, so we can verify shutdown keeps receiving + /// in-flight results after cancellation. + async fn drain_in_flight_result_after_stop( + cancellation_token: CancellationToken, + intents_meta_map: Arc>>, + mut results_rx: broadcast::Receiver, + ) { + let mut draining = false; + loop { + let intent_id = tokio::select! { + biased; + _ = cancellation_token.cancelled(), if !draining => { + draining = true; + continue; + } + execution_result = results_rx.recv() => { + match execution_result { + Ok(result) => result.id, + Err(broadcast::error::RecvError::Closed) => break, + Err(broadcast::error::RecvError::Lagged(_)) => continue, + } + } + }; + + intents_meta_map + .lock() + .expect(POISONED_MUTEX_MSG) + .remove(&intent_id); + } + } + + #[tokio::test] + async fn result_processor_drains_in_flight_after_stop() { + let (results_tx, _) = broadcast::channel(16); + let results_rx = results_tx.subscribe(); + let cancellation_token = CancellationToken::new(); + let intents_meta_map = Arc::new(Mutex::new(HashMap::new())); + intents_meta_map + .lock() + .expect(POISONED_MUTEX_MSG) + .insert(1, ScheduledBaseIntentMeta::new(&test_intent(1))); + + let drain_task = { + let cancellation_token = cancellation_token.clone(); + let intents_meta_map = intents_meta_map.clone(); + tokio::spawn(async move { + drain_in_flight_result_after_stop( + cancellation_token, + intents_meta_map, + results_rx, + ) + .await; + }) + }; + + tokio::task::yield_now().await; + cancellation_token.cancel(); + results_tx + .send(BroadcastedIntentExecutionResult { + id: 1, + inner: Ok(ExecutionOutput::SingleStage( + solana_signature::Signature::new_unique(), + )), + patched_errors: Arc::new(vec![]), + callbacks_report: vec![], + }) + .expect("broadcast result"); + + tokio::time::timeout(Duration::from_secs(5), async { + while ScheduledCommitsProcessorImpl::pending_intent_results_count( + &intents_meta_map, + ) > 0 + { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("timed out waiting for drained intent metadata"); + + drop(results_tx); + tokio::time::timeout(Duration::from_secs(1), drain_task) + .await + .expect("timed out waiting for drain task") + .expect("drain task panicked"); + } + + #[tokio::test] + async fn shutdown_waits_until_pending_results_are_drained() { + let cancellation_token = CancellationToken::new(); + let intents_meta_map = Arc::new(Mutex::new(HashMap::new())); + intents_meta_map + .lock() + .expect(POISONED_MUTEX_MSG) + .insert(1, ScheduledBaseIntentMeta::new(&test_intent(1))); + + let (results_tx, _) = broadcast::channel(16); + let results_rx = results_tx.subscribe(); + + let drain_task = { + let cancellation_token = cancellation_token.clone(); + let intents_meta_map = intents_meta_map.clone(); + tokio::spawn(async move { + drain_in_flight_result_after_stop( + cancellation_token, + intents_meta_map, + results_rx, + ) + .await; + }) + }; + + let shutdown_task = { + let cancellation_token = cancellation_token.clone(); + let intents_meta_map = intents_meta_map.clone(); + tokio::spawn(async move { + if !cancellation_token.is_cancelled() { + cancellation_token.cancel(); + } + while + ScheduledCommitsProcessorImpl::pending_intent_results_count( + &intents_meta_map, + ) > 0 + { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + }; + + tokio::task::yield_now().await; + results_tx + .send(BroadcastedIntentExecutionResult { + id: 1, + inner: Ok(ExecutionOutput::SingleStage( + solana_signature::Signature::new_unique(), + )), + patched_errors: Arc::new(vec![]), + callbacks_report: vec![], + }) + .expect("broadcast result"); + + tokio::time::timeout(Duration::from_secs(5), shutdown_task) + .await + .expect("timed out waiting for shutdown") + .expect("shutdown task panicked"); + assert_eq!( + ScheduledCommitsProcessorImpl::pending_intent_results_count( + &intents_meta_map, + ), + 0 + ); + + drop(results_tx); + let _ = tokio::time::timeout(Duration::from_secs(1), drain_task).await; + } +} diff --git a/magicblock-accounts/src/traits.rs b/magicblock-accounts/src/traits.rs index c039558be..60b8e2778 100644 --- a/magicblock-accounts/src/traits.rs +++ b/magicblock-accounts/src/traits.rs @@ -13,6 +13,9 @@ pub trait ScheduledCommitsProcessor: Send + Sync + 'static { /// Clears all scheduled commits fn clear_scheduled_commits(&self); - /// Stop processor + /// Signals the processor to stop accepting new work fn stop(&self); + + /// Waits until in-flight commit results have been processed + async fn shutdown(&self); } diff --git a/magicblock-api/src/magic_validator.rs b/magicblock-api/src/magic_validator.rs index 7a3df2594..72d7cb8f3 100644 --- a/magicblock-api/src/magic_validator.rs +++ b/magicblock-api/src/magic_validator.rs @@ -9,10 +9,7 @@ use std::{ }; use magicblock_account_cloner::ChainlinkCloner; -use magicblock_accounts::{ - scheduled_commits_processor::ScheduledCommitsProcessorImpl, - ScheduledCommitsProcessor, -}; +use magicblock_accounts::scheduled_commits_processor::ScheduledCommitsProcessorImpl; use magicblock_accounts_db::{traits::AccountsBank, AccountsDb}; use magicblock_aperture::{ initialize_aperture, @@ -1081,25 +1078,32 @@ impl MagicValidator { self.start_unregister_validator_on_chain().await; self.exit.store(true, Ordering::Relaxed); - // Ordering is important here - // Commitor service shall be stopped last - self.token.cancel(); + if let Some(handle) = self.slot_ticker.take() { + let step_start = Instant::now(); + let _ = handle.await; + log_timing("shutdown", "slot_ticker_join", step_start); + } + + if let Some(ref committor_service) = self.committor_service { + let step_start = Instant::now(); + committor_service.stop(); + committor_service.stopped().await; + log_timing("shutdown", "committor_service_stop", step_start); + } + if let Some(ref scheduled_commits_processor) = self.scheduled_commits_processor { let step_start = Instant::now(); - scheduled_commits_processor.stop(); + scheduled_commits_processor.shutdown().await; log_timing( "shutdown", - "scheduled_commits_processor_stop", + "scheduled_commits_processor_shutdown", step_start, ); } - if let Some(ref committor_service) = self.committor_service { - let step_start = Instant::now(); - committor_service.stop(); - log_timing("shutdown", "committor_service_stop", step_start); - } + + self.token.cancel(); let step_start = Instant::now(); self.claim_fees_task.stop().await; @@ -1108,11 +1112,6 @@ impl MagicValidator { let step_start = Instant::now(); let _ = self.rpc_handle.join(); log_timing("shutdown", "rpc_thread_join", step_start); - if let Some(handle) = self.slot_ticker { - let step_start = Instant::now(); - let _ = handle.await; - log_timing("shutdown", "slot_ticker_join", step_start); - } let step_start = Instant::now(); if let Err(err) = self.ledger_truncator.join() { error!(error = ?err, "Ledger truncator did not gracefully exit"); diff --git a/magicblock-committor-service/src/committor_processor.rs b/magicblock-committor-service/src/committor_processor.rs index 5061eb4b0..de3369f3b 100644 --- a/magicblock-committor-service/src/committor_processor.rs +++ b/magicblock-committor-service/src/committor_processor.rs @@ -207,11 +207,23 @@ impl CommittorProcessor { Ok(bundles) } + pub async fn shutdown(&self) { + self.commits_scheduler.shutdown().await; + } + + pub fn is_shutting_down(&self) -> bool { + self.commits_scheduler.is_shutting_down() + } + #[instrument(skip(self, intent_bundles))] pub async fn schedule_intent_bundle( &self, intent_bundles: Vec, ) -> CommittorServiceResult<()> { + if self.is_shutting_down() { + return Err(crate::error::CommittorServiceError::ShuttingDown); + } + if let Err(err) = self.persister.start_base_intents(&intent_bundles) { // We will still try to perform the commits, but the fact that we cannot // persist the intent is very serious and we should probably restart the @@ -234,6 +246,10 @@ impl CommittorProcessor { &self, intent_bundles: Vec, ) -> CommittorServiceResult<()> { + if self.is_shutting_down() { + return Err(crate::error::CommittorServiceError::ShuttingDown); + } + self.commits_scheduler .schedule(intent_bundles) .await @@ -263,6 +279,36 @@ impl CommittorProcessor { } } +#[cfg(test)] +impl CommittorProcessor { + pub(crate) fn new_for_test( + factory: crate::intent_execution_manager::test_support::MockIntentExecutorFactory, + ) -> Self { + let authority = Keypair::new(); + let rpc_client = Arc::new(RpcClient::new("http://127.0.0.1:1".to_string())); + let magicblock_rpc_client = MagicblockRpcClient::new(rpc_client); + let table_mania = TableMania::new( + magicblock_rpc_client.clone(), + &authority, + None, + ); + let persister = IntentPersisterImpl::try_new(":memory:").unwrap(); + let task_info_fetcher = Arc::new(CacheTaskInfoFetcher::new( + RpcTaskInfoFetcher::new(magicblock_rpc_client.clone()), + )); + let commits_scheduler = crate::intent_execution_manager::test_support::new_test_intent_execution_manager(factory); + + Self { + authority, + magicblock_rpc_client, + table_mania, + persister, + commits_scheduler, + task_info_fetcher, + } + } +} + fn pending_rows_to_scheduled_intent_bundles( rows: Vec, payer: Pubkey, @@ -391,10 +437,20 @@ fn committed_account_from_pending_row( #[cfg(test)] mod tests { + use std::{sync::Arc, time::Duration}; + use solana_hash::Hash; use super::*; - use crate::persist::CommitStatus; + use crate::{ + error::CommittorServiceError, + intent_execution_manager::{ + intent_scheduler::create_test_intent, + test_support::MockIntentExecutorFactory, + }, + persist::CommitStatus, + test_utils, + }; fn pending_row( message_id: u64, @@ -448,6 +504,88 @@ mod tests { row } + fn new_test_processor( + factory: MockIntentExecutorFactory, + ) -> CommittorProcessor { + CommittorProcessor::new_for_test(factory) + } + + #[tokio::test] + async fn test_schedule_rejected_after_shutdown() { + test_utils::init_test_logger(); + let processor = + new_test_processor(MockIntentExecutorFactory::new()); + processor.shutdown().await; + + let intent = create_test_intent( + 1, + &[Pubkey::new_unique()], + false, + ); + let err = processor + .schedule_intent_bundle(vec![intent]) + .await + .unwrap_err(); + assert!(matches!(err, CommittorServiceError::ShuttingDown)); + } + + #[tokio::test] + async fn test_recovered_schedule_rejected_after_shutdown() { + test_utils::init_test_logger(); + let processor = + new_test_processor(MockIntentExecutorFactory::new()); + processor.shutdown().await; + + let intent = create_test_intent( + 2, + &[Pubkey::new_unique()], + false, + ); + let err = processor + .schedule_recovered_intent_bundles(vec![intent]) + .await + .unwrap_err(); + assert!(matches!(err, CommittorServiceError::ShuttingDown)); + } + + #[tokio::test] + async fn test_shutdown_waits_for_in_flight_intent() { + test_utils::init_test_logger(); + let processor = Arc::new(new_test_processor( + MockIntentExecutorFactory::new() + .with_execution_delay(Duration::from_millis(200)), + )); + let mut result_receiver = processor.subscribe_for_results(); + + let intent = create_test_intent( + 1, + &[Pubkey::new_unique()], + false, + ); + processor + .schedule_intent_bundle(vec![intent]) + .await + .unwrap(); + + let draining = { + let processor = processor.clone(); + tokio::spawn(async move { processor.shutdown().await }) + }; + + let result = tokio::time::timeout( + Duration::from_secs(5), + result_receiver.recv(), + ) + .await + .expect("timed out waiting for in-flight intent result") + .expect("result channel closed"); + assert!(result.is_ok()); + assert_eq!(result.id, 1); + + draining.await.expect("shutdown task panicked"); + assert!(processor.is_shutting_down()); + } + #[test] fn pending_rows_reconstruct_commit_finalize_bundle() { let payer = Pubkey::new_unique(); diff --git a/magicblock-committor-service/src/error.rs b/magicblock-committor-service/src/error.rs index fa13e49b8..a7afd1302 100644 --- a/magicblock-committor-service/src/error.rs +++ b/magicblock-committor-service/src/error.rs @@ -61,6 +61,9 @@ pub enum CommittorServiceError { #[error("Task {0} failed to create transaction: {1} ({1:?})")] FailedToCreateTransaction(String, solana_signer::SignerError), + #[error("Committor service is shutting down")] + ShuttingDown, + #[error("Could not find commit strategy for bundle {0}")] CouldNotFindCommitStrategyForBundle(u64), diff --git a/magicblock-committor-service/src/intent_execution_manager.rs b/magicblock-committor-service/src/intent_execution_manager.rs index 725224386..92478e973 100644 --- a/magicblock-committor-service/src/intent_execution_manager.rs +++ b/magicblock-committor-service/src/intent_execution_manager.rs @@ -2,6 +2,9 @@ pub(crate) mod db; mod intent_execution_engine; pub mod intent_scheduler; +#[cfg(test)] +pub(crate) mod test_support; + use std::sync::Arc; pub use intent_execution_engine::BroadcastedIntentExecutionResult; @@ -9,7 +12,8 @@ use magicblock_core::traits::ActionsCallbackScheduler; use magicblock_program::magic_scheduled_base_intent::ScheduledIntentBundle; use magicblock_rpc_client::MagicblockRpcClient; use magicblock_table_mania::TableMania; -use tokio::sync::{broadcast, mpsc, mpsc::error::TrySendError}; +use tokio::sync::{broadcast, mpsc, mpsc::error::TrySendError, watch}; +use tokio_util::sync::CancellationToken; use crate::{ intent_execution_manager::{ @@ -27,6 +31,8 @@ pub struct IntentExecutionManager { db: Arc, result_subscriber: ResultSubscriber, intent_sender: mpsc::Sender, + shutdown: CancellationToken, + idle_rx: watch::Receiver, } impl IntentExecutionManager { @@ -54,19 +60,37 @@ impl IntentExecutionManager { }; let (sender, receiver) = mpsc::channel(1000); + let shutdown = CancellationToken::new(); + let (idle_tx, idle_rx) = watch::channel(false); let worker = IntentExecutionEngine::new( db.clone(), executor_factory, intent_persister, receiver, ); - let result_subscriber = worker.spawn(); + let result_subscriber = worker.spawn(shutdown.clone(), idle_tx); Self { db, intent_sender: sender, result_subscriber, + shutdown, + idle_rx, + } + } + + /// Signals the engine to stop accepting new intents and waits until all + /// in-flight executors finish and broadcast their results. + pub async fn shutdown(&self) { + self.shutdown.cancel(); + if *self.idle_rx.borrow() { + return; } + let _ = self.idle_rx.clone().wait_for(|idle| *idle).await; + } + + pub fn is_shutting_down(&self) -> bool { + self.shutdown.is_cancelled() } /// Schedules [`ScheduledBaseIntent`] intent to be executed @@ -76,6 +100,10 @@ impl IntentExecutionManager { &self, intent_bundles: Vec, ) -> Result<(), IntentExecutionManagerError> { + if self.is_shutting_down() { + return Err(IntentExecutionManagerError::ShuttingDown); + } + // If db not empty push el-t there // This means that at some point channel got full // Worker first will clean-up channel, and then DB. @@ -116,6 +144,85 @@ impl IntentExecutionManager { pub enum IntentExecutionManagerError { #[error("Channel was closed")] ChannelClosed, + #[error("Intent execution manager is shutting down")] + ShuttingDown, #[error("DBError: {0}")] DBError(#[from] db::Error), } + +#[cfg(test)] +mod tests { + use std::{sync::Arc, time::Duration}; + + use solana_pubkey::pubkey; + + use super::*; + use crate::intent_execution_manager::{ + db::DummyDB, + intent_scheduler::create_test_intent, + test_support::{self, MockIntentExecutorFactory}, + IntentExecutionManager, + }; + + fn new_for_test(factory: MockIntentExecutorFactory) -> Arc> { + Arc::new(test_support::new_test_intent_execution_manager(factory)) + } + + #[tokio::test] + async fn test_manager_schedule_rejected_after_shutdown() { + let manager = new_for_test(MockIntentExecutorFactory::new()); + manager.shutdown.cancel(); + + let intent = create_test_intent( + 1, + &[pubkey!("1111111111111111111111111111111111111111111")], + false, + ); + let err = manager.schedule(vec![intent]).await.unwrap_err(); + assert!(matches!(err, IntentExecutionManagerError::ShuttingDown)); + } + + #[tokio::test] + async fn test_manager_shutdown_waits_for_in_flight_intent() { + let manager = new_for_test( + MockIntentExecutorFactory::new() + .with_execution_delay(Duration::from_millis(200)), + ); + let mut result_receiver = manager.subscribe_for_results(); + + let intent = create_test_intent( + 1, + &[pubkey!("1111111111111111111111111111111111111111111")], + false, + ); + manager.schedule(vec![intent]).await.unwrap(); + + let shutdown_handle = { + let manager = manager.clone(); + tokio::spawn(async move { manager.shutdown().await }) + }; + + let result = tokio::time::timeout( + Duration::from_secs(5), + result_receiver.recv(), + ) + .await + .expect("timed out waiting for intent result") + .expect("result channel closed"); + assert!(result.is_ok()); + assert_eq!(result.id, 1); + + shutdown_handle.await.expect("shutdown task panicked"); + assert!(manager.is_shutting_down()); + } + + #[tokio::test] + async fn test_manager_shutdown_is_idempotent() { + let manager = new_for_test(MockIntentExecutorFactory::new()); + + manager.shutdown().await; + manager.shutdown().await; + + assert!(manager.is_shutting_down()); + } +} diff --git a/magicblock-committor-service/src/intent_execution_manager/intent_execution_engine.rs b/magicblock-committor-service/src/intent_execution_manager/intent_execution_engine.rs index 7b435a85a..de5fe63d6 100644 --- a/magicblock-committor-service/src/intent_execution_manager/intent_execution_engine.rs +++ b/magicblock-committor-service/src/intent_execution_manager/intent_execution_engine.rs @@ -11,11 +11,12 @@ use magicblock_program::magic_scheduled_base_intent::ScheduledIntentBundle; use solana_signature::Signature; use tokio::{ sync::{ - broadcast, mpsc, mpsc::error::TryRecvError, OwnedSemaphorePermit, - Semaphore, + broadcast, mpsc, mpsc::error::TryRecvError, watch, + OwnedSemaphorePermit, Semaphore, }, task::JoinHandle, }; +use tokio_util::sync::CancellationToken; use tracing::{error, info, instrument, trace, warn}; use crate::{ @@ -126,9 +127,13 @@ where } /// Spawns `main_loop` and return `Receiver` listening to results - pub fn spawn(self) -> ResultSubscriber { + pub fn spawn( + self, + shutdown: CancellationToken, + idle_tx: watch::Sender, + ) -> ResultSubscriber { let (result_sender, _) = broadcast::channel(100); - tokio::spawn(self.main_loop(result_sender.clone())); + tokio::spawn(self.main_loop(result_sender.clone(), shutdown, idle_tx)); ResultSubscriber(result_sender) } @@ -137,18 +142,42 @@ where /// 1. Handles & schedules incoming intents /// 2. Finds available executor /// 3. Spawns execution of scheduled intent - #[instrument(skip(self, result_sender))] + /// 4. On shutdown, drains queued and in-flight work before exiting + #[instrument(skip(self, result_sender, shutdown, idle_tx))] async fn main_loop( mut self, result_sender: broadcast::Sender, + shutdown: CancellationToken, + idle_tx: watch::Sender, ) { + let mut draining = false; + loop { - let intent = match self.next_scheduled_intent().await { + if !draining && shutdown.is_cancelled() { + draining = true; + info!("IntentExecutionEngine entering drain mode"); + } + + if draining && self.is_fully_idle() { + break; + } + + let intent = match self.next_scheduled_intent(draining).await { Ok(value) => value, Err(IntentExecutionManagerError::ChannelClosed) => { + if draining && self.is_fully_idle() { + break; + } + if draining { + trace!("Intent channel closed while draining, waiting for in-flight work"); + continue; + } info!("Channel closed, exiting"); break; } + Err(IntentExecutionManagerError::ShuttingDown) => { + unreachable!("engine does not emit ShuttingDown") + } Err(IntentExecutionManagerError::DBError(err)) => { error!(error = ?err, "Failed to fetch intent"); break; @@ -158,6 +187,7 @@ where // We couldn't pick up intent for execution due to: // 1. All executors are currently busy // 2. All intents are blocked and none could be executed at the moment + // 3. Drain mode with no ingress left to pull trace!("Could not schedule any intents"); continue; }; @@ -189,18 +219,44 @@ where self.running_executors.len() as i64, ); } + + while let Some(result) = self.running_executors.next().await { + if let Err(err) = result { + error!(error = ?err, "Executor failed during shutdown drain"); + } + } + metrics::set_committor_executors_busy_count(0); + + if idle_tx.send(true).is_err() { + trace!("No shutdown observers for intent execution engine"); + } + info!("IntentExecutionEngine shutdown complete"); + } + + fn has_pending_ingress(&self) -> bool { + !self.receiver.is_empty() || !self.db.is_empty() + } + + fn is_fully_idle(&self) -> bool { + self.running_executors.is_empty() + && !self.has_pending_ingress() + && self.inner.lock().expect(POISONED_INNER_MSG).is_idle() } /// Returns [`ScheduledIntentBundle`] or None if all intents are blocked #[instrument(skip(self))] async fn next_scheduled_intent( &mut self, + draining: bool, ) -> Result, IntentExecutionManagerError> { // Limit on number of intents that can be stored in scheduler const SCHEDULER_CAPACITY: usize = 1000; let can_receive = || { + if draining { + return false; + } let num_blocked_intents = self .inner .lock() @@ -214,24 +270,34 @@ where } }; + let can_pull_ingress = + can_receive() || (draining && self.has_pending_ingress()); + let running_executors = &mut self.running_executors; let receiver = &mut self.receiver; let db = &self.db; let intent = tokio::select! { // Notify polled first to prioritize unblocked intents over new one biased; - Some(result) = running_executors.next() => { + Some(result) = running_executors.next(), if !running_executors.is_empty() => { if let Err(err) = result { error!(error = ?err, "Executor failed"); }; trace!("Worker executed intent bundle, fetching new available one"); self.inner.lock().expect(POISONED_INNER_MSG).pop_next_scheduled_intent() }, - result = Self::get_new_intent(receiver, db), if can_receive() => { - let intent = result?; - self.inner.lock().expect(POISONED_INNER_MSG).schedule(intent) + result = Self::get_new_intent(receiver, db, draining), if can_pull_ingress => { + match result? { + Some(intent) => { + self.inner.lock().expect(POISONED_INNER_MSG).schedule(intent) + } + None => return Ok(None), + } }, else => { + if draining { + return Ok(None); + } // Shouldn't be possible: // 1. If no executors spawned -> we can receive // 2. If can't receive -> there are MAX_EXECUTORS running executors @@ -248,18 +314,23 @@ where async fn get_new_intent( receiver: &mut mpsc::Receiver, db: &Arc, - ) -> Result { + draining: bool, + ) -> Result, IntentExecutionManagerError> + { match receiver.try_recv() { - Ok(val) => Ok(val), + Ok(val) => Ok(Some(val)), Err(TryRecvError::Empty) => { // Worker either cleaned-up congested channel and now need to clean-up DB // or we're just waiting on empty channel if let Some(intent_bundle) = db.pop_intent_bundle().await? { - Ok(intent_bundle) + Ok(Some(intent_bundle)) + } else if draining { + Ok(None) } else { receiver .recv() .await + .map(Some) .ok_or(IntentExecutionManagerError::ChannelClosed) } } @@ -379,65 +450,25 @@ mod tests { time::Duration, }; - use async_trait::async_trait; - use magicblock_program::magic_scheduled_base_intent::ScheduledIntentBundle; use solana_pubkey::{pubkey, Pubkey}; - use solana_signature::Signature; - use solana_signer::SignerError; - use solana_transaction_error::TransactionError; - use tokio::{sync::mpsc, time::sleep}; use super::*; - use crate::{ - intent_execution_manager::{ - db::{DummyDB, DB}, - intent_scheduler::{create_test_intent, create_test_intent_bundle}, - }, - intent_executor::{ - error::{IntentExecutorError as ExecutorError, InternalError}, - IntentExecutionResult, - }, - persist::IntentPersisterImpl, - test_utils, - transaction_preparator::delivery_preparator::BufferExecutionError, + use crate::intent_execution_manager::{ + db::DB, + intent_scheduler::{create_test_intent, create_test_intent_bundle}, + test_support::{self, setup_engine, spawn_engine, wait_idle}, }; - type MockIntentExecutionEngine = IntentExecutionEngine< - DummyDB, - IntentPersisterImpl, - MockIntentExecutorFactory, - >; - fn setup_engine( - should_fail: bool, - ) -> ( - mpsc::Sender, - MockIntentExecutionEngine, - ) { - test_utils::init_test_logger(); - - let (sender, receiver) = mpsc::channel(1000); - - let db = Arc::new(DummyDB::new()); - let executor_factory = if !should_fail { - MockIntentExecutorFactory::new() - } else { - MockIntentExecutorFactory::new_failing() - }; - let worker = IntentExecutionEngine::new( - db.clone(), - executor_factory, - None::, - receiver, - ); - - (sender, worker) + fn subscribe( + worker: test_support::TestEngine, + ) -> broadcast::Receiver { + spawn_engine(worker).result_subscriber.subscribe() } #[tokio::test] async fn test_worker_processes_messages() { let (sender, worker) = setup_engine(false); - let result_subscriber = worker.spawn(); - let mut result_receiver = result_subscriber.subscribe(); + let mut result_receiver = subscribe(worker); // Send a test message let msg = create_test_intent( @@ -456,8 +487,7 @@ mod tests { #[tokio::test] async fn test_worker_handles_conflicting_messages() { let (sender, worker) = setup_engine(false); - let result_subscriber = worker.spawn(); - let mut result_receiver = result_subscriber.subscribe(); + let mut result_receiver = subscribe(worker); // Send two conflicting messages let pubkey = pubkey!("1111111111111111111111111111111111111111111"); @@ -481,8 +511,7 @@ mod tests { #[tokio::test] async fn test_worker_handles_conflicting_bundles() { let (sender, worker) = setup_engine(false); - let result_subscriber = worker.spawn(); - let mut result_receiver = result_subscriber.subscribe(); + let mut result_receiver = subscribe(worker); // Send two conflicting messages let a = pubkey!("1111111111111111111111111111111111111111111"); @@ -507,8 +536,7 @@ mod tests { #[tokio::test] async fn test_worker_handles_executor_failure() { let (sender, worker) = setup_engine(true); - let result_subscriber = worker.spawn(); - let mut result_receiver = result_subscriber.subscribe(); + let mut result_receiver = subscribe(worker); // Send a test message that will fail let msg = create_test_intent( @@ -545,8 +573,7 @@ mod tests { worker.db.store_intent_bundle(msg.clone()).await.unwrap(); // Start worker - let result_subscriber = worker.spawn(); - let mut result_receiver = result_subscriber.subscribe(); + let mut result_receiver = subscribe(worker); // Verify the message from DB was processed let result = result_receiver.recv().await.unwrap(); @@ -567,8 +594,7 @@ mod tests { .executor_factory .with_concurrency_tracking(&active_tasks, &max_concurrent); - let result_subscriber = worker.spawn(); - let mut result_receiver = result_subscriber.subscribe(); + let mut result_receiver = subscribe(worker); // Send a flood of messages for i in 0..NUM_MESSAGES { @@ -602,8 +628,7 @@ mod tests { #[tokio::test] async fn test_multiple_failures() { let (sender, worker) = setup_engine(true); // Worker that always fails - let result_subscriber = worker.spawn(); - let mut result_receiver = result_subscriber.subscribe(); + let mut result_receiver = subscribe(worker); // Send several messages that will fail const NUM_FAILURES: usize = 10; @@ -635,8 +660,7 @@ mod tests { .executor_factory .with_concurrency_tracking(&active_tasks, &max_concurrent); - let result_subscriber = worker.spawn(); - let mut result_receiver = result_subscriber.subscribe(); + let mut result_receiver = subscribe(worker); // Send messages with unique keys (non-blocking) let mut received_ids = HashSet::new(); @@ -693,8 +717,7 @@ mod tests { .executor_factory .with_concurrency_tracking(&active_tasks, &max_concurrent); - let result_subscriber = worker.spawn(); - let mut result_receiver = result_subscriber.subscribe(); + let mut result_receiver = subscribe(worker); // Shared key for blocking messages let blocking_key = @@ -729,137 +752,134 @@ mod tests { ); } - // Mock implementations for testing - pub struct MockIntentExecutorFactory { - should_fail: bool, - active_tasks: Option>, - max_concurrent: Option>, - } - - impl MockIntentExecutorFactory { - pub fn new() -> Self { - Self { - should_fail: false, - active_tasks: None, - max_concurrent: None, - } - } + #[tokio::test] + async fn test_shutdown_drains_in_flight_intent() { + let (sender, worker) = test_support::setup_engine_with_factory( + test_support::MockIntentExecutorFactory::new() + .with_execution_delay(Duration::from_millis(200)), + ); + let spawned = spawn_engine(worker); + let mut result_receiver = spawned.result_subscriber.subscribe(); - pub fn new_failing() -> Self { - Self { - should_fail: true, - active_tasks: None, - max_concurrent: None, - } - } + let msg = create_test_intent( + 1, + &[pubkey!("1111111111111111111111111111111111111111111")], + false, + ); + sender.send(msg).await.unwrap(); + spawned.shutdown.cancel(); + + let result = tokio::time::timeout( + Duration::from_secs(5), + result_receiver.recv(), + ) + .await + .expect("timed out waiting for in-flight intent result") + .expect("result channel closed"); + assert!(result.is_ok()); + assert_eq!(result.id, 1); - pub fn with_concurrency_tracking( - &mut self, - active_tasks: &Arc, - max_concurrent: &Arc, - ) { - self.active_tasks = Some(active_tasks.clone()); - self.max_concurrent = Some(max_concurrent.clone()); - } + let mut idle_rx = spawned.idle_rx.clone(); + wait_idle(&mut idle_rx).await; } - impl IntentExecutorFactory for MockIntentExecutorFactory { - type Executor = MockIntentExecutor; + #[tokio::test] + async fn test_shutdown_drains_conflicting_intents() { + let (sender, worker) = setup_engine(false); + let spawned = spawn_engine(worker); + let mut result_receiver = spawned.result_subscriber.subscribe(); - fn create_instance(&self) -> Self::Executor { - MockIntentExecutor { - should_fail: self.should_fail, - active_tasks: self.active_tasks.clone(), - max_concurrent: self.max_concurrent.clone(), - } - } - } + let pubkey = pubkey!("1111111111111111111111111111111111111111111"); + sender + .send(create_test_intent(1, &[pubkey], false)) + .await + .unwrap(); + sender + .send(create_test_intent(2, &[pubkey], false)) + .await + .unwrap(); + spawned.shutdown.cancel(); + + let result1 = tokio::time::timeout( + Duration::from_secs(5), + result_receiver.recv(), + ) + .await + .expect("timed out waiting for first intent result") + .expect("result channel closed"); + assert!(result1.is_ok()); + assert_eq!(result1.id, 1); + + let result2 = tokio::time::timeout( + Duration::from_secs(5), + result_receiver.recv(), + ) + .await + .expect("timed out waiting for second intent result") + .expect("result channel closed"); + assert!(result2.is_ok()); + assert_eq!(result2.id, 2); - pub struct MockIntentExecutor { - should_fail: bool, - active_tasks: Option>, - max_concurrent: Option>, + let mut idle_rx = spawned.idle_rx.clone(); + wait_idle(&mut idle_rx).await; } - impl MockIntentExecutor { - fn on_task_started(&self) { - if let (Some(active), Some(max)) = - (&self.active_tasks, &self.max_concurrent) - { - // Increment active task count - let current = active.fetch_add(1, Ordering::SeqCst) + 1; - - // Update max concurrent if needed - let mut observed_max = max.load(Ordering::SeqCst); - while current > observed_max { - match max.compare_exchange_weak( - observed_max, - current, - Ordering::SeqCst, - Ordering::SeqCst, - ) { - Ok(_) => break, - Err(x) => observed_max = x, - } - } - } - } + #[tokio::test] + async fn test_shutdown_drains_db_backlog() { + let (_sender, worker) = setup_engine(false); + let msg = create_test_intent( + 1, + &[pubkey!("1111111111111111111111111111111111111111111")], + false, + ); + worker.db.store_intent_bundle(msg).await.unwrap(); + + let spawned = spawn_engine(worker); + let mut result_receiver = spawned.result_subscriber.subscribe(); + spawned.shutdown.cancel(); + + let result = tokio::time::timeout( + Duration::from_secs(5), + result_receiver.recv(), + ) + .await + .expect("timed out waiting for db backlog intent result") + .expect("result channel closed"); + assert!(result.is_ok()); + assert_eq!(result.id, 1); - fn on_task_finished(&self) { - if let Some(active) = &self.active_tasks { - active.fetch_sub(1, Ordering::SeqCst); - } - } + let mut idle_rx = spawned.idle_rx.clone(); + wait_idle(&mut idle_rx).await; } - #[async_trait] - impl IntentExecutor for MockIntentExecutor { - async fn execute( - &mut self, - _base_intent: ScheduledIntentBundle, - _persister: Option

, - ) -> IntentExecutionResult { - self.on_task_started(); - - // Simulate some work - sleep(Duration::from_millis(50)).await; - - let result = if self.should_fail { - IntentExecutionResult { - inner: Err(ExecutorError::FailedToCommitError { - err: TransactionStrategyExecutionError::InternalError( - InternalError::SignerError(SignerError::Custom( - "oops".to_string(), - )), - ), - signature: None, - }), - patched_errors: vec![ - TransactionStrategyExecutionError::ActionsError( - TransactionError::AccountNotFound, - None, - ), - ], - callbacks_report: vec![], - } - } else { - IntentExecutionResult { - inner: Ok(ExecutionOutput::TwoStage { - commit_signature: Signature::default(), - finalize_signature: Signature::default(), - }), - patched_errors: vec![], - callbacks_report: vec![], - } - }; - - self.on_task_finished(); + #[tokio::test] + async fn test_shutdown_drains_queued_channel_intents() { + let (sender, worker) = setup_engine(false); + let spawned = spawn_engine(worker); + let mut result_receiver = spawned.result_subscriber.subscribe(); - result + let pubkey = pubkey!("1111111111111111111111111111111111111111111"); + for id in 1..=3 { + sender + .send(create_test_intent(id, &[pubkey], false)) + .await + .unwrap(); } - - async fn cleanup(self) -> Result<(), BufferExecutionError> { - Ok(()) + spawned.shutdown.cancel(); + + for expected_id in 1..=3 { + let result = tokio::time::timeout( + Duration::from_secs(5), + result_receiver.recv(), + ) + .await + .expect("timed out waiting for queued intent result") + .expect("result channel closed"); + assert!(result.is_ok()); + assert_eq!(result.id, expected_id); } + + let mut idle_rx = spawned.idle_rx.clone(); + wait_idle(&mut idle_rx).await; } } diff --git a/magicblock-committor-service/src/intent_execution_manager/intent_scheduler.rs b/magicblock-committor-service/src/intent_execution_manager/intent_scheduler.rs index 580acebf6..bca982b40 100644 --- a/magicblock-committor-service/src/intent_execution_manager/intent_scheduler.rs +++ b/magicblock-committor-service/src/intent_execution_manager/intent_scheduler.rs @@ -283,6 +283,11 @@ impl IntentScheduler { pub fn intents_blocked(&self) -> usize { self.blocked_intents.len() } + + /// Returns true when no intents are queued or executing in the scheduler. + pub fn is_idle(&self) -> bool { + self.blocked_intents.is_empty() && self.blocked_keys.is_empty() + } } #[derive(Error, Debug)] @@ -314,9 +319,35 @@ mod simple_test { setup(); let mut scheduler = IntentScheduler::new(); assert_eq!(scheduler.intents_blocked(), 0); + assert!(scheduler.is_idle()); assert!(scheduler.pop_next_scheduled_intent().is_none()); } + #[test] + fn test_is_idle_tracks_executing_and_blocked_intents() { + setup(); + let mut scheduler = IntentScheduler::new(); + let pubkey = pubkey!("1111111111111111111111111111111111111111111"); + let msg1 = create_test_intent(1, &[pubkey], false); + let msg2 = create_test_intent(2, &[pubkey], false); + + assert!(scheduler.is_idle()); + assert!(scheduler.schedule(msg1.clone()).is_some()); + assert!(!scheduler.is_idle()); + assert!(scheduler.schedule(msg2).is_none()); + assert!(!scheduler.is_idle()); + + assert!(scheduler.complete(&msg1).is_ok()); + assert!(!scheduler.is_idle()); + + let next = scheduler.pop_next_scheduled_intent().unwrap(); + assert_eq!(next.id, 2); + assert!(!scheduler.is_idle()); + + assert!(scheduler.complete(&next).is_ok()); + assert!(scheduler.is_idle()); + } + /// Ensure intents with non-conflicting set of keys can run in parallel #[test] fn test_non_conflicting_intents() { diff --git a/magicblock-committor-service/src/intent_execution_manager/test_support.rs b/magicblock-committor-service/src/intent_execution_manager/test_support.rs new file mode 100644 index 000000000..f0aeab6e5 --- /dev/null +++ b/magicblock-committor-service/src/intent_execution_manager/test_support.rs @@ -0,0 +1,261 @@ +#![cfg(test)] +//! Shared test helpers for intent execution manager/engine tests. + +use std::{ + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, + time::Duration, +}; + +use async_trait::async_trait; +use magicblock_program::magic_scheduled_base_intent::ScheduledIntentBundle; +use solana_signature::Signature; +use solana_signer::SignerError; +use solana_transaction_error::TransactionError; +use tokio::sync::{mpsc, watch}; +use tokio_util::sync::CancellationToken; + +use crate::{ + intent_execution_manager::{ + db::DummyDB, + intent_execution_engine::{IntentExecutionEngine, ResultSubscriber}, + IntentExecutionManager, + }, + intent_executor::{ + error::{ + IntentExecutorError as ExecutorError, InternalError, + TransactionStrategyExecutionError, + }, + intent_executor_factory::IntentExecutorFactory, + ExecutionOutput, IntentExecutionResult, IntentExecutor, + }, + persist::IntentPersisterImpl, + test_utils, + transaction_preparator::delivery_preparator::BufferExecutionError, +}; + +pub type TestEngine = IntentExecutionEngine< + DummyDB, + IntentPersisterImpl, + MockIntentExecutorFactory, +>; + +pub struct SpawnedEngine { + pub shutdown: CancellationToken, + pub idle_rx: watch::Receiver, + pub result_subscriber: ResultSubscriber, +} + +pub fn setup_engine( + should_fail: bool, +) -> (mpsc::Sender, TestEngine) { + setup_engine_with_factory(if should_fail { + MockIntentExecutorFactory::new_failing() + } else { + MockIntentExecutorFactory::new() + }) +} + +pub fn setup_engine_with_factory( + factory: MockIntentExecutorFactory, +) -> (mpsc::Sender, TestEngine) { + test_utils::init_test_logger(); + + let (sender, receiver) = mpsc::channel(1000); + let db = Arc::new(DummyDB::new()); + let worker = IntentExecutionEngine::new( + db, + factory, + None::, + receiver, + ); + + (sender, worker) +} + +pub fn new_test_intent_execution_manager( + factory: MockIntentExecutorFactory, +) -> IntentExecutionManager { + let db = Arc::new(DummyDB::new()); + let (sender, receiver) = mpsc::channel(1000); + let shutdown = CancellationToken::new(); + let (idle_tx, idle_rx) = watch::channel(false); + let worker = IntentExecutionEngine::new( + db.clone(), + factory, + None::, + receiver, + ); + let result_subscriber = worker.spawn(shutdown.clone(), idle_tx); + + IntentExecutionManager { + db, + intent_sender: sender, + result_subscriber, + shutdown, + idle_rx, + } +} + +pub fn spawn_engine(worker: TestEngine) -> SpawnedEngine { + let shutdown = CancellationToken::new(); + let (idle_tx, idle_rx) = watch::channel(false); + let result_subscriber = worker.spawn(shutdown.clone(), idle_tx); + + SpawnedEngine { + shutdown, + idle_rx, + result_subscriber, + } +} + +pub async fn wait_idle(idle_rx: &mut watch::Receiver) { + if *idle_rx.borrow() { + return; + } + idle_rx + .changed() + .await + .expect("engine idle signal dropped before drain completed"); + assert!(*idle_rx.borrow()); +} + +pub struct MockIntentExecutorFactory { + should_fail: bool, + execution_delay: Duration, + active_tasks: Option>, + max_concurrent: Option>, +} + +impl MockIntentExecutorFactory { + pub fn new() -> Self { + Self { + should_fail: false, + execution_delay: Duration::from_millis(50), + active_tasks: None, + max_concurrent: None, + } + } + + pub fn new_failing() -> Self { + Self { + should_fail: true, + execution_delay: Duration::from_millis(50), + active_tasks: None, + max_concurrent: None, + } + } + + pub fn with_execution_delay(mut self, execution_delay: Duration) -> Self { + self.execution_delay = execution_delay; + self + } + + pub fn with_concurrency_tracking( + &mut self, + active_tasks: &Arc, + max_concurrent: &Arc, + ) { + self.active_tasks = Some(active_tasks.clone()); + self.max_concurrent = Some(max_concurrent.clone()); + } +} + +impl IntentExecutorFactory for MockIntentExecutorFactory { + type Executor = MockIntentExecutor; + + fn create_instance(&self) -> Self::Executor { + MockIntentExecutor { + should_fail: self.should_fail, + execution_delay: self.execution_delay, + active_tasks: self.active_tasks.clone(), + max_concurrent: self.max_concurrent.clone(), + } + } +} + +pub struct MockIntentExecutor { + should_fail: bool, + execution_delay: Duration, + active_tasks: Option>, + max_concurrent: Option>, +} + +impl MockIntentExecutor { + fn on_task_started(&self) { + if let (Some(active), Some(max)) = + (&self.active_tasks, &self.max_concurrent) + { + let current = active.fetch_add(1, Ordering::SeqCst) + 1; + + let mut observed_max = max.load(Ordering::SeqCst); + while current > observed_max { + match max.compare_exchange_weak( + observed_max, + current, + Ordering::SeqCst, + Ordering::SeqCst, + ) { + Ok(_) => break, + Err(x) => observed_max = x, + } + } + } + } + + fn on_task_finished(&self) { + if let Some(active) = &self.active_tasks { + active.fetch_sub(1, Ordering::SeqCst); + } + } +} + +#[async_trait] +impl IntentExecutor for MockIntentExecutor { + async fn execute( + &mut self, + _base_intent: ScheduledIntentBundle, + _persister: Option

, + ) -> IntentExecutionResult { + self.on_task_started(); + tokio::time::sleep(self.execution_delay).await; + + let result = if self.should_fail { + IntentExecutionResult { + inner: Err(ExecutorError::FailedToCommitError { + err: TransactionStrategyExecutionError::InternalError( + InternalError::SignerError(SignerError::Custom( + "oops".to_string(), + )), + ), + signature: None, + }), + patched_errors: vec![ + TransactionStrategyExecutionError::ActionsError( + TransactionError::AccountNotFound, + None, + ), + ], + callbacks_report: vec![], + } + } else { + IntentExecutionResult { + inner: Ok(ExecutionOutput::TwoStage { + commit_signature: Signature::default(), + finalize_signature: Signature::default(), + }), + patched_errors: vec![], + callbacks_report: vec![], + } + }; + + self.on_task_finished(); + result + } + + async fn cleanup(self) -> Result<(), BufferExecutionError> { + Ok(()) + } +} diff --git a/magicblock-committor-service/src/service.rs b/magicblock-committor-service/src/service.rs index b1415b8bc..5ba5a8b7b 100644 --- a/magicblock-committor-service/src/service.rs +++ b/magicblock-committor-service/src/service.rs @@ -317,24 +317,91 @@ impl CommittorActor { } } - #[instrument(skip(self, cancel_token))] - pub async fn run(&mut self, cancel_token: CancellationToken) { + fn reject_msg(msg: CommittorMessage) { + use CommittorMessage::*; + fn shutdown_err( + respond_to: oneshot::Sender>, + ) { + let _ = respond_to.send(Err(CommittorServiceError::ShuttingDown)); + } + + match msg { + ReservePubkeysForCommittee { respond_to, .. } => { + shutdown_err(respond_to) + } + ReserveCommonPubkeys { respond_to } => shutdown_err(respond_to), + ScheduleIntentBundle { respond_to, .. } => shutdown_err(respond_to), + GetPendingIntentBundles { respond_to } => shutdown_err(respond_to), + ScheduleRecoveredIntentBundle { respond_to, .. } => { + shutdown_err(respond_to) + } + GetCommitStatuses { respond_to, .. } => shutdown_err(respond_to), + GetCommitSignatures { respond_to, .. } => shutdown_err(respond_to), + GetTransaction { respond_to, .. } => shutdown_err(respond_to), + FetchCurrentCommitNonces { respond_to, .. } => { + shutdown_err(respond_to) + } + FetchCurrentCommitNoncesSync { respond_to, .. } => { + let _ = + respond_to.send(Err(CommittorServiceError::ShuttingDown)); + } + + ReleaseCommonPubkeys { respond_to } => { + let _ = respond_to.send(()); + } + GetLookupTables { respond_to } => { + let _ = respond_to.send(LookupTables { + active: Vec::new(), + released: Vec::new(), + }); + } + SubscribeForResults { respond_to } => { + let (_tx, rx) = broadcast::channel(1); + let _ = respond_to.send(rx); + } + } + } + + fn drain_rejected_messages( + receiver: &mut mpsc::Receiver, + ) { + while let Ok(msg) = receiver.try_recv() { + Self::reject_msg(msg); + } + } + + #[instrument(skip(self, shutdown_token))] + pub async fn run(&mut self, shutdown_token: CancellationToken) { + let mut shutting_down = false; + loop { + if shutting_down { + break; + } + select! { + biased; + _ = shutdown_token.cancelled(), if !shutting_down => { + shutting_down = true; + info!("CommittorActor shutdown initiated"); + Self::drain_rejected_messages(&mut self.receiver); + } msg = self.receiver.recv() => { - if let Some(msg) = msg { - self.handle_msg(msg).await; - } else { - break; + match msg { + Some(msg) if !shutting_down => { + self.handle_msg(msg).await; + } + Some(msg) => { + Self::reject_msg(msg); + } + None => break, } } - _ = cancel_token.cancelled() => { - break; - } } } - info!("Actor shutdown"); + self.processor.shutdown().await; + info!("CommittorActor shutdown complete"); } } @@ -343,7 +410,8 @@ impl CommittorActor { // ----------------- pub struct CommittorService { sender: mpsc::Sender, - cancel_token: CancellationToken, + shutdown_token: CancellationToken, + actor_done_token: CancellationToken, } impl CommittorService { @@ -360,9 +428,11 @@ impl CommittorService { { debug!("Starting committor service"); let (sender, receiver) = mpsc::channel(1_000); - let cancel_token = CancellationToken::new(); + let shutdown_token = CancellationToken::new(); + let actor_done_token = CancellationToken::new(); { - let cancel_token = cancel_token.clone(); + let shutdown_token = shutdown_token.clone(); + let actor_done_token = actor_done_token.clone(); let mut actor = CommittorActor::try_new( receiver, authority, @@ -372,12 +442,14 @@ impl CommittorService { actions_callback_executor, )?; tokio::spawn(async move { - actor.run(cancel_token).await; + actor.run(shutdown_token).await; + actor_done_token.cancel(); }); } Ok(Self { sender, - cancel_token, + shutdown_token, + actor_done_token, }) } @@ -569,11 +641,11 @@ impl BaseIntentCommittor for CommittorService { } fn stop(&self) { - self.cancel_token.cancel(); + self.shutdown_token.cancel(); } fn stopped(&self) -> WaitForCancellationFutureOwned { - self.cancel_token.clone().cancelled_owned() + self.actor_done_token.clone().cancelled_owned() } } @@ -625,6 +697,201 @@ pub trait BaseIntentCommittor: Send + Sync + 'static { /// Stops Committor service fn stop(&self); - /// Returns future which resolves once committor `stop` got called + /// Returns future which resolves once the committor actor has finished + /// draining in-flight intents and exited fn stopped(&self) -> WaitForCancellationFutureOwned; } + +#[cfg(test)] +mod shutdown_tests { + use std::{sync::Arc, time::Duration}; + + use solana_pubkey::pubkey; + + use super::*; + use crate::{ + intent_execution_manager::{ + intent_scheduler::create_test_intent, + test_support::MockIntentExecutorFactory, + }, + test_utils, + }; + + fn try_start_for_test( + factory: MockIntentExecutorFactory, + ) -> Arc { + test_utils::init_test_logger(); + + let (sender, receiver) = mpsc::channel(1000); + let processor = Arc::new(CommittorProcessor::new_for_test(factory)); + let shutdown_token = CancellationToken::new(); + let actor_done_token = CancellationToken::new(); + let mut actor = CommittorActor { + receiver, + processor, + }; + tokio::spawn({ + let shutdown_token = shutdown_token.clone(); + let actor_done_token = actor_done_token.clone(); + async move { + actor.run(shutdown_token).await; + actor_done_token.cancel(); + } + }); + + Arc::new(CommittorService { + sender, + shutdown_token, + actor_done_token, + }) + } + + #[test] + fn reject_msg_responds_with_shutting_down_for_schedule() { + let (respond_to, mut response) = oneshot::channel(); + CommittorActor::reject_msg(CommittorMessage::ScheduleIntentBundle { + intent_bundles: Vec::new(), + respond_to, + }); + + let err = response.try_recv().unwrap().unwrap_err(); + assert!(matches!(err, CommittorServiceError::ShuttingDown)); + } + + #[test] + fn reject_msg_returns_empty_lookup_tables() { + let (respond_to, mut response) = oneshot::channel(); + CommittorActor::reject_msg(CommittorMessage::GetLookupTables { + respond_to, + }); + + let tables = response.try_recv().unwrap(); + assert!(tables.active.is_empty()); + assert!(tables.released.is_empty()); + } + + #[tokio::test] + async fn test_stopped_waits_for_in_flight_intent() { + let service = try_start_for_test( + MockIntentExecutorFactory::new() + .with_execution_delay(Duration::from_millis(200)), + ); + let mut result_receiver = + service.subscribe_for_results().await.unwrap(); + + let intent = create_test_intent( + 1, + &[pubkey!("1111111111111111111111111111111111111111111")], + false, + ); + service + .schedule_intent_bundles(vec![intent]) + .await + .unwrap() + .unwrap(); + + let stopped_task = { + let service = service.clone(); + tokio::spawn(async move { service.stopped().await }) + }; + + tokio::time::sleep(Duration::from_millis(20)).await; + assert!( + !stopped_task.is_finished(), + "stopped() must not complete before shutdown is requested" + ); + + service.stop(); + + let result = tokio::time::timeout( + Duration::from_secs(5), + result_receiver.recv(), + ) + .await + .expect("timed out waiting for in-flight intent result") + .expect("result channel closed"); + assert!(result.is_ok()); + assert_eq!(result.id, 1); + + tokio::time::timeout(Duration::from_secs(5), stopped_task) + .await + .expect("timed out waiting for actor shutdown") + .expect("stopped task panicked"); + } + + #[tokio::test] + async fn test_queued_schedule_rejected_when_shutdown_drains_channel() { + let service = try_start_for_test( + MockIntentExecutorFactory::new() + .with_execution_delay(Duration::from_millis(300)), + ); + + let blocking_intent = create_test_intent( + 1, + &[pubkey!("1111111111111111111111111111111111111111111")], + false, + ); + let queued_intent = create_test_intent( + 2, + &[pubkey!("22222222222222222222222222222222222222222222")], + false, + ); + + service + .schedule_intent_bundles(vec![blocking_intent]) + .await + .unwrap() + .unwrap(); + + let queued_schedule = + service.schedule_intent_bundles(vec![queued_intent]); + service.stop(); + + let err = tokio::time::timeout(Duration::from_secs(5), queued_schedule) + .await + .expect("timed out waiting for queued schedule response") + .unwrap() + .unwrap_err(); + assert!(matches!(err, CommittorServiceError::ShuttingDown)); + + tokio::time::timeout(Duration::from_secs(5), service.stopped()) + .await + .expect("timed out waiting for actor shutdown"); + } + + #[tokio::test] + async fn test_service_ext_waits_for_in_flight_result_during_shutdown() { + use crate::service_ext::{BaseIntentCommittorExt, CommittorServiceExt}; + + let service = try_start_for_test( + MockIntentExecutorFactory::new() + .with_execution_delay(Duration::from_millis(200)), + ); + let ext = CommittorServiceExt::new(service.clone()); + + let intent = create_test_intent( + 1, + &[pubkey!("1111111111111111111111111111111111111111111")], + false, + ); + let waiting = tokio::spawn(async move { + ext.schedule_intent_bundles_waiting(vec![intent]).await + }); + + tokio::time::sleep(Duration::from_millis(20)).await; + service.stop(); + + let results = tokio::time::timeout(Duration::from_secs(5), waiting) + .await + .expect("timed out waiting for in-flight result via CommittorServiceExt") + .expect("waiting task panicked") + .expect("schedule_intent_bundles_waiting failed"); + assert_eq!(results.len(), 1); + assert_eq!(results[0].id, 1); + assert!(results[0].inner.is_ok()); + + tokio::time::timeout(Duration::from_secs(5), service.stopped()) + .await + .expect("timed out waiting for actor shutdown"); + } +} diff --git a/magicblock-committor-service/src/service_ext.rs b/magicblock-committor-service/src/service_ext.rs index 024d0423e..5a55dc9b1 100644 --- a/magicblock-committor-service/src/service_ext.rs +++ b/magicblock-committor-service/src/service_ext.rs @@ -13,7 +13,7 @@ use solana_signature::Signature; use solana_transaction_status_client_types::EncodedConfirmedTransactionWithStatusMeta; use tokio::sync::{broadcast, oneshot, oneshot::error::RecvError}; use tokio_util::sync::WaitForCancellationFutureOwned; -use tracing::{error, info, instrument}; +use tracing::{error, info, instrument, warn}; use crate::{ error::{CommittorServiceError, CommittorServiceResult}, @@ -72,12 +72,14 @@ impl CommittorServiceExt { let mut results_subscription = results_subscription.await.unwrap(); tokio::pin!(committor_stopped); + let mut draining = false; loop { let execution_result = tokio::select! { biased; - _ = &mut committor_stopped => { - info!("Shutting down extension"); - return; + _ = &mut committor_stopped, if !draining => { + draining = true; + info!("CommittorServiceExt draining in-flight results"); + continue; } execution_result = results_subscription.recv() => { match execution_result { @@ -113,6 +115,20 @@ impl CommittorServiceExt { ); } } + + // A poisoned mutex means a prior panic occurred while holding this + // lock; continuing would be unsafe, so we abort. + let pending_count = pending_message + .lock() + .expect(POISONED_MUTEX_MSG) + .drain() + .count(); + if pending_count > 0 { + warn!( + pending_count, + "Dropped pending intent result waiters during CommittorServiceExt shutdown" + ); + } } } @@ -251,3 +267,159 @@ impl From for CommittorServiceExtError { pub type BaseIntentCommitorExtResult = Result; + +#[cfg(test)] +mod shutdown_tests { + use std::{ + collections::HashMap, + sync::Arc, + time::{Duration, Instant}, + }; + + use magicblock_program::magic_scheduled_base_intent::ScheduledIntentBundle; + use solana_pubkey::Pubkey; + use solana_signature::Signature; + use solana_transaction_status_client_types::EncodedConfirmedTransactionWithStatusMeta; + use tokio::sync::{broadcast, oneshot}; + use tokio_util::sync::CancellationToken; + + use super::*; + use crate::{ + intent_execution_manager::intent_scheduler::create_test_intent, + intent_executor::ExecutionOutput, test_utils, + }; + + struct MockCommittor { + actor_done_token: CancellationToken, + results_tx: broadcast::Sender, + } + + impl MockCommittor { + fn new() -> Arc { + let (results_tx, _) = broadcast::channel(16); + Arc::new(Self { + actor_done_token: CancellationToken::new(), + results_tx, + }) + } + + fn signal_stopped(&self) { + self.actor_done_token.cancel(); + } + + fn broadcast_result(&self, result: BroadcastedIntentExecutionResult) { + self.results_tx.send(result).expect("broadcast result"); + } + } + + fn test_result(id: u64) -> BroadcastedIntentExecutionResult { + BroadcastedIntentExecutionResult { + id, + inner: Ok(ExecutionOutput::SingleStage(Signature::new_unique())), + patched_errors: Arc::new(vec![]), + callbacks_report: vec![], + } + } + + fn empty_oneshot() -> oneshot::Receiver { + let (_, rx) = oneshot::channel(); + rx + } + + impl BaseIntentCommittor for MockCommittor { + fn reserve_pubkeys_for_committee( + &self, + _: Pubkey, + _: Pubkey, + ) -> oneshot::Receiver> { + empty_oneshot() + } + + fn schedule_intent_bundles( + &self, + _: Vec, + ) -> oneshot::Receiver> { + let (tx, rx) = oneshot::channel(); + tx.send(Ok(())).expect("schedule response"); + rx + } + + fn subscribe_for_results( + &self, + ) -> oneshot::Receiver< + broadcast::Receiver, + > { + let (tx, rx) = oneshot::channel(); + tx.send(self.results_tx.subscribe()) + .expect("results subscription"); + rx + } + + fn get_commit_statuses( + &self, + _: u64, + ) -> oneshot::Receiver>> + { + empty_oneshot() + } + + fn get_commit_signatures( + &self, + _: u64, + _: Pubkey, + ) -> oneshot::Receiver>> + { + empty_oneshot() + } + + fn get_transaction( + &self, + _: &Signature, + ) -> oneshot::Receiver< + CommittorServiceResult, + > { + empty_oneshot() + } + + fn fetch_current_commit_nonces( + &self, + _: &[Pubkey], + _: u64, + ) -> oneshot::Receiver>> + { + empty_oneshot() + } + + fn stop(&self) {} + + fn stopped(&self) -> WaitForCancellationFutureOwned { + self.actor_done_token.clone().cancelled_owned() + } + } + + #[tokio::test] + async fn dispatcher_delivers_result_after_stopped_signal() { + test_utils::init_test_logger(); + + let mock = MockCommittor::new(); + let ext = CommittorServiceExt::new(mock.clone()); + + let intent = create_test_intent(1, &[Pubkey::new_unique()], false); + let waiting = tokio::spawn(async move { + ext.schedule_intent_bundles_waiting(vec![intent]).await + }); + + tokio::time::sleep(Duration::from_millis(20)).await; + mock.signal_stopped(); + mock.broadcast_result(test_result(1)); + + let results = tokio::time::timeout(Duration::from_secs(5), waiting) + .await + .expect("timed out waiting for result delivery") + .expect("waiting task panicked") + .expect("schedule_intent_bundles_waiting failed"); + assert_eq!(results.len(), 1); + assert_eq!(results[0].id, 1); + assert!(results[0].inner.is_ok()); + } +}