arty 0.5.0

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

use std::any::type_name;
use std::fmt::Debug;
use std::thread::ThreadId;

use arty_executor::remote::Pool as RemotePool;
use performables::arc::Arc;

use crate::runtime::Error;
use crate::runtime::blocking_worker::BlockingWorker;
use crate::runtime::dispatch::{DispatcherCore, WorkerIndex};
use crate::task::Builtins;
use crate::task::join::JoinHandle;

/// Cheap shared access to one runtime's routing and shutdown state.
#[derive(Clone)]
pub(crate) struct DispatcherClient {
    core: Arc<DispatcherCore>,
}

impl DispatcherClient {
    pub(crate) const fn new(core: Arc<DispatcherCore>) -> Self {
        Self { core }
    }

    pub(crate) fn worker_index(&self, thread_id: ThreadId) -> Option<WorkerIndex> {
        self.core.worker_index(thread_id)
    }

    pub(crate) fn blocking_worker(&self, worker_index: WorkerIndex) -> Arc<BlockingWorker> {
        self.core.blocking_worker(worker_index)
    }

    pub(crate) fn wake_worker(&self, worker_index: WorkerIndex) {
        self.core.wake_worker(worker_index);
    }

    pub(crate) fn next_blocking_worker(&self) -> Arc<BlockingWorker> {
        self.core.next_blocking_worker()
    }

    pub(crate) fn owns(&self, thread: &thread_aware::Thread) -> bool {
        self.core.owns(thread)
    }

    pub(crate) fn is_shutting_down(&self) -> bool {
        self.core.is_shutting_down()
    }

    pub(crate) fn shutdown_signal(&self) -> Arc<std::sync::atomic::AtomicBool> {
        self.core.shutdown_signal()
    }

    pub(crate) fn remote_results(&self) -> &RemotePool {
        self.core.remote_results()
    }

    #[cfg_attr(test, mutants::skip)] // Removing this wait guard deadlocks runtime-owner teardown.
    pub(crate) fn is_current_blocking_thread(&self) -> bool {
        self.core.is_current_blocking_thread()
    }

    /// Submits to a previously resolved worker without discarding destination affinity.
    pub(crate) fn spawn_on_worker<FF, F, R>(&self, worker_index: WorkerIndex, future_factory: FF) -> JoinHandle<R>
    where
        FF: FnOnce(Builtins) -> F + Send + 'static,
        F: Future<Output = R> + 'static,
        R: Send + 'static,
    {
        self.core.spawn_on_worker(worker_index, future_factory)
    }

    pub(crate) fn spawn<FF, F, R>(&self, future_factory: FF) -> JoinHandle<R>
    where
        FF: FnOnce(Builtins) -> F + Send + 'static,
        F: Future<Output = R> + 'static,
        R: Send + 'static,
    {
        self.core.spawn(future_factory)
    }

    pub(crate) fn spawn_everywhere<M, FF, F, R>(&self, make_factory: M) -> Vec<JoinHandle<R>>
    where
        M: FnMut() -> FF,
        FF: FnOnce(Builtins) -> F + Send + 'static,
        F: Future<Output = R> + 'static,
        R: Send + 'static,
    {
        self.core.spawn_everywhere(make_factory)
    }

    #[cfg_attr(test, mutants::skip)] // Will cause tests to hang due to runtime never stopping.
    pub(crate) fn stop(&self) {
        self.core.stop();
    }

    // Impractical to test real waiting at this API layer. We test the real waiter implementation
    // but not the API layers that simply call the waiter, as it is hard to prove the wait failed.
    #[cfg_attr(test, mutants::skip)]
    pub(crate) fn wait(&self) -> Result<(), Error> {
        self.core.join()
    }
}

impl Debug for DispatcherClient {
    #[cfg_attr(coverage_nightly, coverage(off))] // Opaque diagnostic formatting only.
    #[cfg_attr(test, mutants::skip)] // Debug formatting not tested
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct(type_name::<Self>()).finish()
    }
}

#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))] // Test scaffolding is not runtime behavior.
mod tests {
    use std::task::Waker;

    use observed::Sink;
    use performables::sync::channel;
    use thread_aware::ThreadAware;

    use super::*;
    use crate::runtime::blocking_worker::BlockingPool;
    use crate::runtime::dispatch::{WorkerEndpoint, test_threads};
    use crate::runtime::thread::waiter::ThreadWaiter;
    use crate::task::scheduler::Scheduler;

    #[test]
    fn clients_observe_the_shared_shutdown_state() {
        let (command_tx, _commands) = channel::unbounded();
        let worker = WorkerEndpoint {
            command_tx,
            waker: Waker::noop().clone(),
            thread: test_threads(1).pop().unwrap(),
            blocking_worker: BlockingWorker::new(BlockingPool::new(None), Sink::noop()),
        };
        let client = DispatcherClient::new(Arc::new(DispatcherCore::new(
            ThreadWaiter::new(vec![]),
            nonempty::NonEmpty::new(worker),
            Sink::noop(),
        )));
        let clone = client.clone();
        assert!(std::ptr::eq(client.remote_results(), client.core.remote_results()));
        assert!(!client.is_shutting_down());
        assert!(!clone.is_shutting_down());

        client.stop();
        assert!(client.is_shutting_down());
        assert!(clone.is_shutting_down());
    }

    #[test]
    fn relocate_updates_task_placement() {
        // Create a dispatcher with 2 workers so we can verify which one receives tasks.
        let (worker0_tx, worker0_rx) = channel::unbounded();
        let (worker1_tx, worker1_rx) = channel::unbounded();
        let wfs = ThreadWaiter::new(vec![]);
        let threads = test_threads(2);
        let blocking_worker = BlockingWorker::new(BlockingPool::new(None), Sink::noop());

        let dispatcher = Arc::new(DispatcherCore::new(
            wfs,
            nonempty::NonEmpty::from_vec(vec![
                WorkerEndpoint {
                    command_tx: worker0_tx,
                    waker: Waker::noop().clone(),
                    thread: threads[0].clone(),
                    blocking_worker: Arc::clone(&blocking_worker),
                },
                WorkerEndpoint {
                    command_tx: worker1_tx,
                    waker: Waker::noop().clone(),
                    thread: threads[1].clone(),
                    blocking_worker: Arc::clone(&blocking_worker),
                },
            ])
            .unwrap(),
            Sink::noop(),
        ));

        let mut scheduler = Scheduler::new(DispatcherClient::new(dispatcher), threads[0].clone());

        // Relocate to worker 1.
        scheduler.relocate(Some(&threads[0]), &threads[1]);

        // Spawn a task on worker 1, not worker 0.
        scheduler.spawn(async |_| {});

        // Worker 0 should have received nothing.
        assert!(
            worker0_rx.try_recv().is_err(),
            "task should not be dispatched to original worker after relocate"
        );
        // Worker 1 should have received the task.
        assert!(worker1_rx.try_recv().is_ok(), "task should be dispatched to the relocated worker");
    }

    #[test]
    fn relocate_switches_blocking_worker() {
        let threads = test_threads(2);

        let source_worker = BlockingWorker::new(BlockingPool::new(None), Sink::noop());
        let destination_worker = BlockingWorker::new(BlockingPool::new(None), Sink::noop());
        destination_worker.shutdown();

        let dispatcher = Arc::new(DispatcherCore::new(
            ThreadWaiter::new(vec![]),
            nonempty::NonEmpty::from_vec(vec![
                WorkerEndpoint {
                    command_tx: channel::unbounded().0,
                    waker: Waker::noop().clone(),
                    thread: threads[0].clone(),
                    blocking_worker: Arc::clone(&source_worker),
                },
                WorkerEndpoint {
                    command_tx: channel::unbounded().0,
                    waker: Waker::noop().clone(),
                    thread: threads[1].clone(),
                    blocking_worker: Arc::clone(&destination_worker),
                },
            ])
            .unwrap(),
            Sink::noop(),
        ));

        let mut scheduler = Scheduler::new(DispatcherClient::new(dispatcher), threads[0].clone());

        let before = (
            Arc::ptr_eq(scheduler.blocking_worker(), &source_worker),
            Arc::ptr_eq(scheduler.blocking_worker(), &destination_worker),
        );

        scheduler.relocate(Some(&threads[0]), &threads[1]);

        let after = (
            Arc::ptr_eq(scheduler.blocking_worker(), &source_worker),
            Arc::ptr_eq(scheduler.blocking_worker(), &destination_worker),
        );

        assert_eq!((before, after), ((true, false), (false, true)));
    }
}