velo 0.12.0

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Filesystem-based peer discovery for Velo distributed systems.
//!
//! This crate provides [`FilesystemPeerDiscovery`], a [`crate::messenger::PeerDiscovery`]
//! backend that stores peer information in a JSON file on disk. It is suitable for
//! development, testing, and lightweight single-host deployments.
//!
//! # File Format
//!
//! The file contains a JSON object with a `peers` array:
//! ```json
//! {
//!   "peers": [
//!     {
//!       "instance_id": "uuid-string",
//!       "worker_id": 123,
//!       "worker_address": "<msgpack bytes>",
//!       "address_checksum": 12345678
//!     }
//!   ]
//! }
//! ```
//!
//! # Concurrency
//!
//! Uses file-based locking (via `fs4`) for safe concurrent access across processes.
//! Within a process, uses `RwLock` + `AsyncMutex` for thread safety.

use anyhow::{Context, Result};
use fs4::fs_std::FileExt;
use futures::future::BoxFuture;
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::Mutex as AsyncMutex;

use velo_ext::{InstanceId, PeerInfo, WorkerAddress, WorkerId};

use super::PeerDiscovery;
use super::PeerRegistrationGuard;

/// Filesystem-based peer discovery backend.
///
/// Stores peer information in a JSON file. Provides simple persistence for
/// testing scenarios where external dependencies (etcd, p2p) are not desired.
///
/// Cache invalidation is done on-demand: the cache is marked invalid when writes
/// occur, and reloaded on the next query. For multi-process use, the cache is
/// always reloaded before writes to ensure consistency.
#[derive(Clone)]
pub struct FilesystemPeerDiscovery {
    file_path: PathBuf,
    inner: Arc<RwLock<FilesystemPeerDiscoveryInner>>,
    /// Mutex to serialize write operations (register/unregister) to prevent race conditions.
    write_mutex: Arc<AsyncMutex<()>>,
}

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

#[derive(Debug)]
struct FilesystemPeerDiscoveryInner {
    by_worker_id: HashMap<WorkerId, InstanceId>,
    by_instance_id: HashMap<InstanceId, PeerInfo>,
}

impl FilesystemPeerDiscoveryInner {
    fn new() -> Self {
        Self {
            by_worker_id: HashMap::new(),
            by_instance_id: HashMap::new(),
        }
    }
}

impl FilesystemPeerDiscovery {
    /// Create a new filesystem-based peer discovery at the specified path.
    ///
    /// If the file exists, it will be loaded lazily on first query. If it doesn't
    /// exist, it will be created on first registration.
    pub fn new(file_path: impl Into<PathBuf>) -> Result<Self> {
        let file_path = file_path.into();
        let inner = Arc::new(RwLock::new(FilesystemPeerDiscoveryInner::new()));
        Ok(Self {
            file_path,
            inner,
            write_mutex: Arc::new(AsyncMutex::new(())),
        })
    }

    /// Create a new filesystem-based peer discovery in a temporary directory.
    ///
    /// Useful for testing.
    pub fn new_temp() -> Result<Self> {
        let temp_dir = std::env::temp_dir();
        let file_name = format!("velo-discovery-{}.json", uuid::Uuid::new_v4());
        let file_path = temp_dir.join(file_name);
        Self::new(file_path)
    }

    /// Register peer info.
    pub async fn register_peer_info(
        &self,
        peer_info: &PeerInfo,
    ) -> Result<FilesystemRegistrationGuard> {
        let instance_id = peer_info.instance_id();
        let worker_address = peer_info.worker_address().clone();
        self.register_instance_async(instance_id, worker_address)
            .await
    }

    /// Register an instance with an address.
    ///
    /// Returns a [`FilesystemRegistrationGuard`] that automatically unregisters the instance when dropped.
    pub async fn register_instance_async(
        &self,
        instance_id: InstanceId,
        worker_address: WorkerAddress,
    ) -> Result<FilesystemRegistrationGuard> {
        let _guard = self.write_mutex.lock().await;

        self.load_from_disk().await?;

        {
            let mut state = self.inner.write();

            let worker_id = instance_id.worker_id();
            if let Some(existing_instance) = state.by_worker_id.get(&worker_id)
                && *existing_instance != instance_id
            {
                anyhow::bail!(
                    "Worker ID collision: worker_id={} is already claimed by instance {} (new: {})",
                    worker_id,
                    existing_instance,
                    instance_id
                );
            }

            if let Some(existing_peer_info) = state.by_instance_id.get(&instance_id) {
                if existing_peer_info.worker_address().checksum() == worker_address.checksum() {
                    anyhow::bail!("Instance {} is already registered", instance_id);
                } else {
                    anyhow::bail!(
                        "Instance {} already registered with a different address (checksum mismatch: {} vs {})",
                        instance_id,
                        existing_peer_info.worker_address().checksum(),
                        worker_address.checksum()
                    );
                }
            }

            let peer_info = PeerInfo::new(instance_id, worker_address);
            state.by_worker_id.insert(worker_id, instance_id);
            state.by_instance_id.insert(instance_id, peer_info);
        }

        self.save_to_disk().await?;
        Ok(FilesystemRegistrationGuard::new(self.clone(), instance_id))
    }

    /// Unregister an instance.
    pub async fn unregister_instance_async(&self, instance_id: InstanceId) -> Result<()> {
        let _guard = self.write_mutex.lock().await;

        self.load_from_disk().await?;

        {
            let mut state = self.inner.write();
            state.by_worker_id.remove(&instance_id.worker_id());
            state.by_instance_id.remove(&instance_id);
        }

        self.save_to_disk().await?;
        Ok(())
    }

    /// Load peer registry from disk with file locking.
    async fn load_from_disk(&self) -> Result<()> {
        if !self.file_path.exists() {
            return Ok(());
        }

        let file_path = self.file_path.clone();
        let content = tokio::task::spawn_blocking(move || -> Result<String> {
            let mut file =
                std::fs::File::open(&file_path).context("Failed to open discovery file")?;
            file.lock_shared()
                .context("Failed to acquire shared lock")?;
            let mut content = String::new();
            std::io::Read::read_to_string(&mut file, &mut content)
                .context("Failed to read discovery file")?;
            // Lock is released when file is dropped
            Ok(content)
        })
        .await
        .context("spawn_blocking failed")??;

        let registry: PeerRegistry =
            serde_json::from_str(&content).context("Failed to parse discovery file")?;

        let mut inner = self.inner.write();
        inner.by_worker_id.clear();
        inner.by_instance_id.clear();

        for serialized in registry.peers {
            let peer_info = serialized.to_peer_info()?;
            let instance_id = peer_info.instance_id();
            let worker_id = peer_info.worker_id();
            inner.by_worker_id.insert(worker_id, instance_id);
            inner.by_instance_id.insert(instance_id, peer_info);
        }

        Ok(())
    }

    /// Save peer registry to disk with file locking and atomic rename.
    async fn save_to_disk(&self) -> Result<()> {
        let content = {
            let inner = self.inner.read();
            let peers: Vec<SerializedPeerInfo> = inner
                .by_instance_id
                .values()
                .map(SerializedPeerInfo::from_peer_info)
                .collect();
            let registry = PeerRegistry { peers };
            serde_json::to_string_pretty(&registry).context("Failed to serialize registry")?
        };

        let file_path = self.file_path.clone();
        let temp_file_name = format!(
            "{}.tmp.{}",
            self.file_path
                .file_name()
                .unwrap_or_default()
                .to_string_lossy(),
            uuid::Uuid::new_v4()
        );
        let temp_path = self
            .file_path
            .parent()
            .map(|p| p.join(&temp_file_name))
            .unwrap_or_else(|| PathBuf::from(&temp_file_name));

        tokio::task::spawn_blocking(move || -> Result<()> {
            if let Some(parent) = file_path.parent() {
                std::fs::create_dir_all(parent).ok();
            }

            std::fs::write(&temp_path, &content).context("Failed to write temp file")?;

            let file = std::fs::OpenOptions::new()
                .write(true)
                .create(true)
                .truncate(false)
                .open(&file_path)
                .context("Failed to open discovery file")?;
            file.lock_exclusive()
                .context("Failed to acquire exclusive lock")?;

            std::fs::rename(&temp_path, &file_path).context("Failed to rename file")?;
            Ok(())
        })
        .await
        .context("spawn_blocking failed")??;

        Ok(())
    }

    async fn discover_by_worker_id_async(&self, worker_id: WorkerId) -> Result<PeerInfo> {
        self.load_from_disk().await?;

        let state = self.inner.read();
        if let Some(instance_id) = state.by_worker_id.get(&worker_id)
            && let Some(peer_info) = state.by_instance_id.get(instance_id)
        {
            return Ok(peer_info.clone());
        }
        anyhow::bail!("Worker ID {} not found", worker_id)
    }

    async fn discover_by_instance_id_async(&self, instance_id: InstanceId) -> Result<PeerInfo> {
        self.load_from_disk().await?;

        let state = self.inner.read();
        if let Some(peer_info) = state.by_instance_id.get(&instance_id) {
            return Ok(peer_info.clone());
        }
        anyhow::bail!("Instance ID {} not found", instance_id)
    }
}

impl PeerDiscovery for FilesystemPeerDiscovery {
    fn discover_by_worker_id(&self, worker_id: WorkerId) -> BoxFuture<'_, Result<PeerInfo>> {
        Box::pin(self.discover_by_worker_id_async(worker_id))
    }

    fn discover_by_instance_id(&self, instance_id: InstanceId) -> BoxFuture<'_, Result<PeerInfo>> {
        Box::pin(self.discover_by_instance_id_async(instance_id))
    }
}

/// RAII guard that unregisters instances when dropped.
///
/// Obtained from [`FilesystemPeerDiscovery::register_instance_async`] or
/// [`FilesystemPeerDiscovery::register_peer_info`].
#[derive(Debug)]
pub struct FilesystemRegistrationGuard {
    discovery: FilesystemPeerDiscovery,
    instances: Vec<InstanceId>,
}

impl FilesystemRegistrationGuard {
    fn new(discovery: FilesystemPeerDiscovery, instance_id: InstanceId) -> Self {
        Self {
            discovery,
            instances: vec![instance_id],
        }
    }
}

impl PeerRegistrationGuard for FilesystemRegistrationGuard {
    fn unregister(&mut self) -> BoxFuture<'_, anyhow::Result<()>> {
        Box::pin(async move {
            let instances = std::mem::take(&mut self.instances);
            for id in instances {
                self.discovery.unregister_instance_async(id).await?;
            }
            Ok(())
        })
    }
}

impl Drop for FilesystemRegistrationGuard {
    fn drop(&mut self) {
        let discovery = self.discovery.clone();
        let instances = std::mem::take(&mut self.instances);
        if let Ok(handle) = tokio::runtime::Handle::try_current() {
            handle.spawn(async move {
                for id in instances {
                    let _ = discovery.unregister_instance_async(id).await;
                }
            });
        }
        // If no runtime is active, instances are left in the file (acceptable for
        // process-exit scenarios where the file will be cleaned up externally).
    }
}

/// Serializable peer registry for filesystem storage.
#[derive(Debug, Clone, Serialize, Deserialize)]
struct PeerRegistry {
    peers: Vec<SerializedPeerInfo>,
}

/// Serialized representation of PeerInfo for JSON storage.
#[derive(Debug, Clone, Serialize, Deserialize)]
struct SerializedPeerInfo {
    instance_id: String,
    worker_id: u64,
    worker_address: WorkerAddress,
    address_checksum: u64,
}

impl SerializedPeerInfo {
    fn from_peer_info(peer_info: &PeerInfo) -> Self {
        Self {
            instance_id: peer_info.instance_id().to_string(),
            worker_id: peer_info.worker_id().as_u64(),
            worker_address: peer_info.worker_address().clone(),
            address_checksum: peer_info.worker_address().checksum(),
        }
    }

    fn to_peer_info(&self) -> Result<PeerInfo> {
        let uuid =
            uuid::Uuid::parse_str(&self.instance_id).context("Failed to parse instance_id")?;
        let instance_id = InstanceId::from(uuid);
        Ok(PeerInfo::new(instance_id, self.worker_address.clone()))
    }
}

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

    fn make_test_address(data: &[u8]) -> WorkerAddress {
        let map: HashMap<String, Vec<u8>> = [("endpoint".to_string(), data.to_vec())]
            .into_iter()
            .collect();
        let encoded = rmp_serde::to_vec(&map).unwrap();
        WorkerAddress::from_encoded(encoded)
    }

    fn create_test_peer_info() -> (InstanceId, WorkerAddress) {
        let instance_id = InstanceId::new_v4();
        let worker_address = make_test_address(b"test-address");
        (instance_id, worker_address)
    }

    #[tokio::test]
    async fn test_register_and_discover() {
        let discovery = FilesystemPeerDiscovery::new_temp().unwrap();
        let (instance_id, worker_address) = create_test_peer_info();
        let worker_id = instance_id.worker_id();

        let _guard = discovery
            .register_instance_async(instance_id, worker_address.clone())
            .await
            .unwrap();

        let peer_info = discovery
            .discover_by_worker_id_async(worker_id)
            .await
            .unwrap();
        assert_eq!(peer_info.instance_id(), instance_id);
        assert_eq!(peer_info.worker_id(), worker_id);

        let peer_info = discovery
            .discover_by_instance_id_async(instance_id)
            .await
            .unwrap();
        assert_eq!(peer_info.instance_id(), instance_id);
    }

    #[tokio::test]
    async fn test_persistence() {
        let temp_dir = std::env::temp_dir();
        let file_path = temp_dir.join(format!("test-velo-discovery-{}.json", uuid::Uuid::new_v4()));

        let (instance_id, worker_address) = create_test_peer_info();
        let worker_id = instance_id.worker_id();

        // Keep the guard alive across both scopes so the registration persists on disk.
        let _guard;
        {
            let discovery = FilesystemPeerDiscovery::new(&file_path).unwrap();
            _guard = discovery
                .register_instance_async(instance_id, worker_address.clone())
                .await
                .unwrap();
        }

        {
            let discovery = FilesystemPeerDiscovery::new(&file_path).unwrap();
            let peer_info = discovery
                .discover_by_worker_id_async(worker_id)
                .await
                .unwrap();
            assert_eq!(peer_info.instance_id(), instance_id);
        }

        let _ = std::fs::remove_file(&file_path);
    }

    #[tokio::test]
    async fn test_unregister() {
        let discovery = FilesystemPeerDiscovery::new_temp().unwrap();
        let (instance_id, worker_address) = create_test_peer_info();
        let worker_id = instance_id.worker_id();

        let _guard = discovery
            .register_instance_async(instance_id, worker_address)
            .await
            .unwrap();

        assert!(
            discovery
                .discover_by_worker_id_async(worker_id)
                .await
                .is_ok()
        );

        discovery
            .unregister_instance_async(instance_id)
            .await
            .unwrap();

        assert!(
            discovery
                .discover_by_worker_id_async(worker_id)
                .await
                .is_err()
        );
    }

    #[tokio::test]
    async fn test_duplicate_registration() {
        let discovery = FilesystemPeerDiscovery::new_temp().unwrap();
        let instance_id = InstanceId::new_v4();
        let address1 = make_test_address(b"address1");
        let address2 = make_test_address(b"address2");

        let _guard = discovery
            .register_instance_async(instance_id, address1.clone())
            .await
            .unwrap();

        let result = discovery
            .register_instance_async(instance_id, address1.clone())
            .await;
        assert!(result.is_err());
        let msg = result.unwrap_err().to_string();
        assert!(
            msg.contains("already registered"),
            "Expected 'already registered', got: {msg}"
        );

        let result = discovery
            .register_instance_async(instance_id, address2.clone())
            .await;
        assert!(result.is_err());
        let msg = result.unwrap_err().to_string();
        assert!(
            msg.contains("checksum mismatch") || msg.contains("different address"),
            "Expected checksum mismatch error, got: {msg}"
        );
    }

    #[tokio::test]
    async fn test_concurrent_access() {
        let temp_dir = std::env::temp_dir();
        let file_path = temp_dir.join(format!("test-velo-discovery-{}.json", uuid::Uuid::new_v4()));

        let discovery = FilesystemPeerDiscovery::new(&file_path).unwrap();

        let mut handles = vec![];
        for _ in 0..5 {
            let disc = discovery.clone();
            let handle = tokio::spawn(async move {
                let (instance_id, worker_address) = create_test_peer_info();
                let _guard = disc
                    .register_instance_async(instance_id, worker_address.clone())
                    .await
                    .unwrap();
                (instance_id, worker_address, _guard)
            });
            handles.push(handle);
        }

        let mut peers = vec![];
        for handle in handles {
            let peer = handle.await.unwrap();
            peers.push(peer);
        }

        for (instance_id, _, _guard) in &peers {
            let worker_id = instance_id.worker_id();
            let peer_info = discovery
                .discover_by_worker_id_async(worker_id)
                .await
                .unwrap();
            assert_eq!(peer_info.instance_id(), *instance_id);
        }

        let _ = std::fs::remove_file(&file_path);
    }

    #[tokio::test]
    async fn test_sequential_multi_instance() {
        let temp_dir = std::env::temp_dir();
        let file_path = temp_dir.join(format!("test-velo-discovery-{}.json", uuid::Uuid::new_v4()));

        let discovery1 = FilesystemPeerDiscovery::new(&file_path).unwrap();
        let (instance_id1, worker_address1) = create_test_peer_info();
        let _guard1 = discovery1
            .register_instance_async(instance_id1, worker_address1.clone())
            .await
            .unwrap();

        // Second instance starts with cache_valid=false, so it will load from disk
        let discovery2 = FilesystemPeerDiscovery::new(&file_path).unwrap();
        let peer_info = discovery2
            .discover_by_worker_id_async(instance_id1.worker_id())
            .await
            .unwrap();
        assert_eq!(peer_info.instance_id(), instance_id1);

        let (instance_id2, worker_address2) = create_test_peer_info();
        let _guard2 = discovery2
            .register_instance_async(instance_id2, worker_address2.clone())
            .await
            .unwrap();

        // discovery1 always re-reads from disk, so it sees discovery2's registration
        let peer_info = discovery1
            .discover_by_worker_id_async(instance_id2.worker_id())
            .await
            .unwrap();
        assert_eq!(peer_info.instance_id(), instance_id2);

        let _ = std::fs::remove_file(&file_path);
    }

    #[tokio::test]
    async fn test_peer_discovery_trait() {
        let discovery = FilesystemPeerDiscovery::new_temp().unwrap();
        let (instance_id, worker_address) = create_test_peer_info();
        let worker_id = instance_id.worker_id();

        let _guard = discovery
            .register_instance_async(instance_id, worker_address.clone())
            .await
            .unwrap();

        let peer_info = discovery.discover_by_worker_id(worker_id).await.unwrap();
        assert_eq!(peer_info.instance_id(), instance_id);

        let peer_info = discovery
            .discover_by_instance_id(instance_id)
            .await
            .unwrap();
        assert_eq!(peer_info.instance_id(), instance_id);
    }

    #[tokio::test]
    async fn test_registration_guard_trait_unregister() {
        use crate::discovery::PeerRegistrationGuard as _;
        let discovery = FilesystemPeerDiscovery::new_temp().unwrap();
        let (instance_id, worker_address) = create_test_peer_info();
        let worker_id = instance_id.worker_id();

        let mut guard = discovery
            .register_instance_async(instance_id, worker_address)
            .await
            .unwrap();

        assert!(
            discovery
                .discover_by_worker_id_async(worker_id)
                .await
                .is_ok()
        );

        guard.unregister().await.unwrap();

        assert!(
            discovery
                .discover_by_worker_id_async(worker_id)
                .await
                .is_err(),
            "Instance should be unregistered after explicit unregister()"
        );
    }
}