Line data    Source code

       1              : use pageserver_api::models::{AuxFilePolicy, TenantConfig, TenantConfigRequest};
       2              : use pageserver_api::shard::TenantShardId;
       3              : use utils::id::TenantTimelineId;
       4              : use utils::lsn::Lsn;
       5              : 
       6              : use std::collections::HashMap;
       7              : use std::sync::Arc;
       8              : use std::time::Instant;
       9              : 
      10              : /// Ingest aux files into the pageserver.
      11            0 : #[derive(clap::Parser)]
      12              : pub(crate) struct Args {
      13              :     #[clap(long, default_value = "http://localhost:9898")]
      14            0 :     mgmt_api_endpoint: String,
      15              :     #[clap(long, default_value = "postgres://postgres@localhost:64000")]
      16            0 :     page_service_connstring: String,
      17              :     #[clap(long)]
      18              :     pageserver_jwt: Option<String>,
      19              : 
      20            0 :     targets: Option<Vec<TenantTimelineId>>,
      21              : }
      22              : 
      23            0 : pub(crate) fn main(args: Args) -> anyhow::Result<()> {
      24            0 :     let rt = tokio::runtime::Builder::new_multi_thread()
      25            0 :         .enable_all()
      26            0 :         .build()
      27            0 :         .unwrap();
      28            0 : 
      29            0 :     let main_task = rt.spawn(main_impl(args));
      30            0 :     rt.block_on(main_task).unwrap()
      31            0 : }
      32              : 
      33            0 : async fn main_impl(args: Args) -> anyhow::Result<()> {
      34            0 :     let args: &'static Args = Box::leak(Box::new(args));
      35            0 : 
      36            0 :     let mgmt_api_client = Arc::new(pageserver_client::mgmt_api::Client::new(
      37            0 :         args.mgmt_api_endpoint.clone(),
      38            0 :         args.pageserver_jwt.as_deref(),
      39            0 :     ));
      40              : 
      41              :     // discover targets
      42            0 :     let timelines: Vec<TenantTimelineId> = crate::util::cli::targets::discover(
      43            0 :         &mgmt_api_client,
      44            0 :         crate::util::cli::targets::Spec {
      45            0 :             limit_to_first_n_targets: None,
      46              :             targets: {
      47            0 :                 if let Some(targets) = &args.targets {
      48            0 :                     if targets.len() != 1 {
      49            0 :                         anyhow::bail!("must specify exactly one target");
      50            0 :                     }
      51            0 :                     Some(targets.clone())
      52              :                 } else {
      53            0 :                     None
      54              :                 }
      55              :             },
      56              :         },
      57              :     )
      58            0 :     .await?;
      59              : 
      60            0 :     let timeline = timelines[0];
      61            0 :     let tenant_shard_id = TenantShardId::unsharded(timeline.tenant_id);
      62            0 :     let timeline_id = timeline.timeline_id;
      63            0 : 
      64            0 :     println!("operating on timeline {}", timeline);
      65            0 : 
      66            0 :     mgmt_api_client
      67            0 :         .tenant_config(&TenantConfigRequest {
      68            0 :             tenant_id: timeline.tenant_id,
      69            0 :             config: TenantConfig {
      70            0 :                 switch_aux_file_policy: Some(AuxFilePolicy::V2),
      71            0 :                 ..Default::default()
      72            0 :             },
      73            0 :         })
      74            0 :         .await?;
      75              : 
      76            0 :     for batch in 0..100 {
      77            0 :         let items = (0..100)
      78            0 :             .map(|id| {
      79            0 :                 (
      80            0 :                     format!("pg_logical/mappings/{:03}.{:03}", batch, id),
      81            0 :                     format!("{:08}", id),
      82            0 :                 )
      83            0 :             })
      84            0 :             .collect::<HashMap<_, _>>();
      85            0 :         let file_cnt = items.len();
      86            0 :         mgmt_api_client
      87            0 :             .ingest_aux_files(tenant_shard_id, timeline_id, items)
      88            0 :             .await?;
      89            0 :         println!("ingested {file_cnt} files");
      90              :     }
      91              : 
      92            0 :     for _ in 0..100 {
      93            0 :         let start = Instant::now();
      94            0 :         let files = mgmt_api_client
      95            0 :             .list_aux_files(tenant_shard_id, timeline_id, Lsn(Lsn::MAX.0 - 1))
      96            0 :             .await?;
      97            0 :         println!(
      98            0 :             "{} files found in {}s",
      99            0 :             files.len(),
     100            0 :             start.elapsed().as_secs_f64()
     101            0 :         );
     102              :     }
     103              : 
     104            0 :     anyhow::Ok(())
     105            0 : }