rstango 0.1.0

//! Asynchronous ORB executor.

use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{Duration, Instant};
use std::collections::VecDeque;
use std::fmt::Debug;
use std::sync::Arc;
use std::thread;

use tracing::{debug, debug_span, trace};

use tokio::sync::oneshot;

use crate::error::{Error, ErrorKind};
use crate::omniorb::runtime_queue::WorkerQueue;

use super::runtime_queue::Queue;
use super::orb;


/// An internal tracker for the async ORB implementation.
/// 
/// Regarding the use of Async Method Invocation (AMI), it's feasible to have such 
/// bindings, but the two methods, namely Callback and Polling, are not made for relying
/// heavily on them like we would like to do. They are relying on (linked list) linear 
/// search and are using internal thread pools, which we will do in this implementation.
/// In turn, the omniORB implementation ends up using synchronous call from within pool
/// threads, so we just do the same and it will prevent omniORB from having to manage
/// these pools.
#[derive(Debug, Clone)]
pub struct AsyncOrb {
    /// The total task count ever created.
    task_count: Arc<AtomicUsize>,
    /// A task sender in case no task sender is available!
    task_queue: Queue<Task>,
    /// For tasks that could not be sent through a queue immediately, this will reach the
    /// balancer thread for creating new worker threads.
    delayed_task_tx: crossbeam_channel::Sender<Task>,
}

impl AsyncOrb {

    /// Initialize this asynchronous ORB.
    pub fn init() -> Result<Self, orb::CORBA_Exception> {
        
        let orb = orb::CORBA_ORB::init()?;
        debug!(?orb, "Initialized ORB");

        let task_queue = Queue::new();

        let (
            delayed_task_tx,
            delayed_task_rx,
        ) = crossbeam_channel::unbounded();

        let balancer_task_queue = task_queue.clone();
        thread::Builder::new()
            .name(format!("orb-balancer"))
            .spawn(move || Self::balancer_thread(orb, balancer_task_queue, delayed_task_rx))
            .unwrap();

        Ok(Self {
            task_count: Arc::new(AtomicUsize::new(0)),
            task_queue,
            delayed_task_tx,
        })

    }

    /// Invoke a blocking ORB operation and execute it in another thread, returning a
    /// future for the result of that operation.
    pub fn dispatch<R, F>(&self, operation: F) -> impl Future<Output = Result<R, Error>> + use<R, F> + Send
    where
        F: FnOnce(&orb::CORBA_ORB) -> Result<R, Error>,
        F: Send + 'static,
        R: Send + 'static,
    {
        
        // We save the time when dispatch started, to measure and trace the execution 
        // latency of our async runtime!
        let dispatch_time = Instant::now();

        // Create a oneshot channel that we use to transfer the result of the operation.
        let (tx, rx) = oneshot::channel::<Result<R, Error>>();

        // Unique task id...
        let task_id = self.task_count.fetch_add(1, Ordering::Relaxed);

        let task = Task {
            func: Box::new(move |orb| {

                let delay = dispatch_time.elapsed();
                trace!(task_id, ?delay, "Executing operation");

                if tx.is_closed() {
                    trace!(task_id, "Cancelling operation because the channel is closed");
                    return;
                }

                let start_time = Instant::now();
                let res = operation(orb);
                trace!(task_id, delay = ?start_time.elapsed(), "Executed operation");

                // Explicitly ignore if the channel was already closed!
                let _ = tx.send(res);
                
            }), 
        };

        if let Err(task) = self.task_queue.send(task) {
            self.delayed_task_tx.send(task).unwrap();
            trace!(?task_id, "Dispatched task, to balancer");
        } else {
            trace!(?task_id, "Dispatched task, to worker");
        }

        // If this future happens to be dropped before being run, then the channel RX 
        // will be dropped and therefore the channel will be closed, thus causing 
        // cancellation of the task when run.
        async move {

            let res = rx.await;
            let return_duration = dispatch_time.elapsed();

            trace!(?task_id, "Returned task in {return_duration:.2?}");

            match res {
                Ok(v) => v,
                Err(_e) => Err(Error::new(ErrorKind::Internal)
                    .with_description("Unable to received async ORB result since the channel is unexpectedly closed")),
            }

        }
        
    }

    /// The single dispatcher thread loop.
    fn balancer_thread(
        orb: orb::CORBA_ORB, 
        task_queue: Queue<Task>,
        delayed_task_rx: crossbeam_channel::Receiver<Task>,
    ) {

        let _span = debug_span!("orb_balancer").entered();

        let mut worker_count = 0usize;
        let mut last_worker_time = Instant::now();

        let mut delayed_tasks = VecDeque::new();
        
        'main: loop {

            let mut new_worker = false;
            if worker_count < 1 {
                debug!("Creating an initial worker");
                new_worker = true;
            } else {

                loop {
                    if delayed_tasks.is_empty() {
                        match delayed_task_rx.recv() {
                            Ok(task) => delayed_tasks.push_back(task),
                            Err(_) => break 'main,
                        }
                    } else {
                        match delayed_task_rx.recv_timeout(Duration::from_micros(20)) {
                            Ok(task) => delayed_tasks.push_back(task),
                            Err(crossbeam_channel::RecvTimeoutError::Disconnected) => break 'main,
                            Err(crossbeam_channel::RecvTimeoutError::Timeout) => break,
                        }
                    }
                }

                // If we land here, we necessarily have delayed tasks. Here we compute a
                // minimum interval to respect between two instantiation of workers.
                // We take the inverse of the number of delayed tasks, and multiply it
                // by the maximum interval we want.
                // Max 1ms between task creation, then 100 us, then 50us, 25, and so on.
                const MAX_INTERVAL: f32 = 0.004;
                let new_worker_interval = MAX_INTERVAL / delayed_tasks.len() as f32;

                let new_worker_interval = Duration::from_secs_f32(new_worker_interval);
                let now = Instant::now();
                let new_worker_elapsed = now - last_worker_time;
                if new_worker_elapsed >= new_worker_interval {
                    debug!(?new_worker_elapsed, ?new_worker_interval, "Creating a new worker");
                    new_worker = true;
                }

            }
            
            if new_worker {

                let worker_id = worker_count;
                worker_count += 1;

                if let Some(worker_task_queue) = task_queue.worker() {

                    let worker_orb = orb.duplicate();
                    thread::Builder::new()
                        .name(format!("orb-worker-{worker_id}"))
                        .spawn(move || Self::worker_thread(worker_id, worker_orb, worker_task_queue))
                        .unwrap();

                    last_worker_time = Instant::now();

                }

            }

            while let Some(task) = delayed_tasks.pop_front() {
                if let Err(task) = task_queue.send(task) {
                    delayed_tasks.push_front(task);
                    break;
                }
            }

        }

    }

    /// Any worker thread loop.
    fn worker_thread(
        worker_id: usize,
        orb: orb::CORBA_ORB, 
        mut task_queue: WorkerQueue<Task>,
    ) {

        let _span = debug_span!("orb_worker", worker_id).entered();
        debug!("Spawned");

        loop {
            let task = task_queue.recv();
            (task.func)(&orb);
        }

    }

}

/// A task to be run in the ORB task threads.
struct Task {
    /// This function is internally implemented to execute the user task and then return
    /// its result through the generic oneshot channel it owns.
    func: Box<dyn FnOnce(&orb::CORBA_ORB) + Send>,
}

impl Debug for Task {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AsyncOrbTask").finish()
    }
}