Line data    Source code

       1              : use std::sync::Arc;
       2              : 
       3              : use humantime::Duration;
       4              : use tokio::task::JoinSet;
       5              : use utils::id::TenantTimelineId;
       6              : 
       7              : use pageserver_client::mgmt_api::ForceAwaitLogicalSize;
       8              : 
       9            0 : #[derive(clap::Parser)]
      10              : pub(crate) struct Args {
      11              :     #[clap(long, default_value = "http://localhost:9898")]
      12            0 :     mgmt_api_endpoint: String,
      13              :     #[clap(long, default_value = "localhost:64000")]
      14            0 :     page_service_host_port: String,
      15              :     #[clap(long)]
      16              :     pageserver_jwt: Option<String>,
      17              :     #[clap(
      18              :         long,
      19              :         help = "if specified, poll mgmt api to check whether init logical size calculation has completed"
      20              :     )]
      21              :     poll_for_completion: Option<Duration>,
      22              :     #[clap(long)]
      23              :     limit_to_first_n_targets: Option<usize>,
      24            0 :     targets: Option<Vec<TenantTimelineId>>,
      25              : }
      26              : 
      27            0 : pub(crate) fn main(args: Args) -> anyhow::Result<()> {
      28            0 :     let rt = tokio::runtime::Builder::new_multi_thread()
      29            0 :         .enable_all()
      30            0 :         .build()
      31            0 :         .unwrap();
      32            0 : 
      33            0 :     let main_task = rt.spawn(main_impl(args));
      34            0 :     rt.block_on(main_task).unwrap()
      35            0 : }
      36              : 
      37            0 : async fn main_impl(args: Args) -> anyhow::Result<()> {
      38            0 :     let args: &'static Args = Box::leak(Box::new(args));
      39            0 : 
      40            0 :     let mgmt_api_client = Arc::new(pageserver_client::mgmt_api::Client::new(
      41            0 :         args.mgmt_api_endpoint.clone(),
      42            0 :         args.pageserver_jwt.as_deref(),
      43            0 :     ));
      44              : 
      45              :     // discover targets
      46            0 :     let timelines: Vec<TenantTimelineId> = crate::util::cli::targets::discover(
      47            0 :         &mgmt_api_client,
      48            0 :         crate::util::cli::targets::Spec {
      49            0 :             limit_to_first_n_targets: args.limit_to_first_n_targets,
      50            0 :             targets: args.targets.clone(),
      51            0 :         },
      52            0 :     )
      53            0 :     .await?;
      54              : 
      55              :     // kick it off
      56              : 
      57            0 :     let mut js = JoinSet::new();
      58            0 :     for tl in timelines {
      59            0 :         let mgmt_api_client = Arc::clone(&mgmt_api_client);
      60            0 :         js.spawn(async move {
      61            0 :             let info = mgmt_api_client
      62            0 :                 .timeline_info(tl.tenant_id, tl.timeline_id, ForceAwaitLogicalSize::Yes)
      63            0 :                 .await
      64            0 :                 .unwrap();
      65              : 
      66              :             // Polling should not be strictly required here since we await
      67              :             // for the initial logical size, however it's possible for the request
      68              :             // to land before the timeline is initialised. This results in an approximate
      69              :             // logical size.
      70            0 :             if let Some(period) = args.poll_for_completion {
      71            0 :                 let mut ticker = tokio::time::interval(period.into());
      72            0 :                 ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
      73            0 :                 let mut info = info;
      74            0 :                 while !info.current_logical_size_is_accurate {
      75            0 :                     ticker.tick().await;
      76            0 :                     info = mgmt_api_client
      77            0 :                         .timeline_info(tl.tenant_id, tl.timeline_id, ForceAwaitLogicalSize::Yes)
      78            0 :                         .await
      79            0 :                         .unwrap();
      80              :                 }
      81            0 :             }
      82            0 :         });
      83            0 :     }
      84            0 :     while let Some(res) = js.join_next().await {
      85            0 :         let _: () = res.unwrap();
      86            0 :     }
      87            0 :     Ok(())
      88            0 : }