Browse Source

fud/download: create the DHT events sub as early as possible

epiphany 8 months ago
parent
commit
d9aa63f0af
2 changed files with 59 additions and 19 deletions
  1. 17 10
      bin/fud/fud/src/download.rs
  2. 42 9
      bin/fud/fud/src/lib.rs

+ 17 - 10
bin/fud/fud/src/download.rs

@@ -30,9 +30,10 @@ use rand::{
 use tracing::{error, info, warn};
 
 use darkfi::{
-    dht::{event::DhtEvent, DhtNode},
+    dht::{event::DhtEvent, DhtHandler, DhtNode},
     geode::{hash_to_string, ChunkedStorage},
     net::ChannelPtr,
+    system::Subscription,
     Error, Result,
 };
 use darkfi_serial::serialize_async;
@@ -47,6 +48,8 @@ use crate::{
     Fud, FudSeeder, ResourceStatus, ResourceType, Scrap,
 };
 
+type FudDhtEvent = DhtEvent<<Fud as DhtHandler>::Node, <Fud as DhtHandler>::Value>;
+
 /// Receive seeders from a DHT events subscription, and execute an async
 /// expression for each deduplicated seeder once (seeder order is random).
 /// It will keep going until the expression returns `Ok(())`, or there are
@@ -54,7 +57,7 @@ use crate::{
 /// It has an optional `favored_seeder` argument that will be tried first if
 /// specified.
 macro_rules! seeders_loop {
-    ($key:expr, $fud:expr, $favored_seeder:expr, $code:expr) => {
+    ($key:expr, $fud:expr, $dht_sub:expr, $favored_seeder:expr, $code:expr) => {
         let mut queried_seeders: HashSet<blake3::Hash> = HashSet::new();
         let mut is_done = false;
 
@@ -68,9 +71,8 @@ macro_rules! seeders_loop {
         }
 
         // Try other seeders using the DHT subscription
-        let dht_sub = $fud.dht.subscribe().await;
         while !is_done {
-            let event = dht_sub.receive().await;
+            let event = $dht_sub.receive().await;
             if event.key() != Some($key) {
                 continue // Ignore this event if it's not about the right key
             }
@@ -102,10 +104,9 @@ macro_rules! seeders_loop {
                 break
             }
         }
-        dht_sub.unsubscribe().await;
     };
-    ($key:expr, $fud:expr, $code:expr) => {
-        seeders_loop!($key, $fud, None, $code)
+    ($key:expr, $fud:expr, $dht_sub:expr, $code:expr) => {
+        seeders_loop!($key, $fud, $dht_sub, None, $code)
     };
 }
 
@@ -127,12 +128,13 @@ pub async fn fetch_chunks(
     fud: &Fud,
     hash: &blake3::Hash,
     chunked: &mut ChunkedStorage,
+    dht_sub: &Subscription<FudDhtEvent>,
     favored_seeder: Option<FudSeeder>,
     chunks: &mut HashSet<blake3::Hash>,
 ) -> Result<()> {
     let mut ctx = ChunkFetchContext { fud, hash, chunked, chunks };
 
-    seeders_loop!(hash, fud, favored_seeder, async |seeder: FudSeeder| -> Result<()> {
+    seeders_loop!(hash, fud, dht_sub, favored_seeder, async |seeder: FudSeeder| -> Result<()> {
         let (channel, _) = match fud.dht.get_channel(&seeder.node).await {
             Ok(channel) => channel,
             Err(e) => {
@@ -323,10 +325,15 @@ enum MetadataFetchReply {
 /// 1. Wait for seeders from the subscription
 /// 2. Request the metadata from the seeders
 /// 3. Insert metadata to geode using the reply
-pub async fn fetch_metadata(fud: &Fud, hash: &blake3::Hash, path: &Path) -> Result<FudSeeder> {
+pub async fn fetch_metadata(
+    fud: &Fud,
+    hash: &blake3::Hash,
+    path: &Path,
+    dht_sub: &Subscription<FudDhtEvent>,
+) -> Result<FudSeeder> {
     let mut result: Option<(FudSeeder, MetadataFetchReply)> = None;
 
-    seeders_loop!(hash, fud, async |seeder: FudSeeder| -> Result<()> {
+    seeders_loop!(hash, fud, dht_sub, async |seeder: FudSeeder| -> Result<()> {
         let (channel, _) = fud.dht.get_channel(&seeder.node).await?;
         let msg_subscriber_chunk = channel.subscribe_msg::<FudChunkReply>().await.unwrap();
         let msg_subscriber_file = channel.subscribe_msg::<FudFileReply>().await.unwrap();

+ 42 - 9
bin/fud/fud/src/lib.rs

@@ -688,10 +688,15 @@ impl Fud {
             Err(Error::GeodeFileNotFound) => {
                 // Find nodes close to the file hash
                 info!(target: "fud::get_metadata()", "Requested metadata {} not found in Geode, triggering fetch", hash_to_string(hash));
-                self.lookup_tx.send(*hash).await?;
+                let dht_sub = self.dht.subscribe().await;
+                if let Err(e) = self.lookup_tx.send(*hash).await {
+                    dht_sub.unsubscribe().await;
+                    return Err(e.into())
+                }
 
                 // Fetch resource metadata
-                let fetch_res = fetch_metadata(self, hash, path).await;
+                let fetch_res = fetch_metadata(self, hash, path, &dht_sub).await;
+                dht_sub.unsubscribe().await;
                 let seeder = fetch_res?;
                 Ok((self.geode.get(hash, path).await?, Some(seeder)))
             }
@@ -764,6 +769,9 @@ impl Fud {
         resources_write.insert(*hash, resource.clone());
         drop(resources_write);
 
+        // Subscribe to DHT events early for `fetch_chunks()`
+        let dht_sub = self.dht.subscribe().await;
+
         // Send a DownloadStarted event
         notify_event!(self, DownloadStarted, resource);
 
@@ -774,6 +782,7 @@ impl Fud {
             // Set resource status to `Incomplete` and send a `MetadataNotFound` event
             let resource = update_resource!(hash, { status = ResourceStatus::Incomplete });
             notify_event!(self, MetadataNotFound, resource);
+            dht_sub.unsubscribe().await;
             return Err(e)
         }
         let (mut chunked, metadata_seeder) = metadata_result.unwrap();
@@ -782,13 +791,20 @@ impl Fud {
         let resources_read = self.resources.read().await;
         let resource = match resources_read.get(hash) {
             Some(resource) => resource,
-            None => return Ok(()), // Resource was removed, abort
+            None => {
+                // Resource was removed, abort
+                dht_sub.unsubscribe().await;
+                return Ok(())
+            }
         };
         let files_vec: Vec<PathBuf> = resource.get_selected_files(&chunked);
         drop(resources_read);
 
         // Create all files (and all necessary directories)
-        create_all_files(&files_vec).await?;
+        if let Err(e) = create_all_files(&files_vec).await {
+            dht_sub.unsubscribe().await;
+            return Err(e)
+        }
 
         // Set resource status to `Verifying` and send a `MetadataDownloadCompleted` event
         let resource = update_resource!(hash, {
@@ -806,11 +822,15 @@ impl Fud {
         let chunk_hashes = resource.get_selected_chunks(&chunked);
 
         // Write all scraps to make sure the data on the filesystem is correct
-        self.write_scraps(&mut chunked, &chunk_hashes).await?;
+        if let Err(e) = self.write_scraps(&mut chunked, &chunk_hashes).await {
+            dht_sub.unsubscribe().await;
+            return Err(e)
+        }
 
         // Mark locally available chunks as such
         let verify_res = self.verify_chunks(&resource, &mut chunked).await;
         if let Err(e) = verify_res {
+            dht_sub.unsubscribe().await;
             error!(target: "fud::fetch_resource()", "Error while verifying chunks: {e}");
             return Err(e);
         }
@@ -827,8 +847,12 @@ impl Fud {
         // This fixes two edge-cases: a file that exactly ends at the end of
         // a chunk, and a file with no chunk.
         if !chunked.is_dir() {
-            let fs_metadata = fs::metadata(&path).await?;
-            if fs_metadata.len() > (chunked.len() * MAX_CHUNK_SIZE) as u64 {
+            let fs_metadata = fs::metadata(&path).await;
+            if let Err(e) = fs_metadata {
+                dht_sub.unsubscribe().await;
+                return Err(e.into());
+            }
+            if fs_metadata.unwrap().len() > (chunked.len() * MAX_CHUNK_SIZE) as u64 {
                 if let Ok(file) = OpenOptions::new().write(true).create(true).open(path).await {
                     let _ = file.set_len((chunked.len() * MAX_CHUNK_SIZE) as u64).await;
                 }
@@ -881,6 +905,7 @@ impl Fud {
 
         // If we don't need to download any chunk
         if missing_chunks.is_empty() {
+            dht_sub.unsubscribe().await;
             return download_completed(&chunked).await;
         }
 
@@ -892,11 +917,19 @@ impl Fud {
 
         // Start looking up seeders if we did not need to do it for the metadata
         if metadata_seeder.is_none() {
-            self.lookup_tx.send(*hash).await?;
+            if let Err(e) = self.lookup_tx.send(*hash).await {
+                dht_sub.unsubscribe().await;
+                return Err(e.into())
+            }
         }
 
         // Fetch missing chunks from seeders
-        let _ = fetch_chunks(self, hash, &mut chunked, metadata_seeder, &mut missing_chunks).await;
+        let _ =
+            fetch_chunks(self, hash, &mut chunked, &dht_sub, metadata_seeder, &mut missing_chunks)
+                .await;
+
+        // We don't need the DHT events sub anymore
+        dht_sub.unsubscribe().await;
 
         // Get chunked file from geode
         let mut chunked = self.geode.get(hash, path).await?;