begin to pass around generic file type

Author: seanaye

src/store/bao_file.rs

@@ -32,7 +32,7 @@ use iroh_io::AsyncSliceReader;
 use super::mutable_mem_storage::{MutableMemStorage, SizeInfo};
 use crate::{
     store::BaoBatchWriter,
-    util::{get_limited_slice, MemOrFile, SparseMemFile},
+    util::{callback_lock::CallbackLock, get_limited_slice, FileAndSize, MemOrFile, SparseMemFile},
     Hash, IROH_BLOCK_SIZE,
 };
 
@@ -81,46 +81,63 @@ struct DataPaths {
 ///
 /// For the memory variant, it does reading in a zero copy way, since storage
 /// is already a `Bytes`.
-#[derive(Default, derive_more::Debug)]
-pub struct CompleteStorage {
+#[derive(derive_more::Debug)]
+#[debug(bound(T: Debug))]
+pub struct CompleteStorage<T> {
     /// data part, which can be in memory or on disk.
     #[debug("{:?}", data.as_ref().map_mem(|x| x.len()))]
-    pub data: MemOrFile<Bytes, (File, u64)>,
+    pub data: MemOrFile<Bytes, FileAndSize<T>>,
     /// outboard part, which can be in memory or on disk.
     #[debug("{:?}", outboard.as_ref().map_mem(|x| x.len()))]
-    pub outboard: MemOrFile<Bytes, (File, u64)>,
+    pub outboard: MemOrFile<Bytes, FileAndSize<T>>,
 }
 
-impl CompleteStorage {
+impl<T> Default for CompleteStorage<T> {
+    fn default() -> Self {
+        Self {
+            data: Default::default(),
+            outboard: Default::default(),
+        }
+    }
+}
+
+impl<T> CompleteStorage<T>
+where
+    T: bao_tree::io::sync::ReadAt,
+{
     /// Read from the data file at the given offset, until end of file or max bytes.
     pub fn read_data_at(&self, offset: u64, len: usize) -> Bytes {
         match &self.data {
             MemOrFile::Mem(mem) => get_limited_slice(mem, offset, len),
-            MemOrFile::File((file, _size)) => read_to_end(file, offset, len).unwrap(),
+            MemOrFile::File(FileAndSize { file, size: _ }) => {
+                read_to_end(file, offset, len).unwrap()
+            }
         }
     }
 
     /// Read from the outboard file at the given offset, until end of file or max bytes.
     pub fn read_outboard_at(&self, offset: u64, len: usize) -> Bytes {
         match &self.outboard {
             MemOrFile::Mem(mem) => get_limited_slice(mem, offset, len),
-            MemOrFile::File((file, _size)) => read_to_end(file, offset, len).unwrap(),
+            MemOrFile::File(FileAndSize { file, size: _ }) => {
+                read_to_end(file, offset, len).unwrap()
+            }
         }
     }
 
     /// The size of the data file.
     pub fn data_size(&self) -> u64 {
         match &self.data {
             MemOrFile::Mem(mem) => mem.len() as u64,
-            MemOrFile::File((_file, size)) => *size,
+            MemOrFile::File(FileAndSize { file: _, size }) => *size,
         }
     }
 
     /// The size of the outboard file.
     pub fn outboard_size(&self) -> u64 {
         match &self.outboard {
             MemOrFile::Mem(mem) => mem.len() as u64,
-            MemOrFile::File((_file, size)) => *size,
+            MemOrFile::File(FileAndSize { file: _, size }) => *size,
         }
     }
 }
@@ -244,7 +261,7 @@ impl FileStorage {
 
 /// The storage for a bao file. This can be either in memory or on disk.
 #[derive(Debug)]
-pub(crate) enum BaoFileStorage {
+pub(crate) enum BaoFileStorage<T> {
     /// The entry is incomplete and in memory.
     ///
     /// Since it is incomplete, it must be writeable.
@@ -261,16 +278,16 @@ pub(crate) enum BaoFileStorage {
     /// (memory or file).
     ///
     /// Writing to this is a no-op, since it is already complete.
-    Complete(CompleteStorage),
+    Complete(CompleteStorage<T>),
 }
 
-impl Default for BaoFileStorage {
+impl<T> Default for BaoFileStorage<T> {
     fn default() -> Self {
         BaoFileStorage::Complete(Default::default())
     }
 }
 
-impl BaoFileStorage {
+impl<T> BaoFileStorage<T> {
     /// Take the storage out, leaving an empty storage in its place.
     ///
     /// Be careful to put something back in its place, or you will lose data.
@@ -310,11 +327,11 @@ impl BaoFileStorage {
 
 /// A weak reference to a bao file handle.
 #[derive(Debug, Clone)]
-pub struct BaoFileHandleWeak(Weak<BaoFileHandleInner>);
+pub struct BaoFileHandleWeak<T>(Weak<BaoFileHandleInner<T>>);
 
-impl BaoFileHandleWeak {
+impl<T> BaoFileHandleWeak<T> {
     /// Upgrade to a strong reference if possible.
-    pub fn upgrade(&self) -> Option<BaoFileHandle> {
+    pub fn upgrade(&self) -> Option<BaoFileHandle<T>> {
         self.0.upgrade().map(BaoFileHandle)
     }
 
@@ -326,15 +343,29 @@ impl BaoFileHandleWeak {
 
 /// The inner part of a bao file handle.
 #[derive(Debug)]
-pub struct BaoFileHandleInner {
-    pub(crate) storage: RwLock<BaoFileStorage>,
+pub struct BaoFileHandleInner<T> {
+    pub(crate) storage: RwLock<BaoFileStorage<T>>,
     config: Arc<BaoFileConfig>,
     hash: Hash,
 }
 
 /// A cheaply cloneable handle to a bao file, including the hash and the configuration.
-#[derive(Debug, Clone, derive_more::Deref)]
-pub struct BaoFileHandle(Arc<BaoFileHandleInner>);
+#[derive(Debug)]
+pub struct BaoFileHandle<T>(Arc<BaoFileHandleInner<T>>);
+
+impl<T> Deref for BaoFileHandle<T> {
+    type Target = Arc<BaoFileHandleInner<T>>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<T> Clone for BaoFileHandle<T> {
+    fn clone(&self) -> Self {
+        Self(self.0.clone())
+    }
+}
 
 pub(crate) type CreateCb = Arc<dyn Fn(&Hash) -> io::Result<()> + Send + Sync>;
 
@@ -375,13 +406,18 @@ impl BaoFileConfig {
 
 /// A reader for a bao file, reading just the data.
 #[derive(Debug)]
-pub struct DataReader(Option<BaoFileHandle>);
+pub struct DataReader<T>(Option<BaoFileHandle<T>>);
 
-async fn with_storage<T, P, F>(opt: &mut Option<BaoFileHandle>, no_io: P, f: F) -> io::Result<T>
+async fn with_storage<T, P, F, H>(
+    opt: &mut Option<BaoFileHandle<H>>,
+    no_io: P,
+    f: F,
+) -> io::Result<T>
 where
-    P: Fn(&BaoFileStorage) -> bool + Send + 'static,
-    F: FnOnce(&BaoFileStorage) -> io::Result<T> + Send + 'static,
+    P: Fn(&BaoFileStorage<H>) -> bool + Send + 'static,
+    F: FnOnce(&BaoFileStorage<H>) -> io::Result<T> + Send + 'static,
     T: Send + 'static,
+    H: Send + Sync + 'static,
 {
     let handle = opt
         .take()
@@ -410,7 +446,10 @@ where
     res
 }
 
-impl AsyncSliceReader for DataReader {
+impl<T> AsyncSliceReader for DataReader<T>
+where
+    T: Send + Sync + bao_tree::io::sync::ReadAt + 'static,
+{
     async fn read_at(&mut self, offset: u64, len: usize) -> io::Result<Bytes> {
         with_storage(
             &mut self.0,
@@ -440,9 +479,12 @@ impl AsyncSliceReader for DataReader {
 
 /// A reader for the outboard part of a bao file.
 #[derive(Debug)]
-pub struct OutboardReader(Option<BaoFileHandle>);
+pub struct OutboardReader<T>(Option<BaoFileHandle<T>>);
 
-impl AsyncSliceReader for OutboardReader {
+impl<T> AsyncSliceReader for OutboardReader<T>
+where
+    T: Send + Sync + bao_tree::io::sync::ReadAt + 'static,
+{
     async fn read_at(&mut self, offset: u64, len: usize) -> io::Result<Bytes> {
         with_storage(
             &mut self.0,
@@ -476,7 +518,10 @@ enum HandleChange {
     // later: size verified
 }
 
-impl BaoFileHandle {
+impl<T> BaoFileHandle<T>
+where
+    T: bao_tree::io::sync::ReadAt,
+{
     /// Create a new bao file handle.
     ///
     /// This will create a new file handle with an empty memory storage.
@@ -509,8 +554,8 @@ impl BaoFileHandle {
     pub fn new_complete(
         config: Arc<BaoFileConfig>,
         hash: Hash,
-        data: MemOrFile<Bytes, (File, u64)>,
-        outboard: MemOrFile<Bytes, (File, u64)>,
+        data: MemOrFile<Bytes, FileAndSize<T>>,
+        outboard: MemOrFile<Bytes, FileAndSize<T>>,
     ) -> Self {
         let storage = BaoFileStorage::Complete(CompleteStorage { data, outboard });
         Self(Arc::new(BaoFileHandleInner {
@@ -525,7 +570,7 @@ impl BaoFileHandle {
     #[cfg(feature = "fs-store")]
     pub(crate) fn transform(
         &self,
-        f: impl FnOnce(BaoFileStorage) -> io::Result<BaoFileStorage>,
+        f: impl FnOnce(BaoFileStorage<T>) -> io::Result<BaoFileStorage<T>>,
     ) -> io::Result<()> {
         let mut lock = self.storage.write().unwrap();
         let storage = lock.take();
@@ -545,15 +590,15 @@ impl BaoFileHandle {
     ///
     /// Caution: this is a reader for the unvalidated data file. Reading this
     /// can produce data that does not match the hash.
-    pub fn data_reader(&self) -> DataReader {
+    pub fn data_reader(&self) -> DataReader<T> {
         DataReader(Some(self.clone()))
     }
 
     /// An AsyncSliceReader for the outboard file.
     ///
     /// The outboard file is used to validate the data file. It is not guaranteed
     /// to be complete.
-    pub fn outboard_reader(&self) -> OutboardReader {
+    pub fn outboard_reader(&self) -> OutboardReader<T> {
         OutboardReader(Some(self.clone()))
     }
 
@@ -567,7 +612,7 @@ impl BaoFileHandle {
     }
 
     /// The outboard for the file.
-    pub fn outboard(&self) -> io::Result<PreOrderOutboard<OutboardReader>> {
+    pub fn outboard(&self) -> io::Result<PreOrderOutboard<OutboardReader<T>>> {
         let root = self.hash.into();
         let tree = BaoTree::new(self.current_size()?, IROH_BLOCK_SIZE);
         let outboard = self.outboard_reader();
@@ -584,7 +629,7 @@ impl BaoFileHandle {
     }
 
     /// Create a new writer from the handle.
-    pub fn writer(&self) -> BaoFileWriter {
+    pub fn writer(&self) -> BaoFileWriter<T> {
         BaoFileWriter(Some(self.clone()))
     }
 
@@ -625,7 +670,7 @@ impl BaoFileHandle {
     }
 
     /// Downgrade to a weak reference.
-    pub fn downgrade(&self) -> BaoFileHandleWeak {
+    pub fn downgrade(&self) -> BaoFileHandleWeak<T> {
         BaoFileHandleWeak(Arc::downgrade(&self.0))
     }
 }
@@ -676,9 +721,12 @@ impl MutableMemStorage {
 /// It is a BaoFileHandle wrapped in an Option, so that we can take it out
 /// in the future.
 #[derive(Debug)]
-pub struct BaoFileWriter(Option<BaoFileHandle>);
+pub struct BaoFileWriter<T>(Option<BaoFileHandle<T>>);
 
-impl BaoBatchWriter for BaoFileWriter {
+impl<T> BaoBatchWriter for BaoFileWriter<T>
+where
+    T: Send + Sync + bao_tree::io::sync::ReadAt + 'static,
+{
     async fn write_batch(&mut self, size: u64, batch: Vec<BaoContentItem>) -> std::io::Result<()> {
         let Some(handle) = self.0.take() else {
             return Err(io::Error::new(io::ErrorKind::Other, "deferred batch busy"));
@@ -828,7 +876,11 @@ pub mod test_support {
         (outboard.root.into(), chunk_ranges, encoded)
     }
 
-    pub async fn validate(handle: &BaoFileHandle, original: &[u8], ranges: &[Range<u64>]) {
+    pub async fn validate(
+        handle: &BaoFileHandle<std::fs::File>,
+        original: &[u8],
+        ranges: &[Range<u64>],
+    ) {
         let mut r = handle.data_reader();
         for range in ranges {
             let start = range.start;