feat(core): add prefetching to ArchiveBlockProvider to download next archive early

Author: 0xdeafbeef

core/src/block_strider/provider/archive_provider.rs

@@ -23,12 +23,14 @@ use crate::overlay_client::{Neighbour, PunishReason};
 #[serde(default)]
 pub struct ArchiveBlockProviderConfig {
     pub max_archive_to_memory_size: ByteSize,
+    pub enable_prefetching: bool,
 }
 
 impl Default for ArchiveBlockProviderConfig {
     fn default() -> Self {
         Self {
             max_archive_to_memory_size: ByteSize::mb(100),
+            enable_prefetching: true,
         }
     }
 }
@@ -61,7 +63,7 @@ impl ArchiveBlockProvider {
     }
 
     async fn get_next_block_impl(&self, block_id: &BlockId) -> OptionalBlockStuff {
-        let this = self.inner.as_ref();
+        let this = &self.inner;
 
         let next_mc_seqno = block_id.seqno + 1;
 
@@ -95,7 +97,7 @@ impl ArchiveBlockProvider {
     }
 
     async fn get_block_impl(&self, block_id_relation: &BlockIdRelation) -> OptionalBlockStuff {
-        let this = self.inner.as_ref();
+        let this = &self.inner;
 
         let block_id = block_id_relation.block_id;
         let mc_block_id = block_id_relation.mc_block_id;
@@ -161,7 +163,7 @@ struct Inner {
 }
 
 impl Inner {
-    async fn get_archive(&self, mc_seqno: u32) -> Option<(u32, ArchiveInfo)> {
+    async fn get_archive(self: &Arc<Self>, mc_seqno: u32) -> Option<(u32, ArchiveInfo)> {
         loop {
             let mut pending = 'pending: {
                 let mut guard = self.known_archives.lock();
@@ -171,6 +173,19 @@ impl Inner {
                     match value {
                         ArchiveSlot::Downloaded(info) => {
                             if info.archive.mc_block_ids.contains_key(&mc_seqno) {
+                                // Prefetch next archive if enabled
+                                if self.config.enable_prefetching {
+                                    if let Some((last_seqno, _)) =
+                                        info.archive.mc_block_ids.last_key_value()
+                                    {
+                                        let next_archive_start_seqno = last_seqno.saturating_add(1);
+                                        let this = self.clone();
+                                        tokio::spawn(async move {
+                                            this.try_prefetch_archive(next_archive_start_seqno)
+                                                .await;
+                                        });
+                                    }
+                                }
                                 return Some((*archive_key, info.clone()));
                             }
                         }
@@ -221,6 +236,21 @@ impl Inner {
             }
 
             if finished {
+                // Prefetch next archive if enabled
+                if self.config.enable_prefetching {
+                    if let Some(info) = res
+                        .as_ref()
+                        .and_then(|i| i.archive.mc_block_ids.last_key_value())
+                    {
+                        let (last_seqno, _) = info;
+                        let next_archive_start_seqno = last_seqno.saturating_add(1);
+                        let this = self.clone();
+                        // Fire and forget prefetch
+                        tokio::spawn(async move {
+                            this.try_prefetch_archive(next_archive_start_seqno).await;
+                        });
+                    }
+                }
                 return res.map(|info| (pending.archive_key, info));
             }
 
@@ -248,6 +278,39 @@ impl Inner {
         }
     }
 
+    /// Try to prefetch the archive containing `prefetch_seqno` if not already present or pending.
+    async fn try_prefetch_archive(&self, prefetch_seqno: u32) {
+        let mut guard = self.known_archives.lock();
+
+        // Check if archive already present or pending
+        for value in guard.values() {
+            match value {
+                ArchiveSlot::Downloaded(info) => {
+                    if info.archive.mc_block_ids.contains_key(&prefetch_seqno) {
+                        tracing::trace!(
+                            prefetch_seqno,
+                            "archive already downloaded, skipping prefetch"
+                        );
+                        return;
+                    }
+                }
+                ArchiveSlot::Pending(task) => {
+                    if task.archive_key == prefetch_seqno {
+                        tracing::trace!(
+                            prefetch_seqno,
+                            "archive download already pending, skipping prefetch"
+                        );
+                        return;
+                    }
+                }
+            }
+        }
+
+        let task = self.make_downloader().spawn(prefetch_seqno);
+        guard.insert(prefetch_seqno, ArchiveSlot::Pending(task));
+        tracing::debug!(prefetch_seqno, "starting archive prefetch");
+    }
+
     fn make_downloader(&self) -> ArchiveDownloader {
         ArchiveDownloader {
             client: self.client.clone(),