Merge pull request #219 from cocool97/improve-storage

feat: add on_piece_completed method on TorrentStorage

crates/librqbit/src/file_ops.rs

@@ -215,6 +215,8 @@ impl<'a> FileOps<'a> {
                 format!("error reading {to_read_in_file} bytes, file_id: {file_idx} (\"{name:?}\")")
             })?;
 
+            self.files.on_piece_completed(piece_index)?;
+
             piece_remaining_bytes -= to_read_in_file;
 
             if piece_remaining_bytes == 0 {

crates/librqbit/src/peer_connection.rs

@@ -14,7 +14,10 @@ use librqbit_core::{
 };
 use parking_lot::RwLock;
 use peer_binary_protocol::{
-    extended::{handshake::{ExtendedHandshake, YourIP}, ExtendedMessage},
+    extended::{
+        handshake::{ExtendedHandshake, YourIP},
+        ExtendedMessage,
+    },
     serialize_piece_preamble, Handshake, Message, MessageOwned, PIECE_MESSAGE_DEFAULT_LEN,
 };
 use serde::{Deserialize, Serialize};

crates/librqbit/src/storage/examples/inmemory.rs

@@ -27,7 +27,7 @@ impl StorageFactory for InMemoryExampleStorageFactory {
 
     fn create(
         &self,
-        info: &crate::torrent_state::ManagedTorrentInfo,
+        info: &crate::torrent_state::ManagedTorrentShared,
     ) -> anyhow::Result<InMemoryExampleStorage> {
         InMemoryExampleStorage::new(info.lengths, info.file_infos.clone())
     }
@@ -111,7 +111,7 @@ impl TorrentStorage for InMemoryExampleStorage {
         }))
     }
 
-    fn init(&mut self, _meta: &crate::ManagedTorrentInfo) -> anyhow::Result<()> {
+    fn init(&mut self, _meta: &crate::ManagedTorrentShared) -> anyhow::Result<()> {
         Ok(())
     }
 

crates/librqbit/src/storage/examples/mmap.rs

@@ -4,7 +4,7 @@ use parking_lot::RwLock;
 
 use crate::{
     storage::{StorageFactory, StorageFactoryExt, TorrentStorage},
-    FileInfos, ManagedTorrentInfo,
+    FileInfos, ManagedTorrentShared,
 };
 
 #[derive(Default, Clone)]
@@ -18,7 +18,7 @@ pub struct MmapStorage {
 impl StorageFactory for MmapStorageFactory {
     type Storage = MmapStorage;
 
-    fn create(&self, info: &ManagedTorrentInfo) -> anyhow::Result<Self::Storage> {
+    fn create(&self, info: &ManagedTorrentShared) -> anyhow::Result<Self::Storage> {
         Ok(MmapStorage {
             mmap: RwLock::new(
                 MmapOptions::new()
@@ -63,11 +63,15 @@ impl TorrentStorage for MmapStorage {
         anyhow::bail!("not implemented")
     }
 
-    fn init(&mut self, _meta: &ManagedTorrentInfo) -> anyhow::Result<()> {
+    fn init(&mut self, _meta: &ManagedTorrentShared) -> anyhow::Result<()> {
         Ok(())
     }
 
     fn remove_directory_if_empty(&self, _path: &std::path::Path) -> anyhow::Result<()> {
         Ok(())
     }
+
+    fn on_piece_completed(&self, _file_id: usize, _offset: u64) -> anyhow::Result<()> {
+        Ok(())
+    }
 }

crates/librqbit/src/storage/mod.rs

@@ -11,6 +11,8 @@ use std::{
     path::Path,
 };
 
+use librqbit_core::lengths::ValidPieceIndex;
+
 use crate::torrent_state::ManagedTorrentShared;
 
 pub trait StorageFactory: Send + Sync + Any {
@@ -97,6 +99,12 @@ pub trait TorrentStorage: Send + Sync {
     /// Replace the current storage with a dummy, and return a new one that should be used instead.
     /// This is used to make the underlying object useless when e.g. pausing the torrent.
     fn take(&self) -> anyhow::Result<Box<dyn TorrentStorage>>;
+
+    /// Callback called every time a piece has completed and has been validated.
+    /// Default implementation does nothing, but can be override in trait implementations.
+    fn on_piece_completed(&self, _piece_index: ValidPieceIndex) -> anyhow::Result<()> {
+        Ok(())
+    }
 }
 
 impl<U: TorrentStorage + ?Sized> TorrentStorage for Box<U> {
@@ -127,4 +135,8 @@ impl<U: TorrentStorage + ?Sized> TorrentStorage for Box<U> {
     fn init(&mut self, meta: &ManagedTorrentShared) -> anyhow::Result<()> {
         (**self).init(meta)
     }
+
+    fn on_piece_completed(&self, piece_id: ValidPieceIndex) -> anyhow::Result<()> {
+        (**self).on_piece_completed(piece_id)
+    }
 }

crates/librqbit/src/tests/test_util.rs

@@ -5,7 +5,7 @@ use std::{
     time::Duration,
 };
 
-use anyhow::{bail, Context};
+use anyhow::bail;
 use librqbit_core::Id20;
 use parking_lot::RwLock;
 use rand::{thread_rng, Rng, RngCore, SeedableRng};
@@ -97,6 +97,7 @@ impl TestPeerMetadata {
 
 #[cfg(feature = "http-api")]
 async fn debug_server() -> anyhow::Result<()> {
+    use anyhow::Context;
     use axum::{response::IntoResponse, routing::get, Router};
     async fn backtraces() -> impl IntoResponse {
         #[cfg(feature = "async-bt")]

crates/peer_binary_protocol/src/extended/handshake.rs

@@ -99,7 +99,7 @@ impl Serialize for YourIP {
             IpAddr::V6(ipv6) => {
                 let buf = ipv6.octets();
                 serializer.serialize_bytes(&buf)
-            },
+            }
         }
     }
 }