Move the files around

crates/musicfs-cache/src/prefetch.rs

@@ -0,0 +1,197 @@
+use crate::patterns::{AccessContext, PatternStore};
+use musicfs_cas::ContentFetcher;
+use musicfs_core::{Event, EventBus, FileId};
+use parking_lot::Mutex as ParkingMutex;
+use std::collections::HashSet;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::sync::Semaphore;
+use tokio::task::JoinHandle;
+use tracing::{debug, info, warn};
+
+const DEFAULT_PREFETCH_LOOKAHEAD: usize = 3;
+const DEFAULT_MAX_CONCURRENT: usize = 2;
+const DEFAULT_COOLDOWN_MS: u64 = 100;
+
+#[derive(Debug, Clone)]
+pub struct PrefetchConfig {
+    pub lookahead: usize,
+    pub max_concurrent: usize,
+    pub cooldown: Duration,
+    pub enabled: bool,
+}
+
+impl Default for PrefetchConfig {
+    fn default() -> Self {
+        Self {
+            lookahead: DEFAULT_PREFETCH_LOOKAHEAD,
+            max_concurrent: DEFAULT_MAX_CONCURRENT,
+            cooldown: Duration::from_millis(DEFAULT_COOLDOWN_MS),
+            enabled: true,
+        }
+    }
+}
+
+pub struct PrefetchEngine {
+    config: PrefetchConfig,
+    fetcher: Arc<ContentFetcher>,
+    in_flight: Arc<ParkingMutex<HashSet<FileId>>>,
+    semaphore: Arc<Semaphore>,
+    running: Arc<AtomicBool>,
+}
+
+pub struct PrefetchHandle {
+    handle: JoinHandle<()>,
+    running: Arc<AtomicBool>,
+}
+
+impl PrefetchHandle {
+    pub async fn stop(self) {
+        self.running.store(false, Ordering::SeqCst);
+        let _ = self.handle.await;
+    }
+}
+
+impl PrefetchEngine {
+    pub fn new(
+        config: PrefetchConfig,
+        _pattern_store: Arc<PatternStore>,
+        fetcher: Arc<ContentFetcher>,
+    ) -> Self {
+        let semaphore = Arc::new(Semaphore::new(config.max_concurrent));
+
+        Self {
+            config,
+            fetcher,
+            in_flight: Arc::new(ParkingMutex::new(HashSet::new())),
+            semaphore,
+            running: Arc::new(AtomicBool::new(false)),
+        }
+    }
+
+    pub fn start(
+        self: Arc<Self>,
+        event_bus: Arc<EventBus>,
+        pattern_store: Arc<PatternStore>,
+    ) -> PrefetchHandle {
+        self.running.store(true, Ordering::SeqCst);
+        let running = self.running.clone();
+
+        let config = self.config.clone();
+        let fetcher = self.fetcher.clone();
+        let in_flight = self.in_flight.clone();
+        let semaphore = self.semaphore.clone();
+        let running_inner = running.clone();
+
+        let handle = tokio::spawn(async move {
+            let mut rx = event_bus.subscribe();
+
+            while running_inner.load(Ordering::SeqCst) {
+                match tokio::time::timeout(Duration::from_secs(1), rx.recv()).await {
+                    Ok(Ok(event)) => {
+                        if let Event::FileAccessed { file_id, .. } = event {
+                            if config.enabled {
+                                let ctx = AccessContext::default();
+                                if let Err(e) = pattern_store.record(file_id, ctx) {
+                                    warn!("Failed to record access pattern: {}", e);
+                                    continue;
+                                }
+
+                                let predictions =
+                                    pattern_store.predict_next(file_id, config.lookahead);
+
+                                for predicted_id in predictions {
+                                    prefetch_file(predicted_id, &fetcher, &in_flight, &semaphore)
+                                        .await;
+                                }
+
+                                tokio::time::sleep(config.cooldown).await;
+                            }
+                        }
+                    }
+                    Ok(Err(_)) => break,
+                    Err(_) => continue,
+                }
+            }
+
+            info!("Prefetch engine stopped");
+        });
+
+        PrefetchHandle { handle, running }
+    }
+
+    pub fn is_running(&self) -> bool {
+        self.running.load(Ordering::SeqCst)
+    }
+
+    pub fn in_flight_count(&self) -> usize {
+        self.in_flight.lock().len()
+    }
+
+    pub fn update_config(&mut self, config: PrefetchConfig) {
+        self.config = config;
+    }
+}
+
+async fn prefetch_file(
+    file_id: FileId,
+    fetcher: &Arc<ContentFetcher>,
+    in_flight: &Arc<ParkingMutex<HashSet<FileId>>>,
+    semaphore: &Arc<Semaphore>,
+) {
+    {
+        let mut guard = in_flight.lock();
+        if guard.contains(&file_id) {
+            debug!("Skipping prefetch for {:?} - already in flight", file_id);
+            return;
+        }
+        guard.insert(file_id);
+    }
+
+    let permit = match semaphore.clone().try_acquire_owned() {
+        Ok(p) => p,
+        Err(_) => {
+            debug!("Skipping prefetch for {:?} - concurrency limit", file_id);
+            in_flight.lock().remove(&file_id);
+            return;
+        }
+    };
+
+    let fetcher = fetcher.clone();
+    let in_flight = in_flight.clone();
+
+    tokio::spawn(async move {
+        debug!("Prefetching file {:?}", file_id);
+
+        match fetcher.ensure_cached(file_id).await {
+            Ok(manifest) => {
+                info!(
+                    "Prefetched {:?}: {} chunks, {} bytes",
+                    file_id,
+                    manifest.chunks.len(),
+                    manifest.total_size
+                );
+            }
+            Err(e) => {
+                debug!("Prefetch failed for {:?}: {}", file_id, e);
+            }
+        }
+
+        in_flight.lock().remove(&file_id);
+        drop(permit);
+    });
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_prefetch_config_defaults() {
+        let config = PrefetchConfig::default();
+        assert_eq!(config.lookahead, 3);
+        assert_eq!(config.max_concurrent, 2);
+        assert!(config.enabled);
+    }
+}