aether-lspd 0.3.4

use lsp_types::Uri;
use std::collections::{HashSet, VecDeque};
use std::sync::{Arc, Mutex, PoisonError};
use std::time::Duration;
use tokio::sync::{Notify, watch};

#[derive(Clone)]
pub(crate) struct RefreshQueue {
    state: Arc<Mutex<RefreshQueueState>>,
    wake: Arc<Notify>,
    completed_generation: Arc<watch::Sender<u64>>,
}

struct RefreshQueueState {
    pending_queue: VecDeque<Uri>,
    pending_set: HashSet<Uri>,
    scheduled_generation: u64,
    bootstrap_in_progress: bool,
    active: bool,
    shutdown: bool,
}

impl RefreshQueue {
    pub(crate) fn new() -> Self {
        Self {
            state: Arc::new(Mutex::new(RefreshQueueState {
                pending_queue: VecDeque::new(),
                pending_set: HashSet::new(),
                scheduled_generation: 1,
                bootstrap_in_progress: true,
                active: false,
                shutdown: false,
            })),
            wake: Arc::new(Notify::new()),
            completed_generation: Arc::new(watch::Sender::new(0)),
        }
    }

    pub(crate) fn enqueue(&self, uris: Vec<Uri>) {
        if uris.is_empty() {
            return;
        }

        let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
        let mut added = false;
        for uri in uris {
            if state.pending_set.insert(uri.clone()) {
                state.pending_queue.push_back(uri);
                added = true;
            }
        }

        if !added {
            return;
        }

        if !state.bootstrap_in_progress {
            state.scheduled_generation += 1;
        }
        drop(state);
        self.wake.notify_one();
    }

    pub(crate) fn complete_bootstrap(&self) {
        {
            let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
            state.bootstrap_in_progress = false;
            if !state.active && state.pending_queue.is_empty() {
                self.completed_generation.send_replace(state.scheduled_generation);
            }
        }
        self.wake.notify_one();
    }

    pub(crate) async fn wait_for_current_generation(&self, timeout: Duration) {
        let target = self.state.lock().unwrap_or_else(PoisonError::into_inner).scheduled_generation;
        let mut completed = self.completed_generation.subscribe();
        let _ = tokio::time::timeout(timeout, completed.wait_for(|generation| *generation >= target)).await;
    }

    pub(crate) async fn recv(&self) -> Option<Uri> {
        loop {
            let wake = self.wake.notified();
            {
                let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
                if state.shutdown {
                    return None;
                }

                if let Some(uri) = state.pending_queue.pop_front() {
                    state.pending_set.remove(&uri);
                    state.active = true;
                    return Some(uri);
                }

                state.active = false;
                if !state.bootstrap_in_progress {
                    self.completed_generation.send_if_modified(|completed| {
                        let changed = *completed != state.scheduled_generation;
                        *completed = state.scheduled_generation;
                        changed
                    });
                }
            }

            wake.await;
        }
    }

    pub(crate) fn shutdown(&self) {
        {
            let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
            state.shutdown = true;
        }
        self.wake.notify_waiters();
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn uri(s: &str) -> Uri {
        s.parse().unwrap()
    }

    #[tokio::test]
    async fn enqueue_deduplicates_uris() {
        let queue = RefreshQueue::new();
        queue.complete_bootstrap();

        let a = uri("file:///a.rs");
        queue.enqueue(vec![a.clone(), a.clone()]);

        assert_eq!(queue.recv().await, Some(a));
    }

    #[tokio::test(start_paused = true)]
    async fn wait_for_current_generation_resolves_when_drained() {
        let queue = RefreshQueue::new();
        queue.complete_bootstrap();
        queue.enqueue(vec![uri("file:///a.rs"), uri("file:///b.rs")]);

        let wait_queue = queue.clone();
        let waiter = tokio::spawn(async move {
            let start = tokio::time::Instant::now();
            wait_queue.wait_for_current_generation(Duration::from_secs(60)).await;
            start.elapsed()
        });

        assert!(queue.recv().await.is_some());
        assert!(queue.recv().await.is_some());
        let idle_queue = queue.clone();
        let idle_receiver = tokio::spawn(async move { idle_queue.recv().await });

        assert!(waiter.await.unwrap() < Duration::from_secs(60), "draining the queue should wake the waiter");
        queue.shutdown();
        assert_eq!(idle_receiver.await.unwrap(), None);
    }

    #[tokio::test]
    async fn bootstrap_completion_without_pending_wakes_waiters() {
        let queue = RefreshQueue::new();

        let wait_queue = queue.clone();
        let waiter = tokio::spawn(async move {
            wait_queue.wait_for_current_generation(Duration::from_secs(2)).await;
        });

        tokio::task::yield_now().await;
        queue.complete_bootstrap();

        waiter.await.unwrap();
    }

    #[tokio::test]
    async fn shutdown_causes_recv_to_return_none() {
        let queue = RefreshQueue::new();
        queue.complete_bootstrap();

        queue.enqueue(vec![uri("file:///a.rs")]);
        queue.shutdown();

        assert_eq!(queue.recv().await, None);
    }

    #[tokio::test]
    async fn shutdown_wakes_waiting_receiver() {
        let queue = RefreshQueue::new();
        queue.complete_bootstrap();

        let wait_queue = queue.clone();
        let receiver = tokio::spawn(async move { wait_queue.recv().await });

        tokio::task::yield_now().await;
        queue.shutdown();

        assert_eq!(receiver.await.unwrap(), None);
    }
}