LCOV - code coverage report
Current view: top level - storage_scrubber/src - lib.rs (source / functions) Coverage Total Hit
Test: e5024a5c05016c30dec7897aca22d1040a340f63.info Lines: 0.0 % 376 0
Test Date: 2024-11-20 11:45:54 Functions: 0.0 % 81 0

            Line data    Source code
       1              : #![deny(unsafe_code)]
       2              : #![deny(clippy::undocumented_unsafe_blocks)]
       3              : pub mod checks;
       4              : pub mod cloud_admin_api;
       5              : pub mod find_large_objects;
       6              : pub mod garbage;
       7              : pub mod metadata_stream;
       8              : pub mod pageserver_physical_gc;
       9              : pub mod scan_pageserver_metadata;
      10              : pub mod scan_safekeeper_metadata;
      11              : pub mod tenant_snapshot;
      12              : 
      13              : use std::env;
      14              : use std::fmt::Display;
      15              : use std::sync::Arc;
      16              : use std::time::Duration;
      17              : 
      18              : use anyhow::Context;
      19              : use aws_config::retry::{RetryConfigBuilder, RetryMode};
      20              : use aws_sdk_s3::config::Region;
      21              : use aws_sdk_s3::error::DisplayErrorContext;
      22              : use aws_sdk_s3::Client;
      23              : 
      24              : use camino::{Utf8Path, Utf8PathBuf};
      25              : use clap::ValueEnum;
      26              : use futures::{Stream, StreamExt};
      27              : use pageserver::tenant::remote_timeline_client::{remote_tenant_path, remote_timeline_path};
      28              : use pageserver::tenant::TENANTS_SEGMENT_NAME;
      29              : use pageserver_api::shard::TenantShardId;
      30              : use remote_storage::{
      31              :     DownloadOpts, GenericRemoteStorage, Listing, ListingMode, RemotePath, RemoteStorageConfig,
      32              :     RemoteStorageKind, S3Config,
      33              : };
      34              : use reqwest::Url;
      35              : use serde::{Deserialize, Serialize};
      36              : use storage_controller_client::control_api;
      37              : use tokio::io::AsyncReadExt;
      38              : use tokio_util::sync::CancellationToken;
      39              : use tracing::{error, warn};
      40              : use tracing_appender::non_blocking::WorkerGuard;
      41              : use tracing_subscriber::{fmt, prelude::*, EnvFilter};
      42              : use utils::fs_ext;
      43              : use utils::id::{TenantId, TenantTimelineId, TimelineId};
      44              : 
      45              : const MAX_RETRIES: usize = 20;
      46              : const CLOUD_ADMIN_API_TOKEN_ENV_VAR: &str = "CLOUD_ADMIN_API_TOKEN";
      47              : 
      48              : #[derive(Debug, Clone)]
      49              : pub struct S3Target {
      50              :     pub desc_str: String,
      51              :     /// This `prefix_in_bucket` is only equal to the PS/SK config of the same
      52              :     /// name for the RootTarget: other instances of S3Target will have prefix_in_bucket
      53              :     /// with extra parts.
      54              :     pub prefix_in_bucket: String,
      55              :     pub delimiter: String,
      56              : }
      57              : 
      58              : /// Convenience for referring to timelines within a particular shard: more ergonomic
      59              : /// than using a 2-tuple.
      60              : ///
      61              : /// This is the shard-aware equivalent of TenantTimelineId.  It's defined here rather
      62              : /// than somewhere more broadly exposed, because this kind of thing is rarely needed
      63              : /// in the pageserver, as all timeline objects existing in the scope of a particular
      64              : /// tenant: the scrubber is different in that it handles collections of data referring to many
      65              : /// TenantShardTimelineIds in on place.
      66            0 : #[derive(Serialize, Deserialize, Debug, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord)]
      67              : pub struct TenantShardTimelineId {
      68              :     tenant_shard_id: TenantShardId,
      69              :     timeline_id: TimelineId,
      70              : }
      71              : 
      72              : impl TenantShardTimelineId {
      73            0 :     fn new(tenant_shard_id: TenantShardId, timeline_id: TimelineId) -> Self {
      74            0 :         Self {
      75            0 :             tenant_shard_id,
      76            0 :             timeline_id,
      77            0 :         }
      78            0 :     }
      79              : 
      80            0 :     fn as_tenant_timeline_id(&self) -> TenantTimelineId {
      81            0 :         TenantTimelineId::new(self.tenant_shard_id.tenant_id, self.timeline_id)
      82            0 :     }
      83              : }
      84              : 
      85              : impl Display for TenantShardTimelineId {
      86            0 :     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
      87            0 :         write!(f, "{}/{}", self.tenant_shard_id, self.timeline_id)
      88            0 :     }
      89              : }
      90              : 
      91            0 : #[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
      92              : pub enum TraversingDepth {
      93              :     Tenant,
      94              :     Timeline,
      95              : }
      96              : 
      97              : impl Display for TraversingDepth {
      98            0 :     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
      99            0 :         f.write_str(match self {
     100            0 :             Self::Tenant => "tenant",
     101            0 :             Self::Timeline => "timeline",
     102              :         })
     103            0 :     }
     104              : }
     105              : 
     106            0 : #[derive(ValueEnum, Clone, Copy, Eq, PartialEq, Debug, Serialize, Deserialize)]
     107              : pub enum NodeKind {
     108              :     Safekeeper,
     109              :     Pageserver,
     110              : }
     111              : 
     112              : impl NodeKind {
     113            0 :     fn as_str(&self) -> &'static str {
     114            0 :         match self {
     115            0 :             Self::Safekeeper => "safekeeper",
     116            0 :             Self::Pageserver => "pageserver",
     117              :         }
     118            0 :     }
     119              : }
     120              : 
     121              : impl Display for NodeKind {
     122            0 :     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
     123            0 :         f.write_str(self.as_str())
     124            0 :     }
     125              : }
     126              : 
     127              : impl S3Target {
     128            0 :     pub fn with_sub_segment(&self, new_segment: &str) -> Self {
     129            0 :         let mut new_self = self.clone();
     130            0 :         if new_self.prefix_in_bucket.is_empty() {
     131            0 :             new_self.prefix_in_bucket = format!("/{}/", new_segment);
     132            0 :         } else {
     133            0 :             if new_self.prefix_in_bucket.ends_with('/') {
     134            0 :                 new_self.prefix_in_bucket.pop();
     135            0 :             }
     136            0 :             new_self.prefix_in_bucket =
     137            0 :                 [&new_self.prefix_in_bucket, new_segment, ""].join(&new_self.delimiter);
     138              :         }
     139            0 :         new_self
     140            0 :     }
     141              : }
     142              : 
     143              : #[derive(Clone)]
     144              : pub enum RootTarget {
     145              :     Pageserver(S3Target),
     146              :     Safekeeper(S3Target),
     147              : }
     148              : 
     149              : impl RootTarget {
     150            0 :     pub fn tenants_root(&self) -> S3Target {
     151            0 :         match self {
     152            0 :             Self::Pageserver(root) => root.with_sub_segment(TENANTS_SEGMENT_NAME),
     153            0 :             Self::Safekeeper(root) => root.clone(),
     154              :         }
     155            0 :     }
     156              : 
     157            0 :     pub fn tenant_root(&self, tenant_id: &TenantShardId) -> S3Target {
     158            0 :         match self {
     159            0 :             Self::Pageserver(_) => self.tenants_root().with_sub_segment(&tenant_id.to_string()),
     160            0 :             Self::Safekeeper(_) => self
     161            0 :                 .tenants_root()
     162            0 :                 .with_sub_segment(&tenant_id.tenant_id.to_string()),
     163              :         }
     164            0 :     }
     165              : 
     166            0 :     pub(crate) fn tenant_shards_prefix(&self, tenant_id: &TenantId) -> S3Target {
     167            0 :         // Only pageserver remote storage contains tenant-shards
     168            0 :         assert!(matches!(self, Self::Pageserver(_)));
     169            0 :         let Self::Pageserver(root) = self else {
     170            0 :             panic!();
     171              :         };
     172              : 
     173            0 :         S3Target {
     174            0 :             desc_str: root.desc_str.clone(),
     175            0 :             prefix_in_bucket: format!(
     176            0 :                 "{}/{TENANTS_SEGMENT_NAME}/{tenant_id}",
     177            0 :                 root.prefix_in_bucket
     178            0 :             ),
     179            0 :             delimiter: root.delimiter.clone(),
     180            0 :         }
     181            0 :     }
     182              : 
     183            0 :     pub fn timelines_root(&self, tenant_id: &TenantShardId) -> S3Target {
     184            0 :         match self {
     185            0 :             Self::Pageserver(_) => self.tenant_root(tenant_id).with_sub_segment("timelines"),
     186            0 :             Self::Safekeeper(_) => self.tenant_root(tenant_id),
     187              :         }
     188            0 :     }
     189              : 
     190            0 :     pub fn timeline_root(&self, id: &TenantShardTimelineId) -> S3Target {
     191            0 :         self.timelines_root(&id.tenant_shard_id)
     192            0 :             .with_sub_segment(&id.timeline_id.to_string())
     193            0 :     }
     194              : 
     195              :     /// Given RemotePath "tenants/foo/timelines/bar/layerxyz", prefix it to a literal
     196              :     /// key in the S3 bucket.
     197            0 :     pub fn absolute_key(&self, key: &RemotePath) -> String {
     198            0 :         let root = match self {
     199            0 :             Self::Pageserver(root) => root,
     200            0 :             Self::Safekeeper(root) => root,
     201              :         };
     202              : 
     203            0 :         let prefix = &root.prefix_in_bucket;
     204            0 :         if prefix.ends_with('/') {
     205            0 :             format!("{prefix}{key}")
     206              :         } else {
     207            0 :             format!("{prefix}/{key}")
     208              :         }
     209            0 :     }
     210              : 
     211            0 :     pub fn desc_str(&self) -> &str {
     212            0 :         match self {
     213            0 :             Self::Pageserver(root) => &root.desc_str,
     214            0 :             Self::Safekeeper(root) => &root.desc_str,
     215              :         }
     216            0 :     }
     217              : 
     218            0 :     pub fn delimiter(&self) -> &str {
     219            0 :         match self {
     220            0 :             Self::Pageserver(root) => &root.delimiter,
     221            0 :             Self::Safekeeper(root) => &root.delimiter,
     222              :         }
     223            0 :     }
     224              : }
     225              : 
     226            0 : pub fn remote_timeline_path_id(id: &TenantShardTimelineId) -> RemotePath {
     227            0 :     remote_timeline_path(&id.tenant_shard_id, &id.timeline_id)
     228            0 : }
     229              : 
     230            0 : #[derive(Debug, Clone, Serialize, Deserialize)]
     231              : #[serde(deny_unknown_fields)]
     232              : pub struct BucketConfig(RemoteStorageConfig);
     233              : 
     234              : impl BucketConfig {
     235            0 :     pub fn from_env() -> anyhow::Result<Self> {
     236            0 :         if let Ok(legacy) = Self::from_env_legacy() {
     237            0 :             return Ok(legacy);
     238            0 :         }
     239            0 :         let config_toml =
     240            0 :             env::var("REMOTE_STORAGE_CONFIG").context("'REMOTE_STORAGE_CONFIG' retrieval")?;
     241            0 :         let remote_config = RemoteStorageConfig::from_toml_str(&config_toml)?;
     242            0 :         Ok(BucketConfig(remote_config))
     243            0 :     }
     244              : 
     245            0 :     fn from_env_legacy() -> anyhow::Result<Self> {
     246            0 :         let bucket_region = env::var("REGION").context("'REGION' param retrieval")?;
     247            0 :         let bucket_name = env::var("BUCKET").context("'BUCKET' param retrieval")?;
     248            0 :         let prefix_in_bucket = env::var("BUCKET_PREFIX").ok();
     249            0 :         let endpoint = env::var("AWS_ENDPOINT_URL").ok();
     250            0 :         // Create a json object which we then deserialize so that we don't
     251            0 :         // have to repeat all of the S3Config fields.
     252            0 :         let s3_config_json = serde_json::json!({
     253            0 :             "bucket_name": bucket_name,
     254            0 :             "bucket_region": bucket_region,
     255            0 :             "prefix_in_bucket": prefix_in_bucket,
     256            0 :             "endpoint": endpoint,
     257            0 :         });
     258            0 :         let config: RemoteStorageConfig = serde_json::from_value(s3_config_json)?;
     259            0 :         Ok(BucketConfig(config))
     260            0 :     }
     261            0 :     pub fn desc_str(&self) -> String {
     262            0 :         match &self.0.storage {
     263            0 :             RemoteStorageKind::LocalFs { local_path } => {
     264            0 :                 format!("local path {local_path}")
     265              :             }
     266            0 :             RemoteStorageKind::AwsS3(config) => format!(
     267            0 :                 "bucket {}, region {}",
     268            0 :                 config.bucket_name, config.bucket_region
     269            0 :             ),
     270            0 :             RemoteStorageKind::AzureContainer(config) => format!(
     271            0 :                 "bucket {}, storage account {:?}, region {}",
     272            0 :                 config.container_name, config.storage_account, config.container_region
     273            0 :             ),
     274              :         }
     275            0 :     }
     276            0 :     pub fn bucket_name(&self) -> Option<&str> {
     277            0 :         self.0.storage.bucket_name()
     278            0 :     }
     279              : }
     280              : 
     281            0 : #[derive(Debug, Clone, Serialize, Deserialize)]
     282              : #[serde(deny_unknown_fields)]
     283              : pub struct BucketConfigLegacy {
     284              :     pub region: String,
     285              :     pub bucket: String,
     286              :     pub prefix_in_bucket: Option<String>,
     287              : }
     288              : 
     289              : pub struct ControllerClientConfig {
     290              :     /// URL to storage controller.  e.g. http://127.0.0.1:1234 when using `neon_local`
     291              :     pub controller_api: Url,
     292              : 
     293              :     /// JWT token for authenticating with storage controller.  Requires scope 'scrubber' or 'admin'.
     294              :     pub controller_jwt: String,
     295              : }
     296              : 
     297              : impl ControllerClientConfig {
     298            0 :     pub fn build_client(self) -> control_api::Client {
     299            0 :         control_api::Client::new(self.controller_api, Some(self.controller_jwt))
     300            0 :     }
     301              : }
     302              : 
     303              : pub struct ConsoleConfig {
     304              :     pub token: String,
     305              :     pub base_url: Url,
     306              : }
     307              : 
     308              : impl ConsoleConfig {
     309            0 :     pub fn from_env() -> anyhow::Result<Self> {
     310            0 :         let base_url: Url = env::var("CLOUD_ADMIN_API_URL")
     311            0 :             .context("'CLOUD_ADMIN_API_URL' param retrieval")?
     312            0 :             .parse()
     313            0 :             .context("'CLOUD_ADMIN_API_URL' param parsing")?;
     314              : 
     315            0 :         let token = env::var(CLOUD_ADMIN_API_TOKEN_ENV_VAR)
     316            0 :             .context("'CLOUD_ADMIN_API_TOKEN' environment variable fetch")?;
     317              : 
     318            0 :         Ok(Self { base_url, token })
     319            0 :     }
     320              : }
     321              : 
     322            0 : pub fn init_logging(file_name: &str) -> Option<WorkerGuard> {
     323            0 :     let stderr_logs = fmt::Layer::new()
     324            0 :         .with_target(false)
     325            0 :         .with_writer(std::io::stderr);
     326              : 
     327            0 :     let disable_file_logging = match std::env::var("PAGESERVER_DISABLE_FILE_LOGGING") {
     328            0 :         Ok(s) => s == "1" || s.to_lowercase() == "true",
     329            0 :         Err(_) => false,
     330              :     };
     331              : 
     332            0 :     if disable_file_logging {
     333            0 :         tracing_subscriber::registry()
     334            0 :             .with(EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")))
     335            0 :             .with(stderr_logs)
     336            0 :             .init();
     337            0 :         None
     338              :     } else {
     339            0 :         let (file_writer, guard) =
     340            0 :             tracing_appender::non_blocking(tracing_appender::rolling::never("./logs/", file_name));
     341            0 :         let file_logs = fmt::Layer::new()
     342            0 :             .with_target(false)
     343            0 :             .with_ansi(false)
     344            0 :             .with_writer(file_writer);
     345            0 :         tracing_subscriber::registry()
     346            0 :             .with(EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")))
     347            0 :             .with(stderr_logs)
     348            0 :             .with(file_logs)
     349            0 :             .init();
     350            0 :         Some(guard)
     351              :     }
     352            0 : }
     353              : 
     354            0 : async fn init_s3_client(bucket_region: Region) -> Client {
     355            0 :     let mut retry_config_builder = RetryConfigBuilder::new();
     356            0 : 
     357            0 :     retry_config_builder
     358            0 :         .set_max_attempts(Some(3))
     359            0 :         .set_mode(Some(RetryMode::Adaptive));
     360              : 
     361            0 :     let config = aws_config::defaults(aws_config::BehaviorVersion::v2024_03_28())
     362            0 :         .region(bucket_region)
     363            0 :         .retry_config(retry_config_builder.build())
     364            0 :         .load()
     365            0 :         .await;
     366            0 :     Client::new(&config)
     367            0 : }
     368              : 
     369            0 : fn default_prefix_in_bucket(node_kind: NodeKind) -> &'static str {
     370            0 :     match node_kind {
     371            0 :         NodeKind::Pageserver => "pageserver/v1/",
     372            0 :         NodeKind::Safekeeper => "wal/",
     373              :     }
     374            0 : }
     375              : 
     376            0 : fn make_root_target(desc_str: String, prefix_in_bucket: String, node_kind: NodeKind) -> RootTarget {
     377            0 :     let s3_target = S3Target {
     378            0 :         desc_str,
     379            0 :         prefix_in_bucket,
     380            0 :         delimiter: "/".to_string(),
     381            0 :     };
     382            0 :     match node_kind {
     383            0 :         NodeKind::Pageserver => RootTarget::Pageserver(s3_target),
     384            0 :         NodeKind::Safekeeper => RootTarget::Safekeeper(s3_target),
     385              :     }
     386            0 : }
     387              : 
     388            0 : async fn init_remote_s3(
     389            0 :     bucket_config: S3Config,
     390            0 :     node_kind: NodeKind,
     391            0 : ) -> anyhow::Result<(Arc<Client>, RootTarget)> {
     392            0 :     let bucket_region = Region::new(bucket_config.bucket_region);
     393            0 :     let s3_client = Arc::new(init_s3_client(bucket_region).await);
     394            0 :     let default_prefix = default_prefix_in_bucket(node_kind).to_string();
     395            0 : 
     396            0 :     let s3_root = make_root_target(
     397            0 :         bucket_config.bucket_name,
     398            0 :         bucket_config.prefix_in_bucket.unwrap_or(default_prefix),
     399            0 :         node_kind,
     400            0 :     );
     401            0 : 
     402            0 :     Ok((s3_client, s3_root))
     403            0 : }
     404              : 
     405            0 : async fn init_remote(
     406            0 :     mut storage_config: BucketConfig,
     407            0 :     node_kind: NodeKind,
     408            0 : ) -> anyhow::Result<(GenericRemoteStorage, RootTarget)> {
     409            0 :     let desc_str = storage_config.desc_str();
     410            0 : 
     411            0 :     let default_prefix = default_prefix_in_bucket(node_kind).to_string();
     412            0 : 
     413            0 :     match &mut storage_config.0.storage {
     414            0 :         RemoteStorageKind::AwsS3(ref mut config) => {
     415            0 :             config.prefix_in_bucket.get_or_insert(default_prefix);
     416            0 :         }
     417            0 :         RemoteStorageKind::AzureContainer(ref mut config) => {
     418            0 :             config.prefix_in_container.get_or_insert(default_prefix);
     419            0 :         }
     420            0 :         RemoteStorageKind::LocalFs { .. } => (),
     421              :     }
     422              : 
     423              :     // We already pass the prefix to the remote client above
     424            0 :     let prefix_in_root_target = String::new();
     425            0 :     let root_target = make_root_target(desc_str, prefix_in_root_target, node_kind);
     426              : 
     427            0 :     let client = GenericRemoteStorage::from_config(&storage_config.0).await?;
     428            0 :     Ok((client, root_target))
     429            0 : }
     430              : 
     431              : /// Listing possibly large amounts of keys in a streaming fashion.
     432            0 : fn stream_objects_with_retries<'a>(
     433            0 :     storage_client: &'a GenericRemoteStorage,
     434            0 :     listing_mode: ListingMode,
     435            0 :     s3_target: &'a S3Target,
     436            0 : ) -> impl Stream<Item = Result<Listing, anyhow::Error>> + 'a {
     437            0 :     async_stream::stream! {
     438            0 :         let mut trial = 0;
     439            0 :         let cancel = CancellationToken::new();
     440            0 :         let prefix_str = &s3_target
     441            0 :             .prefix_in_bucket
     442            0 :             .strip_prefix("/")
     443            0 :             .unwrap_or(&s3_target.prefix_in_bucket);
     444            0 :         let prefix = RemotePath::from_string(prefix_str)?;
     445            0 :         let mut list_stream =
     446            0 :             storage_client.list_streaming(Some(&prefix), listing_mode, None, &cancel);
     447            0 :         while let Some(res) = list_stream.next().await {
     448            0 :             match res {
     449            0 :                 Err(err) => {
     450            0 :                     let yield_err = if err.is_permanent() {
     451            0 :                         true
     452            0 :                     } else {
     453            0 :                         let backoff_time = 1 << trial.min(5);
     454            0 :                         tokio::time::sleep(Duration::from_secs(backoff_time)).await;
     455            0 :                         trial += 1;
     456            0 :                         trial == MAX_RETRIES - 1
     457            0 :                     };
     458            0 :                     if yield_err {
     459            0 :                         yield Err(err)
     460            0 :                             .with_context(|| format!("Failed to list objects {MAX_RETRIES} times"));
     461            0 :                         break;
     462            0 :                     }
     463            0 :                 }
     464            0 :                 Ok(res) => {
     465            0 :                     trial = 0;
     466            0 :                     yield Ok(res);
     467            0 :                 }
     468            0 :             }
     469            0 :         }
     470            0 :     }
     471            0 : }
     472              : 
     473              : /// If you want to list a bounded amount of prefixes or keys. For larger numbers of keys/prefixes,
     474              : /// use [`stream_objects_with_retries`] instead.
     475            0 : async fn list_objects_with_retries(
     476            0 :     remote_client: &GenericRemoteStorage,
     477            0 :     listing_mode: ListingMode,
     478            0 :     s3_target: &S3Target,
     479            0 : ) -> anyhow::Result<Listing> {
     480            0 :     let cancel = CancellationToken::new();
     481            0 :     let prefix_str = &s3_target
     482            0 :         .prefix_in_bucket
     483            0 :         .strip_prefix("/")
     484            0 :         .unwrap_or(&s3_target.prefix_in_bucket);
     485            0 :     let prefix = RemotePath::from_string(prefix_str)?;
     486            0 :     for trial in 0..MAX_RETRIES {
     487            0 :         match remote_client
     488            0 :             .list(Some(&prefix), listing_mode, None, &cancel)
     489            0 :             .await
     490              :         {
     491            0 :             Ok(response) => return Ok(response),
     492            0 :             Err(e) => {
     493            0 :                 if trial == MAX_RETRIES - 1 {
     494            0 :                     return Err(e)
     495            0 :                         .with_context(|| format!("Failed to list objects {MAX_RETRIES} times"));
     496            0 :                 }
     497            0 :                 warn!(
     498            0 :                     "list_objects_v2 query failed: bucket_name={}, prefix={}, delimiter={}, error={}",
     499            0 :                     remote_client.bucket_name().unwrap_or_default(),
     500            0 :                     s3_target.prefix_in_bucket,
     501            0 :                     s3_target.delimiter,
     502            0 :                     DisplayErrorContext(e),
     503              :                 );
     504            0 :                 let backoff_time = 1 << trial.min(5);
     505            0 :                 tokio::time::sleep(Duration::from_secs(backoff_time)).await;
     506              :             }
     507              :         }
     508              :     }
     509            0 :     panic!("MAX_RETRIES is not allowed to be 0");
     510            0 : }
     511              : 
     512            0 : async fn download_object_with_retries(
     513            0 :     remote_client: &GenericRemoteStorage,
     514            0 :     key: &RemotePath,
     515            0 : ) -> anyhow::Result<Vec<u8>> {
     516            0 :     let cancel = CancellationToken::new();
     517            0 :     for trial in 0..MAX_RETRIES {
     518            0 :         let mut buf = Vec::new();
     519            0 :         let download = match remote_client
     520            0 :             .download(key, &DownloadOpts::default(), &cancel)
     521            0 :             .await
     522              :         {
     523            0 :             Ok(response) => response,
     524            0 :             Err(e) => {
     525            0 :                 error!("Failed to download object for key {key}: {e}");
     526            0 :                 let backoff_time = 1 << trial.min(5);
     527            0 :                 tokio::time::sleep(Duration::from_secs(backoff_time)).await;
     528            0 :                 continue;
     529              :             }
     530              :         };
     531              : 
     532            0 :         match tokio_util::io::StreamReader::new(download.download_stream)
     533            0 :             .read_to_end(&mut buf)
     534            0 :             .await
     535              :         {
     536            0 :             Ok(bytes_read) => {
     537            0 :                 tracing::debug!("Downloaded {bytes_read} bytes for object {key}");
     538            0 :                 return Ok(buf);
     539              :             }
     540            0 :             Err(e) => {
     541            0 :                 error!("Failed to stream object body for key {key}: {e}");
     542            0 :                 let backoff_time = 1 << trial.min(5);
     543            0 :                 tokio::time::sleep(Duration::from_secs(backoff_time)).await;
     544              :             }
     545              :         }
     546              :     }
     547              : 
     548            0 :     anyhow::bail!("Failed to download objects with key {key} {MAX_RETRIES} times")
     549            0 : }
     550              : 
     551            0 : async fn download_object_to_file_s3(
     552            0 :     s3_client: &Client,
     553            0 :     bucket_name: &str,
     554            0 :     key: &str,
     555            0 :     version_id: Option<&str>,
     556            0 :     local_path: &Utf8Path,
     557            0 : ) -> anyhow::Result<()> {
     558            0 :     let tmp_path = Utf8PathBuf::from(format!("{local_path}.tmp"));
     559            0 :     for _ in 0..MAX_RETRIES {
     560            0 :         tokio::fs::remove_file(&tmp_path)
     561            0 :             .await
     562            0 :             .or_else(fs_ext::ignore_not_found)?;
     563              : 
     564            0 :         let mut file = tokio::fs::File::create(&tmp_path)
     565            0 :             .await
     566            0 :             .context("Opening output file")?;
     567              : 
     568            0 :         let request = s3_client.get_object().bucket(bucket_name).key(key);
     569              : 
     570            0 :         let request = match version_id {
     571            0 :             Some(version_id) => request.version_id(version_id),
     572            0 :             None => request,
     573              :         };
     574              : 
     575            0 :         let response_stream = match request.send().await {
     576            0 :             Ok(response) => response,
     577            0 :             Err(e) => {
     578            0 :                 error!(
     579            0 :                     "Failed to download object for key {key} version {}: {e:#}",
     580            0 :                     version_id.unwrap_or("")
     581              :                 );
     582            0 :                 tokio::time::sleep(Duration::from_secs(1)).await;
     583            0 :                 continue;
     584              :             }
     585              :         };
     586              : 
     587            0 :         let mut read_stream = response_stream.body.into_async_read();
     588            0 : 
     589            0 :         tokio::io::copy(&mut read_stream, &mut file).await?;
     590              : 
     591            0 :         tokio::fs::rename(&tmp_path, local_path).await?;
     592            0 :         return Ok(());
     593              :     }
     594              : 
     595            0 :     anyhow::bail!("Failed to download objects with key {key} {MAX_RETRIES} times")
     596            0 : }
        

Generated by: LCOV version 2.1-beta