MusicFS/crates/musicfs-client/src/origins/network/mod.rs

pub mod transport;
pub mod watcher;

use std::{
    collections::BTreeMap,
    io,
    path::{Path, PathBuf},
    sync::{Arc, RwLock},
    time::{Duration, SystemTime},
};

use fuser::INodeNo;
use sea_orm::EntityTrait;

use crate::db::cache::{delete_cached_bytes_for, get_cached_bytes, put_cached_bytes};
use crate::db::entities as item_entities;
use crate::db::sync::{run_db_blocking, sync_items_to_db};
use crate::item::{FileType, Item};
use crate::music::db::restore_music_metadata_from_db;
use crate::music::metadata::MusicMetadata;
use crate::origins::attrs::FileAttrs;
use crate::origins::{ByteSource, FileWatcher, Origin};
use crate::proto::ManifestEntry as ProtoManifestEntry;
use crate::virtual_dirs::{ensure_virtual_dirs, restore_virtual_paths};
use tracing::{debug, error, info, trace};

use self::transport::NetworkTransport;

pub struct NetworkOrigin {
    pub(crate) endpoint: String,
    pub(crate) destination: PathBuf,
    handle: tokio::runtime::Handle,
    transport: NetworkTransport,
    client: sea_orm::DatabaseConnection,
    latest_manifest: Arc<RwLock<BTreeMap<u64, ProtoManifestEntry>>>,
    server_status: crate::health::ServerStatus,
}

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

impl NetworkOrigin {
    pub fn new(
        endpoint: String,
        destination: String,
        client: sea_orm::DatabaseConnection,
    ) -> io::Result<Self> {
        let handle = tokio::runtime::Handle::current();
        let transport = NetworkTransport::new(endpoint.clone()).map_err(io_err)?;
        let server_status =
            crate::health::ServerStatus::new_network(destination.clone(), endpoint.clone());
        return Ok(NetworkOrigin {
            endpoint,
            destination: destination.into(),
            handle,
            transport,
            client,
            latest_manifest: Arc::new(RwLock::new(BTreeMap::new())),
            server_status,
        });
    }

    pub fn runtime_handle(&self) -> tokio::runtime::Handle {
        return self.handle.clone();
    }

    pub fn server_status(&self) -> crate::health::ServerStatus {
        self.server_status.clone()
    }
}

impl Origin for NetworkOrigin {
    fn snapshot(&self) -> io::Result<BTreeMap<INodeNo, Item>> {
        return run_db_blocking(self.snapshot_async()).map_err(io_err);
    }

    fn byte_source(&self) -> Arc<dyn ByteSource> {
        return Arc::new(NetworkByteSource {
            transport: self.transport.clone(),
            runtime_handle: self.runtime_handle(),
            client: self.client.clone(),
        });
    }

    fn watcher(&self) -> Box<dyn FileWatcher> {
        return Box::new(watcher::NetworkOriginFileWatcher::new(
            self.transport.clone(),
            self.runtime_handle(),
            self.client.clone(),
            self.destination.clone(),
            self.latest_manifest.clone(),
            self.server_status.clone(),
        ));
    }
}

impl NetworkOrigin {
    /// Async snapshot driven directly on the caller's runtime. `main` (already
    /// `#[tokio::main]`) awaits this; the sync `Origin::snapshot()` wrapper is
    /// only for non-async callers and must not be invoked from within a runtime
    /// (it builds and `block_on`s a throwaway one).
    pub async fn snapshot_async(&self) -> io::Result<BTreeMap<INodeNo, Item>> {
        // 1. Pull client's current (inode, hash) pairs from the DB.
        let client_entries: Vec<(u64, u64)> = item_entities::Entity::find()
            .all(&self.client)
            .await
            .map_err(|e| {
                error!(error = %e, "network snapshot: client entries DB read failed");
                io_err(e)
            })?
            .into_iter()
            .filter(|m| m.file_type == "file")
            .map(|m| (m.inode as u64, m.hash as u64))
            .collect();

        // 2. Ask server what changed.
        let response = self
            .transport
            .reconcile(client_entries)
            .await
            .map_err(|e| {
                error!(error = %e, "network snapshot: reconcile with server failed");
                io_err(e)
            })?;

        // 3. Invalidate cached bytes for changed + deleted inodes — they are
        //    stale by definition.
        let mut changed_or_deleted: Vec<i64> =
            response.changed.iter().map(|ih| ih.inode as i64).collect();
        changed_or_deleted.extend(response.deleted.iter().map(|i| *i as i64));
        delete_cached_bytes_for(&changed_or_deleted, &self.client).await;

        // latest_manifest is in-memory; on restart it's empty so the reconcile
        // delta misses unchanged files. Fetch full manifest then, delta otherwise.
        let mut current_manifest = self.latest_manifest.read().unwrap().clone();
        if current_manifest.is_empty() {
            let entries = self.transport.get_manifest().await.map_err(|e| {
                error!(error = %e, "network snapshot: get_manifest from server failed");
                io_err(e)
            })?;
            current_manifest = entries.into_iter().map(|e| (e.id, e)).collect();
        } else {
            let wanted: Vec<u64> = response.changed.iter().map(|ih| ih.inode).collect();
            if !wanted.is_empty() {
                let wanted_count = wanted.len();
                let entries = self.transport.get_metadata(wanted).await.map_err(|e| {
                    error!(
                        wanted = wanted_count,
                        error = %e,
                        "network snapshot: get_metadata from server failed"
                    );
                    io_err(e)
                })?;
                for entry in entries {
                    current_manifest.insert(entry.id, entry);
                }
            }
        }
        for inode in &response.deleted {
            current_manifest.remove(inode);
        }
        *self.latest_manifest.write().unwrap() = current_manifest.clone();

        let snapshot =
            build_snapshot_from_manifest(&current_manifest, &self.destination, &self.client)
                .await?;

        info!(
            changed = response.changed.len(),
            deleted = response.deleted.len(),
            files = snapshot.len(),
            "network snapshot complete"
        );
        return Ok(snapshot);
    }
}

/// Build a complete Items snapshot from the server manifest, sync it to the
/// DB, and restore music metadata + virtual paths from existing DB rows.
///
/// Shared between `snapshot_async` (initial mount) and the watcher's
/// `reconcile_once` (runtime updates) so both paths produce identical
/// snapshots and keep the DB in sync.
pub(crate) async fn build_snapshot_from_manifest(
    manifest: &BTreeMap<u64, ProtoManifestEntry>,
    destination: &Path,
    client: &sea_orm::DatabaseConnection,
) -> io::Result<BTreeMap<INodeNo, Item>> {
    let source_root = PathBuf::from("/");
    let mut snapshot = BTreeMap::new();
    let root_attrs = FileAttrs {
        size: 0,
        blocks: 0,
        atime: SystemTime::UNIX_EPOCH,
        mtime: SystemTime::UNIX_EPOCH,
        ctime: SystemTime::UNIX_EPOCH,
        crtime: SystemTime::UNIX_EPOCH,
        perm: 0o755,
        nlink: 2,
        uid: 0,
        gid: 0,
        rdev: 0,
        blksize: 4096,
    };
    snapshot.insert(
        INodeNo::ROOT,
        Item::new(
            INodeNo::ROOT,
            INodeNo::ROOT,
            "/".to_string(),
            destination.to_path_buf(),
            destination.to_path_buf(),
            FileType::Directory,
            root_attrs,
            None,
        ),
    );

    for (_id, entry) in manifest {
        let item = manifest_entry_to_item(entry, &source_root, &mut snapshot);
        snapshot.insert(item.inode, item);
    }

    let db_items: std::collections::HashMap<i64, item_entities::Model> =
        item_entities::Entity::find()
            .all(client)
            .await
            .map_err(|e| {
                error!(error = %e, "build_snapshot_from_manifest: DB read failed");
                io_err(e)
            })?
            .into_iter()
            .map(|e| (e.inode, e))
            .collect();
    sync_items_to_db(&snapshot, &db_items, client).await;
    restore_music_metadata_from_db(&mut snapshot, &db_items, client).await;
    restore_virtual_paths(&mut snapshot, &db_items, &source_root);

    return Ok(snapshot);
}

fn manifest_entry_to_item(
    entry: &ProtoManifestEntry,
    source_root: &Path,
    snapshot: &mut BTreeMap<INodeNo, Item>,
) -> Item {
    let inode = INodeNo(entry.id);
    let original_path = PathBuf::from(&entry.rel_path);
    let attrs = FileAttrs {
        size: entry.size,
        blocks: 0,
        atime: SystemTime::UNIX_EPOCH + Duration::from_secs(entry.mtime),
        mtime: SystemTime::UNIX_EPOCH + Duration::from_secs(entry.mtime),
        ctime: SystemTime::UNIX_EPOCH + Duration::from_secs(entry.ctime),
        crtime: SystemTime::UNIX_EPOCH + Duration::from_secs(entry.crtime),
        perm: 0o644,
        nlink: 1,
        uid: 0,
        gid: 0,
        rdev: 0,
        blksize: 4096,
    };
    let music_metadata: Option<MusicMetadata> =
        entry.music_metadata.clone().map(music_metadata_from_proto);

    let mut local_path = PathBuf::new();
    if let Some(mm) = &music_metadata {
        let joined;
        let artist_dir = match mm.album_artist.as_deref() {
            Some(a) => a,
            None => {
                joined = mm.artist.join("-");
                &joined
            }
        };
        local_path.push(artist_dir);
        local_path.push(&mm.album);
    }
    let name = Path::new(&entry.rel_path)
        .file_name()
        .map(|n| n.to_string_lossy().into_owned())
        .unwrap_or_else(|| entry.rel_path.clone());
    local_path.push(&name);

    let parent_inode = ensure_virtual_dirs(&local_path, source_root, snapshot);
    return Item::new(
        inode,
        parent_inode,
        name,
        original_path,
        local_path,
        FileType::File,
        attrs,
        music_metadata,
    );
}

fn music_metadata_from_proto(mm: crate::proto::MusicMetadata) -> MusicMetadata {
    MusicMetadata {
        artist: mm.artist,
        album_artist: mm.album_artist,
        album: mm.album,
        track_number: mm.track_number,
        track_title: mm.track_title,
        other_tags: mm.other_tags,
        header: mm.header,
        picture_block_headers: mm.picture_block_headers,
        picture_data_ranges: mm
            .picture_data_ranges
            .into_iter()
            .map(|p| (p.offset, p.length))
            .collect(),
        real_audio_start: mm.real_audio_start,
        vorbis_comment_offset: mm.vorbis_comment_offset,
        vorbis_comment_length: mm.vorbis_comment_length,
    }
}

fn io_err<E: std::fmt::Display>(e: E) -> io::Error {
    return io::Error::new(io::ErrorKind::Other, e.to_string());
}

pub struct NetworkByteSource {
    transport: NetworkTransport,
    runtime_handle: tokio::runtime::Handle,
    client: sea_orm::DatabaseConnection,
}

impl ByteSource for NetworkByteSource {
    fn read_at(
        &self,
        inode: INodeNo,
        _locator: &Path,
        offset: u64,
        len: usize,
    ) -> io::Result<Vec<u8>> {
        let inode_i64 = inode.0 as i64;
        trace!(%inode, offset, len, "network read_at");

        let cached = self
            .runtime_handle
            .block_on(get_cached_bytes(inode_i64, &self.client));
        if let Some(data) = cached {
            debug!(%inode, "read_at cache hit");
            return slice_range(&data, offset, len);
        }

        debug!(%inode, "read_at cache miss; fetching from server");
        let (data, _total_size) = self
            .runtime_handle
            .block_on(self.transport.fetch_file_range(inode.0, 0, 0))
            .map_err(|e| {
                error!(%inode, error = %e, "read_at: fetch_file_range failed");
                io_err(e)
            })?;
        self.runtime_handle
            .block_on(put_cached_bytes(inode_i64, data.clone(), &self.client));

        return slice_range(&data, offset, len);
    }
}

fn slice_range(data: &[u8], offset: u64, len: usize) -> io::Result<Vec<u8>> {
    let start = (offset as usize).min(data.len());
    let end = (start + len).min(data.len());
    return Ok(data[start..end].to_vec());
}