zippa-db 0.1.1

//! Tokio runtime used for database work.
//!
//! sqlx needs a Tokio reactor, which GPUI's executor is not. Every database
//! future therefore runs on this dedicated runtime and its result travels back
//! to the UI over a oneshot channel that GPUI can await.

use std::future::Future;
use std::pin::Pin;
use std::sync::OnceLock;
use std::task::{Context, Poll};

use tokio::runtime::{Builder, Runtime};
use tokio::sync::oneshot;

fn runtime() -> &'static Runtime {
    static RUNTIME: OnceLock<Runtime> = OnceLock::new();
    RUNTIME.get_or_init(|| {
        Builder::new_multi_thread()
            .worker_threads(2)
            .enable_all()
            .thread_name("zippa-db")
            .build()
            .expect("failed to start the database runtime")
    })
}

/// A database future that is on its way back to the UI.
///
/// Awaiting it yields what the future returned. Giving up on it — the user
/// cancelling a long query — drops the work on the runtime, which is what
/// closes the connection sqlx was using and lets the pool open a fresh one.
pub struct Task<T> {
    receiver: oneshot::Receiver<T>,
    /// `None` under test, where the work has already run to completion.
    running: Option<tokio::task::JoinHandle<()>>,
    /// Whether dropping this stops the work. Off by default: plenty of work
    /// is handed to the runtime and never awaited (a pool closing, a rollback)
    /// and must finish whether or not anyone waits for it.
    abort_on_drop: bool,
}

impl<T> Drop for Task<T> {
    fn drop(&mut self) {
        if self.abort_on_drop
            && let Some(running) = &self.running
        {
            running.abort();
        }
    }
}

impl<T> Task<T> {
    /// A handle that can stop the work after the task itself has been moved
    /// into whatever is awaiting it.
    ///
    /// `None` under test, where the work has already run.
    pub fn abort_handle(&self) -> Option<tokio::task::AbortHandle> {
        self.running.as_ref().map(|running| running.abort_handle())
    }

    /// The work's answer, if it has arrived — without waiting for it. `None`
    /// while the runtime is still working. For fire-and-forget work whose
    /// result is picked up later, like a schema fetched for the completion
    /// menu, rather than awaited.
    pub fn try_recv(&mut self) -> Option<T> {
        self.receiver.try_recv().ok()
    }

    /// Stop the work when this is dropped unfinished — for work only its
    /// awaiter wants, like a connection test the user has walked away from.
    pub fn abort_on_drop(mut self) -> Self {
        self.abort_on_drop = true;
        self
    }
}

impl<T> Future for Task<T> {
    type Output = Result<T, oneshot::error::RecvError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        Pin::new(&mut self.receiver).poll(cx)
    }
}

/// Run `future` on the database runtime.
///
/// Await the returned receiver from a GPUI task:
///
/// ```ignore
/// let task = db::runtime::spawn(async move { db::connect(&config, None).await });
/// cx.spawn(async move |this, cx| {
///     let result = task.await;
/// })
/// .detach();
/// ```
#[cfg(not(test))]
pub fn spawn<F>(future: F) -> Task<F::Output>
where
    F: Future + Send + 'static,
    F::Output: Send + 'static,
{
    let (tx, rx) = oneshot::channel();
    let running = runtime().spawn(async move {
        let output = future.await;
        // The receiver is dropped when the UI stops caring about the result.
        let _ = tx.send(output);
    });

    Task {
        receiver: rx,
        running: Some(running),
        abort_on_drop: false,
    }
}

/// Under test the future runs to completion on the calling thread.
///
/// GPUI's test scheduler is single-threaded and aborts the process if one of
/// its tasks is woken from another thread, which is exactly what a background
/// runtime does when its result arrives. Running the work inline keeps every
/// wake-up on the test thread, and makes database results deterministic.
#[cfg(test)]
pub fn spawn<F>(future: F) -> Task<F::Output>
where
    F: Future + Send + 'static,
    F::Output: Send + 'static,
{
    let (tx, rx) = oneshot::channel();
    let _ = tx.send(runtime().block_on(future));

    Task {
        receiver: rx,
        running: None,
        abort_on_drop: false,
    }
}

/// Block the current thread on `future`. Tests only: the UI never blocks.
#[cfg(test)]
pub fn block_on<F: Future>(future: F) -> F::Output {
    runtime().block_on(future)
}