lanspread/crates/lanspread-peer/src/context.rs

//! Shared context types for the peer system.

use std::{collections::HashMap, net::SocketAddr, path::PathBuf, sync::Arc};

use lanspread_db::db::{GameCatalog, GameDB};
use tokio::sync::{RwLock, mpsc::UnboundedSender};
use tokio_util::{sync::CancellationToken, task::TaskTracker};

use crate::{PeerEvent, Unpacker, events, library::LocalLibraryState, peer_db::PeerGameDB};

/// Thread-safe map of active outbound file transfers grouped by game ID.
pub type OutboundTransfers = Arc<RwLock<HashMap<String, Vec<(u64, CancellationToken)>>>>;

/// Mutating filesystem operation currently in flight for a game root.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OperationKind {
    /// Downloading or replacing archive files.
    Downloading,
    /// Extracting into a previously uninstalled game root.
    Installing,
    /// Replacing an existing `local/` install.
    Updating,
    /// Removing an existing `local/` install.
    Uninstalling,
    /// Removing downloaded archive files for an uninstalled game.
    RemovingDownload,
}

/// Main context for the peer system.
#[derive(Clone)]
pub struct Ctx {
    pub game_dir: Arc<RwLock<PathBuf>>,
    pub state_dir: Arc<PathBuf>,
    pub local_game_db: Arc<RwLock<Option<GameDB>>>,
    pub local_library: Arc<RwLock<LocalLibraryState>>,
    pub peer_game_db: Arc<RwLock<PeerGameDB>>,
    pub local_peer_addr: Arc<RwLock<Option<SocketAddr>>>,
    pub active_operations: Arc<RwLock<HashMap<String, OperationKind>>>,
    pub active_downloads: Arc<RwLock<HashMap<String, CancellationToken>>>,
    pub unpacker: Arc<dyn Unpacker>,
    pub catalog: Arc<RwLock<GameCatalog>>,
    pub peer_id: Arc<String>,
    pub shutdown: CancellationToken,
    pub task_tracker: TaskTracker,
    pub active_outbound_transfers: OutboundTransfers,
}

/// Context for peer connection handling.
#[derive(Clone)]
pub struct PeerCtx {
    pub game_dir: Arc<RwLock<PathBuf>>,
    pub local_game_db: Arc<RwLock<Option<GameDB>>>,
    pub local_library: Arc<RwLock<LocalLibraryState>>,
    pub local_peer_addr: Arc<RwLock<Option<SocketAddr>>>,
    pub active_operations: Arc<RwLock<HashMap<String, OperationKind>>>,
    pub peer_game_db: Arc<RwLock<PeerGameDB>>,
    pub catalog: Arc<RwLock<GameCatalog>>,
    pub peer_id: Arc<String>,
    pub tx_notify_ui: tokio::sync::mpsc::UnboundedSender<PeerEvent>,
    pub shutdown: CancellationToken,
    pub task_tracker: TaskTracker,
    pub active_outbound_transfers: OutboundTransfers,
}

impl std::fmt::Debug for PeerCtx {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PeerCtx")
            .field("game_dir", &"...")
            .field("local_game_db", &"...")
            .field("local_peer_addr", &"...")
            .field("active_operations", &"...")
            .finish()
    }
}

impl Ctx {
    /// Creates a new context with the given peer game database.
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        peer_game_db: Arc<RwLock<PeerGameDB>>,
        peer_id: String,
        game_dir: PathBuf,
        state_dir: PathBuf,
        unpacker: Arc<dyn Unpacker>,
        shutdown: CancellationToken,
        task_tracker: TaskTracker,
        catalog: Arc<RwLock<GameCatalog>>,
        active_outbound_transfers: OutboundTransfers,
    ) -> Self {
        Self {
            game_dir: Arc::new(RwLock::new(game_dir)),
            state_dir: Arc::new(state_dir),
            local_game_db: Arc::new(RwLock::new(None)),
            local_library: Arc::new(RwLock::new(LocalLibraryState::empty())),
            peer_game_db,
            local_peer_addr: Arc::new(RwLock::new(None)),
            active_operations: Arc::new(RwLock::new(HashMap::new())),
            active_downloads: Arc::new(RwLock::new(HashMap::new())),
            unpacker,
            catalog,
            peer_id: Arc::new(peer_id),
            shutdown,
            task_tracker,
            active_outbound_transfers,
        }
    }

    /// Creates a `PeerCtx` from this context.
    pub fn to_peer_ctx(
        &self,
        tx_notify_ui: tokio::sync::mpsc::UnboundedSender<PeerEvent>,
    ) -> PeerCtx {
        PeerCtx {
            game_dir: self.game_dir.clone(),
            local_game_db: self.local_game_db.clone(),
            local_library: self.local_library.clone(),
            local_peer_addr: self.local_peer_addr.clone(),
            active_operations: self.active_operations.clone(),
            peer_game_db: self.peer_game_db.clone(),
            catalog: self.catalog.clone(),
            peer_id: self.peer_id.clone(),
            tx_notify_ui,
            shutdown: self.shutdown.clone(),
            task_tracker: self.task_tracker.clone(),
            active_outbound_transfers: self.active_outbound_transfers.clone(),
        }
    }
}

/// Removes operation tracking no matter how a task exits.
pub(crate) struct OperationGuard {
    id: String,
    active_operations: Arc<RwLock<HashMap<String, OperationKind>>>,
    active_downloads: Arc<RwLock<HashMap<String, CancellationToken>>>,
    tx_notify_ui: UnboundedSender<PeerEvent>,
    clears_download: bool,
    armed: bool,
}

impl OperationGuard {
    pub(crate) fn new(
        id: String,
        active_operations: Arc<RwLock<HashMap<String, OperationKind>>>,
        tx_notify_ui: UnboundedSender<PeerEvent>,
    ) -> Self {
        Self {
            id,
            active_operations,
            active_downloads: Arc::new(RwLock::new(HashMap::new())),
            tx_notify_ui,
            clears_download: false,
            armed: true,
        }
    }

    pub(crate) fn download(
        id: String,
        active_operations: Arc<RwLock<HashMap<String, OperationKind>>>,
        active_downloads: Arc<RwLock<HashMap<String, CancellationToken>>>,
        tx_notify_ui: UnboundedSender<PeerEvent>,
    ) -> Self {
        Self {
            id,
            active_operations,
            active_downloads,
            tx_notify_ui,
            clears_download: true,
            armed: true,
        }
    }

    pub(crate) fn disarm(mut self) {
        self.armed = false;
    }
}

impl Drop for OperationGuard {
    fn drop(&mut self) {
        if !self.armed {
            return;
        }

        let id = self.id.clone();
        log::error!(
            "Operation guard is cleaning up {id}; operation ended without explicit state cleanup"
        );

        if let Ok(mut guard) = self.active_operations.try_write() {
            if guard.remove(&id).is_some() {
                events::send_active_operations_snapshot(&self.tx_notify_ui, &guard);
            }
        } else if let Ok(handle) = tokio::runtime::Handle::try_current() {
            let active_operations = self.active_operations.clone();
            let tx_notify_ui = self.tx_notify_ui.clone();
            handle.spawn({
                let id = id.clone();
                async move {
                    let mut active_operations = active_operations.write().await;
                    if active_operations.remove(&id).is_some() {
                        events::send_active_operations_snapshot(&tx_notify_ui, &active_operations);
                    }
                }
            });
        } else {
            log::error!("Failed to clean operation state for {id}: no Tokio runtime");
        }

        if !self.clears_download {
            return;
        }

        if let Ok(mut guard) = self.active_downloads.try_write() {
            guard.remove(&id);
        } else if let Ok(handle) = tokio::runtime::Handle::try_current() {
            let active_downloads = self.active_downloads.clone();
            handle.spawn({
                let id = id.clone();
                async move {
                    active_downloads.write().await.remove(&id);
                }
            });
        } else {
            log::error!("Failed to clean active download state for {id}: no Tokio runtime");
        }
    }
}

#[cfg(test)]
mod tests {
    use std::{collections::HashMap, sync::Arc, time::Duration};

    use tokio::sync::{RwLock, mpsc};
    use tokio_util::sync::CancellationToken;

    use super::{OperationGuard, OperationKind};
    use crate::{ActiveOperation, ActiveOperationKind, PeerEvent};

    type OperationTracking = (
        Arc<RwLock<HashMap<String, OperationKind>>>,
        Arc<RwLock<HashMap<String, CancellationToken>>>,
        CancellationToken,
    );

    async fn wait_for_tracking_clear(
        id: &str,
        active_operations: &Arc<RwLock<HashMap<String, OperationKind>>>,
        active_downloads: &Arc<RwLock<HashMap<String, CancellationToken>>>,
    ) {
        tokio::time::timeout(Duration::from_secs(1), async {
            loop {
                let operation_contains = active_operations.read().await.contains_key(id);
                let active_contains = active_downloads.read().await.contains_key(id);
                if !operation_contains && !active_contains {
                    break;
                }
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("download tracking should be cleared");
    }

    fn tracked_download_state(id: &str) -> OperationTracking {
        let active_operations = Arc::new(RwLock::new(HashMap::from([(
            id.to_string(),
            OperationKind::Downloading,
        )])));
        let cancel = CancellationToken::new();
        let active_downloads = Arc::new(RwLock::new(HashMap::from([(
            id.to_string(),
            cancel.clone(),
        )])));
        (active_operations, active_downloads, cancel)
    }

    async fn recv_active_operations(
        rx: &mut mpsc::UnboundedReceiver<PeerEvent>,
    ) -> Vec<ActiveOperation> {
        let event = tokio::time::timeout(Duration::from_secs(1), rx.recv())
            .await
            .expect("active operation event should arrive")
            .expect("event channel should remain open");
        let PeerEvent::ActiveOperationsChanged { active_operations } = event else {
            panic!("expected ActiveOperationsChanged");
        };
        active_operations
    }

    #[tokio::test]
    async fn operation_guard_cleans_tracking_when_not_disarmed() {
        let id = "game-complete";
        let (active_operations, active_downloads, _) = tracked_download_state(id);
        let (tx, mut rx) = mpsc::unbounded_channel();

        drop(OperationGuard::download(
            id.to_string(),
            active_operations.clone(),
            active_downloads.clone(),
            tx,
        ));

        wait_for_tracking_clear(id, &active_operations, &active_downloads).await;
        assert!(recv_active_operations(&mut rx).await.is_empty());
    }

    #[tokio::test]
    async fn operation_guard_cleans_tracking_after_cancellation() {
        let id = "game-cancelled";
        let (active_operations, active_downloads, cancel) = tracked_download_state(id);
        cancel.cancel();
        let (tx, mut rx) = mpsc::unbounded_channel();

        drop(OperationGuard::download(
            id.to_string(),
            active_operations.clone(),
            active_downloads.clone(),
            tx,
        ));

        wait_for_tracking_clear(id, &active_operations, &active_downloads).await;
        assert!(recv_active_operations(&mut rx).await.is_empty());
    }

    #[tokio::test]
    async fn disarmed_operation_guard_does_not_clean_tracking() {
        let id = "game-finished";
        let (active_operations, active_downloads, _) = tracked_download_state(id);
        let (tx, _rx) = mpsc::unbounded_channel();

        OperationGuard::download(
            id.to_string(),
            active_operations.clone(),
            active_downloads.clone(),
            tx,
        )
        .disarm();

        assert!(active_operations.read().await.contains_key(id));
        assert!(active_downloads.read().await.contains_key(id));
    }

    #[tokio::test]
    async fn operation_guard_cleans_tracking_when_task_is_dropped() {
        let id = "game-aborted";
        let (active_operations, active_downloads, _) = tracked_download_state(id);
        let (ready_tx, ready_rx) = tokio::sync::oneshot::channel();
        let (tx, mut rx) = mpsc::unbounded_channel();

        let handle = tokio::spawn({
            let active_operations = active_operations.clone();
            let active_downloads = active_downloads.clone();
            let tx = tx.clone();
            async move {
                let _guard = OperationGuard::download(
                    id.to_string(),
                    active_operations,
                    active_downloads,
                    tx,
                );
                let _ = ready_tx.send(());
                std::future::pending::<()>().await;
            }
        });

        ready_rx.await.expect("download guard should be created");
        handle.abort();
        let _ = handle.await;

        wait_for_tracking_clear(id, &active_operations, &active_downloads).await;
        assert_eq!(
            recv_active_operations(&mut rx).await,
            Vec::<ActiveOperation>::new()
        );
    }

    #[tokio::test]
    async fn operation_guard_cleanup_snapshot_keeps_other_operations() {
        let active_operations = Arc::new(RwLock::new(HashMap::from([
            ("aborted".to_string(), OperationKind::Downloading),
            ("other".to_string(), OperationKind::Installing),
        ])));
        let active_downloads = Arc::new(RwLock::new(HashMap::new()));
        let (tx, mut rx) = mpsc::unbounded_channel();

        drop(OperationGuard::download(
            "aborted".to_string(),
            active_operations,
            active_downloads,
            tx,
        ));

        assert_eq!(
            recv_active_operations(&mut rx).await,
            vec![ActiveOperation {
                id: "other".to_string(),
                operation: ActiveOperationKind::Installing,
            }]
        );
    }
}