blvm_sdk/module/
runner.rs

1//! Module runner and invocation context.
2//!
3//! Provides `InvocationContext` for CLI/RPC handlers, `run_async` for sync-over-async CLI,
4//! and `run_module` for the unified connect/dispatch/loop lifecycle.
5
6use blvm_node::module::integration::ModuleIntegration;
7use blvm_node::module::ipc::protocol::{
8    CliSpec, InvocationMessage, InvocationResultMessage, InvocationResultPayload, InvocationType,
9    ModuleMessage,
10};
11use blvm_node::module::traits::{ModuleError, NodeAPI};
12use blvm_node::storage::database::Database;
13use std::collections::VecDeque;
14use std::path::Path;
15use std::pin::Pin;
16use std::sync::Arc;
17use tokio::sync::oneshot;
18use tokio::time::{Duration, sleep};
19use tracing::info;
20
21use crate::module::timers::{Exec, TimerRegistry, sdk_tick_spec};
22
23use crate::module::storage::{DatabaseStorageAdapter, ModuleStorage, ModuleStorageDatabaseBridge};
24
25/// Core RPC allowlist: registered with [`NodeAPI::register_core_rpc_override`] in module setup,
26/// not via [`NodeAPI::register_rpc_endpoint`] (the RPC server rejects extension registration for them).
27fn is_overrideable_core_rpc_method(method: &str) -> bool {
28    blvm_node::rpc::methods::OVERRIDABLE_CORE_RPC_METHODS.contains(&method)
29}
30
31/// Run an async future from a sync context (e.g. CLI handler).
32/// Blocks the current thread and executes the future on the current runtime.
33/// Use when `#[command]` methods need to call async APIs.
34///
35/// When the future only returns `Ok(_)` with no error path, use `Ok::<_, String>(...)` to fix inference.
36pub fn run_async<F, T, E>(f: F) -> Result<T, ModuleError>
37where
38    F: std::future::Future<Output = Result<T, E>>,
39    E: std::fmt::Display,
40{
41    tokio::task::block_in_place(|| tokio::runtime::Handle::current().block_on(f))
42        .map_err(|e| ModuleError::Other(e.to_string()))
43}
44
45/// Context passed to CLI handlers for database and config access.
46///
47/// Uses `ModuleStorage` internally; `ctx.db()` returns the database interface for compatibility.
48/// When connected to a node, `node_api()` provides access to blockchain data (e.g. get_transaction).
49#[derive(Clone)]
50pub struct InvocationContext {
51    db: Arc<dyn Database>,
52    node_api: Option<Arc<dyn NodeAPI>>,
53    timers: Option<Arc<TimerRegistry>>,
54}
55
56impl InvocationContext {
57    /// Create a new invocation context from module storage.
58    pub fn from_storage(storage: Arc<dyn ModuleStorage>) -> Self {
59        let db = Arc::new(ModuleStorageDatabaseBridge::new(storage));
60        Self {
61            db,
62            node_api: None,
63            timers: None,
64        }
65    }
66
67    /// Create a new invocation context from a database (legacy; wraps in ModuleStorage).
68    pub fn new(db: Arc<dyn Database>) -> Self {
69        let storage = Arc::new(DatabaseStorageAdapter::new(db));
70        Self::from_storage(storage)
71    }
72
73    /// Create invocation context with NodeAPI for CLI commands that need blockchain access.
74    pub fn with_node_api(db: Arc<dyn Database>, node_api: Arc<dyn NodeAPI>) -> Self {
75        let storage = Arc::new(DatabaseStorageAdapter::new(db));
76        Self {
77            db: Arc::new(ModuleStorageDatabaseBridge::new(storage)),
78            node_api: Some(node_api),
79            timers: None,
80        }
81    }
82
83    pub(crate) fn with_registry(mut self, timers: Arc<TimerRegistry>) -> Self {
84        self.timers = Some(timers);
85        self
86    }
87
88    /// Builder bound to this module's node connection and timer registry.
89    pub fn timers(&self) -> crate::module::timers::TimerBuilder {
90        match (&self.node_api, &self.timers) {
91            (Some(api), Some(registry)) => crate::module::timers::TimerBuilder::with_runtime(
92                Arc::clone(api),
93                Arc::clone(registry),
94            ),
95            _ => crate::module::timers::TimerBuilder::new(),
96        }
97    }
98
99    /// Get the module's database.
100    pub fn db(&self) -> &Arc<dyn Database> {
101        &self.db
102    }
103
104    /// Get NodeAPI when connected to node (for fetch-by-txid, etc.).
105    pub fn node_api(&self) -> Option<Arc<dyn NodeAPI>> {
106        self.node_api.clone()
107    }
108}
109
110/// Run a module with automatic connect, registration, event subscription, and dispatch.
111///
112/// Handles the full lifecycle: connect → register CLI/RPC/events → loop (invocations + events) → unload on disconnect.
113#[allow(clippy::too_many_arguments)] // Explicit wiring for embedders; splitting would obscure the lifecycle.
114pub async fn run_module<M, C, F, FE, Fut>(
115    socket_path: impl AsRef<Path>,
116    module_id: &str,
117    module_name: &str,
118    version: &str,
119    cli_spec: CliSpec,
120    rpc_methods: &[&str],
121    event_types: Vec<blvm_node::module::traits::EventType>,
122    dispatch: F,
123    on_event: FE,
124    module: M,
125    cli: C,
126    db: Arc<dyn Database>,
127) -> Result<(), ModuleError>
128where
129    F: Fn(InvocationMessage, InvocationContext, &M, &C) -> InvocationResultMessage,
130    FE: Fn(blvm_node::module::ipc::protocol::EventMessage, &M, &InvocationContext) -> Fut,
131    Fut: std::future::Future<Output = Result<(), ModuleError>> + Send,
132{
133    let socket_path = socket_path.as_ref().to_path_buf();
134
135    match ModuleIntegration::connect(
136        socket_path.clone(),
137        module_id.to_string(),
138        module_name.to_string(),
139        version.to_string(),
140        Some(cli_spec),
141    )
142    .await
143    {
144        Ok(mut integration) => {
145            info!("Connected to node");
146
147            let node_api = integration.node_api();
148            for method in rpc_methods {
149                if is_overrideable_core_rpc_method(method) {
150                    continue;
151                }
152                node_api
153                    .register_rpc_endpoint((*method).to_string(), String::new())
154                    .await?;
155            }
156
157            integration.subscribe_events(event_types).await?;
158
159            let event_rx = integration.event_receiver();
160            let invocation_rx = integration.invocation_receiver().ok_or_else(|| {
161                ModuleError::IpcError(
162                    "Invocation receiver not available for this module integration".to_string(),
163                )
164            })?;
165            let timers = TimerRegistry::new();
166            let ctx =
167                InvocationContext::with_node_api(db, node_api).with_registry(Arc::clone(&timers));
168            module_loop(
169                event_rx,
170                invocation_rx,
171                &module,
172                &cli,
173                ctx,
174                dispatch,
175                on_event,
176                None,
177                timers,
178            )
179            .await;
180        }
181        Err(e) => {
182            info!("Node not running, standalone mode: {}", e);
183            loop {
184                sleep(Duration::from_secs(5)).await;
185            }
186        }
187    }
188
189    Ok(())
190}
191
192/// Run a module where (module, cli) are created after connect.
193///
194/// Use when the module depends on NodeAPI (e.g. datum creates DatumServer with node_api).
195/// The setup receives (node_api, db, data_dir) and returns (module, cli).
196#[allow(clippy::too_many_arguments)]
197pub async fn run_module_with_setup<M, C, F, FE, Fut, FSetup, FutSetup>(
198    socket_path: impl AsRef<Path>,
199    module_id: &str,
200    module_name: &str,
201    version: &str,
202    cli_spec: CliSpec,
203    rpc_methods: &[&str],
204    event_types: Vec<blvm_node::module::traits::EventType>,
205    dispatch: F,
206    on_event: FE,
207    setup: FSetup,
208    db: Arc<dyn Database>,
209    data_dir: &Path,
210) -> Result<(), ModuleError>
211where
212    F: Fn(InvocationMessage, InvocationContext, &M, &C) -> InvocationResultMessage,
213    FE: Fn(blvm_node::module::ipc::protocol::EventMessage, &M, &InvocationContext) -> Fut,
214    Fut: std::future::Future<Output = Result<(), ModuleError>> + Send,
215    FSetup: Fn(Arc<dyn NodeAPI>, Arc<dyn Database>, &Path) -> FutSetup,
216    FutSetup: std::future::Future<Output = Result<(M, C), ModuleError>> + Send,
217{
218    let socket_path = socket_path.as_ref().to_path_buf();
219
220    match ModuleIntegration::connect(
221        socket_path.clone(),
222        module_id.to_string(),
223        module_name.to_string(),
224        version.to_string(),
225        Some(cli_spec),
226    )
227    .await
228    {
229        Ok(mut integration) => {
230            info!("Connected to node");
231
232            let node_api = integration.node_api();
233            for method in rpc_methods {
234                if is_overrideable_core_rpc_method(method) {
235                    continue;
236                }
237                node_api
238                    .register_rpc_endpoint((*method).to_string(), String::new())
239                    .await?;
240            }
241
242            integration.subscribe_events(event_types).await?;
243
244            let (module, cli) = setup(node_api.clone(), Arc::clone(&db), data_dir).await?;
245            let module = Arc::new(module);
246
247            let event_rx = integration.event_receiver();
248            let invocation_rx = integration.invocation_receiver().ok_or_else(|| {
249                ModuleError::IpcError(
250                    "Invocation receiver not available for this module integration".to_string(),
251                )
252            })?;
253            let timers = TimerRegistry::new();
254            let ctx = InvocationContext::with_node_api(Arc::clone(&db), node_api)
255                .with_registry(Arc::clone(&timers));
256            module_loop(
257                event_rx,
258                invocation_rx,
259                &*module,
260                &cli,
261                ctx,
262                dispatch,
263                on_event,
264                None,
265                timers,
266            )
267            .await;
268        }
269        Err(e) => {
270            info!("Node not running, standalone mode: {}", e);
271            loop {
272                sleep(Duration::from_secs(5)).await;
273            }
274        }
275    }
276
277    Ok(())
278}
279
280/// Like [`run_module_with_setup`], but setup also returns a `ModuleAPI` for IPC forwarding.
281#[allow(clippy::too_many_arguments)]
282pub async fn run_module_with_setup_and_api<M, C, F, FE, Fut, FSetup, FutSetup>(
283    socket_path: impl AsRef<Path>,
284    module_id: &str,
285    module_name: &str,
286    version: &str,
287    cli_spec: CliSpec,
288    rpc_methods: &[&str],
289    event_types: Vec<blvm_node::module::traits::EventType>,
290    dispatch: F,
291    on_event: FE,
292    setup: FSetup,
293    db: Arc<dyn Database>,
294    data_dir: &Path,
295) -> Result<(), ModuleError>
296where
297    F: Fn(InvocationMessage, InvocationContext, &M, &C) -> InvocationResultMessage,
298    FE: Fn(blvm_node::module::ipc::protocol::EventMessage, &M, &InvocationContext) -> Fut,
299    Fut: std::future::Future<Output = Result<(), ModuleError>> + Send,
300    FSetup: Fn(Arc<dyn NodeAPI>, Arc<dyn Database>, &Path) -> FutSetup,
301    FutSetup: std::future::Future<
302            Output = Result<
303                (
304                    M,
305                    C,
306                    Arc<dyn blvm_node::module::inter_module::api::ModuleAPI>,
307                ),
308                ModuleError,
309            >,
310        > + Send,
311{
312    let socket_path = socket_path.as_ref().to_path_buf();
313
314    match ModuleIntegration::connect(
315        socket_path.clone(),
316        module_id.to_string(),
317        module_name.to_string(),
318        version.to_string(),
319        Some(cli_spec),
320    )
321    .await
322    {
323        Ok(mut integration) => {
324            info!("Connected to node");
325
326            let node_api = integration.node_api();
327            for method in rpc_methods {
328                if is_overrideable_core_rpc_method(method) {
329                    continue;
330                }
331                node_api
332                    .register_rpc_endpoint((*method).to_string(), String::new())
333                    .await?;
334            }
335
336            integration.subscribe_events(event_types).await?;
337
338            let (module, cli, module_api) =
339                setup(node_api.clone(), Arc::clone(&db), data_dir).await?;
340            let module = Arc::new(module);
341            let module_api = Arc::clone(&module_api);
342
343            if let Err(e) = node_api.register_module_api(module_api.clone()).await {
344                return Err(ModuleError::Other(format!(
345                    "Failed to register module API descriptor: {e}"
346                )));
347            }
348            info!("Module API descriptor registered with node");
349
350            let event_rx = integration.event_receiver();
351            let invocation_rx = integration.invocation_receiver().ok_or_else(|| {
352                ModuleError::IpcError(
353                    "Invocation receiver not available for this module integration".to_string(),
354                )
355            })?;
356            let timers = TimerRegistry::new();
357            let ctx = InvocationContext::with_node_api(Arc::clone(&db), node_api)
358                .with_registry(Arc::clone(&timers));
359            module_loop(
360                event_rx,
361                invocation_rx,
362                &*module,
363                &cli,
364                ctx,
365                dispatch,
366                on_event,
367                Some(module_api),
368                timers,
369            )
370            .await;
371        }
372        Err(e) => {
373            info!("Node not running, standalone mode: {}", e);
374            loop {
375                sleep(Duration::from_secs(5)).await;
376            }
377        }
378    }
379
380    Ok(())
381}
382
383/// Run a module with optional on_connect (setup) and on_tick (periodic) callbacks.
384#[allow(clippy::too_many_arguments)]
385pub async fn run_module_with_tick<M, C, F, FE, Fut, FConnect, FutConnect, FTick, FutTick>(
386    socket_path: impl AsRef<Path>,
387    module_id: &str,
388    module_name: &str,
389    version: &str,
390    cli_spec: CliSpec,
391    rpc_methods: &[&str],
392    event_types: Vec<blvm_node::module::traits::EventType>,
393    dispatch: F,
394    on_event: FE,
395    on_connect: Option<FConnect>,
396    on_tick: Option<FTick>,
397    module: M,
398    cli: C,
399    db: Arc<dyn Database>,
400) -> Result<(), ModuleError>
401where
402    F: Fn(InvocationMessage, InvocationContext, &M, &C) -> InvocationResultMessage,
403    FE: Fn(blvm_node::module::ipc::protocol::EventMessage, &M, &InvocationContext) -> Fut,
404    Fut: std::future::Future<Output = Result<(), ModuleError>> + Send,
405    FConnect: Fn(Arc<dyn NodeAPI>, Arc<dyn Database>) -> FutConnect,
406    FutConnect: std::future::Future<Output = Result<(), ModuleError>> + Send,
407    FTick: Fn(Arc<dyn NodeAPI>, Arc<dyn Database>) -> FutTick + Send + Sync + 'static,
408    FutTick: std::future::Future<Output = ()> + Send + 'static,
409{
410    let socket_path = socket_path.as_ref().to_path_buf();
411
412    match ModuleIntegration::connect(
413        socket_path.clone(),
414        module_id.to_string(),
415        module_name.to_string(),
416        version.to_string(),
417        Some(cli_spec),
418    )
419    .await
420    {
421        Ok(mut integration) => {
422            info!("Connected to node");
423
424            let node_api = integration.node_api();
425            for method in rpc_methods {
426                if is_overrideable_core_rpc_method(method) {
427                    continue;
428                }
429                node_api
430                    .register_rpc_endpoint((*method).to_string(), String::new())
431                    .await?;
432            }
433
434            integration.subscribe_events(event_types).await?;
435
436            if let Some(ref connect) = on_connect {
437                connect(node_api.clone(), Arc::clone(&db)).await?;
438            }
439
440            let event_rx = integration.event_receiver();
441            let invocation_rx = integration.invocation_receiver().ok_or_else(|| {
442                ModuleError::IpcError(
443                    "Invocation receiver not available for this module integration".to_string(),
444                )
445            })?;
446            let timers = TimerRegistry::new();
447            if let Some(on_tick) = on_tick {
448                register_sdk_tick(&node_api, &db, &timers, on_tick).await;
449            }
450            let ctx = InvocationContext::with_node_api(Arc::clone(&db), Arc::clone(&node_api))
451                .with_registry(Arc::clone(&timers));
452            module_loop(
453                event_rx,
454                invocation_rx,
455                &module,
456                &cli,
457                ctx,
458                dispatch,
459                on_event,
460                None,
461                timers,
462            )
463            .await;
464        }
465        Err(e) => {
466            info!("Node not running, standalone mode: {}", e);
467            loop {
468                sleep(Duration::from_secs(5)).await;
469            }
470        }
471    }
472
473    Ok(())
474}
475
476async fn register_sdk_tick<FTick, FutTick>(
477    node_api: &Arc<dyn NodeAPI>,
478    db: &Arc<dyn Database>,
479    timers: &TimerRegistry,
480    on_tick: FTick,
481) where
482    FTick: Fn(Arc<dyn NodeAPI>, Arc<dyn Database>) -> FutTick + Send + Sync + 'static,
483    FutTick: std::future::Future<Output = ()> + Send + 'static,
484{
485    use blvm_node::module::timers::types::{AckMode, CreateMode, TimerAck};
486
487    let info = match node_api
488        .timer_create(sdk_tick_spec(), CreateMode::Ensure)
489        .await
490    {
491        Ok(info) => info,
492        Err(err) => {
493            info!("sdk tick was not registered: {err}");
494            return;
495        }
496    };
497    let api = Arc::clone(node_api);
498    let db = Arc::clone(db);
499    let tick = Arc::new(on_tick);
500    timers.register(
501        info.timer_id,
502        Exec::Serial,
503        AckMode::Completion,
504        Arc::new(move |_fire| {
505            let api = Arc::clone(&api);
506            let db = Arc::clone(&db);
507            let tick = Arc::clone(&tick);
508            Box::pin(async move {
509                tick(api, db).await;
510                info!("Module running");
511                Ok(TimerAck::Continue)
512            })
513        }),
514    );
515}
516
517enum LoopWake {
518    SerialDone,
519    Event(Box<Result<ModuleMessage, tokio::sync::broadcast::error::RecvError>>),
520    Invocation(Option<(InvocationMessage, oneshot::Sender<InvocationResultMessage>)>),
521}
522
523/// One module loop. A `Serial` timer is awaited here. `Task`, `Blocking`, and
524/// `Dedicated` fires are started without waiting, including while a serial handler runs.
525/// Events that arrive during a serial handler wait until it returns.
526#[allow(clippy::too_many_arguments)]
527pub(crate) async fn module_loop<M, C, F, FE, Fut>(
528    mut event_rx: tokio::sync::broadcast::Receiver<ModuleMessage>,
529    invocation_rx: &mut tokio::sync::mpsc::Receiver<(
530        InvocationMessage,
531        oneshot::Sender<InvocationResultMessage>,
532    )>,
533    module: &M,
534    cli: &C,
535    ctx: InvocationContext,
536    dispatch: F,
537    on_event: FE,
538    module_api: Option<Arc<dyn blvm_node::module::inter_module::api::ModuleAPI>>,
539    timers: Arc<TimerRegistry>,
540) where
541    F: Fn(InvocationMessage, InvocationContext, &M, &C) -> InvocationResultMessage,
542    FE: Fn(blvm_node::module::ipc::protocol::EventMessage, &M, &InvocationContext) -> Fut,
543    Fut: std::future::Future<Output = Result<(), ModuleError>> + Send,
544{
545    let mut queued_events = VecDeque::new();
546    let mut queued_invocations = VecDeque::new();
547    let mut serial: Option<Pin<Box<dyn std::future::Future<Output = ()> + Send>>> = None;
548    let mut closing = false;
549
550    loop {
551        if serial.is_none() {
552            if let Some(event) = queued_events.pop_front() {
553                let _ = on_event(event, module, &ctx).await;
554                continue;
555            }
556            if let Some((invocation, result_tx)) = queued_invocations.pop_front() {
557                if let Some(job) = dispatch_invocation(
558                    invocation,
559                    result_tx,
560                    module,
561                    cli,
562                    &ctx,
563                    &dispatch,
564                    module_api.as_ref(),
565                    &timers,
566                )
567                .await
568                {
569                    serial = Some(job);
570                }
571                continue;
572            }
573            if closing {
574                info!("Invocation channel closed, module unloading");
575                break;
576            }
577        }
578
579        if let Some(mut job) = serial.take() {
580            let wake = tokio::select! {
581                _ = &mut job => LoopWake::SerialDone,
582                msg = event_rx.recv() => LoopWake::Event(Box::new(msg)),
583                inv = invocation_rx.recv() => LoopWake::Invocation(inv),
584            };
585            if !matches!(wake, LoopWake::SerialDone) {
586                serial = Some(job);
587            }
588            match wake {
589                LoopWake::SerialDone => {}
590                LoopWake::Event(msg) => {
591                    if let Ok(ModuleMessage::Event(event)) = *msg {
592                        queued_events.push_back(event);
593                    }
594                }
595                LoopWake::Invocation(Some((invocation, result_tx))) => {
596                    if detached_timer(&timers, &invocation) {
597                        if let InvocationType::Timer(fire) = invocation.invocation_type {
598                            timers.start_detached(fire, invocation.correlation_id, result_tx);
599                        }
600                    } else {
601                        queued_invocations.push_back((invocation, result_tx));
602                    }
603                }
604                LoopWake::Invocation(None) => closing = true,
605            }
606            continue;
607        }
608
609        tokio::select! {
610            msg = event_rx.recv() => {
611                if let Ok(ModuleMessage::Event(event)) = msg {
612                    let _ = on_event(event, module, &ctx).await;
613                }
614            }
615            inv = invocation_rx.recv() => {
616                match inv {
617                    Some((invocation, result_tx)) => {
618                        if let Some(job) = dispatch_invocation(
619                            invocation,
620                            result_tx,
621                            module,
622                            cli,
623                            &ctx,
624                            &dispatch,
625                            module_api.as_ref(),
626                            &timers,
627                        )
628                        .await
629                        {
630                            serial = Some(job);
631                        }
632                    }
633                    None => {
634                        info!("Invocation channel closed, module unloading");
635                        break;
636                    }
637                }
638            }
639        }
640    }
641}
642
643fn detached_timer(timers: &TimerRegistry, invocation: &InvocationMessage) -> bool {
644    match &invocation.invocation_type {
645        InvocationType::Timer(fire) => timers.is_detached(&fire.timer_id),
646        _ => false,
647    }
648}
649
650#[allow(clippy::too_many_arguments)]
651async fn dispatch_invocation<M, C, F>(
652    invocation: InvocationMessage,
653    result_tx: oneshot::Sender<InvocationResultMessage>,
654    module: &M,
655    cli: &C,
656    ctx: &InvocationContext,
657    dispatch: &F,
658    module_api: Option<&Arc<dyn blvm_node::module::inter_module::api::ModuleAPI>>,
659    timers: &Arc<TimerRegistry>,
660) -> Option<Pin<Box<dyn std::future::Future<Output = ()> + Send>>>
661where
662    F: Fn(InvocationMessage, InvocationContext, &M, &C) -> InvocationResultMessage,
663{
664    if let InvocationType::Timer(fire) = invocation.invocation_type.clone() {
665        if timers.is_detached(&fire.timer_id) {
666            timers.start_detached(fire, invocation.correlation_id, result_tx);
667            return None;
668        }
669        let correlation_id = invocation.correlation_id;
670        let timers = Arc::clone(timers);
671        return Some(Box::pin(async move {
672            timers.run_serial(correlation_id, &fire, result_tx).await;
673        }));
674    }
675
676    if let Some(module_api) = module_api {
677        if let InvocationType::ModuleApi {
678            method,
679            params,
680            caller_module_id,
681        } = &invocation.invocation_type
682        {
683            let result = match module_api
684                .handle_request(method, params, caller_module_id)
685                .await
686            {
687                Ok(data) => InvocationResultMessage {
688                    correlation_id: invocation.correlation_id,
689                    success: true,
690                    payload: Some(InvocationResultPayload::ModuleApi(data)),
691                    error: None,
692                },
693                Err(err) => InvocationResultMessage {
694                    correlation_id: invocation.correlation_id,
695                    success: false,
696                    payload: None,
697                    error: Some(err.to_string()),
698                },
699            };
700            let _ = result_tx.send(result);
701            return None;
702        }
703    }
704
705    let result = dispatch(invocation, ctx.clone(), module, cli);
706    let _ = result_tx.send(result);
707    None
708}
709
710#[cfg(test)]
711mod tests {
712    use super::*;
713    use crate::module::storage::ModuleStorage;
714    use crate::module::timers::Exec;
715    use blvm_node::module::ipc::protocol::{EventMessage, EventPayload};
716    use blvm_node::module::timers::types::{AckMode, Clock, TimerAck, TimerFire, TimerId};
717    use blvm_node::module::traits::EventType;
718    use std::sync::Mutex;
719    use std::sync::atomic::{AtomicUsize, Ordering};
720    use tokio::sync::{Notify, broadcast, mpsc};
721
722    struct NoStorage;
723
724    impl ModuleStorage for NoStorage {
725        fn open_tree(
726            &self,
727            _: &str,
728        ) -> anyhow::Result<Arc<dyn crate::module::storage::ModuleTree>> {
729            anyhow::bail!("unused")
730        }
731    }
732
733    fn context(timers: Arc<TimerRegistry>) -> InvocationContext {
734        InvocationContext::from_storage(Arc::new(NoStorage)).with_registry(timers)
735    }
736
737    fn fire(id: TimerId) -> TimerFire {
738        TimerFire {
739            timer_id: id,
740            key: None,
741            payload: Vec::new(),
742            seq: 1,
743            attempt: 1,
744            clock: Clock::Monotonic,
745            scheduled_unix_ms: 0,
746            planned_unix_ms: 0,
747            fired_unix_ms: 0,
748            lateness_ms: 0,
749            missed: 0,
750            coalesced: 0,
751            skipped_overlap: 0,
752            clock_step_ms: 0,
753            tip_height: None,
754            tip_hash: None,
755            tip_mtp: None,
756            is_last: false,
757        }
758    }
759
760    fn timer_message(id: TimerId, correlation_id: u64) -> InvocationMessage {
761        InvocationMessage {
762            correlation_id,
763            invocation_type: InvocationType::Timer(fire(id)),
764        }
765    }
766
767    fn event_message() -> ModuleMessage {
768        ModuleMessage::Event(EventMessage {
769            event_type: EventType::NewTransaction,
770            payload: EventPayload::NewTransaction {
771                tx_hash: blvm_node::Hash::default(),
772            },
773        })
774    }
775
776    fn dummy_result(invocation: InvocationMessage) -> InvocationResultMessage {
777        InvocationResultMessage {
778            correlation_id: invocation.correlation_id,
779            success: true,
780            payload: None,
781            error: None,
782        }
783    }
784
785    async fn until(order: &Mutex<Vec<&'static str>>, needle: &'static str) {
786        tokio::time::timeout(Duration::from_secs(2), async {
787            loop {
788                if order.lock().unwrap().contains(&needle) {
789                    break;
790                }
791                tokio::task::yield_now().await;
792            }
793        })
794        .await
795        .expect(needle);
796    }
797
798    #[tokio::test]
799    async fn serial_blocks_events_until_the_handler_returns() {
800        let timers = TimerRegistry::new();
801        let serial_id = TimerId([1; 16]);
802        let order = Arc::new(Mutex::new(Vec::new()));
803        let release = Arc::new(Notify::new());
804        let order_s = Arc::clone(&order);
805        let release_s = Arc::clone(&release);
806        timers.register(
807            serial_id,
808            Exec::Serial,
809            AckMode::Completion,
810            Arc::new(move |_fire| {
811                let order_s = Arc::clone(&order_s);
812                let release_s = Arc::clone(&release_s);
813                Box::pin(async move {
814                    order_s.lock().unwrap().push("serial-start");
815                    tokio::time::sleep(Duration::from_millis(20)).await;
816                    release_s.notified().await;
817                    order_s.lock().unwrap().push("serial-end");
818                    Ok(TimerAck::Continue)
819                })
820            }),
821        );
822        let (event_tx, event_rx) = broadcast::channel(32);
823        let (inv_tx, mut inv_rx) = mpsc::channel(8);
824        let order_e = Arc::clone(&order);
825        let timers_loop = Arc::clone(&timers);
826        let ctx = context(Arc::clone(&timers));
827        let running = tokio::spawn(async move {
828            module_loop(
829                event_rx,
830                &mut inv_rx,
831                &(),
832                &(),
833                ctx,
834                |invocation, _, _, _| dummy_result(invocation),
835                move |_event, _, _| {
836                    let order_e = Arc::clone(&order_e);
837                    async move {
838                        order_e.lock().unwrap().push("event");
839                        Ok(())
840                    }
841                },
842                None,
843                timers_loop,
844            )
845            .await;
846        });
847        inv_tx
848            .send((timer_message(serial_id, 1), oneshot::channel().0))
849            .await
850            .unwrap();
851        until(&order, "serial-start").await;
852        event_tx.send(event_message()).unwrap();
853        tokio::time::sleep(Duration::from_millis(30)).await;
854        assert_eq!(order.lock().unwrap().as_slice(), ["serial-start"]);
855        release.notify_one();
856        until(&order, "event").await;
857        assert_eq!(
858            order.lock().unwrap().as_slice(),
859            ["serial-start", "serial-end", "event"]
860        );
861        drop(inv_tx);
862        running.await.unwrap();
863    }
864
865    #[tokio::test]
866    async fn task_does_not_block_events() {
867        let timers = TimerRegistry::new();
868        let task_id = TimerId([2; 16]);
869        let order = Arc::new(Mutex::new(Vec::new()));
870        let order_t = Arc::clone(&order);
871        timers.register(
872            task_id,
873            Exec::Task,
874            AckMode::Completion,
875            Arc::new(move |_fire| {
876                let order_t = Arc::clone(&order_t);
877                Box::pin(async move {
878                    order_t.lock().unwrap().push("task-start");
879                    tokio::time::sleep(Duration::from_millis(40)).await;
880                    order_t.lock().unwrap().push("task-end");
881                    Ok(TimerAck::Continue)
882                })
883            }),
884        );
885        let (event_tx, event_rx) = broadcast::channel(32);
886        let (inv_tx, mut inv_rx) = mpsc::channel(8);
887        let order_e = Arc::clone(&order);
888        let timers_loop = Arc::clone(&timers);
889        let ctx = context(Arc::clone(&timers));
890        let running = tokio::spawn(async move {
891            module_loop(
892                event_rx,
893                &mut inv_rx,
894                &(),
895                &(),
896                ctx,
897                |invocation, _, _, _| dummy_result(invocation),
898                move |_event, _, _| {
899                    let order_e = Arc::clone(&order_e);
900                    async move {
901                        order_e.lock().unwrap().push("event");
902                        Ok(())
903                    }
904                },
905                None,
906                timers_loop,
907            )
908            .await;
909        });
910        let (result_tx, result_rx) = oneshot::channel();
911        inv_tx
912            .send((timer_message(task_id, 2), result_tx))
913            .await
914            .unwrap();
915        until(&order, "task-start").await;
916        event_tx.send(event_message()).unwrap();
917        until(&order, "event").await;
918        until(&order, "task-end").await;
919        let order = order.lock().unwrap().clone();
920        let event_at = order.iter().position(|item| *item == "event").unwrap();
921        let end_at = order.iter().position(|item| *item == "task-end").unwrap();
922        assert!(event_at < end_at);
923        assert!(result_rx.await.unwrap().success);
924        drop(inv_tx);
925        running.await.unwrap();
926    }
927
928    #[tokio::test]
929    async fn dedicated_is_not_starved_by_a_serial_handler() {
930        let timers = TimerRegistry::new();
931        let serial_id = TimerId([3; 16]);
932        let dedicated_id = TimerId([4; 16]);
933        let release = Arc::new(Notify::new());
934        let saw = Arc::new(AtomicUsize::new(0));
935        let release_s = Arc::clone(&release);
936        timers.register(
937            serial_id,
938            Exec::Serial,
939            AckMode::Completion,
940            Arc::new(move |_fire| {
941                let release_s = Arc::clone(&release_s);
942                Box::pin(async move {
943                    release_s.notified().await;
944                    Ok(TimerAck::Continue)
945                })
946            }),
947        );
948        let saw_d = Arc::clone(&saw);
949        timers.register(
950            dedicated_id,
951            Exec::Dedicated("poll".into()),
952            AckMode::Completion,
953            Arc::new(move |_fire| {
954                let saw_d = Arc::clone(&saw_d);
955                Box::pin(async move {
956                    saw_d.fetch_add(1, Ordering::SeqCst);
957                    Ok(TimerAck::Continue)
958                })
959            }),
960        );
961        let (_event_tx, event_rx) = broadcast::channel(8);
962        let (inv_tx, mut inv_rx) = mpsc::channel(8);
963        let timers_loop = Arc::clone(&timers);
964        let ctx = context(Arc::clone(&timers));
965        let running = tokio::spawn(async move {
966            module_loop(
967                event_rx,
968                &mut inv_rx,
969                &(),
970                &(),
971                ctx,
972                |invocation, _, _, _| dummy_result(invocation),
973                |_event, _, _| async { Ok(()) },
974                None,
975                timers_loop,
976            )
977            .await;
978        });
979        inv_tx
980            .send((timer_message(serial_id, 3), oneshot::channel().0))
981            .await
982            .unwrap();
983        tokio::time::sleep(Duration::from_millis(20)).await;
984        let (result_tx, result_rx) = oneshot::channel();
985        inv_tx
986            .send((timer_message(dedicated_id, 4), result_tx))
987            .await
988            .unwrap();
989        tokio::time::timeout(Duration::from_secs(2), async {
990            while saw.load(Ordering::SeqCst) == 0 {
991                tokio::task::yield_now().await;
992            }
993        })
994        .await
995        .expect("dedicated handler waited for the serial handler");
996        release.notify_one();
997        assert!(result_rx.await.unwrap().success);
998        drop(inv_tx);
999        running.await.unwrap();
1000    }
1001
1002    #[tokio::test]
1003    async fn a_registered_tick_runs_under_constant_event_traffic() {
1004        let timers = TimerRegistry::new();
1005        let tick_id = TimerId([5; 16]);
1006        let ticks = Arc::new(AtomicUsize::new(0));
1007        let events = Arc::new(AtomicUsize::new(0));
1008        let ticks_h = Arc::clone(&ticks);
1009        timers.register(
1010            tick_id,
1011            Exec::Serial,
1012            AckMode::Completion,
1013            Arc::new(move |_fire| {
1014                let ticks_h = Arc::clone(&ticks_h);
1015                Box::pin(async move {
1016                    ticks_h.fetch_add(1, Ordering::SeqCst);
1017                    Ok(TimerAck::Continue)
1018                })
1019            }),
1020        );
1021        let spec = sdk_tick_spec();
1022        assert_eq!(
1023            spec.clock,
1024            blvm_node::module::timers::types::Clock::Monotonic
1025        );
1026        assert_eq!(
1027            spec.first,
1028            blvm_node::module::timers::types::First::Delay(Duration::from_secs(30))
1029        );
1030        assert_eq!(
1031            spec.overlap,
1032            blvm_node::module::timers::types::Overlap::Skip
1033        );
1034        assert_eq!(spec.key.as_deref(), Some("sdk.tick"));
1035
1036        let (event_tx, event_rx) = broadcast::channel(64);
1037        let (inv_tx, mut inv_rx) = mpsc::channel(8);
1038        let events_h = Arc::clone(&events);
1039        let timers_loop = Arc::clone(&timers);
1040        let ctx = context(Arc::clone(&timers));
1041        let running = tokio::spawn(async move {
1042            module_loop(
1043                event_rx,
1044                &mut inv_rx,
1045                &(),
1046                &(),
1047                ctx,
1048                |invocation, _, _, _| dummy_result(invocation),
1049                move |_event, _, _| {
1050                    let events_h = Arc::clone(&events_h);
1051                    async move {
1052                        events_h.fetch_add(1, Ordering::SeqCst);
1053                        Ok(())
1054                    }
1055                },
1056                None,
1057                timers_loop,
1058            )
1059            .await;
1060        });
1061        for _ in 0..20 {
1062            event_tx.send(event_message()).unwrap();
1063        }
1064        let (result_tx, result_rx) = oneshot::channel();
1065        inv_tx
1066            .send((timer_message(tick_id, 9), result_tx))
1067            .await
1068            .unwrap();
1069        for _ in 0..20 {
1070            event_tx.send(event_message()).unwrap();
1071        }
1072        assert!(result_rx.await.unwrap().success);
1073        tokio::time::timeout(Duration::from_secs(2), async {
1074            while events.load(Ordering::SeqCst) < 40 || ticks.load(Ordering::SeqCst) < 1 {
1075                tokio::task::yield_now().await;
1076            }
1077        })
1078        .await
1079        .expect("tick or events stalled under traffic");
1080        assert_eq!(ticks.load(Ordering::SeqCst), 1);
1081        assert_eq!(events.load(Ordering::SeqCst), 40);
1082        drop(inv_tx);
1083        running.await.unwrap();
1084    }
1085}
blvm_sdk/module/runner.rs

blvm_sdk/module/
runner.rs
