OpenPXE/crates/core/src/queue.rs

//! Queued Deployment queue.
//!
//! When a client selects "Queued Deployment" at the PXE menu, iPXE POSTs to
//! `/api/queue/join` and receives a queue position. It then enters a poll
//! loop hitting `/api/queue/poll/<id>`; the server holds the request open
//! until either (a) the operator assigns an ISO from the WebUI, in which
//! case the poll returns an iPXE `chain` URL, or (b) the poll times out
//! (iPXE's HTTP client has its own timeout), in which case iPXE re-POSTs.
//!
//! The WebUI shows the queue (`GET /api/queue`) and issues
//! `POST /api/queue/assign { iso_id, entry_ids: [...] }` to launch a single
//! ISO across many queued clients at once. Every waiting machine receives
//! the assignment without operator visits at the rack.

use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::Arc;
use time::OffsetDateTime;
use tokio::sync::Notify;
use uuid::Uuid;

use crate::ClientArch;

/// Per-client queue state visible to the WebUI.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueueEntry {
    pub id: String,
    /// 1-based queue position — position 1 is whoever got there first.
    pub position: u32,
    pub mac: String,
    pub ip: Option<IpAddr>,
    pub arch: Option<ClientArch>,
    #[serde(with = "time::serde::rfc3339")]
    pub joined_at: OffsetDateTime,
    #[serde(with = "time::serde::rfc3339")]
    pub last_poll_at: OffsetDateTime,
    pub assigned_target: Option<String>,
}

#[derive(Debug)]
struct QueueEntryInner {
    id: String,
    position: u32,
    mac: String,
    ip: Option<IpAddr>,
    arch: Option<ClientArch>,
    joined_at: OffsetDateTime,
    last_poll_at: OffsetDateTime,
    assigned_target: Option<String>,
    /// Broadcast primitive that wakes the long-poll as soon as an
    /// assignment lands — no polling on our side, no sleep-loops.
    notify: Arc<Notify>,
}

impl QueueEntryInner {
    fn snapshot(&self) -> QueueEntry {
        QueueEntry {
            id: self.id.clone(),
            position: self.position,
            mac: self.mac.clone(),
            ip: self.ip,
            arch: self.arch,
            joined_at: self.joined_at,
            last_poll_at: self.last_poll_at,
            assigned_target: self.assigned_target.clone(),
        }
    }
}

#[derive(Debug, Default)]
pub struct DeploymentQueue {
    inner: RwLock<HashMap<String, QueueEntryInner>>,
}

impl DeploymentQueue {
    #[must_use]
    pub fn new() -> Arc<Self> {
        Arc::new(Self::default())
    }

    /// Add a client to the queue. Returns the current queue snapshot. If the
    /// MAC is already queued, the existing entry is returned unchanged —
    /// retrying iPXE clients don't duplicate their slot.
    pub fn join(&self, mac: &str, ip: Option<IpAddr>, arch: Option<ClientArch>) -> QueueEntry {
        let now = OffsetDateTime::now_utc();
        let mut guard = self.inner.write();

        if let Some(existing) = guard.values_mut().find(|g| g.mac == mac) {
            existing.last_poll_at = now;
            if ip.is_some() {
                existing.ip = ip;
            }
            if arch.is_some() {
                existing.arch = arch;
            }
            return existing.snapshot();
        }

        // Queue position = max(position) + 1, or 1 if empty.
        let next_pos = guard.values().map(|g| g.position).max().unwrap_or(0) + 1;
        let id = Uuid::new_v4().to_string();
        let inner = QueueEntryInner {
            id: id.clone(),
            position: next_pos,
            mac: mac.to_string(),
            ip,
            arch,
            joined_at: now,
            last_poll_at: now,
            assigned_target: None,
            notify: Arc::new(Notify::new()),
        };
        let snap = inner.snapshot();
        guard.insert(id, inner);
        snap
    }

    /// Look up the `Notify` primitive for a given queue entry id, for long-polling.
    #[must_use]
    pub fn notifier(&self, entry_id: &str) -> Option<Arc<Notify>> {
        self.inner.read().get(entry_id).map(|g| g.notify.clone())
    }

    /// Update the last-poll timestamp (keeps the queue's "live" indicator
    /// fresh in the UI) and return the current snapshot. Returns None if
    /// the entry was released/expired between requests.
    pub fn touch(&self, entry_id: &str) -> Option<QueueEntry> {
        let mut guard = self.inner.write();
        let g = guard.get_mut(entry_id)?;
        g.last_poll_at = OffsetDateTime::now_utc();
        Some(g.snapshot())
    }

    /// Operator assigns an ISO entry (boot_entry id) to one or more clients.
    /// Returns the number of queue entries that were updated. Entries not in the
    /// queue are silently skipped.
    pub fn assign(&self, entry_ids: &[String], target: &str) -> usize {
        let mut guard = self.inner.write();
        let mut updated = 0;
        for id in entry_ids {
            if let Some(g) = guard.get_mut(id) {
                g.assigned_target = Some(target.to_string());
                g.notify.notify_waiters();
                updated += 1;
            }
        }
        updated
    }

    /// Remove a queue entry and return its final snapshot. Called after the client
    /// has successfully chained onto its assignment.
    pub fn release(&self, entry_id: &str) -> Option<QueueEntry> {
        let mut guard = self.inner.write();
        let g = guard.remove(entry_id)?;
        g.notify.notify_waiters();
        // Renumber positions so the display stays contiguous (1..N). This
        // is O(N) but the queue is expected to be small (dozens of hosts).
        let mut remaining: Vec<_> = guard.values_mut().collect();
        remaining.sort_by_key(|g| g.position);
        for (i, g) in remaining.iter_mut().enumerate() {
            g.position = (i + 1) as u32;
        }
        Some(g.snapshot())
    }

    #[must_use]
    pub fn list(&self) -> Vec<QueueEntry> {
        let guard = self.inner.read();
        let mut v: Vec<_> = guard.values().map(QueueEntryInner::snapshot).collect();
        v.sort_by_key(|g| g.position);
        v
    }

    #[must_use]
    pub fn len(&self) -> usize {
        self.inner.read().len()
    }

    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

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

    #[test]
    fn join_assigns_sequential_positions() {
        let q = DeploymentQueue::new();
        let g1 = q.join("aa:bb:cc:00:00:01", None, None);
        let g2 = q.join("aa:bb:cc:00:00:02", None, None);
        let g3 = q.join("aa:bb:cc:00:00:03", None, None);
        assert_eq!(g1.position, 1);
        assert_eq!(g2.position, 2);
        assert_eq!(g3.position, 3);
    }

    #[test]
    fn rejoining_same_mac_is_idempotent() {
        let q = DeploymentQueue::new();
        let g1 = q.join("aa:bb:cc:00:00:01", None, None);
        let g2 = q.join("aa:bb:cc:00:00:01", None, None);
        assert_eq!(g1.id, g2.id);
        assert_eq!(g1.position, g2.position);
        assert_eq!(q.len(), 1);
    }

    #[test]
    fn assign_broadcasts_target() {
        let q = DeploymentQueue::new();
        let g1 = q.join("aa:bb:cc:00:00:01", None, None);
        let g2 = q.join("aa:bb:cc:00:00:02", None, None);
        let n = q.assign(&[g1.id.clone(), g2.id.clone()], "ubuntu-24-04-linux");
        assert_eq!(n, 2);
        for g in q.list() {
            assert_eq!(g.assigned_target.as_deref(), Some("ubuntu-24-04-linux"));
        }
    }

    #[test]
    fn release_renumbers() {
        let q = DeploymentQueue::new();
        let a = q.join("aa:00:00:00:00:01", None, None);
        let _b = q.join("aa:00:00:00:00:02", None, None);
        let c = q.join("aa:00:00:00:00:03", None, None);
        q.release(&a.id);
        let list = q.list();
        assert_eq!(list.len(), 2);
        assert_eq!(list[0].position, 1);
        assert_eq!(list[1].position, 2);
        // c had position 3, now renumbered to 2.
        assert_eq!(list[1].id, c.id);
    }

    #[tokio::test]
    async fn assign_wakes_waiter() {
        let q = DeploymentQueue::new();
        let g = q.join("aa:00:00:00:00:01", None, None);
        let notify = q.notifier(&g.id).unwrap();

        let q2 = Arc::new(q);
        let q3 = q2.clone();
        let id = g.id.clone();
        let fut = tokio::spawn(async move {
            notify.notified().await;
            q3.touch(&id)
        });
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
        q2.assign(std::slice::from_ref(&g.id), "x");
        let result = fut.await.unwrap();
        assert!(result.is_some());
        assert_eq!(result.unwrap().assigned_target.as_deref(), Some("x"));
    }
}