create-sia-app 0.1.1

package/template/rust/sia-sdk-rs/indexd/src/download.rs

@@ -0,0 +1,379 @@
+use std::collections::VecDeque;
+use std::fmt::Debug;
+use std::sync::Arc;
+use std::time::Duration;
+
+use log::debug;
+use sia::encryption::{EncryptionKey, encrypt_shard};
+use sia::erasure_coding::{self, ErasureCoder};
+use sia::rhp::SEGMENT_SIZE;
+use sia::signing::PrivateKey;
+use thiserror::Error;
+use tokio::io::AsyncWriteExt;
+use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc};
+use tokio::task::JoinSet;
+#[cfg(not(target_arch = "wasm32"))]
+use tokio::task::spawn_blocking;
+#[cfg(not(target_arch = "wasm32"))]
+use tokio::time::error::Elapsed;
+#[cfg(not(target_arch = "wasm32"))]
+use tokio::time::sleep;
+#[cfg(target_arch = "wasm32")]
+use crate::wasm_time::sleep;
+
+use crate::rhp4::RHP4Client;
+use crate::{Hosts, Object, Sector};
+
+/// Spawns a task on a [`JoinSet`]. Uses `spawn` on native (requires `Send`)
+/// and `spawn_local` on WASM (runs on the current [`tokio::task::LocalSet`]).
+macro_rules! join_set_spawn {
+    ($set:expr, $fut:expr) => {{
+        #[cfg(not(target_arch = "wasm32"))]
+        $set.spawn($fut);
+        #[cfg(target_arch = "wasm32")]
+        $set.spawn_local($fut);
+    }};
+}
+
+#[derive(Debug, Error)]
+pub enum DownloadError {
+    #[error("I/O error: {0}")]
+    Io(#[from] std::io::Error),
+
+    #[error("encoder error: {0}")]
+    Encoder(#[from] erasure_coding::Error),
+
+    #[error("not enough shards: {0}/{1}")]
+    NotEnoughShards(u8, u8),
+
+    #[error("invalid range: {0}-{1}")]
+    OutOfRange(usize, usize),
+
+    #[cfg(not(target_arch = "wasm32"))]
+    #[error("timeout error: {0}")]
+    Timeout(#[from] Elapsed),
+
+    #[error("semaphore error: {0}")]
+    SemaphoreError(#[from] tokio::sync::AcquireError),
+
+    #[error("join error: {0}")]
+    JoinError(#[from] tokio::task::JoinError),
+
+    #[error("invalid slab: {0}")]
+    InvalidSlab(String),
+
+    #[error("rhp4 error: {0}")]
+    Rhp4(#[from] crate::rhp4::Error),
+
+    #[error("custom error: {0}")]
+    Custom(String),
+}
+
+pub struct DownloadOptions {
+    /// Maximum number of concurrent sector downloads.
+    pub max_inflight: usize,
+    pub offset: u64,
+    pub length: Option<u64>,
+
+    /// Optional channel to notify when each slab is downloaded.
+    /// This can be used to implement progress reporting.
+    pub slab_downloaded: Option<mpsc::UnboundedSender<()>>,
+}
+
+impl Default for DownloadOptions {
+    fn default() -> Self {
+        Self {
+            #[cfg(target_arch = "wasm32")]
+            max_inflight: 2, // Browsers can't handle many concurrent WebTransport connections (no connection pooling)
+            #[cfg(not(target_arch = "wasm32"))]
+            max_inflight: 20,
+            offset: 0,
+            length: None,
+            slab_downloaded: None,
+        }
+    }
+}
+
+#[derive(Clone)]
+pub(crate) struct Downloader {
+    account_key: Arc<PrivateKey>,
+    hosts: Hosts,
+    transport: Arc<dyn RHP4Client>,
+    #[cfg(target_arch = "wasm32")]
+    default_max_inflight: usize,
+}
+
+struct SectorDownloadTask {
+    sector: Sector,
+    offset: u64,
+    length: u64,
+    index: usize,
+}
+
+impl Downloader {
+    // helper to pair a sector with its erasure-coded index.
+    // Required because [FuturesUnordered.push] does not
+    // preserve ordering and doesn't play nice with closures.
+    async fn try_download_sector(
+        _permit: OwnedSemaphorePermit,
+        transport: Arc<dyn RHP4Client>,
+        account_key: Arc<PrivateKey>,
+        task: SectorDownloadTask,
+    ) -> Result<(usize, Vec<u8>), DownloadError> {
+        let data = transport
+            .read_sector(
+                task.sector.host_key,
+                &account_key,
+                task.sector.root,
+                task.offset as usize,
+                task.length as usize,
+            )
+            .await?;
+        Ok((task.index, data.to_vec()))
+    }
+
+    #[cfg(not(target_arch = "wasm32"))]
+    pub fn new(hosts: Hosts, transport: Arc<dyn RHP4Client>, account_key: Arc<PrivateKey>) -> Self {
+        Self {
+            account_key,
+            hosts,
+            transport,
+        }
+    }
+
+    #[cfg(target_arch = "wasm32")]
+    pub fn new(
+        hosts: Hosts,
+        transport: Arc<dyn RHP4Client>,
+        account_key: Arc<PrivateKey>,
+        default_max_inflight: usize,
+    ) -> Self {
+        Self {
+            account_key,
+            hosts,
+            transport,
+            default_max_inflight,
+        }
+    }
+
+    /// Returns default download options with the configured max_inflight.
+    #[cfg(target_arch = "wasm32")]
+    pub fn default_options(&self) -> DownloadOptions {
+        DownloadOptions {
+            max_inflight: self.default_max_inflight,
+            ..Default::default()
+        }
+    }
+
+    /// Downloads the shards of an erasure-coded slab.
+    /// Successful shards will be decrypted using the
+    /// encryption_key.
+    ///
+    /// offset and limit are the byte range to download
+    /// from each sector.
+    async fn download_slab_shards(
+        &self,
+        encryption_key: &EncryptionKey,
+        sectors: &[Sector],
+        min_shards: u8,
+        offset: u64,
+        length: u64,
+        max_inflight: usize,
+    ) -> Result<Vec<Option<Vec<u8>>>, DownloadError> {
+        if sectors.len() < min_shards as usize {
+            return Err(DownloadError::InvalidSlab(format!(
+                "not enough sectors: have {}, need at least {}",
+                sectors.len(),
+                min_shards
+            )));
+        }
+
+        let semaphore = Arc::new(Semaphore::new(max_inflight));
+        let mut sectors = sectors
+            .iter()
+            .enumerate()
+            .map(|(index, s)| SectorDownloadTask {
+                sector: s.clone(),
+                offset,
+                length,
+                index,
+            })
+            .collect::<Vec<_>>();
+        self.hosts
+            .prioritize(&mut sectors, |task| &task.sector.host_key);
+        let total_shards = sectors.len();
+        let mut sectors = VecDeque::from(sectors);
+        let mut download_tasks = JoinSet::new();
+        for _ in 0..min_shards {
+            match sectors.pop_front() {
+                Some(task) => {
+                    let permit = semaphore.clone().acquire_owned().await?;
+                    let transport = self.transport.clone();
+                    let account_key = self.account_key.clone();
+                    join_set_spawn!(download_tasks, Self::try_download_sector(
+                        permit,
+                        transport,
+                        account_key,
+                        task,
+                    ));
+                }
+                None => panic!("not enough sectors to satisfy min_shards"), // should be unreachable
+            };
+        }
+
+        let mut successful: u8 = 0;
+        let mut shards = vec![None; total_shards];
+        let mut total_failures: usize = 0;
+        const MAX_TOTAL_FAILURES: usize = 30; // Give up after 30 failed attempts total
+
+        loop {
+            tokio::select! {
+                biased;
+                Some(res) = download_tasks.join_next() => {
+                    match res? { // safe because tasks are never cancelled
+                        Ok((index, mut data)) => {
+                            let encryption_key = encryption_key.clone();
+                            #[cfg(not(target_arch = "wasm32"))]
+                            let data = spawn_blocking(move || {
+                                encrypt_shard(&encryption_key, index as u8, offset as usize, &mut data);
+                                data
+                            }).await?;
+                            #[cfg(target_arch = "wasm32")]
+                            let data = {
+                                encrypt_shard(&encryption_key, index as u8, offset as usize, &mut data);
+                                data
+                            };
+                            shards[index] = Some(data);
+                            successful += 1;
+                            if successful >= min_shards {
+                               return Ok(shards);
+                            }
+                        }
+                        Err(e) => {
+                            total_failures += 1;
+                            debug!("sector download failed ({total_failures}/{MAX_TOTAL_FAILURES}): {:?}", e);
+
+                            if total_failures >= MAX_TOTAL_FAILURES {
+                                return Err(DownloadError::NotEnoughShards(successful, min_shards));
+                            }
+
+                            let rem = min_shards.saturating_sub(successful);
+                            if rem == 0 {
+                                return Ok(shards); // sanity check
+                            } else if download_tasks.len() + sectors.len() < rem as usize {
+                                return Err(DownloadError::NotEnoughShards(successful, min_shards));
+                            } else if download_tasks.len() <= rem as usize && let Some(task) = sectors.pop_front() {
+                                let transport = self.transport.clone();
+                                let account_key = self.account_key.clone();
+                                let permit = semaphore.clone().acquire_owned().await?;
+                                // only spawn additional download tasks if there
+                                // are not enough to satisfy the required number
+                                // of shards. The sleep arm will handle slow
+                                // hosts.
+                                join_set_spawn!(download_tasks, Self::try_download_sector(
+                                    permit,
+                                    transport,
+                                    account_key,
+                                    task,
+                                ));
+                            }
+                        }
+                    }
+                },
+                _ = sleep(Duration::from_secs(1)) => {
+                    if let Ok(racer_permit) = semaphore.clone().try_acquire_owned()
+                        && let Some(task) = sectors.pop_front() {
+                            let transport = self.transport.clone();
+                            let account_key = self.account_key.clone();
+                            join_set_spawn!(download_tasks, Self::try_download_sector(
+                                racer_permit,
+                                transport,
+                                account_key,
+                                task,
+                            ));
+                        }
+                }
+            }
+        }
+    }
+
+    /// Downloads the provided slabs and writes the decrypted data to the
+    /// provided writer.
+    pub async fn download<W: AsyncWriteExt + Unpin>(
+        &self,
+        w: &mut W,
+        object: &Object,
+        options: DownloadOptions,
+    ) -> Result<(), DownloadError> {
+        let mut w = object.writer(w, options.offset as usize);
+        let mut offset = options.offset;
+        let max_length = object.size();
+        let mut length = options.length.unwrap_or(max_length);
+        if offset > max_length || length == 0 {
+            return Ok(());
+        }
+
+        for slab in object.slabs() {
+            if length == 0 {
+                break;
+            }
+
+            let slab_length = slab.length as u64;
+            if offset >= slab_length {
+                offset -= slab_length;
+                continue;
+            }
+
+            // adjust slab range based on offset and length
+            let slab_offset = slab.offset as u64 + offset;
+            let slab_length = (slab_length - offset).min(length);
+            offset = 0;
+
+            // compute the sector aligned region to download
+            let chunk_size = SEGMENT_SIZE as u64 * slab.min_shards as u64;
+            let start = (slab_offset / chunk_size) * SEGMENT_SIZE as u64;
+            let end = (slab_offset + slab_length).div_ceil(chunk_size) * SEGMENT_SIZE as u64;
+            let shard_offset = start;
+            let shard_length = end - start;
+
+            let mut shards = self
+                .download_slab_shards(
+                    &slab.encryption_key,
+                    &slab.sectors,
+                    slab.min_shards,
+                    shard_offset,
+                    shard_length,
+                    options.max_inflight,
+                )
+                .await?;
+            let data_shards = slab.min_shards as usize;
+            let parity_shards = slab.sectors.len() - slab.min_shards as usize;
+            #[cfg(not(target_arch = "wasm32"))]
+            let shards = spawn_blocking(move || -> Result<Vec<Option<Vec<u8>>>, DownloadError> {
+                let rs = ErasureCoder::new(data_shards, parity_shards)?;
+                rs.reconstruct_data_shards(&mut shards)?;
+                Ok(shards)
+            })
+            .await??;
+            #[cfg(target_arch = "wasm32")]
+            let shards = {
+                let rs = ErasureCoder::new(data_shards, parity_shards)?;
+                rs.reconstruct_data_shards(&mut shards)?;
+                shards
+            };
+            ErasureCoder::write_data_shards(
+                &mut w,
+                &shards[..data_shards],
+                slab_offset as usize % (SEGMENT_SIZE * slab.min_shards as usize),
+                slab_length as usize,
+            )
+            .await?;
+            length -= slab_length;
+            if let Some(ref tx) = options.slab_downloaded {
+                let _ = tx.send(());
+            }
+        }
+        w.flush().await?;
+        Ok(())
+    }
+}