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(crate) fn with_local_owner(&self) -> Self {
111 match &self.inner {
112 FileHandleInner::Inline { data, offset, len } => Self {
113 inner: FileHandleInner::Inline {
114 data: data.clone().with_local_owner(),
115 offset: *offset,
116 len: *len,
117 },
118 },
119 FileHandleInner::Lazy { .. } => self.clone(),
120 }
121 }
122
123 pub fn from_bytes(data: OwnedBytes) -> Self {
126 let len = data.len() as u64;
127 Self {
128 inner: FileHandleInner::Inline {
129 data,
130 offset: 0,
131 len,
132 },
133 }
134 }
135
136 pub fn empty() -> Self {
138 Self::from_bytes(OwnedBytes::empty())
139 }
140
141 pub fn lazy(len: u64, read_fn: RangeReadFn) -> Self {
146 Self::lazy_labeled(len, read_fn, Arc::from("unknown"))
147 }
148
149 pub fn lazy_labeled(len: u64, read_fn: RangeReadFn, label: Arc<str>) -> Self {
151 Self {
152 inner: FileHandleInner::Lazy {
153 read_fn,
154 offset: 0,
155 len,
156 label,
157 },
158 }
159 }
160
161 #[inline]
163 pub fn len(&self) -> u64 {
164 match &self.inner {
165 FileHandleInner::Inline { len, .. } => *len,
166 FileHandleInner::Lazy { len, .. } => *len,
167 }
168 }
169
170 #[inline]
172 pub fn is_empty(&self) -> bool {
173 self.len() == 0
174 }
175
176 #[inline]
178 pub fn is_sync(&self) -> bool {
179 matches!(&self.inner, FileHandleInner::Inline { .. })
180 }
181
182 pub fn slice(&self, range: Range<u64>) -> Self {
184 match &self.inner {
185 FileHandleInner::Inline { data, offset, len } => {
186 let new_offset = offset + range.start;
187 let new_len = range.end - range.start;
188 debug_assert!(
189 new_offset + new_len <= offset + len,
190 "slice out of bounds: {}+{} > {}+{}",
191 new_offset,
192 new_len,
193 offset,
194 len
195 );
196 Self {
197 inner: FileHandleInner::Inline {
198 data: data.clone(),
199 offset: new_offset,
200 len: new_len,
201 },
202 }
203 }
204 FileHandleInner::Lazy {
205 read_fn,
206 offset,
207 len,
208 label,
209 } => {
210 let new_offset = offset + range.start;
211 let new_len = range.end - range.start;
212 debug_assert!(
213 new_offset + new_len <= offset + len,
214 "slice out of bounds: {}+{} > {}+{}",
215 new_offset,
216 new_len,
217 offset,
218 len
219 );
220 Self {
221 inner: FileHandleInner::Lazy {
222 read_fn: Arc::clone(read_fn),
223 offset: new_offset,
224 len: new_len,
225 label: Arc::clone(label),
226 },
227 }
228 }
229 }
230 }
231
232 #[cfg(feature = "native")]
237 pub fn madvise_range(&self, range: Range<u64>, advice: libc::c_int) {
238 if let FileHandleInner::Inline { data, offset, len } = &self.inner {
239 let end = range.end.min(*len);
240 if range.start >= end {
241 return;
242 }
243 let start = (*offset + range.start) as usize;
244 let end = (*offset + end) as usize;
245 data.madvise_range(start..end, advice);
246 }
247 }
248
249 pub async fn read_bytes_range(&self, range: Range<u64>) -> io::Result<OwnedBytes> {
251 match &self.inner {
252 FileHandleInner::Inline { data, offset, len } => {
253 if range.end > *len {
254 return Err(io::Error::new(
255 io::ErrorKind::InvalidInput,
256 format!("Range {:?} out of bounds (len: {})", range, len),
257 ));
258 }
259 let start = (*offset + range.start) as usize;
260 let end = (*offset + range.end) as usize;
261 Ok(data.slice(start..end))
262 }
263 FileHandleInner::Lazy {
264 read_fn,
265 offset,
266 len,
267 label,
268 } => {
269 if range.end > *len {
270 return Err(io::Error::new(
271 io::ErrorKind::InvalidInput,
272 format!("Range {:?} out of bounds (len: {})", range, len),
273 ));
274 }
275 let abs_start = offset + range.start;
276 let abs_end = offset + range.end;
277 let t = crate::observe::Timer::start();
281 let result = (read_fn)(abs_start..abs_end).await;
282 if let Ok(bytes) = &result {
283 crate::observe::directory_read(label, "lazy_range", t.secs(), bytes.len());
284 }
285 result
286 }
287 }
288 }
289
290 pub async fn read_bytes(&self) -> io::Result<OwnedBytes> {
292 self.read_bytes_range(0..self.len()).await
293 }
294
295 #[inline]
298 pub fn read_bytes_range_sync(&self, range: Range<u64>) -> io::Result<OwnedBytes> {
299 match &self.inner {
300 FileHandleInner::Inline { data, offset, len } => {
301 if range.end > *len {
302 return Err(io::Error::new(
303 io::ErrorKind::InvalidInput,
304 format!("Range {:?} out of bounds (len: {})", range, len),
305 ));
306 }
307 let start = (*offset + range.start) as usize;
308 let end = (*offset + range.end) as usize;
309 Ok(data.slice(start..end))
310 }
311 FileHandleInner::Lazy { .. } => Err(io::Error::new(
312 io::ErrorKind::Unsupported,
313 "Synchronous read not available on lazy file handle",
314 )),
315 }
316 }
317
318 #[inline]
320 pub fn read_bytes_sync(&self) -> io::Result<OwnedBytes> {
321 self.read_bytes_range_sync(0..self.len())
322 }
323}
324
325#[derive(Clone)]
327enum SharedBytes {
328 Vec(Arc<Vec<u8>>),
329 #[cfg(feature = "native")]
330 Mmap(Arc<memmap2::Mmap>),
331 Local(Arc<SharedBytes>),
332}
333
334impl SharedBytes {
335 #[inline]
336 fn as_bytes(&self) -> &[u8] {
337 match self {
338 SharedBytes::Vec(v) => v.as_slice(),
339 #[cfg(feature = "native")]
340 SharedBytes::Mmap(m) => m.as_ref(),
341 SharedBytes::Local(owner) => owner.as_bytes(),
342 }
343 }
344}
345
346impl std::fmt::Debug for SharedBytes {
347 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
348 match self {
349 SharedBytes::Vec(v) => write!(f, "Vec(len={})", v.len()),
350 #[cfg(feature = "native")]
351 SharedBytes::Mmap(m) => write!(f, "Mmap(len={})", m.len()),
352 SharedBytes::Local(owner) => owner.fmt(f),
353 }
354 }
355}
356
357#[derive(Clone)]
363pub struct OwnedBytes {
364 data: SharedBytes,
365 view: std::ptr::NonNull<[u8]>,
368}
369
370unsafe impl Send for OwnedBytes {}
374unsafe impl Sync for OwnedBytes {}
376
377impl std::fmt::Debug for OwnedBytes {
378 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
379 f.debug_struct("OwnedBytes")
380 .field("data", &self.data)
381 .field("len", &self.len())
382 .finish()
383 }
384}
385
386impl OwnedBytes {
387 fn with_local_owner(mut self) -> Self {
388 if !matches!(self.data, SharedBytes::Local(_)) {
389 self.data = SharedBytes::Local(Arc::new(self.data));
390 }
391 self
392 }
393
394 fn with_range(data: SharedBytes, range: Range<usize>) -> Self {
397 let view = std::ptr::NonNull::from(&data.as_bytes()[range]);
398 Self { data, view }
399 }
400
401 pub fn new(data: Vec<u8>) -> Self {
402 let len = data.len();
403 Self::with_range(SharedBytes::Vec(Arc::new(data)), 0..len)
404 }
405
406 pub fn empty() -> Self {
407 Self::new(Vec::new())
408 }
409
410 pub(crate) fn from_arc_vec(data: Arc<Vec<u8>>, range: Range<usize>) -> Self {
413 Self::with_range(SharedBytes::Vec(data), range)
414 }
415
416 #[cfg(feature = "native")]
418 pub(crate) fn from_mmap(mmap: Arc<memmap2::Mmap>) -> Self {
419 let len = mmap.len();
420 Self::with_range(SharedBytes::Mmap(mmap), 0..len)
421 }
422
423 #[cfg(feature = "native")]
425 pub(crate) fn from_mmap_range(mmap: Arc<memmap2::Mmap>, range: Range<usize>) -> Self {
426 Self::with_range(SharedBytes::Mmap(mmap), range)
427 }
428
429 #[inline]
430 pub fn len(&self) -> usize {
431 self.view.len()
432 }
433
434 #[inline]
435 pub fn is_empty(&self) -> bool {
436 self.len() == 0
437 }
438
439 pub fn slice(&self, range: Range<usize>) -> Self {
441 let view = std::ptr::NonNull::from(&self.as_slice()[range]);
442 Self {
443 data: self.data.clone(),
444 view,
445 }
446 }
447
448 #[inline]
449 pub fn as_slice(&self) -> &[u8] {
450 unsafe { self.view.as_ref() }
454 }
455
456 #[cfg(feature = "native")]
461 #[inline]
462 pub fn is_mmap(&self) -> bool {
463 match &self.data {
464 SharedBytes::Mmap(_) => true,
465 SharedBytes::Local(owner) => matches!(owner.as_ref(), SharedBytes::Mmap(_)),
466 SharedBytes::Vec(_) => false,
467 }
468 }
469
470 #[cfg(feature = "native")]
476 pub fn madvise(&self, advice: libc::c_int) {
477 self.madvise_range(0..self.len(), advice);
478 }
479
480 #[cfg(feature = "native")]
485 pub fn mlock(&self) -> bool {
486 if !self.is_mmap() {
487 return false;
488 }
489 let slice = self.as_slice();
490 if slice.is_empty() {
491 return true;
492 }
493 let ptr = slice.as_ptr();
494 let len = slice.len();
495 let page_size = 4096usize;
496 let aligned_ptr = (ptr as usize) & !(page_size - 1);
497 let aligned_len = len + (ptr as usize - aligned_ptr);
498 unsafe { libc::mlock(aligned_ptr as *const libc::c_void, aligned_len) == 0 }
499 }
500
501 #[cfg(feature = "native")]
507 pub fn madvise_range(&self, range: Range<usize>, advice: libc::c_int) {
508 if !self.is_mmap() {
509 return;
510 }
511 let slice = &self.as_slice()[range];
512 if slice.is_empty() {
513 return;
514 }
515 let ptr = slice.as_ptr();
516 let len = slice.len();
517 let page_size = 4096usize;
518 let aligned_ptr = (ptr as usize) & !(page_size - 1);
519 let aligned_len = len + (ptr as usize - aligned_ptr);
520 unsafe {
521 libc::madvise(aligned_ptr as *mut libc::c_void, aligned_len, advice);
522 }
523 }
524
525 pub fn to_vec(&self) -> Vec<u8> {
526 self.as_slice().to_vec()
527 }
528}
529
530impl AsRef<[u8]> for OwnedBytes {
531 fn as_ref(&self) -> &[u8] {
532 self.as_slice()
533 }
534}
535
536impl std::ops::Deref for OwnedBytes {
537 type Target = [u8];
538
539 fn deref(&self) -> &Self::Target {
540 self.as_slice()
541 }
542}
543
544#[cfg(not(target_arch = "wasm32"))]
546#[async_trait]
547pub trait Directory: Send + Sync + 'static {
548 async fn exists(&self, path: &Path) -> io::Result<bool>;
550
551 async fn file_size(&self, path: &Path) -> io::Result<u64>;
553
554 async fn open_read(&self, path: &Path) -> io::Result<FileHandle>;
556
557 async fn read_range(&self, path: &Path, range: Range<u64>) -> io::Result<OwnedBytes>;
559
560 async fn list_files(&self, prefix: &Path) -> io::Result<Vec<PathBuf>>;
562
563 async fn open_lazy(&self, path: &Path) -> io::Result<FileHandle>;
567
568 fn set_index_label(&self, _label: &str) {}
574
575 fn local_path(&self, _path: &Path) -> Option<PathBuf> {
582 None
583 }
584}
585
586#[cfg(target_arch = "wasm32")]
588#[async_trait(?Send)]
589pub trait Directory: 'static {
590 async fn exists(&self, path: &Path) -> io::Result<bool>;
592
593 async fn file_size(&self, path: &Path) -> io::Result<u64>;
595
596 async fn open_read(&self, path: &Path) -> io::Result<FileHandle>;
598
599 async fn read_range(&self, path: &Path, range: Range<u64>) -> io::Result<OwnedBytes>;
601
602 async fn list_files(&self, prefix: &Path) -> io::Result<Vec<PathBuf>>;
604
605 async fn open_lazy(&self, path: &Path) -> io::Result<FileHandle>;
607
608 fn set_index_label(&self, _label: &str) {}
611
612 fn local_path(&self, _path: &Path) -> Option<PathBuf> {
614 None
615 }
616}
617
618pub trait StreamingWriter: io::Write + Send {
623 fn finish(self: Box<Self>) -> io::Result<()>;
625
626 fn bytes_written(&self) -> u64;
628
629 #[cfg(feature = "native")]
636 fn copy_from_file_range(
637 &mut self,
638 _source: &std::fs::File,
639 _source_offset: &mut u64,
640 _len: usize,
641 ) -> io::Result<usize> {
642 Err(io::Error::new(
643 io::ErrorKind::Unsupported,
644 "streaming writer does not support kernel-assisted range copies",
645 ))
646 }
647}
648
649struct BufferedStreamingWriter {
652 path: PathBuf,
653 buffer: Vec<u8>,
654 files: Arc<RwLock<HashMap<PathBuf, Arc<Vec<u8>>>>>,
657}
658
659impl io::Write for BufferedStreamingWriter {
660 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
661 self.buffer.extend_from_slice(buf);
662 Ok(buf.len())
663 }
664
665 fn flush(&mut self) -> io::Result<()> {
666 Ok(())
667 }
668}
669
670impl StreamingWriter for BufferedStreamingWriter {
671 fn finish(self: Box<Self>) -> io::Result<()> {
672 self.files.write().insert(self.path, Arc::new(self.buffer));
673 Ok(())
674 }
675
676 fn bytes_written(&self) -> u64 {
677 self.buffer.len() as u64
678 }
679}
680
681#[cfg(feature = "native")]
685const FILE_STREAMING_BUF_SIZE: usize = 8 * 1024 * 1024;
686
687#[cfg(feature = "native")]
689pub(crate) struct FileStreamingWriter {
690 pub(crate) file: io::BufWriter<std::fs::File>,
691 pub(crate) written: u64,
692}
693
694#[cfg(feature = "native")]
695impl FileStreamingWriter {
696 pub(crate) fn new(file: std::fs::File) -> Self {
697 Self {
698 file: io::BufWriter::with_capacity(FILE_STREAMING_BUF_SIZE, file),
699 written: 0,
700 }
701 }
702}
703
704#[cfg(feature = "native")]
705impl io::Write for FileStreamingWriter {
706 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
707 let n = self.file.write(buf)?;
708 self.written += n as u64;
709 Ok(n)
710 }
711
712 fn flush(&mut self) -> io::Result<()> {
713 self.file.flush()
714 }
715}
716
717#[cfg(feature = "native")]
718impl StreamingWriter for FileStreamingWriter {
719 fn finish(self: Box<Self>) -> io::Result<()> {
720 let file = self.file.into_inner().map_err(|e| e.into_error())?;
721 file.sync_all()?;
722 Ok(())
723 }
724
725 fn bytes_written(&self) -> u64 {
726 self.written
727 }
728
729 fn copy_from_file_range(
730 &mut self,
731 source: &std::fs::File,
732 source_offset: &mut u64,
733 len: usize,
734 ) -> io::Result<usize> {
735 io::Write::flush(&mut self.file)?;
736 let copied = copy_file_range_once(source, source_offset, self.file.get_ref(), len)?;
737 self.written = self
738 .written
739 .checked_add(copied as u64)
740 .ok_or_else(|| io::Error::other("streaming-writer byte count overflow"))?;
741 Ok(copied)
742 }
743}
744
745#[cfg(feature = "native")]
746pub(crate) fn copy_file_range_once(
747 source: &std::fs::File,
748 source_offset: &mut u64,
749 destination: &std::fs::File,
750 len: usize,
751) -> io::Result<usize> {
752 #[cfg(target_os = "linux")]
753 {
754 use std::os::fd::AsRawFd;
755
756 let mut offset = libc::loff_t::try_from(*source_offset).map_err(|_| {
757 io::Error::new(io::ErrorKind::InvalidInput, "source offset exceeds i64")
758 })?;
759 let copied = unsafe {
760 libc::copy_file_range(
761 source.as_raw_fd(),
762 &mut offset,
763 destination.as_raw_fd(),
764 std::ptr::null_mut(),
765 len,
766 0,
767 )
768 };
769 if copied < 0 {
770 let error = io::Error::last_os_error();
771 let unsupported = error.raw_os_error().is_some_and(|code| {
772 code == libc::ENOSYS
773 || code == libc::EXDEV
774 || code == libc::EOPNOTSUPP
775 || code == libc::EINVAL
776 });
777 return if unsupported {
778 Err(io::Error::new(io::ErrorKind::Unsupported, error))
779 } else {
780 Err(error)
781 };
782 }
783 *source_offset = u64::try_from(offset)
784 .map_err(|_| io::Error::other("copy_file_range returned a negative source offset"))?;
785 Ok(copied as usize)
786 }
787 #[cfg(not(target_os = "linux"))]
788 {
789 let _ = (source, source_offset, destination, len);
790 Err(io::Error::new(
791 io::ErrorKind::Unsupported,
792 "kernel-assisted range copies are only available on Linux",
793 ))
794 }
795}
796
797#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
799#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
800pub trait DirectoryWriter: Directory {
801 async fn write(&self, path: &Path, data: &[u8]) -> io::Result<()>;
803
804 async fn write_durable(&self, path: &Path, data: &[u8]) -> io::Result<()> {
813 use io::Write as _;
814 let mut writer = self.streaming_writer(path).await?;
815 writer.write_all(data)?;
816 writer.finish()
817 }
818
819 async fn delete(&self, path: &Path) -> io::Result<()>;
821
822 async fn rename(&self, from: &Path, to: &Path) -> io::Result<()>;
824
825 async fn link(&self, _from: &Path, _to: &Path) -> io::Result<()> {
831 Err(io::Error::new(
832 io::ErrorKind::Unsupported,
833 "directory backend does not support immutable file links",
834 ))
835 }
836
837 async fn sync(&self) -> io::Result<()>;
839
840 async fn streaming_writer(&self, path: &Path) -> io::Result<Box<dyn StreamingWriter>>;
843
844 async fn streaming_writer_cold(&self, path: &Path) -> io::Result<Box<dyn StreamingWriter>> {
851 self.streaming_writer(path).await
852 }
853
854 async fn streaming_writer_cold_with_capacity(
859 &self,
860 path: &Path,
861 _buffer_capacity: usize,
862 ) -> io::Result<Box<dyn StreamingWriter>> {
863 self.streaming_writer_cold(path).await
864 }
865}
866
867#[derive(Debug, Default)]
869pub struct RamDirectory {
870 files: Arc<RwLock<HashMap<PathBuf, Arc<Vec<u8>>>>>,
871}
872
873impl Clone for RamDirectory {
874 fn clone(&self) -> Self {
875 Self {
876 files: Arc::clone(&self.files),
877 }
878 }
879}
880
881impl RamDirectory {
882 pub fn new() -> Self {
883 Self::default()
884 }
885
886 pub fn list_files_sync(&self, prefix: &Path) -> io::Result<Vec<PathBuf>> {
888 let files = self.files.read();
889 Ok(files
890 .keys()
891 .filter(|p| p.starts_with(prefix))
892 .cloned()
893 .collect())
894 }
895
896 pub fn read_file_sync(&self, path: &Path) -> io::Result<Vec<u8>> {
898 let files = self.files.read();
899 files
900 .get(path)
901 .map(|data| data.as_ref().clone())
902 .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "File not found"))
903 }
904
905 pub fn write_sync(&self, path: &Path, data: &[u8]) -> io::Result<()> {
907 self.files
908 .write()
909 .insert(path.to_path_buf(), Arc::new(data.to_vec()));
910 Ok(())
911 }
912}
913
914#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
915#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
916impl Directory for RamDirectory {
917 async fn exists(&self, path: &Path) -> io::Result<bool> {
918 Ok(self.files.read().contains_key(path))
919 }
920
921 async fn file_size(&self, path: &Path) -> io::Result<u64> {
922 self.files
923 .read()
924 .get(path)
925 .map(|data| data.len() as u64)
926 .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "File not found"))
927 }
928
929 async fn open_read(&self, path: &Path) -> io::Result<FileHandle> {
930 let files = self.files.read();
931 let data = files
932 .get(path)
933 .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "File not found"))?;
934
935 Ok(FileHandle::from_bytes(OwnedBytes::from_arc_vec(
936 Arc::clone(data),
937 0..data.len(),
938 )))
939 }
940
941 async fn read_range(&self, path: &Path, range: Range<u64>) -> io::Result<OwnedBytes> {
942 let files = self.files.read();
943 let data = files
944 .get(path)
945 .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "File not found"))?;
946
947 let start = range.start as usize;
948 let end = range.end as usize;
949
950 if end > data.len() {
951 return Err(io::Error::new(
952 io::ErrorKind::InvalidInput,
953 "Range out of bounds",
954 ));
955 }
956
957 Ok(OwnedBytes::from_arc_vec(Arc::clone(data), start..end))
958 }
959
960 async fn list_files(&self, prefix: &Path) -> io::Result<Vec<PathBuf>> {
961 let files = self.files.read();
962 Ok(files
963 .keys()
964 .filter(|p| p.starts_with(prefix))
965 .cloned()
966 .collect())
967 }
968
969 async fn open_lazy(&self, path: &Path) -> io::Result<FileHandle> {
970 self.open_read(path).await
972 }
973}
974
975#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
976#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
977impl DirectoryWriter for RamDirectory {
978 async fn write(&self, path: &Path, data: &[u8]) -> io::Result<()> {
979 self.files
980 .write()
981 .insert(path.to_path_buf(), Arc::new(data.to_vec()));
982 Ok(())
983 }
984
985 async fn delete(&self, path: &Path) -> io::Result<()> {
986 self.files.write().remove(path);
987 Ok(())
988 }
989
990 async fn rename(&self, from: &Path, to: &Path) -> io::Result<()> {
991 let mut files = self.files.write();
992 if let Some(data) = files.remove(from) {
993 files.insert(to.to_path_buf(), data);
994 }
995 Ok(())
996 }
997
998 async fn link(&self, from: &Path, to: &Path) -> io::Result<()> {
999 let mut files = self.files.write();
1000 let data = files.get(from).cloned().ok_or_else(|| {
1001 io::Error::new(
1002 io::ErrorKind::NotFound,
1003 format!("source file {from:?} does not exist"),
1004 )
1005 })?;
1006 files.insert(to.to_path_buf(), data);
1007 Ok(())
1008 }
1009
1010 async fn sync(&self) -> io::Result<()> {
1011 Ok(())
1012 }
1013
1014 async fn streaming_writer(&self, path: &Path) -> io::Result<Box<dyn StreamingWriter>> {
1015 Ok(Box::new(BufferedStreamingWriter {
1016 path: path.to_path_buf(),
1017 buffer: Vec::new(),
1018 files: Arc::clone(&self.files),
1019 }))
1020 }
1021}
1022
1023#[cfg(feature = "native")]
1025#[derive(Debug, Clone)]
1026pub struct FsDirectory {
1027 root: PathBuf,
1028 label: IndexLabel,
1029}
1030
1031#[cfg(all(feature = "native", unix))]
1034fn read_exact_at(file: &std::fs::File, buffer: &mut [u8], offset: u64) -> io::Result<()> {
1035 use std::os::unix::fs::FileExt;
1036 file.read_exact_at(buffer, offset)
1037}
1038
1039#[cfg(all(feature = "native", windows))]
1040fn read_exact_at(file: &std::fs::File, mut buffer: &mut [u8], mut offset: u64) -> io::Result<()> {
1041 use std::os::windows::fs::FileExt;
1042 while !buffer.is_empty() {
1043 match file.seek_read(buffer, offset) {
1044 Ok(0) => {
1045 return Err(io::Error::new(
1046 io::ErrorKind::UnexpectedEof,
1047 "failed to fill whole buffer",
1048 ));
1049 }
1050 Ok(read) => {
1051 buffer = &mut buffer[read..];
1052 offset += read as u64;
1053 }
1054 Err(error) if error.kind() == io::ErrorKind::Interrupted => {}
1055 Err(error) => return Err(error),
1056 }
1057 }
1058 Ok(())
1059}
1060
1061#[cfg(feature = "native")]
1062impl FsDirectory {
1063 pub fn new(root: impl AsRef<Path>) -> Self {
1064 Self {
1065 root: root.as_ref().to_path_buf(),
1066 label: IndexLabel::default(),
1067 }
1068 }
1069
1070 fn resolve(&self, path: &Path) -> PathBuf {
1071 self.root.join(path)
1072 }
1073}
1074
1075#[cfg(feature = "native")]
1076#[async_trait]
1077impl Directory for FsDirectory {
1078 async fn exists(&self, path: &Path) -> io::Result<bool> {
1079 let full_path = self.resolve(path);
1080 tokio::fs::try_exists(&full_path).await
1086 }
1087
1088 async fn file_size(&self, path: &Path) -> io::Result<u64> {
1089 let full_path = self.resolve(path);
1090 let metadata = tokio::fs::metadata(&full_path).await?;
1091 Ok(metadata.len())
1092 }
1093
1094 async fn open_read(&self, path: &Path) -> io::Result<FileHandle> {
1095 let full_path = self.resolve(path);
1096 let data = tokio::fs::read(&full_path).await?;
1097 Ok(FileHandle::from_bytes(OwnedBytes::new(data)))
1098 }
1099
1100 async fn read_range(&self, path: &Path, range: Range<u64>) -> io::Result<OwnedBytes> {
1101 use tokio::io::{AsyncReadExt, AsyncSeekExt};
1102
1103 let full_path = self.resolve(path);
1104 let mut file = tokio::fs::File::open(&full_path).await?;
1105
1106 file.seek(std::io::SeekFrom::Start(range.start)).await?;
1107
1108 let len = (range.end - range.start) as usize;
1109 let mut buffer = vec![0u8; len];
1110 file.read_exact(&mut buffer).await?;
1111
1112 Ok(OwnedBytes::new(buffer))
1113 }
1114
1115 async fn list_files(&self, prefix: &Path) -> io::Result<Vec<PathBuf>> {
1116 super::local::list_files(&self.root, prefix).await
1117 }
1118
1119 async fn open_lazy(&self, path: &Path) -> io::Result<FileHandle> {
1120 let full_path = self.resolve(path);
1124 let (file, file_size) = tokio::task::spawn_blocking(move || {
1125 let file = std::fs::File::open(&full_path)?;
1126 let file_size = file.metadata()?.len();
1127 Ok::<_, io::Error>((file, file_size))
1128 })
1129 .await
1130 .map_err(io::Error::other)??;
1131 let file = Arc::new(file);
1132
1133 let read_fn: RangeReadFn = Arc::new(move |range: Range<u64>| {
1134 let file = Arc::clone(&file);
1135 Box::pin(async move {
1136 tokio::task::spawn_blocking(move || {
1137 let len = (range.end - range.start) as usize;
1138 let mut buffer = vec![0u8; len];
1139 read_exact_at(&file, &mut buffer, range.start)?;
1140 Ok(OwnedBytes::new(buffer))
1141 })
1142 .await
1143 .map_err(io::Error::other)?
1144 })
1145 });
1146
1147 Ok(FileHandle::lazy_labeled(
1148 file_size,
1149 read_fn,
1150 self.label.get(),
1151 ))
1152 }
1153
1154 fn set_index_label(&self, label: &str) {
1155 self.label.set(label);
1156 }
1157
1158 fn local_path(&self, path: &Path) -> Option<PathBuf> {
1159 Some(self.resolve(path))
1160 }
1161}
1162
1163#[cfg(feature = "native")]
1164#[async_trait]
1165impl DirectoryWriter for FsDirectory {
1166 async fn write(&self, path: &Path, data: &[u8]) -> io::Result<()> {
1167 let full_path = self.resolve(path);
1168
1169 if let Some(parent) = full_path.parent() {
1171 tokio::fs::create_dir_all(parent).await?;
1172 }
1173
1174 tokio::fs::write(&full_path, data).await
1175 }
1176
1177 async fn delete(&self, path: &Path) -> io::Result<()> {
1178 let full_path = self.resolve(path);
1179 tokio::fs::remove_file(&full_path).await
1180 }
1181
1182 async fn rename(&self, from: &Path, to: &Path) -> io::Result<()> {
1183 let from_path = self.resolve(from);
1184 let to_path = self.resolve(to);
1185 std::fs::rename(&from_path, &to_path)
1192 }
1193
1194 async fn link(&self, from: &Path, to: &Path) -> io::Result<()> {
1195 std::fs::hard_link(self.resolve(from), self.resolve(to))
1196 }
1197
1198 async fn sync(&self) -> io::Result<()> {
1199 let dir = std::fs::File::open(&self.root)?;
1201 dir.sync_all()?;
1202 Ok(())
1203 }
1204
1205 async fn streaming_writer(&self, path: &Path) -> io::Result<Box<dyn StreamingWriter>> {
1206 super::local::streaming_writer(&self.resolve(path)).await
1207 }
1208
1209 async fn streaming_writer_cold(&self, path: &Path) -> io::Result<Box<dyn StreamingWriter>> {
1210 super::local::streaming_writer_cold(&self.resolve(path), self.label.get(), None).await
1211 }
1212
1213 async fn streaming_writer_cold_with_capacity(
1214 &self,
1215 path: &Path,
1216 buffer_capacity: usize,
1217 ) -> io::Result<Box<dyn StreamingWriter>> {
1218 super::local::streaming_writer_cold(
1219 &self.resolve(path),
1220 self.label.get(),
1221 Some(buffer_capacity),
1222 )
1223 .await
1224 }
1225}
1226
1227pub struct CachingDirectory<D: Directory> {
1229 inner: D,
1230 cache: RwLock<HashMap<PathBuf, Arc<Vec<u8>>>>,
1231 max_cached_bytes: usize,
1232 current_bytes: RwLock<usize>,
1233}
1234
1235impl<D: Directory> CachingDirectory<D> {
1236 pub fn new(inner: D, max_cached_bytes: usize) -> Self {
1237 Self {
1238 inner,
1239 cache: RwLock::new(HashMap::new()),
1240 max_cached_bytes,
1241 current_bytes: RwLock::new(0),
1242 }
1243 }
1244
1245 fn try_cache(&self, path: &Path, data: &[u8]) {
1246 let mut current = self.current_bytes.write();
1247 if *current + data.len() <= self.max_cached_bytes {
1248 self.cache
1249 .write()
1250 .insert(path.to_path_buf(), Arc::new(data.to_vec()));
1251 *current += data.len();
1252 }
1253 }
1254}
1255
1256#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
1257#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
1258impl<D: Directory> Directory for CachingDirectory<D> {
1259 async fn exists(&self, path: &Path) -> io::Result<bool> {
1260 if self.cache.read().contains_key(path) {
1261 return Ok(true);
1262 }
1263 self.inner.exists(path).await
1264 }
1265
1266 async fn file_size(&self, path: &Path) -> io::Result<u64> {
1267 if let Some(data) = self.cache.read().get(path) {
1268 return Ok(data.len() as u64);
1269 }
1270 self.inner.file_size(path).await
1271 }
1272
1273 async fn open_read(&self, path: &Path) -> io::Result<FileHandle> {
1274 if let Some(data) = self.cache.read().get(path) {
1276 return Ok(FileHandle::from_bytes(OwnedBytes::from_arc_vec(
1277 Arc::clone(data),
1278 0..data.len(),
1279 )));
1280 }
1281
1282 let handle = self.inner.open_read(path).await?;
1284 let bytes = handle.read_bytes().await?;
1285
1286 self.try_cache(path, bytes.as_slice());
1287
1288 Ok(FileHandle::from_bytes(bytes))
1289 }
1290
1291 async fn read_range(&self, path: &Path, range: Range<u64>) -> io::Result<OwnedBytes> {
1292 if let Some(data) = self.cache.read().get(path) {
1294 let start = range.start as usize;
1295 let end = range.end as usize;
1296 return Ok(OwnedBytes::from_arc_vec(Arc::clone(data), start..end));
1297 }
1298
1299 self.inner.read_range(path, range).await
1300 }
1301
1302 async fn list_files(&self, prefix: &Path) -> io::Result<Vec<PathBuf>> {
1303 self.inner.list_files(prefix).await
1304 }
1305
1306 async fn open_lazy(&self, path: &Path) -> io::Result<FileHandle> {
1307 self.inner.open_lazy(path).await
1309 }
1310
1311 fn set_index_label(&self, label: &str) {
1312 self.inner.set_index_label(label);
1313 }
1314
1315 fn local_path(&self, path: &Path) -> Option<PathBuf> {
1316 self.inner.local_path(path)
1317 }
1318}
1319
1320#[cfg(test)]
1321mod tests {
1322 use super::*;
1323
1324 #[tokio::test]
1325 async fn test_ram_directory() {
1326 let dir = RamDirectory::new();
1327
1328 dir.write(Path::new("test.txt"), b"hello world")
1330 .await
1331 .unwrap();
1332
1333 assert!(dir.exists(Path::new("test.txt")).await.unwrap());
1335 assert!(!dir.exists(Path::new("nonexistent.txt")).await.unwrap());
1336
1337 let slice = dir.open_read(Path::new("test.txt")).await.unwrap();
1339 let data = slice.read_bytes().await.unwrap();
1340 assert_eq!(data.as_slice(), b"hello world");
1341
1342 let range_data = dir.read_range(Path::new("test.txt"), 0..5).await.unwrap();
1344 assert_eq!(range_data.as_slice(), b"hello");
1345
1346 dir.delete(Path::new("test.txt")).await.unwrap();
1348 assert!(!dir.exists(Path::new("test.txt")).await.unwrap());
1349 }
1350
1351 #[cfg(all(unix, feature = "native"))]
1356 #[tokio::test]
1357 async fn test_fs_exists_propagates_stat_errors_instead_of_reporting_missing() {
1358 use std::os::unix::fs::PermissionsExt;
1359
1360 let temp_dir = tempfile::TempDir::new().unwrap();
1361 let dir = FsDirectory::new(temp_dir.path());
1362 dir.write(Path::new("locked/seg.meta"), b"data")
1363 .await
1364 .unwrap();
1365
1366 let locked = temp_dir.path().join("locked");
1369 let original = std::fs::metadata(&locked).unwrap().permissions();
1370 std::fs::set_permissions(&locked, std::fs::Permissions::from_mode(0o000)).unwrap();
1371 if std::fs::metadata(locked.join("seg.meta")).is_ok() {
1372 std::fs::set_permissions(&locked, original).unwrap();
1375 return;
1376 }
1377 let result = dir.exists(Path::new("locked/seg.meta")).await;
1378 std::fs::set_permissions(&locked, original).unwrap();
1379
1380 let error =
1381 result.expect_err("stat failure must propagate as Err, not be misreported as missing");
1382 assert_ne!(error.kind(), io::ErrorKind::NotFound);
1383 assert!(dir.exists(Path::new("locked/seg.meta")).await.unwrap());
1385 }
1386
1387 #[tokio::test]
1388 async fn test_file_handle() {
1389 let data = OwnedBytes::new(b"hello world".to_vec());
1390 let handle = FileHandle::from_bytes(data);
1391
1392 assert_eq!(handle.len(), 11);
1393 assert!(handle.is_sync());
1394
1395 let sub = handle.slice(0..5);
1396 let bytes = sub.read_bytes().await.unwrap();
1397 assert_eq!(bytes.as_slice(), b"hello");
1398
1399 let sub2 = handle.slice(6..11);
1400 let bytes2 = sub2.read_bytes().await.unwrap();
1401 assert_eq!(bytes2.as_slice(), b"world");
1402
1403 let sync_bytes = handle.read_bytes_range_sync(0..5).unwrap();
1405 assert_eq!(sync_bytes.as_slice(), b"hello");
1406 }
1407
1408 #[test]
1409 fn local_byte_owners_share_storage_without_repeated_global_refcounts() {
1410 let backing = Arc::new((0u8..64).collect::<Vec<_>>());
1411 let handle = FileHandle::from_bytes(OwnedBytes::from_arc_vec(backing.clone(), 3..61));
1412 let local = handle.slice(2..40).with_local_owner().with_local_owner();
1413 let count = Arc::strong_count(&backing);
1414 let views: Vec<_> = (0..20)
1415 .map(|i| local.read_bytes_range_sync(i..i + 3).unwrap())
1416 .collect();
1417 assert_eq!(Arc::strong_count(&backing), count);
1418 drop(local);
1419 drop(handle);
1420 let survivor = std::thread::spawn(move || {
1421 for (i, view) in views.iter().enumerate() {
1422 assert_eq!(
1423 view.as_slice(),
1424 &[(i + 5) as u8, (i + 6) as u8, (i + 7) as u8]
1425 );
1426 }
1427 views[7].clone()
1428 })
1429 .join()
1430 .unwrap();
1431 assert_eq!(survivor.as_slice(), &[12, 13, 14]);
1432 drop(survivor);
1433 assert_eq!(Arc::strong_count(&backing), 1);
1434 }
1435
1436 #[test]
1437 fn owned_byte_views_retain_heap_storage_across_moves_clones_and_empty_slices() {
1438 let backing = Arc::new((0u8..64).collect::<Vec<_>>());
1439 let weak = Arc::downgrade(&backing);
1440 let bytes = OwnedBytes::from_arc_vec(backing.clone(), 3..61);
1441 let nested = bytes.slice(1..57).slice(2..53);
1442 let expected = (6u8..57).collect::<Vec<_>>();
1443 let pointer = nested.as_slice().as_ptr();
1444 let empty = nested.slice(nested.len()..nested.len());
1445 assert!(empty.is_empty());
1446 assert!(OwnedBytes::empty().slice(0..0).as_slice().is_empty());
1447 let cloned = nested.clone();
1448 drop(backing);
1449 drop(bytes);
1450 drop(nested);
1451 assert_eq!(cloned.as_slice().as_ptr(), pointer);
1452 assert_eq!(cloned.as_slice(), expected);
1453 drop(cloned);
1454 assert!(
1455 weak.upgrade().is_some(),
1456 "empty views also retain their owner"
1457 );
1458 drop(empty);
1459 assert!(weak.upgrade().is_none());
1460 assert_eq!(
1461 std::mem::size_of::<OwnedBytes>(),
1462 std::mem::size_of::<super::SharedBytes>() + 2 * std::mem::size_of::<usize>()
1463 );
1464 }
1465
1466 #[cfg(feature = "native")]
1467 #[test]
1468 fn owned_byte_views_keep_heap_and_mmap_owners_alive_across_threads() {
1469 fn check(bytes: OwnedBytes, mapped: bool) {
1470 assert_eq!(bytes.is_mmap(), mapped);
1471 let survivor = bytes.slice(3..61).slice(1..56);
1472 let copied = survivor.clone();
1473 drop(bytes);
1474 let thread = std::thread::spawn(move || {
1475 assert_eq!(survivor.as_slice(), &(4u8..59).collect::<Vec<_>>());
1476 assert_eq!(survivor.is_mmap(), mapped);
1477 survivor.slice(2..9)
1478 });
1479 assert_eq!(copied.as_slice(), &(4u8..59).collect::<Vec<_>>());
1480 drop(copied);
1481 let final_view = thread.join().unwrap();
1482 assert_eq!(final_view.as_slice(), &[6, 7, 8, 9, 10, 11, 12]);
1483 assert_eq!(final_view.is_mmap(), mapped);
1484 }
1485 check(OwnedBytes::new((0u8..64).collect()), false);
1486 let mut mapping = memmap2::MmapMut::map_anon(64).unwrap();
1487 mapping.copy_from_slice(&(0u8..64).collect::<Vec<_>>());
1488 let mapping = Arc::new(mapping.make_read_only().unwrap());
1489 let weak = Arc::downgrade(&mapping);
1490 check(OwnedBytes::from_mmap_range(mapping.clone(), 0..64), true);
1491 check(
1492 OwnedBytes::from_mmap_range(mapping.clone(), 0..64).with_local_owner(),
1493 true,
1494 );
1495 check(
1496 OwnedBytes::new((0u8..64).collect()).with_local_owner(),
1497 false,
1498 );
1499 assert_eq!(Arc::strong_count(&mapping), 1);
1500 drop(mapping);
1501 assert!(weak.upgrade().is_none());
1502 }
1503
1504 #[test]
1505 fn owned_byte_subslices_reject_access_outside_the_parent_view() {
1506 let bytes = OwnedBytes::new(vec![1, 2, 3, 4, 5]);
1507 let parent = bytes.slice(1..3);
1508 assert_eq!(parent.slice(0..2).as_slice(), &[2, 3]);
1509 assert!(std::panic::catch_unwind(|| parent.slice(0..3).to_vec()).is_err());
1510 assert!(
1511 std::panic::catch_unwind(|| parent.slice(Range { start: 2, end: 1 }).to_vec()).is_err()
1512 );
1513 assert!(
1514 std::panic::catch_unwind(|| parent.slice(usize::MAX..usize::MAX).to_vec()).is_err()
1515 );
1516 }
1517
1518 #[tokio::test]
1519 async fn test_owned_bytes() {
1520 let bytes = OwnedBytes::new(vec![1, 2, 3, 4, 5]);
1521
1522 assert_eq!(bytes.len(), 5);
1523 assert_eq!(bytes.as_slice(), &[1, 2, 3, 4, 5]);
1524
1525 let sliced = bytes.slice(1..4);
1526 assert_eq!(sliced.as_slice(), &[2, 3, 4]);
1527
1528 assert_eq!(bytes.as_slice(), &[1, 2, 3, 4, 5]);
1530 }
1531}