TLA  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 UBC           0 : #[derive(clap::Parser)]
       8                 : pub(crate) struct Args {
       9                 :     #[clap(long, default_value = "http://localhost:9898")]
      10               0 :     mgmt_api_endpoint: String,
      11                 :     #[clap(long, default_value = "localhost:64000")]
      12               0 :     page_service_host_port: String,
      13                 :     #[clap(long)]
      14                 :     pageserver_jwt: Option<String>,
      15                 :     #[clap(
      16                 :         long,
      17                 :         help = "if specified, poll mgmt api to check whether init logical size calculation has completed"
      18                 :     )]
      19                 :     poll_for_completion: Option<Duration>,
      20                 :     #[clap(long)]
      21                 :     limit_to_first_n_targets: Option<usize>,
      22               0 :     targets: Option<Vec<TenantTimelineId>>,
      23                 : }
      24                 : 
      25               0 : pub(crate) fn main(args: Args) -> anyhow::Result<()> {
      26               0 :     let rt = tokio::runtime::Builder::new_multi_thread()
      27               0 :         .enable_all()
      28               0 :         .build()
      29               0 :         .unwrap();
      30               0 : 
      31               0 :     let main_task = rt.spawn(main_impl(args));
      32               0 :     rt.block_on(main_task).unwrap()
      33               0 : }
      34                 : 
      35               0 : async fn main_impl(args: Args) -> anyhow::Result<()> {
      36               0 :     let args: &'static Args = Box::leak(Box::new(args));
      37               0 : 
      38               0 :     let mgmt_api_client = Arc::new(pageserver_client::mgmt_api::Client::new(
      39               0 :         args.mgmt_api_endpoint.clone(),
      40               0 :         args.pageserver_jwt.as_deref(),
      41               0 :     ));
      42                 : 
      43                 :     // discover targets
      44               0 :     let timelines: Vec<TenantTimelineId> = crate::util::cli::targets::discover(
      45               0 :         &mgmt_api_client,
      46               0 :         crate::util::cli::targets::Spec {
      47               0 :             limit_to_first_n_targets: args.limit_to_first_n_targets,
      48               0 :             targets: args.targets.clone(),
      49               0 :         },
      50               0 :     )
      51               0 :     .await?;
      52                 : 
      53                 :     // kick it off
      54                 : 
      55               0 :     let mut js = JoinSet::new();
      56               0 :     for tl in timelines {
      57               0 :         let mgmt_api_client = Arc::clone(&mgmt_api_client);
      58               0 :         js.spawn(async move {
      59                 :             // TODO: API to explicitly trigger initial logical size computation.
      60                 :             // Should probably also avoid making it a side effect of timeline details to trigger initial logical size calculation.
      61                 :             // => https://github.com/neondatabase/neon/issues/6168
      62               0 :             let info = mgmt_api_client
      63               0 :                 .timeline_info(tl.tenant_id, tl.timeline_id)
      64               0 :                 .await
      65               0 :                 .unwrap();
      66                 : 
      67               0 :             if let Some(period) = args.poll_for_completion {
      68               0 :                 let mut ticker = tokio::time::interval(period.into());
      69               0 :                 ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
      70               0 :                 let mut info = info;
      71               0 :                 while !info.current_logical_size_is_accurate {
      72               0 :                     ticker.tick().await;
      73               0 :                     info = mgmt_api_client
      74               0 :                         .timeline_info(tl.tenant_id, tl.timeline_id)
      75               0 :                         .await
      76               0 :                         .unwrap();
      77                 :                 }
      78               0 :             }
      79               0 :         });
      80               0 :     }
      81               0 :     while let Some(res) = js.join_next().await {
      82               0 :         let _: () = res.unwrap();
      83               0 :     }
      84               0 :     Ok(())
      85               0 : }