TLA  Line data    Source code

       1                 : use anyhow::Context;
       2                 : use async_stream::{stream, try_stream};
       3                 : use aws_sdk_s3::{types::ObjectIdentifier, Client};
       4                 : use tokio_stream::Stream;
       5                 : 
       6                 : use crate::{list_objects_with_retries, RootTarget, S3Target, TenantShardTimelineId};
       7                 : use pageserver_api::shard::TenantShardId;
       8                 : use utils::id::TimelineId;
       9                 : 
      10                 : /// Given an S3 bucket, output a stream of TenantIds discovered via ListObjectsv2
      11 UBC           0 : pub fn stream_tenants<'a>(
      12               0 :     s3_client: &'a Client,
      13               0 :     target: &'a RootTarget,
      14               0 : ) -> impl Stream<Item = anyhow::Result<TenantShardId>> + 'a {
      15               0 :     try_stream! {
      16               0 :         let mut continuation_token = None;
      17               0 :         let tenants_target = target.tenants_root();
      18                 :         loop {
      19               0 :             let fetch_response =
      20               0 :                 list_objects_with_retries(s3_client, &tenants_target, continuation_token.clone()).await?;
      21                 : 
      22               0 :             let new_entry_ids = fetch_response
      23               0 :                 .common_prefixes()
      24               0 :                 .iter()
      25               0 :                 .filter_map(|prefix| prefix.prefix())
      26               0 :                 .filter_map(|prefix| -> Option<&str> {
      27               0 :                     prefix
      28               0 :                         .strip_prefix(&tenants_target.prefix_in_bucket)?
      29               0 :                         .strip_suffix('/')
      30               0 :                 }).map(|entry_id_str| {
      31               0 :                 entry_id_str
      32               0 :                     .parse()
      33               0 :                     .with_context(|| format!("Incorrect entry id str: {entry_id_str}"))
      34               0 :             });
      35                 : 
      36               0 :             for i in new_entry_ids {
      37               0 :                 yield i?;
      38                 :             }
      39                 : 
      40               0 :             match fetch_response.next_continuation_token {
      41               0 :                 Some(new_token) => continuation_token = Some(new_token),
      42               0 :                 None => break,
      43                 :             }
      44                 :         }
      45                 :     }
      46               0 : }
      47                 : 
      48                 : /// Given a TenantShardId, output a stream of the timelines within that tenant, discovered
      49                 : /// using ListObjectsv2.  The listing is done before the stream is built, so that this
      50                 : /// function can be used to generate concurrency on a stream using buffer_unordered.
      51               0 : pub async fn stream_tenant_timelines<'a>(
      52               0 :     s3_client: &'a Client,
      53               0 :     target: &'a RootTarget,
      54               0 :     tenant: TenantShardId,
      55               0 : ) -> anyhow::Result<impl Stream<Item = Result<TenantShardTimelineId, anyhow::Error>> + 'a> {
      56               0 :     let mut timeline_ids: Vec<Result<TimelineId, anyhow::Error>> = Vec::new();
      57               0 :     let mut continuation_token = None;
      58               0 :     let timelines_target = target.timelines_root(&tenant);
      59                 : 
      60                 :     loop {
      61               0 :         tracing::info!("Listing in {}", tenant);
      62               0 :         let fetch_response =
      63               0 :             list_objects_with_retries(s3_client, &timelines_target, continuation_token.clone())
      64               0 :                 .await;
      65               0 :         let fetch_response = match fetch_response {
      66               0 :             Err(e) => {
      67               0 :                 timeline_ids.push(Err(e));
      68               0 :                 break;
      69                 :             }
      70               0 :             Ok(r) => r,
      71               0 :         };
      72               0 : 
      73               0 :         let new_entry_ids = fetch_response
      74               0 :             .common_prefixes()
      75               0 :             .iter()
      76               0 :             .filter_map(|prefix| prefix.prefix())
      77               0 :             .filter_map(|prefix| -> Option<&str> {
      78               0 :                 prefix
      79               0 :                     .strip_prefix(&timelines_target.prefix_in_bucket)?
      80               0 :                     .strip_suffix('/')
      81               0 :             })
      82               0 :             .map(|entry_id_str| {
      83               0 :                 entry_id_str
      84               0 :                     .parse::<TimelineId>()
      85               0 :                     .with_context(|| format!("Incorrect entry id str: {entry_id_str}"))
      86               0 :             });
      87                 : 
      88               0 :         for i in new_entry_ids {
      89               0 :             timeline_ids.push(i);
      90               0 :         }
      91                 : 
      92               0 :         match fetch_response.next_continuation_token {
      93               0 :             Some(new_token) => continuation_token = Some(new_token),
      94               0 :             None => break,
      95                 :         }
      96                 :     }
      97                 : 
      98               0 :     tracing::info!("Yielding for {}", tenant);
      99               0 :     Ok(stream! {
     100               0 :         for i in timeline_ids {
     101               0 :             let id = i?;
     102               0 :             yield Ok(TenantShardTimelineId::new(tenant, id));
     103                 :         }
     104               0 :     })
     105               0 : }
     106                 : 
     107               0 : pub(crate) fn stream_listing<'a>(
     108               0 :     s3_client: &'a Client,
     109               0 :     target: &'a S3Target,
     110               0 : ) -> impl Stream<Item = anyhow::Result<ObjectIdentifier>> + 'a {
     111               0 :     try_stream! {
     112               0 :         let mut continuation_token = None;
     113                 :         loop {
     114               0 :             let fetch_response =
     115               0 :                 list_objects_with_retries(s3_client, target, continuation_token.clone()).await?;
     116                 : 
     117               0 :             if target.delimiter.is_empty() {
     118               0 :                 for object_key in fetch_response.contents().iter().filter_map(|object| object.key())
     119                 :                 {
     120               0 :                     let object_id = ObjectIdentifier::builder().key(object_key).build()?;
     121               0 :                     yield object_id;
     122                 :                 }
     123                 :             } else {
     124               0 :                 for prefix in fetch_response.common_prefixes().iter().filter_map(|p| p.prefix()) {
     125               0 :                     let object_id = ObjectIdentifier::builder().key(prefix).build()?;
     126               0 :                     yield object_id;
     127                 :                 }
     128                 :             }
     129                 : 
     130               0 :             match fetch_response.next_continuation_token {
     131               0 :                 Some(new_token) => continuation_token = Some(new_token),
     132               0 :                 None => break,
     133                 :             }
     134                 :         }
     135                 :     }
     136               0 : }