1use 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#[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#[derive(Clone)]
40pub struct FileHandle {
41 inner: FileHandleInner,
42}
43
44#[derive(Clone)]
45enum FileHandleInner {
46 Inline {
48 data: OwnedBytes,
49 offset: u64,
50 len: u64,
51 },
52 Lazy {
54 read_fn: RangeReadFn,
55 offset: u64,
56 len: u64,
57 label: Arc<str>,
59 },
60}
61
62#[derive(Clone, Debug)]
70pub struct IndexLabel(Arc<std::sync::RwLock<Arc<str>>>);
71
72impl Default for IndexLabel {
73 fn default() -> Self {
74 Self(Arc::new(std::sync::RwLock::new(Arc::from("unknown"))))
75 }
76}
77
78impl IndexLabel {
79 pub fn get(&self) -> Arc<str> {
81 self.0.read().expect("IndexLabel lock poisoned").clone()
82 }
83
84 pub fn set(&self, label: &str) {
86 *self.0.write().expect("IndexLabel lock poisoned") = Arc::from(label);
87 }
88}
89
90impl std::fmt::Debug for FileHandle {
91 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
92 match &self.inner {
93 FileHandleInner::Inline { len, offset, .. } => f
94 .debug_struct("FileHandle::Inline")
95 .field("offset", offset)
96 .field("len", len)
97 .finish(),
98 FileHandleInner::Lazy { len, offset, .. } => f
99 .debug_struct("FileHandle::Lazy")
100 .field("offset", offset)
101 .field("len", len)
102 .finish(),
103 }
104 }
105}
106
107impl FileHandle {
108 pub fn from_bytes(data: OwnedBytes) -> Self {
111 let len = data.len() as u64;
112 Self {
113 inner: FileHandleInner::Inline {
114 data,
115 offset: 0,
116 len,
117 },
118 }
119 }
120
121 pub fn empty() -> Self {
123 Self::from_bytes(OwnedBytes::empty())
124 }
125
126 pub fn lazy(len: u64, read_fn: RangeReadFn) -> Self {
131 Self::lazy_labeled(len, read_fn, Arc::from("unknown"))
132 }
133
134 pub fn lazy_labeled(len: u64, read_fn: RangeReadFn, label: Arc<str>) -> Self {
136 Self {
137 inner: FileHandleInner::Lazy {
138 read_fn,
139 offset: 0,
140 len,
141 label,
142 },
143 }
144 }
145
146 #[inline]
148 pub fn len(&self) -> u64 {
149 match &self.inner {
150 FileHandleInner::Inline { len, .. } => *len,
151 FileHandleInner::Lazy { len, .. } => *len,
152 }
153 }
154
155 #[inline]
157 pub fn is_empty(&self) -> bool {
158 self.len() == 0
159 }
160
161 #[inline]
163 pub fn is_sync(&self) -> bool {
164 matches!(&self.inner, FileHandleInner::Inline { .. })
165 }
166
167 pub fn slice(&self, range: Range<u64>) -> Self {
169 match &self.inner {
170 FileHandleInner::Inline { data, offset, len } => {
171 let new_offset = offset + range.start;
172 let new_len = range.end - range.start;
173 debug_assert!(
174 new_offset + new_len <= offset + len,
175 "slice out of bounds: {}+{} > {}+{}",
176 new_offset,
177 new_len,
178 offset,
179 len
180 );
181 Self {
182 inner: FileHandleInner::Inline {
183 data: data.clone(),
184 offset: new_offset,
185 len: new_len,
186 },
187 }
188 }
189 FileHandleInner::Lazy {
190 read_fn,
191 offset,
192 len,
193 label,
194 } => {
195 let new_offset = offset + range.start;
196 let new_len = range.end - range.start;
197 debug_assert!(
198 new_offset + new_len <= offset + len,
199 "slice out of bounds: {}+{} > {}+{}",
200 new_offset,
201 new_len,
202 offset,
203 len
204 );
205 Self {
206 inner: FileHandleInner::Lazy {
207 read_fn: Arc::clone(read_fn),
208 offset: new_offset,
209 len: new_len,
210 label: Arc::clone(label),
211 },
212 }
213 }
214 }
215 }
216
217 #[cfg(feature = "native")]
222 pub fn madvise_range(&self, range: Range<u64>, advice: libc::c_int) {
223 if let FileHandleInner::Inline { data, offset, len } = &self.inner {
224 let end = range.end.min(*len);
225 if range.start >= end {
226 return;
227 }
228 let start = (*offset + range.start) as usize;
229 let end = (*offset + end) as usize;
230 data.madvise_range(start..end, advice);
231 }
232 }
233
234 pub async fn read_bytes_range(&self, range: Range<u64>) -> io::Result<OwnedBytes> {
236 match &self.inner {
237 FileHandleInner::Inline { data, offset, len } => {
238 if range.end > *len {
239 return Err(io::Error::new(
240 io::ErrorKind::InvalidInput,
241 format!("Range {:?} out of bounds (len: {})", range, len),
242 ));
243 }
244 let start = (*offset + range.start) as usize;
245 let end = (*offset + range.end) as usize;
246 Ok(data.slice(start..end))
247 }
248 FileHandleInner::Lazy {
249 read_fn,
250 offset,
251 len,
252 label,
253 } => {
254 if range.end > *len {
255 return Err(io::Error::new(
256 io::ErrorKind::InvalidInput,
257 format!("Range {:?} out of bounds (len: {})", range, len),
258 ));
259 }
260 let abs_start = offset + range.start;
261 let abs_end = offset + range.end;
262 let t = crate::observe::Timer::start();
266 let result = (read_fn)(abs_start..abs_end).await;
267 if let Ok(bytes) = &result {
268 crate::observe::directory_read(label, "lazy_range", t.secs(), bytes.len());
269 }
270 result
271 }
272 }
273 }
274
275 pub async fn read_bytes(&self) -> io::Result<OwnedBytes> {
277 self.read_bytes_range(0..self.len()).await
278 }
279
280 #[inline]
283 pub fn read_bytes_range_sync(&self, range: Range<u64>) -> io::Result<OwnedBytes> {
284 match &self.inner {
285 FileHandleInner::Inline { data, offset, len } => {
286 if range.end > *len {
287 return Err(io::Error::new(
288 io::ErrorKind::InvalidInput,
289 format!("Range {:?} out of bounds (len: {})", range, len),
290 ));
291 }
292 let start = (*offset + range.start) as usize;
293 let end = (*offset + range.end) as usize;
294 Ok(data.slice(start..end))
295 }
296 FileHandleInner::Lazy { .. } => Err(io::Error::new(
297 io::ErrorKind::Unsupported,
298 "Synchronous read not available on lazy file handle",
299 )),
300 }
301 }
302
303 #[inline]
305 pub fn read_bytes_sync(&self) -> io::Result<OwnedBytes> {
306 self.read_bytes_range_sync(0..self.len())
307 }
308}
309
310#[derive(Clone)]
312enum SharedBytes {
313 Vec(Arc<Vec<u8>>),
314 #[cfg(feature = "native")]
315 Mmap(Arc<memmap2::Mmap>),
316}
317
318impl SharedBytes {
319 #[inline]
320 fn as_bytes(&self) -> &[u8] {
321 match self {
322 SharedBytes::Vec(v) => v.as_slice(),
323 #[cfg(feature = "native")]
324 SharedBytes::Mmap(m) => m.as_ref(),
325 }
326 }
327}
328
329impl std::fmt::Debug for SharedBytes {
330 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
331 match self {
332 SharedBytes::Vec(v) => write!(f, "Vec(len={})", v.len()),
333 #[cfg(feature = "native")]
334 SharedBytes::Mmap(m) => write!(f, "Mmap(len={})", m.len()),
335 }
336 }
337}
338
339#[derive(Debug, Clone)]
345pub struct OwnedBytes {
346 data: SharedBytes,
347 range: Range<usize>,
348}
349
350impl OwnedBytes {
351 pub fn new(data: Vec<u8>) -> Self {
352 let len = data.len();
353 Self {
354 data: SharedBytes::Vec(Arc::new(data)),
355 range: 0..len,
356 }
357 }
358
359 pub fn empty() -> Self {
360 Self {
361 data: SharedBytes::Vec(Arc::new(Vec::new())),
362 range: 0..0,
363 }
364 }
365
366 pub(crate) fn from_arc_vec(data: Arc<Vec<u8>>, range: Range<usize>) -> Self {
369 Self {
370 data: SharedBytes::Vec(data),
371 range,
372 }
373 }
374
375 #[cfg(feature = "native")]
377 pub(crate) fn from_mmap(mmap: Arc<memmap2::Mmap>) -> Self {
378 let len = mmap.len();
379 Self {
380 data: SharedBytes::Mmap(mmap),
381 range: 0..len,
382 }
383 }
384
385 #[cfg(feature = "native")]
387 pub(crate) fn from_mmap_range(mmap: Arc<memmap2::Mmap>, range: Range<usize>) -> Self {
388 Self {
389 data: SharedBytes::Mmap(mmap),
390 range,
391 }
392 }
393
394 pub fn len(&self) -> usize {
395 self.range.len()
396 }
397
398 pub fn is_empty(&self) -> bool {
399 self.range.is_empty()
400 }
401
402 pub fn slice(&self, range: Range<usize>) -> Self {
403 let start = self.range.start + range.start;
404 let end = self.range.start + range.end;
405 Self {
406 data: self.data.clone(),
407 range: start..end,
408 }
409 }
410
411 pub fn as_slice(&self) -> &[u8] {
412 &self.data.as_bytes()[self.range.clone()]
413 }
414
415 #[cfg(feature = "native")]
420 #[inline]
421 pub fn is_mmap(&self) -> bool {
422 matches!(self.data, SharedBytes::Mmap(_))
423 }
424
425 #[cfg(feature = "native")]
431 pub fn madvise(&self, advice: libc::c_int) {
432 self.madvise_range(0..self.len(), advice);
433 }
434
435 #[cfg(feature = "native")]
440 pub fn mlock(&self) -> bool {
441 if !self.is_mmap() {
442 return false;
443 }
444 let slice = self.as_slice();
445 if slice.is_empty() {
446 return true;
447 }
448 let ptr = slice.as_ptr();
449 let len = slice.len();
450 let page_size = 4096usize;
451 let aligned_ptr = (ptr as usize) & !(page_size - 1);
452 let aligned_len = len + (ptr as usize - aligned_ptr);
453 unsafe { libc::mlock(aligned_ptr as *const libc::c_void, aligned_len) == 0 }
454 }
455
456 #[cfg(feature = "native")]
462 pub fn madvise_range(&self, range: Range<usize>, advice: libc::c_int) {
463 if !self.is_mmap() {
464 return;
465 }
466 let slice = &self.as_slice()[range];
467 if slice.is_empty() {
468 return;
469 }
470 let ptr = slice.as_ptr();
471 let len = slice.len();
472 let page_size = 4096usize;
473 let aligned_ptr = (ptr as usize) & !(page_size - 1);
474 let aligned_len = len + (ptr as usize - aligned_ptr);
475 unsafe {
476 libc::madvise(aligned_ptr as *mut libc::c_void, aligned_len, advice);
477 }
478 }
479
480 pub fn to_vec(&self) -> Vec<u8> {
481 self.as_slice().to_vec()
482 }
483}
484
485impl AsRef<[u8]> for OwnedBytes {
486 fn as_ref(&self) -> &[u8] {
487 self.as_slice()
488 }
489}
490
491impl std::ops::Deref for OwnedBytes {
492 type Target = [u8];
493
494 fn deref(&self) -> &Self::Target {
495 self.as_slice()
496 }
497}
498
499#[cfg(not(target_arch = "wasm32"))]
501#[async_trait]
502pub trait Directory: Send + Sync + 'static {
503 async fn exists(&self, path: &Path) -> io::Result<bool>;
505
506 async fn file_size(&self, path: &Path) -> io::Result<u64>;
508
509 async fn open_read(&self, path: &Path) -> io::Result<FileHandle>;
511
512 async fn read_range(&self, path: &Path, range: Range<u64>) -> io::Result<OwnedBytes>;
514
515 async fn list_files(&self, prefix: &Path) -> io::Result<Vec<PathBuf>>;
517
518 async fn open_lazy(&self, path: &Path) -> io::Result<FileHandle>;
522
523 fn set_index_label(&self, _label: &str) {}
529}
530
531#[cfg(target_arch = "wasm32")]
533#[async_trait(?Send)]
534pub trait Directory: 'static {
535 async fn exists(&self, path: &Path) -> io::Result<bool>;
537
538 async fn file_size(&self, path: &Path) -> io::Result<u64>;
540
541 async fn open_read(&self, path: &Path) -> io::Result<FileHandle>;
543
544 async fn read_range(&self, path: &Path, range: Range<u64>) -> io::Result<OwnedBytes>;
546
547 async fn list_files(&self, prefix: &Path) -> io::Result<Vec<PathBuf>>;
549
550 async fn open_lazy(&self, path: &Path) -> io::Result<FileHandle>;
552
553 fn set_index_label(&self, _label: &str) {}
556}
557
558pub trait StreamingWriter: io::Write + Send {
563 fn finish(self: Box<Self>) -> io::Result<()>;
565
566 fn bytes_written(&self) -> u64;
568}
569
570struct BufferedStreamingWriter {
573 path: PathBuf,
574 buffer: Vec<u8>,
575 files: Arc<RwLock<HashMap<PathBuf, Arc<Vec<u8>>>>>,
578}
579
580impl io::Write for BufferedStreamingWriter {
581 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
582 self.buffer.extend_from_slice(buf);
583 Ok(buf.len())
584 }
585
586 fn flush(&mut self) -> io::Result<()> {
587 Ok(())
588 }
589}
590
591impl StreamingWriter for BufferedStreamingWriter {
592 fn finish(self: Box<Self>) -> io::Result<()> {
593 self.files.write().insert(self.path, Arc::new(self.buffer));
594 Ok(())
595 }
596
597 fn bytes_written(&self) -> u64 {
598 self.buffer.len() as u64
599 }
600}
601
602#[cfg(feature = "native")]
606const FILE_STREAMING_BUF_SIZE: usize = 8 * 1024 * 1024;
607
608#[cfg(feature = "native")]
610pub(crate) struct FileStreamingWriter {
611 pub(crate) file: io::BufWriter<std::fs::File>,
612 pub(crate) written: u64,
613}
614
615#[cfg(feature = "native")]
616impl FileStreamingWriter {
617 pub(crate) fn new(file: std::fs::File) -> Self {
618 Self {
619 file: io::BufWriter::with_capacity(FILE_STREAMING_BUF_SIZE, file),
620 written: 0,
621 }
622 }
623}
624
625#[cfg(feature = "native")]
626impl io::Write for FileStreamingWriter {
627 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
628 let n = self.file.write(buf)?;
629 self.written += n as u64;
630 Ok(n)
631 }
632
633 fn flush(&mut self) -> io::Result<()> {
634 self.file.flush()
635 }
636}
637
638#[cfg(feature = "native")]
639impl StreamingWriter for FileStreamingWriter {
640 fn finish(self: Box<Self>) -> io::Result<()> {
641 let file = self.file.into_inner().map_err(|e| e.into_error())?;
642 file.sync_all()?;
643 Ok(())
644 }
645
646 fn bytes_written(&self) -> u64 {
647 self.written
648 }
649}
650
651#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
653#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
654pub trait DirectoryWriter: Directory {
655 async fn write(&self, path: &Path, data: &[u8]) -> io::Result<()>;
657
658 async fn delete(&self, path: &Path) -> io::Result<()>;
660
661 async fn rename(&self, from: &Path, to: &Path) -> io::Result<()>;
663
664 async fn sync(&self) -> io::Result<()>;
666
667 async fn streaming_writer(&self, path: &Path) -> io::Result<Box<dyn StreamingWriter>>;
670
671 async fn streaming_writer_cold(&self, path: &Path) -> io::Result<Box<dyn StreamingWriter>> {
678 self.streaming_writer(path).await
679 }
680}
681
682#[derive(Debug, Default)]
684pub struct RamDirectory {
685 files: Arc<RwLock<HashMap<PathBuf, Arc<Vec<u8>>>>>,
686}
687
688impl Clone for RamDirectory {
689 fn clone(&self) -> Self {
690 Self {
691 files: Arc::clone(&self.files),
692 }
693 }
694}
695
696impl RamDirectory {
697 pub fn new() -> Self {
698 Self::default()
699 }
700
701 pub fn list_files_sync(&self, prefix: &Path) -> io::Result<Vec<PathBuf>> {
703 let files = self.files.read();
704 Ok(files
705 .keys()
706 .filter(|p| p.starts_with(prefix))
707 .cloned()
708 .collect())
709 }
710
711 pub fn read_file_sync(&self, path: &Path) -> io::Result<Vec<u8>> {
713 let files = self.files.read();
714 files
715 .get(path)
716 .map(|data| data.as_ref().clone())
717 .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "File not found"))
718 }
719
720 pub fn write_sync(&self, path: &Path, data: &[u8]) -> io::Result<()> {
722 self.files
723 .write()
724 .insert(path.to_path_buf(), Arc::new(data.to_vec()));
725 Ok(())
726 }
727}
728
729#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
730#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
731impl Directory for RamDirectory {
732 async fn exists(&self, path: &Path) -> io::Result<bool> {
733 Ok(self.files.read().contains_key(path))
734 }
735
736 async fn file_size(&self, path: &Path) -> io::Result<u64> {
737 self.files
738 .read()
739 .get(path)
740 .map(|data| data.len() as u64)
741 .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "File not found"))
742 }
743
744 async fn open_read(&self, path: &Path) -> io::Result<FileHandle> {
745 let files = self.files.read();
746 let data = files
747 .get(path)
748 .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "File not found"))?;
749
750 Ok(FileHandle::from_bytes(OwnedBytes::from_arc_vec(
751 Arc::clone(data),
752 0..data.len(),
753 )))
754 }
755
756 async fn read_range(&self, path: &Path, range: Range<u64>) -> io::Result<OwnedBytes> {
757 let files = self.files.read();
758 let data = files
759 .get(path)
760 .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "File not found"))?;
761
762 let start = range.start as usize;
763 let end = range.end as usize;
764
765 if end > data.len() {
766 return Err(io::Error::new(
767 io::ErrorKind::InvalidInput,
768 "Range out of bounds",
769 ));
770 }
771
772 Ok(OwnedBytes::from_arc_vec(Arc::clone(data), start..end))
773 }
774
775 async fn list_files(&self, prefix: &Path) -> io::Result<Vec<PathBuf>> {
776 let files = self.files.read();
777 Ok(files
778 .keys()
779 .filter(|p| p.starts_with(prefix))
780 .cloned()
781 .collect())
782 }
783
784 async fn open_lazy(&self, path: &Path) -> io::Result<FileHandle> {
785 self.open_read(path).await
787 }
788}
789
790#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
791#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
792impl DirectoryWriter for RamDirectory {
793 async fn write(&self, path: &Path, data: &[u8]) -> io::Result<()> {
794 self.files
795 .write()
796 .insert(path.to_path_buf(), Arc::new(data.to_vec()));
797 Ok(())
798 }
799
800 async fn delete(&self, path: &Path) -> io::Result<()> {
801 self.files.write().remove(path);
802 Ok(())
803 }
804
805 async fn rename(&self, from: &Path, to: &Path) -> io::Result<()> {
806 let mut files = self.files.write();
807 if let Some(data) = files.remove(from) {
808 files.insert(to.to_path_buf(), data);
809 }
810 Ok(())
811 }
812
813 async fn sync(&self) -> io::Result<()> {
814 Ok(())
815 }
816
817 async fn streaming_writer(&self, path: &Path) -> io::Result<Box<dyn StreamingWriter>> {
818 Ok(Box::new(BufferedStreamingWriter {
819 path: path.to_path_buf(),
820 buffer: Vec::new(),
821 files: Arc::clone(&self.files),
822 }))
823 }
824}
825
826#[cfg(feature = "native")]
828#[derive(Debug, Clone)]
829pub struct FsDirectory {
830 root: PathBuf,
831 label: IndexLabel,
832}
833
834#[cfg(feature = "native")]
835impl FsDirectory {
836 pub fn new(root: impl AsRef<Path>) -> Self {
837 Self {
838 root: root.as_ref().to_path_buf(),
839 label: IndexLabel::default(),
840 }
841 }
842
843 fn resolve(&self, path: &Path) -> PathBuf {
844 self.root.join(path)
845 }
846}
847
848#[cfg(feature = "native")]
849#[async_trait]
850impl Directory for FsDirectory {
851 async fn exists(&self, path: &Path) -> io::Result<bool> {
852 let full_path = self.resolve(path);
853 Ok(tokio::fs::try_exists(&full_path).await.unwrap_or(false))
854 }
855
856 async fn file_size(&self, path: &Path) -> io::Result<u64> {
857 let full_path = self.resolve(path);
858 let metadata = tokio::fs::metadata(&full_path).await?;
859 Ok(metadata.len())
860 }
861
862 async fn open_read(&self, path: &Path) -> io::Result<FileHandle> {
863 let full_path = self.resolve(path);
864 let data = tokio::fs::read(&full_path).await?;
865 Ok(FileHandle::from_bytes(OwnedBytes::new(data)))
866 }
867
868 async fn read_range(&self, path: &Path, range: Range<u64>) -> io::Result<OwnedBytes> {
869 use tokio::io::{AsyncReadExt, AsyncSeekExt};
870
871 let full_path = self.resolve(path);
872 let mut file = tokio::fs::File::open(&full_path).await?;
873
874 file.seek(std::io::SeekFrom::Start(range.start)).await?;
875
876 let len = (range.end - range.start) as usize;
877 let mut buffer = vec![0u8; len];
878 file.read_exact(&mut buffer).await?;
879
880 Ok(OwnedBytes::new(buffer))
881 }
882
883 async fn list_files(&self, prefix: &Path) -> io::Result<Vec<PathBuf>> {
884 let full_path = self.resolve(prefix);
885 let mut entries = tokio::fs::read_dir(&full_path).await?;
886 let mut files = Vec::new();
887
888 while let Some(entry) = entries.next_entry().await? {
889 if entry.file_type().await?.is_file() {
890 files.push(entry.path().strip_prefix(&self.root).unwrap().to_path_buf());
891 }
892 }
893
894 Ok(files)
895 }
896
897 async fn open_lazy(&self, path: &Path) -> io::Result<FileHandle> {
898 let full_path = self.resolve(path);
899 let metadata = tokio::fs::metadata(&full_path).await?;
900 let file_size = metadata.len();
901
902 let read_fn: RangeReadFn = Arc::new(move |range: Range<u64>| {
903 let full_path = full_path.clone();
904 Box::pin(async move {
905 use tokio::io::{AsyncReadExt, AsyncSeekExt};
906
907 let mut file = tokio::fs::File::open(&full_path).await?;
908 file.seek(std::io::SeekFrom::Start(range.start)).await?;
909
910 let len = (range.end - range.start) as usize;
911 let mut buffer = vec![0u8; len];
912 file.read_exact(&mut buffer).await?;
913
914 Ok(OwnedBytes::new(buffer))
915 })
916 });
917
918 Ok(FileHandle::lazy_labeled(
919 file_size,
920 read_fn,
921 self.label.get(),
922 ))
923 }
924
925 fn set_index_label(&self, label: &str) {
926 self.label.set(label);
927 }
928}
929
930#[cfg(feature = "native")]
931#[async_trait]
932impl DirectoryWriter for FsDirectory {
933 async fn write(&self, path: &Path, data: &[u8]) -> io::Result<()> {
934 let full_path = self.resolve(path);
935
936 if let Some(parent) = full_path.parent() {
938 tokio::fs::create_dir_all(parent).await?;
939 }
940
941 tokio::fs::write(&full_path, data).await
942 }
943
944 async fn delete(&self, path: &Path) -> io::Result<()> {
945 let full_path = self.resolve(path);
946 tokio::fs::remove_file(&full_path).await
947 }
948
949 async fn rename(&self, from: &Path, to: &Path) -> io::Result<()> {
950 let from_path = self.resolve(from);
951 let to_path = self.resolve(to);
952 tokio::fs::rename(&from_path, &to_path).await
953 }
954
955 async fn sync(&self) -> io::Result<()> {
956 let dir = std::fs::File::open(&self.root)?;
958 dir.sync_all()?;
959 Ok(())
960 }
961
962 async fn streaming_writer(&self, path: &Path) -> io::Result<Box<dyn StreamingWriter>> {
963 let full_path = self.resolve(path);
964 if let Some(parent) = full_path.parent() {
965 tokio::fs::create_dir_all(parent).await?;
966 }
967 let file = std::fs::File::create(&full_path)?;
968 Ok(Box::new(FileStreamingWriter::new(file)))
969 }
970
971 async fn streaming_writer_cold(&self, path: &Path) -> io::Result<Box<dyn StreamingWriter>> {
972 let full_path = self.resolve(path);
973 if let Some(parent) = full_path.parent() {
974 tokio::fs::create_dir_all(parent).await?;
975 }
976 let file = std::fs::File::create(&full_path)?;
977 Ok(Box::new(super::ColdStreamingWriter::new(
978 file,
979 self.label.get(),
980 )))
981 }
982}
983
984pub struct CachingDirectory<D: Directory> {
986 inner: D,
987 cache: RwLock<HashMap<PathBuf, Arc<Vec<u8>>>>,
988 max_cached_bytes: usize,
989 current_bytes: RwLock<usize>,
990}
991
992impl<D: Directory> CachingDirectory<D> {
993 pub fn new(inner: D, max_cached_bytes: usize) -> Self {
994 Self {
995 inner,
996 cache: RwLock::new(HashMap::new()),
997 max_cached_bytes,
998 current_bytes: RwLock::new(0),
999 }
1000 }
1001
1002 fn try_cache(&self, path: &Path, data: &[u8]) {
1003 let mut current = self.current_bytes.write();
1004 if *current + data.len() <= self.max_cached_bytes {
1005 self.cache
1006 .write()
1007 .insert(path.to_path_buf(), Arc::new(data.to_vec()));
1008 *current += data.len();
1009 }
1010 }
1011}
1012
1013#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
1014#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
1015impl<D: Directory> Directory for CachingDirectory<D> {
1016 async fn exists(&self, path: &Path) -> io::Result<bool> {
1017 if self.cache.read().contains_key(path) {
1018 return Ok(true);
1019 }
1020 self.inner.exists(path).await
1021 }
1022
1023 async fn file_size(&self, path: &Path) -> io::Result<u64> {
1024 if let Some(data) = self.cache.read().get(path) {
1025 return Ok(data.len() as u64);
1026 }
1027 self.inner.file_size(path).await
1028 }
1029
1030 async fn open_read(&self, path: &Path) -> io::Result<FileHandle> {
1031 if let Some(data) = self.cache.read().get(path) {
1033 return Ok(FileHandle::from_bytes(OwnedBytes::from_arc_vec(
1034 Arc::clone(data),
1035 0..data.len(),
1036 )));
1037 }
1038
1039 let handle = self.inner.open_read(path).await?;
1041 let bytes = handle.read_bytes().await?;
1042
1043 self.try_cache(path, bytes.as_slice());
1044
1045 Ok(FileHandle::from_bytes(bytes))
1046 }
1047
1048 async fn read_range(&self, path: &Path, range: Range<u64>) -> io::Result<OwnedBytes> {
1049 if let Some(data) = self.cache.read().get(path) {
1051 let start = range.start as usize;
1052 let end = range.end as usize;
1053 return Ok(OwnedBytes::from_arc_vec(Arc::clone(data), start..end));
1054 }
1055
1056 self.inner.read_range(path, range).await
1057 }
1058
1059 async fn list_files(&self, prefix: &Path) -> io::Result<Vec<PathBuf>> {
1060 self.inner.list_files(prefix).await
1061 }
1062
1063 async fn open_lazy(&self, path: &Path) -> io::Result<FileHandle> {
1064 self.inner.open_lazy(path).await
1066 }
1067
1068 fn set_index_label(&self, label: &str) {
1069 self.inner.set_index_label(label);
1070 }
1071}
1072
1073#[cfg(test)]
1074mod tests {
1075 use super::*;
1076
1077 #[tokio::test]
1078 async fn test_ram_directory() {
1079 let dir = RamDirectory::new();
1080
1081 dir.write(Path::new("test.txt"), b"hello world")
1083 .await
1084 .unwrap();
1085
1086 assert!(dir.exists(Path::new("test.txt")).await.unwrap());
1088 assert!(!dir.exists(Path::new("nonexistent.txt")).await.unwrap());
1089
1090 let slice = dir.open_read(Path::new("test.txt")).await.unwrap();
1092 let data = slice.read_bytes().await.unwrap();
1093 assert_eq!(data.as_slice(), b"hello world");
1094
1095 let range_data = dir.read_range(Path::new("test.txt"), 0..5).await.unwrap();
1097 assert_eq!(range_data.as_slice(), b"hello");
1098
1099 dir.delete(Path::new("test.txt")).await.unwrap();
1101 assert!(!dir.exists(Path::new("test.txt")).await.unwrap());
1102 }
1103
1104 #[tokio::test]
1105 async fn test_file_handle() {
1106 let data = OwnedBytes::new(b"hello world".to_vec());
1107 let handle = FileHandle::from_bytes(data);
1108
1109 assert_eq!(handle.len(), 11);
1110 assert!(handle.is_sync());
1111
1112 let sub = handle.slice(0..5);
1113 let bytes = sub.read_bytes().await.unwrap();
1114 assert_eq!(bytes.as_slice(), b"hello");
1115
1116 let sub2 = handle.slice(6..11);
1117 let bytes2 = sub2.read_bytes().await.unwrap();
1118 assert_eq!(bytes2.as_slice(), b"world");
1119
1120 let sync_bytes = handle.read_bytes_range_sync(0..5).unwrap();
1122 assert_eq!(sync_bytes.as_slice(), b"hello");
1123 }
1124
1125 #[tokio::test]
1126 async fn test_owned_bytes() {
1127 let bytes = OwnedBytes::new(vec![1, 2, 3, 4, 5]);
1128
1129 assert_eq!(bytes.len(), 5);
1130 assert_eq!(bytes.as_slice(), &[1, 2, 3, 4, 5]);
1131
1132 let sliced = bytes.slice(1..4);
1133 assert_eq!(sliced.as_slice(), &[2, 3, 4]);
1134
1135 assert_eq!(bytes.as_slice(), &[1, 2, 3, 4, 5]);
1137 }
1138}