hermes_core/directories/
directory.rs

1//! Async Directory abstraction for IO operations
2//!
3//! Supports network, local filesystem, and in-memory storage.
4//! All reads are async to minimize blocking on network latency.
5
6use async_trait::async_trait;
7use parking_lot::RwLock;
8use std::collections::HashMap;
9use std::io;
10use std::ops::Range;
11use std::path::{Path, PathBuf};
12use std::sync::Arc;
13
14/// Callback type for lazy range reading
15#[cfg(not(target_arch = "wasm32"))]
16pub type RangeReadFn = Arc<
17    dyn Fn(
18            Range<u64>,
19        )
20            -> std::pin::Pin<Box<dyn std::future::Future<Output = io::Result<OwnedBytes>> + Send>>
21        + Send
22        + Sync,
23>;
24
25#[cfg(target_arch = "wasm32")]
26pub type RangeReadFn = Arc<
27    dyn Fn(
28        Range<u64>,
29    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = io::Result<OwnedBytes>>>>,
30>;
31
32/// Unified file handle for both inline (mmap/RAM) and lazy (HTTP/filesystem) access.
33///
34/// Replaces the previous `FileSlice`, `LazyFileHandle`, and `LazyFileSlice` types.
35/// - **Inline**: data is available synchronously (mmap, RAM). Sync reads via `read_bytes_range_sync`.
36/// - **Lazy**: data is fetched on-demand via async callback (HTTP, filesystem).
37///
38/// Use `.slice()` to create sub-range views (zero-copy for Inline, offset-adjusted for Lazy).
39#[derive(Clone)]
40pub struct FileHandle {
41    inner: FileHandleInner,
42}
43
44#[derive(Clone)]
45enum FileHandleInner {
46    /// Data available inline — sync reads possible (mmap, RAM)
47    Inline {
48        data: OwnedBytes,
49        offset: u64,
50        len: u64,
51    },
52    /// Data fetched on-demand via async callback (HTTP, filesystem)
53    Lazy {
54        read_fn: RangeReadFn,
55        offset: u64,
56        len: u64,
57    },
58}
59
60impl std::fmt::Debug for FileHandle {
61    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62        match &self.inner {
63            FileHandleInner::Inline { len, offset, .. } => f
64                .debug_struct("FileHandle::Inline")
65                .field("offset", offset)
66                .field("len", len)
67                .finish(),
68            FileHandleInner::Lazy { len, offset, .. } => f
69                .debug_struct("FileHandle::Lazy")
70                .field("offset", offset)
71                .field("len", len)
72                .finish(),
73        }
74    }
75}
76
77impl FileHandle {
78    /// Create an inline file handle from owned bytes (mmap, RAM).
79    /// Sync reads are available.
80    pub fn from_bytes(data: OwnedBytes) -> Self {
81        let len = data.len() as u64;
82        Self {
83            inner: FileHandleInner::Inline {
84                data,
85                offset: 0,
86                len,
87            },
88        }
89    }
90
91    /// Create an empty file handle.
92    pub fn empty() -> Self {
93        Self::from_bytes(OwnedBytes::empty())
94    }
95
96    /// Create a lazy file handle from an async range-read callback.
97    /// Only async reads are available.
98    pub fn lazy(len: u64, read_fn: RangeReadFn) -> Self {
99        Self {
100            inner: FileHandleInner::Lazy {
101                read_fn,
102                offset: 0,
103                len,
104            },
105        }
106    }
107
108    /// Total length in bytes.
109    #[inline]
110    pub fn len(&self) -> u64 {
111        match &self.inner {
112            FileHandleInner::Inline { len, .. } => *len,
113            FileHandleInner::Lazy { len, .. } => *len,
114        }
115    }
116
117    /// Check if empty.
118    #[inline]
119    pub fn is_empty(&self) -> bool {
120        self.len() == 0
121    }
122
123    /// Whether synchronous reads are available (inline/mmap data).
124    #[inline]
125    pub fn is_sync(&self) -> bool {
126        matches!(&self.inner, FileHandleInner::Inline { .. })
127    }
128
129    /// Create a sub-range view. Zero-copy for Inline, offset-adjusted for Lazy.
130    pub fn slice(&self, range: Range<u64>) -> Self {
131        match &self.inner {
132            FileHandleInner::Inline { data, offset, len } => {
133                let new_offset = offset + range.start;
134                let new_len = range.end - range.start;
135                debug_assert!(
136                    new_offset + new_len <= offset + len,
137                    "slice out of bounds: {}+{} > {}+{}",
138                    new_offset,
139                    new_len,
140                    offset,
141                    len
142                );
143                Self {
144                    inner: FileHandleInner::Inline {
145                        data: data.clone(),
146                        offset: new_offset,
147                        len: new_len,
148                    },
149                }
150            }
151            FileHandleInner::Lazy {
152                read_fn,
153                offset,
154                len,
155            } => {
156                let new_offset = offset + range.start;
157                let new_len = range.end - range.start;
158                debug_assert!(
159                    new_offset + new_len <= offset + len,
160                    "slice out of bounds: {}+{} > {}+{}",
161                    new_offset,
162                    new_len,
163                    offset,
164                    len
165                );
166                Self {
167                    inner: FileHandleInner::Lazy {
168                        read_fn: Arc::clone(read_fn),
169                        offset: new_offset,
170                        len: new_len,
171                    },
172                }
173            }
174        }
175    }
176
177    /// Advise the kernel about the access pattern for a byte range of this handle.
178    ///
179    /// Only effective for Inline handles backed by mmap; no-op for Lazy
180    /// handles (HTTP, filesystem callbacks) and heap-backed data.
181    #[cfg(feature = "native")]
182    pub fn madvise_range(&self, range: Range<u64>, advice: libc::c_int) {
183        if let FileHandleInner::Inline { data, offset, len } = &self.inner {
184            let end = range.end.min(*len);
185            if range.start >= end {
186                return;
187            }
188            let start = (*offset + range.start) as usize;
189            let end = (*offset + end) as usize;
190            data.madvise_range(start..end, advice);
191        }
192    }
193
194    /// Async range read — works for both Inline and Lazy.
195    pub async fn read_bytes_range(&self, range: Range<u64>) -> io::Result<OwnedBytes> {
196        match &self.inner {
197            FileHandleInner::Inline { data, offset, len } => {
198                if range.end > *len {
199                    return Err(io::Error::new(
200                        io::ErrorKind::InvalidInput,
201                        format!("Range {:?} out of bounds (len: {})", range, len),
202                    ));
203                }
204                let start = (*offset + range.start) as usize;
205                let end = (*offset + range.end) as usize;
206                Ok(data.slice(start..end))
207            }
208            FileHandleInner::Lazy {
209                read_fn,
210                offset,
211                len,
212            } => {
213                if range.end > *len {
214                    return Err(io::Error::new(
215                        io::ErrorKind::InvalidInput,
216                        format!("Range {:?} out of bounds (len: {})", range, len),
217                    ));
218                }
219                let abs_start = offset + range.start;
220                let abs_end = offset + range.end;
221                // Real IO (HTTP / custom read_fn) — mmap-backed Inline handles
222                // above are zero-copy slices whose latency materializes as
223                // page faults inside the query-phase histograms instead.
224                let t = crate::observe::Timer::start();
225                let result = (read_fn)(abs_start..abs_end).await;
226                if let Ok(bytes) = &result {
227                    crate::observe::directory_read("lazy_range", t.secs(), bytes.len());
228                }
229                result
230            }
231        }
232    }
233
234    /// Read all bytes.
235    pub async fn read_bytes(&self) -> io::Result<OwnedBytes> {
236        self.read_bytes_range(0..self.len()).await
237    }
238
239    /// Synchronous range read — only works for Inline handles.
240    /// Returns `Err` if the handle is Lazy.
241    #[inline]
242    pub fn read_bytes_range_sync(&self, range: Range<u64>) -> io::Result<OwnedBytes> {
243        match &self.inner {
244            FileHandleInner::Inline { data, offset, len } => {
245                if range.end > *len {
246                    return Err(io::Error::new(
247                        io::ErrorKind::InvalidInput,
248                        format!("Range {:?} out of bounds (len: {})", range, len),
249                    ));
250                }
251                let start = (*offset + range.start) as usize;
252                let end = (*offset + range.end) as usize;
253                Ok(data.slice(start..end))
254            }
255            FileHandleInner::Lazy { .. } => Err(io::Error::new(
256                io::ErrorKind::Unsupported,
257                "Synchronous read not available on lazy file handle",
258            )),
259        }
260    }
261
262    /// Synchronous read of all bytes — only works for Inline handles.
263    #[inline]
264    pub fn read_bytes_sync(&self) -> io::Result<OwnedBytes> {
265        self.read_bytes_range_sync(0..self.len())
266    }
267}
268
269/// Backing store for OwnedBytes — supports both heap Vec and mmap.
270#[derive(Clone)]
271enum SharedBytes {
272    Vec(Arc<Vec<u8>>),
273    #[cfg(feature = "native")]
274    Mmap(Arc<memmap2::Mmap>),
275}
276
277impl SharedBytes {
278    #[inline]
279    fn as_bytes(&self) -> &[u8] {
280        match self {
281            SharedBytes::Vec(v) => v.as_slice(),
282            #[cfg(feature = "native")]
283            SharedBytes::Mmap(m) => m.as_ref(),
284        }
285    }
286}
287
288impl std::fmt::Debug for SharedBytes {
289    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
290        match self {
291            SharedBytes::Vec(v) => write!(f, "Vec(len={})", v.len()),
292            #[cfg(feature = "native")]
293            SharedBytes::Mmap(m) => write!(f, "Mmap(len={})", m.len()),
294        }
295    }
296}
297
298/// Owned bytes with cheap cloning (Arc-backed)
299///
300/// Supports two backing stores:
301/// - `Vec<u8>` for owned data (RamDirectory, FsDirectory, decompressed blocks)
302/// - `Mmap` for zero-copy memory-mapped files (MmapDirectory, native only)
303#[derive(Debug, Clone)]
304pub struct OwnedBytes {
305    data: SharedBytes,
306    range: Range<usize>,
307}
308
309impl OwnedBytes {
310    pub fn new(data: Vec<u8>) -> Self {
311        let len = data.len();
312        Self {
313            data: SharedBytes::Vec(Arc::new(data)),
314            range: 0..len,
315        }
316    }
317
318    pub fn empty() -> Self {
319        Self {
320            data: SharedBytes::Vec(Arc::new(Vec::new())),
321            range: 0..0,
322        }
323    }
324
325    /// Create from a pre-existing Arc<Vec<u8>> with a sub-range.
326    /// Used by RamDirectory and CachingDirectory to share data without copying.
327    pub(crate) fn from_arc_vec(data: Arc<Vec<u8>>, range: Range<usize>) -> Self {
328        Self {
329            data: SharedBytes::Vec(data),
330            range,
331        }
332    }
333
334    /// Create from a memory-mapped file (zero-copy).
335    #[cfg(feature = "native")]
336    pub(crate) fn from_mmap(mmap: Arc<memmap2::Mmap>) -> Self {
337        let len = mmap.len();
338        Self {
339            data: SharedBytes::Mmap(mmap),
340            range: 0..len,
341        }
342    }
343
344    /// Create from a memory-mapped file with a sub-range (zero-copy).
345    #[cfg(feature = "native")]
346    pub(crate) fn from_mmap_range(mmap: Arc<memmap2::Mmap>, range: Range<usize>) -> Self {
347        Self {
348            data: SharedBytes::Mmap(mmap),
349            range,
350        }
351    }
352
353    pub fn len(&self) -> usize {
354        self.range.len()
355    }
356
357    pub fn is_empty(&self) -> bool {
358        self.range.is_empty()
359    }
360
361    pub fn slice(&self, range: Range<usize>) -> Self {
362        let start = self.range.start + range.start;
363        let end = self.range.start + range.end;
364        Self {
365            data: self.data.clone(),
366            range: start..end,
367        }
368    }
369
370    pub fn as_slice(&self) -> &[u8] {
371        &self.data.as_bytes()[self.range.clone()]
372    }
373
374    /// Returns `true` if the backing store is a memory-mapped file.
375    ///
376    /// Used to guard `madvise` calls: `MADV_DONTNEED` on heap memory
377    /// zeroes pages on Linux and corrupts allocator metadata.
378    #[cfg(feature = "native")]
379    #[inline]
380    pub fn is_mmap(&self) -> bool {
381        matches!(self.data, SharedBytes::Mmap(_))
382    }
383
384    /// Advise the kernel about the access pattern for these bytes.
385    ///
386    /// No-op unless the backing store is mmap (heap memory must never be
387    /// madvised: `MADV_DONTNEED` on heap zeroes pages and corrupts allocator
388    /// metadata) or the range is empty.
389    #[cfg(feature = "native")]
390    pub fn madvise(&self, advice: libc::c_int) {
391        self.madvise_range(0..self.len(), advice);
392    }
393
394    /// Pin these bytes in physical memory (`mlock`). mmap-backed only —
395    /// heap memory is not evictable by the page cache. Returns whether the
396    /// lock succeeded; failure (e.g. RLIMIT_MEMLOCK) is not fatal.
397    /// Locks are released automatically when the mapping is unmapped.
398    #[cfg(feature = "native")]
399    pub fn mlock(&self) -> bool {
400        if !self.is_mmap() {
401            return false;
402        }
403        let slice = self.as_slice();
404        if slice.is_empty() {
405            return true;
406        }
407        let ptr = slice.as_ptr();
408        let len = slice.len();
409        let page_size = 4096usize;
410        let aligned_ptr = (ptr as usize) & !(page_size - 1);
411        let aligned_len = len + (ptr as usize - aligned_ptr);
412        unsafe { libc::mlock(aligned_ptr as *const libc::c_void, aligned_len) == 0 }
413    }
414
415    /// Advise the kernel about the access pattern for a sub-range.
416    ///
417    /// The range is relative to these bytes. Same mmap-only guard as
418    /// [`Self::madvise`]. The pointer is aligned down to a page boundary
419    /// as required by `madvise`.
420    #[cfg(feature = "native")]
421    pub fn madvise_range(&self, range: Range<usize>, advice: libc::c_int) {
422        if !self.is_mmap() {
423            return;
424        }
425        let slice = &self.as_slice()[range];
426        if slice.is_empty() {
427            return;
428        }
429        let ptr = slice.as_ptr();
430        let len = slice.len();
431        let page_size = 4096usize;
432        let aligned_ptr = (ptr as usize) & !(page_size - 1);
433        let aligned_len = len + (ptr as usize - aligned_ptr);
434        unsafe {
435            libc::madvise(aligned_ptr as *mut libc::c_void, aligned_len, advice);
436        }
437    }
438
439    pub fn to_vec(&self) -> Vec<u8> {
440        self.as_slice().to_vec()
441    }
442}
443
444impl AsRef<[u8]> for OwnedBytes {
445    fn as_ref(&self) -> &[u8] {
446        self.as_slice()
447    }
448}
449
450impl std::ops::Deref for OwnedBytes {
451    type Target = [u8];
452
453    fn deref(&self) -> &Self::Target {
454        self.as_slice()
455    }
456}
457
458/// Async directory trait for reading index files
459#[cfg(not(target_arch = "wasm32"))]
460#[async_trait]
461pub trait Directory: Send + Sync + 'static {
462    /// Check if a file exists
463    async fn exists(&self, path: &Path) -> io::Result<bool>;
464
465    /// Get file size
466    async fn file_size(&self, path: &Path) -> io::Result<u64>;
467
468    /// Open a file for reading (loads entire file into an inline FileHandle)
469    async fn open_read(&self, path: &Path) -> io::Result<FileHandle>;
470
471    /// Read a specific byte range from a file (optimized for network)
472    async fn read_range(&self, path: &Path, range: Range<u64>) -> io::Result<OwnedBytes>;
473
474    /// List files in directory
475    async fn list_files(&self, prefix: &Path) -> io::Result<Vec<PathBuf>>;
476
477    /// Open a file handle that fetches ranges on demand.
478    /// For mmap directories this returns an Inline handle (sync-capable).
479    /// For HTTP/filesystem directories this returns a Lazy handle.
480    async fn open_lazy(&self, path: &Path) -> io::Result<FileHandle>;
481}
482
483/// Async directory trait for reading index files (WASM version - no Send requirement)
484#[cfg(target_arch = "wasm32")]
485#[async_trait(?Send)]
486pub trait Directory: 'static {
487    /// Check if a file exists
488    async fn exists(&self, path: &Path) -> io::Result<bool>;
489
490    /// Get file size
491    async fn file_size(&self, path: &Path) -> io::Result<u64>;
492
493    /// Open a file for reading (loads entire file into an inline FileHandle)
494    async fn open_read(&self, path: &Path) -> io::Result<FileHandle>;
495
496    /// Read a specific byte range from a file (optimized for network)
497    async fn read_range(&self, path: &Path, range: Range<u64>) -> io::Result<OwnedBytes>;
498
499    /// List files in directory
500    async fn list_files(&self, prefix: &Path) -> io::Result<Vec<PathBuf>>;
501
502    /// Open a file handle that fetches ranges on demand.
503    async fn open_lazy(&self, path: &Path) -> io::Result<FileHandle>;
504}
505
506/// A writer for incrementally writing data to a directory file.
507///
508/// Avoids buffering entire files in memory during merge. File-backed
509/// directories write directly to disk; memory directories collect to Vec.
510pub trait StreamingWriter: io::Write + Send {
511    /// Finalize the write, making data available for reading.
512    fn finish(self: Box<Self>) -> io::Result<()>;
513
514    /// Bytes written so far.
515    fn bytes_written(&self) -> u64;
516}
517
518/// StreamingWriter backed by Vec<u8>, finalized via DirectoryWriter::write.
519/// Used as default/fallback and for RamDirectory.
520struct BufferedStreamingWriter {
521    path: PathBuf,
522    buffer: Vec<u8>,
523    /// Callback to write the buffer to the directory on finish.
524    /// We store the files Arc directly for RamDirectory.
525    files: Arc<RwLock<HashMap<PathBuf, Arc<Vec<u8>>>>>,
526}
527
528impl io::Write for BufferedStreamingWriter {
529    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
530        self.buffer.extend_from_slice(buf);
531        Ok(buf.len())
532    }
533
534    fn flush(&mut self) -> io::Result<()> {
535        Ok(())
536    }
537}
538
539impl StreamingWriter for BufferedStreamingWriter {
540    fn finish(self: Box<Self>) -> io::Result<()> {
541        self.files.write().insert(self.path, Arc::new(self.buffer));
542        Ok(())
543    }
544
545    fn bytes_written(&self) -> u64 {
546        self.buffer.len() as u64
547    }
548}
549
550/// Buffer size for FileStreamingWriter (8 MB).
551/// Large enough to coalesce millions of tiny writes (e.g. per-vector doc_id writes)
552/// into efficient sequential I/O.
553#[cfg(feature = "native")]
554const FILE_STREAMING_BUF_SIZE: usize = 8 * 1024 * 1024;
555
556/// StreamingWriter backed by a buffered std::fs::File for filesystem directories.
557#[cfg(feature = "native")]
558pub(crate) struct FileStreamingWriter {
559    pub(crate) file: io::BufWriter<std::fs::File>,
560    pub(crate) written: u64,
561}
562
563#[cfg(feature = "native")]
564impl FileStreamingWriter {
565    pub(crate) fn new(file: std::fs::File) -> Self {
566        Self {
567            file: io::BufWriter::with_capacity(FILE_STREAMING_BUF_SIZE, file),
568            written: 0,
569        }
570    }
571}
572
573#[cfg(feature = "native")]
574impl io::Write for FileStreamingWriter {
575    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
576        let n = self.file.write(buf)?;
577        self.written += n as u64;
578        Ok(n)
579    }
580
581    fn flush(&mut self) -> io::Result<()> {
582        self.file.flush()
583    }
584}
585
586#[cfg(feature = "native")]
587impl StreamingWriter for FileStreamingWriter {
588    fn finish(self: Box<Self>) -> io::Result<()> {
589        let file = self.file.into_inner().map_err(|e| e.into_error())?;
590        file.sync_all()?;
591        Ok(())
592    }
593
594    fn bytes_written(&self) -> u64 {
595        self.written
596    }
597}
598
599/// Async directory trait for writing index files
600#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
601#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
602pub trait DirectoryWriter: Directory {
603    /// Create/overwrite a file with data
604    async fn write(&self, path: &Path, data: &[u8]) -> io::Result<()>;
605
606    /// Delete a file
607    async fn delete(&self, path: &Path) -> io::Result<()>;
608
609    /// Atomic rename
610    async fn rename(&self, from: &Path, to: &Path) -> io::Result<()>;
611
612    /// Sync all pending writes
613    async fn sync(&self) -> io::Result<()>;
614
615    /// Create a streaming writer for incremental file writes.
616    /// Call finish() on the returned writer to finalize.
617    async fn streaming_writer(&self, path: &Path) -> io::Result<Box<dyn StreamingWriter>>;
618
619    /// Streaming writer for **bulk one-shot data** (merge/reorder outputs).
620    ///
621    /// Filesystem directories return a page-cache-dropping writer (see
622    /// `docs/cold-io.md`) so multi-GB merge writes cannot evict the serving
623    /// segments' warm pages. Output is byte-identical to the buffered
624    /// writer. Default impl delegates to [`Self::streaming_writer`].
625    async fn streaming_writer_cold(&self, path: &Path) -> io::Result<Box<dyn StreamingWriter>> {
626        self.streaming_writer(path).await
627    }
628}
629
630/// In-memory directory for testing and small indexes
631#[derive(Debug, Default)]
632pub struct RamDirectory {
633    files: Arc<RwLock<HashMap<PathBuf, Arc<Vec<u8>>>>>,
634}
635
636impl Clone for RamDirectory {
637    fn clone(&self) -> Self {
638        Self {
639            files: Arc::clone(&self.files),
640        }
641    }
642}
643
644impl RamDirectory {
645    pub fn new() -> Self {
646        Self::default()
647    }
648
649    /// Synchronous file listing (for serialization).
650    pub fn list_files_sync(&self, prefix: &Path) -> io::Result<Vec<PathBuf>> {
651        let files = self.files.read();
652        Ok(files
653            .keys()
654            .filter(|p| p.starts_with(prefix))
655            .cloned()
656            .collect())
657    }
658
659    /// Synchronous file read (for serialization).
660    pub fn read_file_sync(&self, path: &Path) -> io::Result<Vec<u8>> {
661        let files = self.files.read();
662        files
663            .get(path)
664            .map(|data| data.as_ref().clone())
665            .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "File not found"))
666    }
667
668    /// Synchronous file write (for deserialization).
669    pub fn write_sync(&self, path: &Path, data: &[u8]) -> io::Result<()> {
670        self.files
671            .write()
672            .insert(path.to_path_buf(), Arc::new(data.to_vec()));
673        Ok(())
674    }
675}
676
677#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
678#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
679impl Directory for RamDirectory {
680    async fn exists(&self, path: &Path) -> io::Result<bool> {
681        Ok(self.files.read().contains_key(path))
682    }
683
684    async fn file_size(&self, path: &Path) -> io::Result<u64> {
685        self.files
686            .read()
687            .get(path)
688            .map(|data| data.len() as u64)
689            .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "File not found"))
690    }
691
692    async fn open_read(&self, path: &Path) -> io::Result<FileHandle> {
693        let files = self.files.read();
694        let data = files
695            .get(path)
696            .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "File not found"))?;
697
698        Ok(FileHandle::from_bytes(OwnedBytes::from_arc_vec(
699            Arc::clone(data),
700            0..data.len(),
701        )))
702    }
703
704    async fn read_range(&self, path: &Path, range: Range<u64>) -> io::Result<OwnedBytes> {
705        let files = self.files.read();
706        let data = files
707            .get(path)
708            .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "File not found"))?;
709
710        let start = range.start as usize;
711        let end = range.end as usize;
712
713        if end > data.len() {
714            return Err(io::Error::new(
715                io::ErrorKind::InvalidInput,
716                "Range out of bounds",
717            ));
718        }
719
720        Ok(OwnedBytes::from_arc_vec(Arc::clone(data), start..end))
721    }
722
723    async fn list_files(&self, prefix: &Path) -> io::Result<Vec<PathBuf>> {
724        let files = self.files.read();
725        Ok(files
726            .keys()
727            .filter(|p| p.starts_with(prefix))
728            .cloned()
729            .collect())
730    }
731
732    async fn open_lazy(&self, path: &Path) -> io::Result<FileHandle> {
733        // RAM data is always available synchronously — return Inline handle
734        self.open_read(path).await
735    }
736}
737
738#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
739#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
740impl DirectoryWriter for RamDirectory {
741    async fn write(&self, path: &Path, data: &[u8]) -> io::Result<()> {
742        self.files
743            .write()
744            .insert(path.to_path_buf(), Arc::new(data.to_vec()));
745        Ok(())
746    }
747
748    async fn delete(&self, path: &Path) -> io::Result<()> {
749        self.files.write().remove(path);
750        Ok(())
751    }
752
753    async fn rename(&self, from: &Path, to: &Path) -> io::Result<()> {
754        let mut files = self.files.write();
755        if let Some(data) = files.remove(from) {
756            files.insert(to.to_path_buf(), data);
757        }
758        Ok(())
759    }
760
761    async fn sync(&self) -> io::Result<()> {
762        Ok(())
763    }
764
765    async fn streaming_writer(&self, path: &Path) -> io::Result<Box<dyn StreamingWriter>> {
766        Ok(Box::new(BufferedStreamingWriter {
767            path: path.to_path_buf(),
768            buffer: Vec::new(),
769            files: Arc::clone(&self.files),
770        }))
771    }
772}
773
774/// Local filesystem directory with async IO via tokio
775#[cfg(feature = "native")]
776#[derive(Debug, Clone)]
777pub struct FsDirectory {
778    root: PathBuf,
779}
780
781#[cfg(feature = "native")]
782impl FsDirectory {
783    pub fn new(root: impl AsRef<Path>) -> Self {
784        Self {
785            root: root.as_ref().to_path_buf(),
786        }
787    }
788
789    fn resolve(&self, path: &Path) -> PathBuf {
790        self.root.join(path)
791    }
792}
793
794#[cfg(feature = "native")]
795#[async_trait]
796impl Directory for FsDirectory {
797    async fn exists(&self, path: &Path) -> io::Result<bool> {
798        let full_path = self.resolve(path);
799        Ok(tokio::fs::try_exists(&full_path).await.unwrap_or(false))
800    }
801
802    async fn file_size(&self, path: &Path) -> io::Result<u64> {
803        let full_path = self.resolve(path);
804        let metadata = tokio::fs::metadata(&full_path).await?;
805        Ok(metadata.len())
806    }
807
808    async fn open_read(&self, path: &Path) -> io::Result<FileHandle> {
809        let full_path = self.resolve(path);
810        let data = tokio::fs::read(&full_path).await?;
811        Ok(FileHandle::from_bytes(OwnedBytes::new(data)))
812    }
813
814    async fn read_range(&self, path: &Path, range: Range<u64>) -> io::Result<OwnedBytes> {
815        use tokio::io::{AsyncReadExt, AsyncSeekExt};
816
817        let full_path = self.resolve(path);
818        let mut file = tokio::fs::File::open(&full_path).await?;
819
820        file.seek(std::io::SeekFrom::Start(range.start)).await?;
821
822        let len = (range.end - range.start) as usize;
823        let mut buffer = vec![0u8; len];
824        file.read_exact(&mut buffer).await?;
825
826        Ok(OwnedBytes::new(buffer))
827    }
828
829    async fn list_files(&self, prefix: &Path) -> io::Result<Vec<PathBuf>> {
830        let full_path = self.resolve(prefix);
831        let mut entries = tokio::fs::read_dir(&full_path).await?;
832        let mut files = Vec::new();
833
834        while let Some(entry) = entries.next_entry().await? {
835            if entry.file_type().await?.is_file() {
836                files.push(entry.path().strip_prefix(&self.root).unwrap().to_path_buf());
837            }
838        }
839
840        Ok(files)
841    }
842
843    async fn open_lazy(&self, path: &Path) -> io::Result<FileHandle> {
844        let full_path = self.resolve(path);
845        let metadata = tokio::fs::metadata(&full_path).await?;
846        let file_size = metadata.len();
847
848        let read_fn: RangeReadFn = Arc::new(move |range: Range<u64>| {
849            let full_path = full_path.clone();
850            Box::pin(async move {
851                use tokio::io::{AsyncReadExt, AsyncSeekExt};
852
853                let mut file = tokio::fs::File::open(&full_path).await?;
854                file.seek(std::io::SeekFrom::Start(range.start)).await?;
855
856                let len = (range.end - range.start) as usize;
857                let mut buffer = vec![0u8; len];
858                file.read_exact(&mut buffer).await?;
859
860                Ok(OwnedBytes::new(buffer))
861            })
862        });
863
864        Ok(FileHandle::lazy(file_size, read_fn))
865    }
866}
867
868#[cfg(feature = "native")]
869#[async_trait]
870impl DirectoryWriter for FsDirectory {
871    async fn write(&self, path: &Path, data: &[u8]) -> io::Result<()> {
872        let full_path = self.resolve(path);
873
874        // Ensure parent directory exists
875        if let Some(parent) = full_path.parent() {
876            tokio::fs::create_dir_all(parent).await?;
877        }
878
879        tokio::fs::write(&full_path, data).await
880    }
881
882    async fn delete(&self, path: &Path) -> io::Result<()> {
883        let full_path = self.resolve(path);
884        tokio::fs::remove_file(&full_path).await
885    }
886
887    async fn rename(&self, from: &Path, to: &Path) -> io::Result<()> {
888        let from_path = self.resolve(from);
889        let to_path = self.resolve(to);
890        tokio::fs::rename(&from_path, &to_path).await
891    }
892
893    async fn sync(&self) -> io::Result<()> {
894        // fsync the directory
895        let dir = std::fs::File::open(&self.root)?;
896        dir.sync_all()?;
897        Ok(())
898    }
899
900    async fn streaming_writer(&self, path: &Path) -> io::Result<Box<dyn StreamingWriter>> {
901        let full_path = self.resolve(path);
902        if let Some(parent) = full_path.parent() {
903            tokio::fs::create_dir_all(parent).await?;
904        }
905        let file = std::fs::File::create(&full_path)?;
906        Ok(Box::new(FileStreamingWriter::new(file)))
907    }
908
909    async fn streaming_writer_cold(&self, path: &Path) -> io::Result<Box<dyn StreamingWriter>> {
910        let full_path = self.resolve(path);
911        if let Some(parent) = full_path.parent() {
912            tokio::fs::create_dir_all(parent).await?;
913        }
914        let file = std::fs::File::create(&full_path)?;
915        Ok(Box::new(super::ColdStreamingWriter::new(file)))
916    }
917}
918
919/// Caching wrapper for any Directory - caches file reads
920pub struct CachingDirectory<D: Directory> {
921    inner: D,
922    cache: RwLock<HashMap<PathBuf, Arc<Vec<u8>>>>,
923    max_cached_bytes: usize,
924    current_bytes: RwLock<usize>,
925}
926
927impl<D: Directory> CachingDirectory<D> {
928    pub fn new(inner: D, max_cached_bytes: usize) -> Self {
929        Self {
930            inner,
931            cache: RwLock::new(HashMap::new()),
932            max_cached_bytes,
933            current_bytes: RwLock::new(0),
934        }
935    }
936
937    fn try_cache(&self, path: &Path, data: &[u8]) {
938        let mut current = self.current_bytes.write();
939        if *current + data.len() <= self.max_cached_bytes {
940            self.cache
941                .write()
942                .insert(path.to_path_buf(), Arc::new(data.to_vec()));
943            *current += data.len();
944        }
945    }
946}
947
948#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
949#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
950impl<D: Directory> Directory for CachingDirectory<D> {
951    async fn exists(&self, path: &Path) -> io::Result<bool> {
952        if self.cache.read().contains_key(path) {
953            return Ok(true);
954        }
955        self.inner.exists(path).await
956    }
957
958    async fn file_size(&self, path: &Path) -> io::Result<u64> {
959        if let Some(data) = self.cache.read().get(path) {
960            return Ok(data.len() as u64);
961        }
962        self.inner.file_size(path).await
963    }
964
965    async fn open_read(&self, path: &Path) -> io::Result<FileHandle> {
966        // Check cache first
967        if let Some(data) = self.cache.read().get(path) {
968            return Ok(FileHandle::from_bytes(OwnedBytes::from_arc_vec(
969                Arc::clone(data),
970                0..data.len(),
971            )));
972        }
973
974        // Read from inner and potentially cache
975        let handle = self.inner.open_read(path).await?;
976        let bytes = handle.read_bytes().await?;
977
978        self.try_cache(path, bytes.as_slice());
979
980        Ok(FileHandle::from_bytes(bytes))
981    }
982
983    async fn read_range(&self, path: &Path, range: Range<u64>) -> io::Result<OwnedBytes> {
984        // Check cache first
985        if let Some(data) = self.cache.read().get(path) {
986            let start = range.start as usize;
987            let end = range.end as usize;
988            return Ok(OwnedBytes::from_arc_vec(Arc::clone(data), start..end));
989        }
990
991        self.inner.read_range(path, range).await
992    }
993
994    async fn list_files(&self, prefix: &Path) -> io::Result<Vec<PathBuf>> {
995        self.inner.list_files(prefix).await
996    }
997
998    async fn open_lazy(&self, path: &Path) -> io::Result<FileHandle> {
999        // For caching directory, delegate to inner - caching happens at read_range level
1000        self.inner.open_lazy(path).await
1001    }
1002}
1003
1004#[cfg(test)]
1005mod tests {
1006    use super::*;
1007
1008    #[tokio::test]
1009    async fn test_ram_directory() {
1010        let dir = RamDirectory::new();
1011
1012        // Write file
1013        dir.write(Path::new("test.txt"), b"hello world")
1014            .await
1015            .unwrap();
1016
1017        // Check exists
1018        assert!(dir.exists(Path::new("test.txt")).await.unwrap());
1019        assert!(!dir.exists(Path::new("nonexistent.txt")).await.unwrap());
1020
1021        // Read file
1022        let slice = dir.open_read(Path::new("test.txt")).await.unwrap();
1023        let data = slice.read_bytes().await.unwrap();
1024        assert_eq!(data.as_slice(), b"hello world");
1025
1026        // Read range
1027        let range_data = dir.read_range(Path::new("test.txt"), 0..5).await.unwrap();
1028        assert_eq!(range_data.as_slice(), b"hello");
1029
1030        // Delete
1031        dir.delete(Path::new("test.txt")).await.unwrap();
1032        assert!(!dir.exists(Path::new("test.txt")).await.unwrap());
1033    }
1034
1035    #[tokio::test]
1036    async fn test_file_handle() {
1037        let data = OwnedBytes::new(b"hello world".to_vec());
1038        let handle = FileHandle::from_bytes(data);
1039
1040        assert_eq!(handle.len(), 11);
1041        assert!(handle.is_sync());
1042
1043        let sub = handle.slice(0..5);
1044        let bytes = sub.read_bytes().await.unwrap();
1045        assert_eq!(bytes.as_slice(), b"hello");
1046
1047        let sub2 = handle.slice(6..11);
1048        let bytes2 = sub2.read_bytes().await.unwrap();
1049        assert_eq!(bytes2.as_slice(), b"world");
1050
1051        // Sync reads work on inline handles
1052        let sync_bytes = handle.read_bytes_range_sync(0..5).unwrap();
1053        assert_eq!(sync_bytes.as_slice(), b"hello");
1054    }
1055
1056    #[tokio::test]
1057    async fn test_owned_bytes() {
1058        let bytes = OwnedBytes::new(vec![1, 2, 3, 4, 5]);
1059
1060        assert_eq!(bytes.len(), 5);
1061        assert_eq!(bytes.as_slice(), &[1, 2, 3, 4, 5]);
1062
1063        let sliced = bytes.slice(1..4);
1064        assert_eq!(sliced.as_slice(), &[2, 3, 4]);
1065
1066        // Original unchanged
1067        assert_eq!(bytes.as_slice(), &[1, 2, 3, 4, 5]);
1068    }
1069}
hermes_core/directories/directory.rs

hermes_core/directories/
directory.rs
