futuresdr 0.6.0

use async_lock::Mutex;
#[cfg(all(not(target_arch = "wasm32"), feature = "ctrl_port"))]
use axum::Router;
use futures::prelude::*;
use std::collections::BTreeMap;
use std::fmt;
use std::sync::Arc;

use crate::runtime;
#[cfg(all(not(target_arch = "wasm32"), feature = "ctrl_port"))]
use crate::runtime::ControlPort;
use crate::runtime::Error;
use crate::runtime::Flowgraph;
use crate::runtime::FlowgraphHandle;
use crate::runtime::FlowgraphId;
use crate::runtime::FlowgraphMessage;
use crate::runtime::FlowgraphTask;
use crate::runtime::RunningFlowgraph;
use crate::runtime::TerminatedFlowgraph;
use crate::runtime::channel::mpsc::channel;
use crate::runtime::channel::oneshot;
use crate::runtime::config;
use crate::runtime::flowgraph::run_flowgraph;
use crate::runtime::flowgraph_handle::RunningFlowgraphRegistry;
use crate::runtime::scheduler::Scheduler;
#[cfg(not(target_arch = "wasm32"))]
use crate::runtime::scheduler::SmolScheduler;
use crate::runtime::scheduler::Task;
#[cfg(target_arch = "wasm32")]
use crate::runtime::scheduler::WasmMainScheduler;

#[cfg(not(target_arch = "wasm32"))]
/// Default scheduler used by [`Runtime`] and [`RuntimeHandle`] on native targets.
pub type DefaultScheduler = SmolScheduler;

#[cfg(target_arch = "wasm32")]
/// Default scheduler used by [`Runtime`] and [`RuntimeHandle`] on WASM targets.
pub type DefaultScheduler = WasmMainScheduler;

/// Executor and control-plane owner for [`Flowgraph`]s and async tasks.
///
/// A [`Runtime`] owns a scheduler, starts flowgraphs, and provides a control
/// port on native targets when the `ctrl_port` feature is enabled. It is generic
/// over the scheduler implementation, but most applications should use
/// [`Runtime::new`] with the default scheduler.
///
/// Use [`Runtime::run`] or [`Runtime::run_async`] when the caller should wait
/// until a flowgraph finishes. Use [`Runtime::start`] or
/// [`Runtime::start_async`] when the caller needs a [`RunningFlowgraph`] handle
/// for live message calls, descriptions, or shutdown.
pub struct Runtime<S = DefaultScheduler> {
    handle: RuntimeHandle<S>,
    #[cfg(all(not(target_arch = "wasm32"), feature = "ctrl_port"))]
    // Kept alive so Drop shuts down the control-port server task.
    #[allow(dead_code)]
    control_port: ControlPort,
}

#[cfg(not(target_arch = "wasm32"))]
impl Runtime<DefaultScheduler> {
    /// Construct a new [`Runtime`] using [`DefaultScheduler::default()`].
    ///
    /// On native targets this also initializes logging and, with the `ctrl_port`
    /// feature, starts the integrated control-port server when the runtime
    /// configuration enables it.
    pub fn new() -> Self {
        Self::with_scheduler(DefaultScheduler::default())
    }

    /// Construct a runtime with additional routes for the integrated web server.
    ///
    /// The routes are merged into the native control-port server. Use this for
    /// application-specific HTTP APIs or UI assets that should be served by the
    /// same process.
    #[cfg(feature = "ctrl_port")]
    pub fn with_custom_routes(routes: Router) -> Self {
        Self::with_config(DefaultScheduler::default(), routes)
    }
}

#[cfg(target_arch = "wasm32")]
impl Runtime<DefaultScheduler> {
    /// Construct a runtime using the default main-thread WASM scheduler.
    ///
    /// Normal blocks run on the browser main thread without requiring shared
    /// memory or a worker script. Use
    /// [`WasmScheduler`](crate::runtime::scheduler::WasmScheduler) explicitly
    /// for worker-backed execution.
    pub fn new() -> Self {
        Self::with_scheduler(DefaultScheduler::default())
    }
}

impl Default for Runtime<DefaultScheduler> {
    fn default() -> Self {
        Self::new()
    }
}

impl<S: Scheduler> Runtime<S> {
    /// Spawn an async task on the runtime scheduler and return its task handle.
    ///
    /// The task is unrelated to any particular flowgraph. Dropping the returned
    /// task cancels or detaches according to the underlying scheduler task type.
    pub fn spawn<T: Send + 'static>(
        &self,
        future: impl Future<Output = T> + Send + 'static,
    ) -> Task<T> {
        self.handle.scheduler.spawn(future)
    }

    /// Spawn an async task on the runtime scheduler and detach it immediately.
    pub fn spawn_background<T: Send + 'static>(
        &self,
        future: impl Future<Output = T> + Send + 'static,
    ) {
        self.handle.scheduler.spawn(future).detach();
    }

    /// Start a [`Flowgraph`] on the [`Runtime`] and await initialization.
    ///
    /// Returns once the flowgraph is initialized and running. The returned
    /// [`RunningFlowgraph`] can be used to send messages, stop the graph, or
    /// wait for completion.
    pub async fn start_async(&self, fg: Flowgraph) -> Result<RunningFlowgraph, Error> {
        self.handle.start(fg).await
    }

    /// Start a [`Flowgraph`] on the [`Runtime`] and await its termination.
    ///
    /// This consumes the input flowgraph, runs it until every block finishes or
    /// an error stops execution, and returns a [`TerminatedFlowgraph`] so final
    /// block state can be inspected.
    pub async fn run_async(&self, fg: Flowgraph) -> Result<TerminatedFlowgraph, Error> {
        self.start_async(fg).await?.wait_async().await
    }

    /// Get the [`Scheduler`] that is associated with the [`Runtime`].
    pub fn scheduler(&self) -> &S {
        &self.handle.scheduler
    }

    /// Create a clonable [`RuntimeHandle`] for starting and querying flowgraphs.
    ///
    /// Handles share the same scheduler and control-plane registry as this
    /// runtime. They are intended for web handlers, callbacks, and other async
    /// tasks that cannot borrow the runtime directly.
    pub fn handle(&self) -> RuntimeHandle<S> {
        self.handle.clone()
    }
}

#[cfg(not(target_arch = "wasm32"))]
impl<S: Scheduler> Runtime<S> {
    /// Start a [`Flowgraph`] on the [`Runtime`].
    ///
    /// Blocks until the flowgraph is initialized and running.
    pub fn start(&self, fg: Flowgraph) -> Result<RunningFlowgraph, Error> {
        runtime::block_on(self.start_async(fg))
    }

    /// Start a [`Flowgraph`] on the [`Runtime`] and block until it terminates.
    ///
    /// This is the synchronous counterpart of [`Runtime::run_async`].
    pub fn run(&self, fg: Flowgraph) -> Result<TerminatedFlowgraph, Error> {
        let running = runtime::block_on(self.start_async(fg))?;
        running.wait()
    }
}

#[cfg(all(not(target_arch = "wasm32"), feature = "ctrl_port"))]
impl<S: Scheduler + Sync> Runtime<S> {
    /// Construct a [`Runtime`] with a custom [`Scheduler`].
    ///
    /// This uses the normal native control-port routes without adding
    /// application-specific routes.
    pub fn with_scheduler(scheduler: S) -> Self {
        Self::with_config(scheduler, Router::new())
    }

    /// Construct a runtime with a custom scheduler and web server routes.
    pub fn with_config(scheduler: S, routes: Router) -> Self {
        runtime::init();

        let handle = RuntimeHandle::new(scheduler.clone());
        let control_port = ControlPort::new(handle.clone(), scheduler, routes);

        Runtime {
            handle,
            control_port,
        }
    }
}

#[cfg(all(not(target_arch = "wasm32"), not(feature = "ctrl_port")))]
impl<S: Scheduler> Runtime<S> {
    /// Construct a [`Runtime`] with a custom [`Scheduler`].
    pub fn with_scheduler(scheduler: S) -> Self {
        runtime::init();

        Runtime {
            handle: RuntimeHandle::new(scheduler),
        }
    }
}

#[cfg(target_arch = "wasm32")]
impl<S: Scheduler> Runtime<S> {
    /// Construct a [`Runtime`] with a custom [`Scheduler`].
    pub fn with_scheduler(scheduler: S) -> Self {
        runtime::init();

        Runtime {
            handle: RuntimeHandle::new(scheduler),
        }
    }
}

/// Clonable runtime control handle used by web handlers and external control code.
///
/// A `RuntimeHandle` can start new flowgraphs on the same scheduler as the
/// owning [`Runtime`] and look up flowgraphs that have been registered with the
/// control plane.
pub struct RuntimeHandle<S = DefaultScheduler> {
    scheduler: S,
    flowgraphs: Arc<Mutex<BTreeMap<FlowgraphId, FlowgraphHandle>>>,
}

impl<S> RuntimeHandle<S> {
    fn new(scheduler: S) -> Self {
        Self {
            scheduler,
            flowgraphs: Arc::new(Mutex::new(BTreeMap::new())),
        }
    }
}

impl<S: Clone> Clone for RuntimeHandle<S> {
    fn clone(&self) -> Self {
        Self {
            scheduler: self.scheduler.clone(),
            flowgraphs: self.flowgraphs.clone(),
        }
    }
}

impl<S> fmt::Debug for RuntimeHandle<S> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("RuntimeHandle")
            .field("flowgraphs", &self.flowgraphs)
            .finish()
    }
}

impl<S> PartialEq for RuntimeHandle<S> {
    fn eq(&self, other: &Self) -> bool {
        Arc::ptr_eq(&self.flowgraphs, &other.flowgraphs)
    }
}

impl<S: Scheduler> RuntimeHandle<S> {
    /// Start a [`Flowgraph`] on the runtime and register it with the control plane.
    ///
    /// This has the same startup semantics as [`Runtime::start_async`]. The
    /// returned flowgraph is available through [`RuntimeHandle::get_flowgraph`]
    /// and the native control-port API until it terminates.
    pub async fn start(&self, fg: Flowgraph) -> Result<RunningFlowgraph, Error> {
        start_flowgraph(self.scheduler.clone(), self.flowgraphs.clone(), fg).await
    }

    /// Get the control handle for a flowgraph by stable flowgraph id.
    ///
    /// Flowgraphs are removed from the registry when their runtime task exits.
    /// A graph may still terminate between listing ids and looking up a handle.
    pub async fn get_flowgraph(&self, id: FlowgraphId) -> Option<FlowgraphHandle> {
        self.flowgraphs.lock().await.get(&id).cloned()
    }

    /// Get the stable ids of flowgraphs currently registered with this runtime handle.
    pub async fn get_flowgraphs(&self) -> Vec<FlowgraphId> {
        self.flowgraphs.lock().await.keys().copied().collect()
    }
}

async fn start_flowgraph<S: Scheduler>(
    scheduler: S,
    flowgraphs: Arc<Mutex<BTreeMap<FlowgraphId, FlowgraphHandle>>>,
    fg: Flowgraph,
) -> Result<RunningFlowgraph, Error> {
    let id = fg.id();
    let queue_size = config::config().queue_size;
    let (fg_inbox, fg_inbox_rx) = channel::<FlowgraphMessage>(queue_size);

    let (startup_tx, startup_rx) =
        oneshot::channel::<Result<Arc<RunningFlowgraphRegistry>, Error>>();
    let (commit_tx, commit_rx) = oneshot::channel::<()>();
    let (completion_tx, completion_rx) = oneshot::channel::<Result<TerminatedFlowgraph, Error>>();
    let cleanup_flowgraphs = flowgraphs.clone();
    let main_channel = fg_inbox.clone();
    let scheduler_clone = scheduler.clone();
    let supervisor = async move {
        let result = run_flowgraph(
            fg,
            scheduler_clone,
            main_channel,
            fg_inbox_rx,
            startup_tx,
            commit_rx,
        )
        .await;
        cleanup_flowgraphs.lock().await.remove(&id);
        let _ = completion_tx.send(result);
    };

    #[cfg(not(target_arch = "wasm32"))]
    scheduler.spawn(supervisor).detach();

    #[cfg(target_arch = "wasm32")]
    wasm_bindgen_futures::spawn_local(supervisor);

    let registry = startup_rx
        .await
        .map_err(|_| Error::RuntimeError("run_flowgraph panicked".to_string()))??;

    let handle = FlowgraphHandle::new(id, fg_inbox, registry);
    flowgraphs.lock().await.insert(id, handle.clone());
    let _ = commit_tx.send(());
    Ok(RunningFlowgraph::new(
        handle,
        FlowgraphTask::new(completion_rx),
    ))
}

#[cfg(all(test, not(target_arch = "wasm32")))]
mod tests {
    use super::*;
    use crate::blocks::MessageSourceBuilder;
    use crate::runtime::Pmt;
    use crate::runtime::Timer;
    use crate::runtime::dev::prelude::*;
    use futures::FutureExt;
    use futures::select;
    use std::sync::Arc;
    use std::sync::atomic::AtomicUsize;
    use std::sync::atomic::Ordering;
    use std::time::Duration;
    use std::time::Instant;

    #[derive(Clone, Default)]
    struct StartupCounters {
        init: Arc<AtomicUsize>,
        deinit: Arc<AtomicUsize>,
    }

    impl StartupCounters {
        fn init(&self) -> usize {
            self.init.load(Ordering::SeqCst)
        }

        fn deinit(&self) -> usize {
            self.deinit.load(Ordering::SeqCst)
        }
    }

    #[derive(Block)]
    struct StartupBlock {
        counters: StartupCounters,
        init_entered: Option<oneshot::Sender<()>>,
        release_init: Option<oneshot::Receiver<()>>,
    }

    impl StartupBlock {
        fn new(
            counters: StartupCounters,
            init_entered: oneshot::Sender<()>,
            release_init: oneshot::Receiver<()>,
        ) -> Self {
            Self {
                counters,
                init_entered: Some(init_entered),
                release_init: Some(release_init),
            }
        }
    }

    impl Kernel for StartupBlock {
        async fn init(
            &mut self,
            _mo: &mut MessageOutputs,
            _meta: &BlockMeta,
        ) -> crate::runtime::Result<()> {
            self.counters.init.fetch_add(1, Ordering::SeqCst);
            if let Some(init_entered) = self.init_entered.take() {
                let _ = init_entered.send(());
            }
            if let Some(release_init) = self.release_init.take() {
                let _ = release_init.await;
            }
            Ok(())
        }

        async fn work(
            &mut self,
            _io: &mut WorkIo,
            _mo: &mut MessageOutputs,
            _meta: &BlockMeta,
        ) -> crate::runtime::Result<()> {
            Ok(())
        }

        async fn deinit(
            &mut self,
            _mo: &mut MessageOutputs,
            _meta: &BlockMeta,
        ) -> crate::runtime::Result<()> {
            self.counters.deinit.fetch_add(1, Ordering::SeqCst);
            Ok(())
        }
    }

    fn message_source_flowgraph(n_messages: Option<usize>) -> Flowgraph {
        let mut fg = Flowgraph::new();
        let builder = MessageSourceBuilder::new(Pmt::Null, Duration::from_millis(1));
        let builder = match n_messages {
            Some(n) => builder.n_messages(n),
            None => builder,
        };
        fg.add(builder.build()).unwrap();
        fg
    }

    fn startup_block_flowgraph(
        counters: StartupCounters,
    ) -> (Flowgraph, oneshot::Receiver<()>, oneshot::Sender<()>) {
        let mut fg = Flowgraph::new();
        let (init_entered_tx, init_entered_rx) = oneshot::channel();
        let (release_init_tx, release_init_rx) = oneshot::channel();
        fg.add(StartupBlock::new(
            counters,
            init_entered_tx,
            release_init_rx,
        ))
        .unwrap();
        (fg, init_entered_rx, release_init_tx)
    }

    async fn wait_for_deinit(counters: &StartupCounters) {
        let deadline = Instant::now() + Duration::from_secs(1);
        loop {
            if counters.deinit() == 1 {
                return;
            }
            assert!(
                Instant::now() < deadline,
                "flowgraph startup cancellation did not deinit block"
            );
            Timer::after(Duration::from_millis(10)).await;
        }
    }

    #[test]
    fn dropped_startup_future_cleans_up_initializing_flowgraph() {
        let scheduler = DefaultScheduler::default();
        let flowgraphs = Arc::new(Mutex::new(BTreeMap::new()));
        let counters = StartupCounters::default();
        let (fg, init_entered_rx, release_init_tx) = startup_block_flowgraph(counters.clone());

        runtime::block_on(async {
            let mut startup = Box::pin(start_flowgraph(scheduler, flowgraphs.clone(), fg).fuse());
            let init_entered = init_entered_rx.fuse();
            futures::pin_mut!(init_entered);

            select! {
                result = startup.as_mut() => match result {
                    Ok(_) => panic!("startup completed unexpectedly"),
                    Err(e) => panic!("startup failed unexpectedly: {e}"),
                },
                result = init_entered => result.unwrap(),
            }

            drop(startup);
            let _ = release_init_tx.send(());
            wait_for_deinit(&counters).await;

            let registry = flowgraphs.lock().await;
            assert!(registry.is_empty());
        });

        assert_eq!(counters.init(), 1);
        assert_eq!(counters.deinit(), 1);
    }

    #[test]
    fn dropped_startup_commit_cleans_up_initialized_flowgraph() {
        let scheduler = DefaultScheduler::default();
        let counters = StartupCounters::default();
        let (fg, init_entered_rx, release_init_tx) = startup_block_flowgraph(counters.clone());
        let queue_size = config::config().queue_size;
        let (fg_inbox, fg_inbox_rx) = channel::<FlowgraphMessage>(queue_size);
        let (startup_tx, startup_rx) =
            oneshot::channel::<Result<Arc<RunningFlowgraphRegistry>, Error>>();
        let (commit_tx, commit_rx) = oneshot::channel::<()>();

        runtime::block_on(async {
            let task = scheduler.spawn(run_flowgraph(
                fg,
                scheduler.clone(),
                fg_inbox,
                fg_inbox_rx,
                startup_tx,
                commit_rx,
            ));

            init_entered_rx.await.unwrap();
            let _ = release_init_tx.send(());
            startup_rx.await.unwrap().unwrap();

            drop(commit_tx);
            let result = task.await;
            assert!(matches!(
                result,
                Err(Error::RuntimeError(msg))
                    if msg == "main thread dropped flowgraph startup before registration"
            ));
        });

        assert_eq!(counters.init(), 1);
        assert_eq!(counters.deinit(), 1);
    }

    #[test]
    fn terminated_flowgraphs_are_not_returned_by_registry() {
        let scheduler = DefaultScheduler::default();
        let handle = RuntimeHandle {
            scheduler,
            flowgraphs: Arc::new(Mutex::new(BTreeMap::new())),
        };

        runtime::block_on(async {
            let running = handle.start(message_source_flowgraph(None)).await.unwrap();
            let id = running.id();

            assert_eq!(handle.get_flowgraphs().await, vec![id]);

            running.stop_and_wait().await.unwrap();

            assert!(handle.get_flowgraph(id).await.is_none());
            assert!(handle.get_flowgraphs().await.is_empty());
        });
    }

    #[test]
    fn completed_flowgraph_removes_registry_entry_without_query() {
        let scheduler = DefaultScheduler::default();
        let flowgraphs = Arc::new(Mutex::new(BTreeMap::new()));
        let handle = RuntimeHandle {
            scheduler,
            flowgraphs: flowgraphs.clone(),
        };

        runtime::block_on(async {
            let running = handle.start(message_source_flowgraph(None)).await.unwrap();
            let id = running.id();

            assert!(flowgraphs.lock().await.contains_key(&id));

            running.stop_and_wait().await.unwrap();

            assert!(!flowgraphs.lock().await.contains_key(&id));
        });
    }

    #[test]
    fn terminated_flowgraph_cleanup_does_not_change_other_ids() {
        let scheduler = DefaultScheduler::default();
        let handle = RuntimeHandle {
            scheduler,
            flowgraphs: Arc::new(Mutex::new(BTreeMap::new())),
        };

        runtime::block_on(async {
            let first = handle
                .start(message_source_flowgraph(Some(1)))
                .await
                .unwrap();
            let first_id = first.id();
            let second = handle.start(message_source_flowgraph(None)).await.unwrap();
            let second_id = second.id();

            first.wait_async().await.unwrap();

            assert!(handle.get_flowgraph(first_id).await.is_none());
            assert_eq!(
                handle.get_flowgraph(second_id).await.unwrap().id(),
                second_id
            );
            assert_eq!(handle.get_flowgraphs().await, vec![second_id]);

            let third = handle.start(message_source_flowgraph(None)).await.unwrap();
            let third_id = third.id();

            assert_ne!(third_id, first_id);
            assert_ne!(third_id, second_id);
            assert_eq!(
                handle.get_flowgraph(second_id).await.unwrap().id(),
                second_id
            );
            assert_eq!(handle.get_flowgraphs().await, vec![second_id, third_id]);

            second.stop_and_wait().await.unwrap();
            third.stop_and_wait().await.unwrap();
        });
    }
}