app_window 0.3.5

// SPDX-License-Identifier: MPL-2.0

/*!
A mini-executor for running futures on the main thread.

This module provides an executor specifically designed to run futures on the application's
main thread while cooperating with the native event loop. This is essential for platforms
that require certain operations (like UI updates) to happen on the main thread.

The executor uses a task queue that's processed whenever the main thread's event loop
allows it, ensuring that futures can yield control back to the native event loop for
smooth operation.

# Thread Safety

All futures submitted to this executor will run on the main thread. The executor
provides two main entry points:
- [`on_main_thread_async`](crate::executor::on_main_thread_async): Can be called from any thread to run a future on the main thread
- [`already_on_main_thread_submit`](crate::executor::already_on_main_thread_submit): Must be called from the main thread

# Integration with `some_executor`

When the `some_executor` feature is enabled, this executor can be wrapped with
`crate::some_executor::MainThreadExecutor` to provide a `some_executor::SomeExecutor`
implementation.
*/
use crate::application::submit_to_main_thread;
use crate::sys;
use std::cell::Cell;
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::task::{Context, RawWaker, RawWakerVTable};

/// Static counter for generating unique task IDs.
static NEXT_TASK_ID: AtomicUsize = AtomicUsize::new(1);

/// Internal state shared between a task and its waker.
///
/// This struct tracks the task ID for waker operations.
struct Inner {
    task_id: usize,
}

impl Inner {
    fn new(task_id: usize) -> Self {
        Inner { task_id }
    }
}

/// A waker implementation that integrates with the main thread executor.
///
/// When woken, this sets a flag indicating the associated task needs polling
/// and submits a request to pump the task queue on the main thread.
struct Waker {
    inner: Arc<Inner>,
}

/// Virtual function table for our custom waker implementation.
///
/// This defines how to clone, wake, wake_by_ref, and drop our waker type.
const WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(
    |data| {
        let w = unsafe { Arc::from_raw(data as *const Waker) };
        let w2 = w.clone();
        _ = Arc::into_raw(w); //leave original arc unchanged
        RawWaker::new(Arc::into_raw(w2) as *const (), &WAKER_VTABLE)
    },
    |data| {
        let w = unsafe { Arc::from_raw(data as *const Waker) };
        wake_task(w.inner.task_id);
    },
    |data| {
        let w = unsafe { Arc::from_raw(data as *const Waker) };
        wake_task(w.inner.task_id);
        std::mem::forget(w);
    },
    |data| {
        let w = unsafe { Arc::from_raw(data as *const Waker) };
        drop(w);
    },
);
impl Waker {
    fn into_waker(self) -> std::task::Waker {
        let arc_waker = Arc::into_raw(Arc::new(self));
        unsafe { std::task::Waker::from_raw(RawWaker::new(arc_waker as *const (), &WAKER_VTABLE)) }
    }
}
/// A task in the executor's queue.
///
/// Each task contains a pinned future, unique ID, and shared state for wake notifications.
struct Task {
    /// The task's durable causal context.
    ///
    /// A token, not a guard: the token outlives every poll and is entered only
    /// while the task is actually being polled. A `ContextGuard` is `!Send` and
    /// restores thread-local state on drop, so storing one in a task that
    /// outlives its poll would be wrong even here, where the executor is
    /// main-thread-only.
    context: logwise::ContextToken,
    our_task_id: usize,
    future: Pin<Box<dyn Future<Output = ()> + 'static>>,
    wake_inner: Arc<Inner>,
}

/// Wakes a task by moving it from RUNNING to POLLABLE and scheduling executor iteration.
///
/// This function handles the wake notification for a specific task ID.
fn wake_task(task_id: usize) {
    // Schedule main executor iteration on the main thread
    crate::application::submit_to_main_thread("wake_task".to_string(), move || {
        // Add the task to the pollable queue
        let mut pollable = POLLABLE.take();
        pollable.push(task_id);
        POLLABLE.replace(pollable);
        main_executor_iter();
    });
}

thread_local! {
    // Thread-local storage for tasks that are running but not currently pollable.
    static RUNNING: Cell<Option<HashMap<usize, Task>>> = const { Cell::new(None) };
    // Thread-local storage for task IDs that are ready to be polled.
    static POLLABLE: Cell<Vec<usize>> = const { Cell::new(Vec::new()) };
    // Task IDs currently checked out of RUNNING and being polled somewhere up the
    // stack.  A wake for one of these must be deferred, not dropped.
    static IN_FLIGHT: Cell<Vec<usize>> = const { Cell::new(Vec::new()) };
}

/// Runs the specified future on the main thread and returns its result.
///
/// This function can be called from any thread. It submits the given future to the
/// main thread executor and waits for its completion. While the future is executing,
/// the main thread can still process other events, allowing for cooperative multitasking.
///
/// # Type Parameters
///
/// * `R` - The return type of the future, must be `Send + 'static`
/// * `F` - The future type, must be `Send + 'static` and return `R`
///
/// # Examples
///
/// ```
/// # use std::future::Future;
/// # fn test() -> impl Future<Output = ()> {
/// # async {
/// // Call from any thread to compute on the main thread
/// let result = app_window::executor::on_main_thread_async("ex".to_owned(),async {
///     // This code runs on the main thread
///     // Perform computation that needs main thread access
///     2 + 2
/// }).await;
///
/// assert_eq!(result, 4);
/// # }
/// # }
/// ```
///
/// # Platform Behavior
///
/// On all supported platforms, this ensures the future runs on the thread that owns
/// the native event loop, which is required for UI operations.
pub async fn on_main_thread_async<R: Send + 'static, F: Future<Output = R> + Send + 'static>(
    debug_label: String,
    future: F,
) -> R {
    let (sender, fut) = r#continue::continuation();
    crate::application::submit_to_main_thread(debug_label.clone(), || {
        already_on_main_thread_submit(debug_label, async move {
            let r = future.await;
            sender.send(r);
        })
    });
    fut.await
}

/// Submits a future to the main thread executor for execution.
///
/// This function must be called from the main thread. It adds the future to the
/// executor's task queue where it will be polled during the main event loop.
///
/// # Panics
///
/// This function will panic if not called from the main thread.
///
/// # Examples
///
/// ```
/// # use std::future::Future;
/// # fn setup_main_thread() {
/// // This code must run on the main thread
/// app_window::executor::already_on_main_thread_submit("ex".to_owned(), async {
///     println!("Running async task on main thread");
///     // Perform async operations that yield to the event loop
/// });
/// # }
/// ```
///
/// # Use Case
///
/// This is primarily used internally by the `crate::some_executor::MainThreadExecutor`
/// (when the `some_executor` feature is enabled) when spawning tasks, but can be used
/// directly when you're already on the main thread and want to submit work to be
/// executed asynchronously.
pub fn already_on_main_thread_submit<F: Future<Output = ()> + 'static>(
    debug_label: String,
    future: F,
) {
    assert!(sys::is_main_thread());

    // Generate unique task ID
    let task_id = NEXT_TASK_ID.fetch_add(1, Ordering::Relaxed);

    // Create task with unique ID
    let wake_inner = Arc::new(Inner::new(task_id));
    // Minted from whatever is entered at the submit site, so the task's lineage
    // records who asked for it rather than which thread happened to run it.
    let new_context = logwise::context::child(logwise::context::capture(), "app_window.task");

    logwise::log!(
        "app_window: creating task {id} {label}",
        id = task_id,
        label = debug_label
    );
    let task = Task {
        our_task_id: task_id,
        context: new_context,
        future: Box::pin(future),
        wake_inner,
    };

    // Add task to POLLABLE queue
    let mut pollable = POLLABLE.take();
    // logwise::info_sync!("Submitting task {id} to main executor", id = task_id);
    pollable.push(task_id);
    POLLABLE.replace(pollable);

    // Add task to RUNNING collection
    let mut running = RUNNING.take().unwrap_or_default();
    running.insert(task_id, task);
    RUNNING.replace(Some(running));

    // Execute the tasks
    main_executor_iter();
}

/// Polls all tasks that need attention.
///
/// This function loops while there are pollable tasks, handling new tasks
/// that may be added during polling without losing them.
fn main_executor_iter() {
    let begin_iter = crate::application::time::Instant::now();
    // Pop off a pollable task
    // let iter = perfwarn_begin!("main_executor_iter");
    let mut swap_pollable = POLLABLE.take();
    //FIFO so a frequently-woken task can't starve older pollable tasks
    let poll = if swap_pollable.is_empty() {
        None
    } else {
        Some(swap_pollable.remove(0))
    };
    POLLABLE.replace(swap_pollable);
    match poll {
        None => {
            //No more pollable tasks, nothing to do.
        }
        Some(woke_task_id) => {
            //If this task is being polled higher up the stack (it woke itself inline
            //during its own poll), defer the wake: put it back for the iteration
            //scheduled when that poll completes.
            let in_flight = IN_FLIGHT.take();
            let is_in_flight = in_flight.contains(&woke_task_id);
            IN_FLIGHT.replace(in_flight);
            if is_in_flight {
                let mut pollable = POLLABLE.take();
                pollable.push(woke_task_id);
                POLLABLE.replace(pollable);
                return;
            }
            // Get the task from RUNNING
            let mut running = RUNNING.take().unwrap_or_default();
            let task = running.remove(&woke_task_id);
            RUNNING.replace(Some(running));
            let mut task = match task {
                Some(task) => task,
                None => {
                    //A wake of a task that already completed.  The Waker contract
                    //requires this to be a no-op, not a panic.
                    return;
                }
            };
            let task_id = task.our_task_id;
            let mut in_flight = IN_FLIGHT.take();
            in_flight.push(task.our_task_id);
            IN_FLIGHT.replace(in_flight);

            //with that out of the way, we can poll the task
            let waker = Waker {
                inner: task.wake_inner.clone(),
            };
            let into_waker = waker.into_waker();
            let mut context = Context::from_waker(&into_waker);
            // Entered only for the poll. The guard restores the previous context
            // when it drops, so work that runs after this task -- including the
            // rest of this iteration -- is not attributed to it.
            let poll_result = {
                let _entered = logwise::context::enter(task.context);
                task.future.as_mut().poll(&mut context)
            };
            let mut in_flight = IN_FLIGHT.take();
            in_flight.retain(|&id| id != task.our_task_id);
            IN_FLIGHT.replace(in_flight);
            match poll_result {
                std::task::Poll::Ready(()) => {
                    // Task completed, don't put it back
                }
                std::task::Poll::Pending => {
                    // Task is still running, put it back in RUNNING
                    let mut running = RUNNING.take().unwrap_or_default();
                    running.insert(task.our_task_id, task);
                    RUNNING.replace(Some(running));
                }
            }
            //there MAY be more pollable tasks.  However, we want to yield here
            submit_to_main_thread("main_executor_iter".to_string(), main_executor_iter);
            if begin_iter.elapsed() > crate::application::time::Duration::from_millis(10) {
                // A main-thread poll that overruns is an operational problem for
                // the whole application, not a developer aside: everything else
                // on this thread was blocked for that long.
                logwise::event!(
                    class: performance,
                    severity: warn,
                    name: "app_window.main_executor.poll_overran",
                    task = support(task_id as u64),
                    duration_ms = support(begin_iter.elapsed().as_millis() as u64),
                );
            }
        }
    }
    // drop(iter);
}