graphile_worker 0.14.0

use std::future::Future;
use std::pin::Pin;
use std::time::Duration;

use futures::FutureExt;
use graphile_worker_runtime as runtime;
use graphile_worker_shutdown_signal::ShutdownSignal;

pub(super) struct RecvOrShutdown<T> {
    pub(super) item: Option<T>,
    pub(super) shutdown: bool,
}

pub(super) async fn recv_or_shutdown<T>(
    rx: &runtime::Receiver<T>,
    shutdown_signal: &mut ShutdownSignal,
) -> RecvOrShutdown<T> {
    let recv = rx.recv().fuse();
    let shutdown = shutdown_signal.fuse();
    futures::pin_mut!(recv, shutdown);

    futures::select_biased! {
        _ = shutdown => RecvOrShutdown { item: None, shutdown: true },
        item = recv => RecvOrShutdown { item: item.ok(), shutdown: false },
    }
}

pub(super) trait BatchProcessor<T>: Send + Sync + 'static {
    fn flush<'a>(&'a self, batch: &'a [T]) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
}

pub(super) async fn run_batcher_task<T, P>(
    rx: runtime::Receiver<T>,
    delay: Duration,
    processor: P,
    mut shutdown_signal: ShutdownSignal,
) where
    T: Send + 'static,
    P: BatchProcessor<T>,
{
    let mut batch = Vec::new();

    loop {
        let first = recv_or_shutdown(&rx, &mut shutdown_signal).await;

        if first.shutdown {
            drain_and_flush(&rx, &mut batch, &processor).await;
            return;
        }

        let Some(first) = first.item else {
            processor.flush(&batch).await;
            return;
        };

        batch.push(first);

        let timeout = runtime::sleep(delay).fuse();
        futures::pin_mut!(timeout);

        loop {
            let recv = rx.recv().fuse();
            let shutdown = (&mut shutdown_signal).fuse();
            futures::pin_mut!(recv, shutdown);

            let result = futures::select_biased! {
                _ = shutdown => {
                    drain_and_flush(&rx, &mut batch, &processor).await;
                    return;
                }
                _ = timeout => break,
                result = recv => result,
            };

            match result {
                Ok(item) => batch.push(item),
                Err(_) => {
                    processor.flush(&batch).await;
                    return;
                }
            }
        }

        processor.flush(&batch).await;
        batch.clear();
    }
}

async fn drain_and_flush<T, P>(rx: &runtime::Receiver<T>, batch: &mut Vec<T>, processor: &P)
where
    T: Send + 'static,
    P: BatchProcessor<T>,
{
    // Seal acceptance before draining. Sends during the final database flush
    // must use direct fallback rather than enter a buffer that will be dropped.
    rx.close();
    while let Ok(item) = rx.try_recv() {
        batch.push(item);
    }
    processor.flush(batch).await;
}

#[cfg(all(test, feature = "runtime-tokio"))]
mod tests {
    use std::sync::{Arc, Mutex};

    use super::*;

    struct GatedProcessor {
        started: Arc<runtime::Notify>,
        finish: Arc<runtime::Notify>,
        flushed: Arc<Mutex<Vec<u32>>>,
    }

    impl BatchProcessor<u32> for GatedProcessor {
        fn flush<'a>(&'a self, batch: &'a [u32]) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> {
            Box::pin(async move {
                self.flushed.lock().unwrap().extend_from_slice(batch);
                self.started.notify_one();
                self.finish.notified().await;
            })
        }
    }

    #[tokio::test]
    async fn shutdown_rejects_late_requests_before_flushing_buffered_work() {
        let (tx, rx) = runtime::channel(8);
        tx.send(1u32).await.unwrap();
        let started = Arc::new(runtime::Notify::new());
        let finish = Arc::new(runtime::Notify::new());
        let flushed = Arc::new(Mutex::new(Vec::new()));
        let processor = GatedProcessor {
            started: started.clone(),
            finish: finish.clone(),
            flushed: flushed.clone(),
        };
        let run = runtime::spawn(run_batcher_task(
            rx,
            Duration::from_secs(60),
            processor,
            futures::future::ready(()).boxed().shared(),
        ));
        tokio::time::timeout(Duration::from_secs(5), started.notified())
            .await
            .unwrap();
        // The shutdown drain has finished, but its database flush is still pending.
        let late_request = tx.send(2).await;
        finish.notify_one();
        tokio::time::timeout(Duration::from_secs(5), run)
            .await
            .unwrap()
            .unwrap();
        assert_eq!(*flushed.lock().unwrap(), [1]);
        assert!(
            late_request.is_err(),
            "a late send must use direct fallback, not disappear after the final drain"
        );
    }
}