LCOV - code coverage report
Current view: top level - control_plane/storcon_cli/src - main.rs (source / functions) Coverage Total Hit
Test: 1b0a6a0c05cee5a7de360813c8034804e105ce1c.info Lines: 0.0 % 1103 0
Test Date: 2025-03-12 00:01:28 Functions: 0.0 % 139 0

            Line data    Source code
       1              : use std::collections::{HashMap, HashSet};
       2              : use std::path::PathBuf;
       3              : use std::str::FromStr;
       4              : use std::time::Duration;
       5              : 
       6              : use clap::{Parser, Subcommand};
       7              : use futures::StreamExt;
       8              : use pageserver_api::controller_api::{
       9              :     AvailabilityZone, MigrationConfig, NodeAvailabilityWrapper, NodeConfigureRequest,
      10              :     NodeDescribeResponse, NodeRegisterRequest, NodeSchedulingPolicy, NodeShardResponse,
      11              :     PlacementPolicy, SafekeeperDescribeResponse, SafekeeperSchedulingPolicyRequest,
      12              :     ShardSchedulingPolicy, ShardsPreferredAzsRequest, ShardsPreferredAzsResponse,
      13              :     SkSchedulingPolicy, TenantCreateRequest, TenantDescribeResponse, TenantPolicyRequest,
      14              :     TenantShardMigrateRequest, TenantShardMigrateResponse,
      15              : };
      16              : use pageserver_api::models::{
      17              :     EvictionPolicy, EvictionPolicyLayerAccessThreshold, LocationConfigSecondary, ShardParameters,
      18              :     TenantConfig, TenantConfigPatchRequest, TenantConfigRequest, TenantShardSplitRequest,
      19              :     TenantShardSplitResponse,
      20              : };
      21              : use pageserver_api::shard::{ShardStripeSize, TenantShardId};
      22              : use pageserver_client::mgmt_api::{self};
      23              : use reqwest::{Method, StatusCode, Url};
      24              : use storage_controller_client::control_api::Client;
      25              : use utils::id::{NodeId, TenantId, TimelineId};
      26              : 
      27              : #[derive(Subcommand, Debug)]
      28              : enum Command {
      29              :     /// Register a pageserver with the storage controller.  This shouldn't usually be necessary,
      30              :     /// since pageservers auto-register when they start up
      31              :     NodeRegister {
      32              :         #[arg(long)]
      33            0 :         node_id: NodeId,
      34              : 
      35              :         #[arg(long)]
      36            0 :         listen_pg_addr: String,
      37              :         #[arg(long)]
      38            0 :         listen_pg_port: u16,
      39              : 
      40              :         #[arg(long)]
      41            0 :         listen_http_addr: String,
      42              :         #[arg(long)]
      43            0 :         listen_http_port: u16,
      44              :         #[arg(long)]
      45              :         listen_https_port: Option<u16>,
      46              : 
      47              :         #[arg(long)]
      48            0 :         availability_zone_id: String,
      49              :     },
      50              : 
      51              :     /// Modify a node's configuration in the storage controller
      52              :     NodeConfigure {
      53              :         #[arg(long)]
      54            0 :         node_id: NodeId,
      55              : 
      56              :         /// Availability is usually auto-detected based on heartbeats.  Set 'offline' here to
      57              :         /// manually mark a node offline
      58              :         #[arg(long)]
      59              :         availability: Option<NodeAvailabilityArg>,
      60              :         /// Scheduling policy controls whether tenant shards may be scheduled onto this node.
      61              :         #[arg(long)]
      62              :         scheduling: Option<NodeSchedulingPolicy>,
      63              :     },
      64              :     NodeDelete {
      65              :         #[arg(long)]
      66            0 :         node_id: NodeId,
      67              :     },
      68              :     /// Modify a tenant's policies in the storage controller
      69              :     TenantPolicy {
      70              :         #[arg(long)]
      71            0 :         tenant_id: TenantId,
      72              :         /// Placement policy controls whether a tenant is `detached`, has only a secondary location (`secondary`),
      73              :         /// or is in the normal attached state with N secondary locations (`attached:N`)
      74              :         #[arg(long)]
      75              :         placement: Option<PlacementPolicyArg>,
      76              :         /// Scheduling policy enables pausing the controller's scheduling activity involving this tenant.  `active` is normal,
      77              :         /// `essential` disables optimization scheduling changes, `pause` disables all scheduling changes, and `stop` prevents
      78              :         /// all reconciliation activity including for scheduling changes already made.  `pause` and `stop` can make a tenant
      79              :         /// unavailable, and are only for use in emergencies.
      80              :         #[arg(long)]
      81              :         scheduling: Option<ShardSchedulingPolicyArg>,
      82              :     },
      83              :     /// List nodes known to the storage controller
      84              :     Nodes {},
      85              :     /// List tenants known to the storage controller
      86              :     Tenants {
      87              :         /// If this field is set, it will list the tenants on a specific node
      88              :         node_id: Option<NodeId>,
      89              :     },
      90              :     /// Create a new tenant in the storage controller, and by extension on pageservers.
      91              :     TenantCreate {
      92              :         #[arg(long)]
      93            0 :         tenant_id: TenantId,
      94              :     },
      95              :     /// Delete a tenant in the storage controller, and by extension on pageservers.
      96              :     TenantDelete {
      97              :         #[arg(long)]
      98            0 :         tenant_id: TenantId,
      99              :     },
     100              :     /// Split an existing tenant into a higher number of shards than its current shard count.
     101              :     TenantShardSplit {
     102              :         #[arg(long)]
     103            0 :         tenant_id: TenantId,
     104              :         #[arg(long)]
     105            0 :         shard_count: u8,
     106              :         /// Optional, in 8kiB pages.  e.g. set 2048 for 16MB stripes.
     107              :         #[arg(long)]
     108              :         stripe_size: Option<u32>,
     109              :     },
     110              :     /// Migrate the attached location for a tenant shard to a specific pageserver.
     111              :     TenantShardMigrate {
     112              :         #[arg(long)]
     113            0 :         tenant_shard_id: TenantShardId,
     114              :         #[arg(long)]
     115            0 :         node: NodeId,
     116            0 :         #[arg(long, default_value_t = true, action = clap::ArgAction::Set)]
     117            0 :         prewarm: bool,
     118            0 :         #[arg(long, default_value_t = false, action = clap::ArgAction::Set)]
     119            0 :         override_scheduler: bool,
     120              :     },
     121              :     /// Watch the location of a tenant shard evolve, e.g. while expecting it to migrate
     122              :     TenantShardWatch {
     123              :         #[arg(long)]
     124            0 :         tenant_shard_id: TenantShardId,
     125              :     },
     126              :     /// Migrate the secondary location for a tenant shard to a specific pageserver.
     127              :     TenantShardMigrateSecondary {
     128              :         #[arg(long)]
     129            0 :         tenant_shard_id: TenantShardId,
     130              :         #[arg(long)]
     131            0 :         node: NodeId,
     132              :     },
     133              :     /// Cancel any ongoing reconciliation for this shard
     134              :     TenantShardCancelReconcile {
     135              :         #[arg(long)]
     136            0 :         tenant_shard_id: TenantShardId,
     137              :     },
     138              :     /// Set the pageserver tenant configuration of a tenant: this is the configuration structure
     139              :     /// that is passed through to pageservers, and does not affect storage controller behavior.
     140              :     /// Any previous tenant configs are overwritten.
     141              :     SetTenantConfig {
     142              :         #[arg(long)]
     143            0 :         tenant_id: TenantId,
     144              :         #[arg(long)]
     145            0 :         config: String,
     146              :     },
     147              :     /// Patch the pageserver tenant configuration of a tenant. Any fields with null values in the
     148              :     /// provided JSON are unset from the tenant config and all fields with non-null values are set.
     149              :     /// Unspecified fields are not changed.
     150              :     PatchTenantConfig {
     151              :         #[arg(long)]
     152            0 :         tenant_id: TenantId,
     153              :         #[arg(long)]
     154            0 :         config: String,
     155              :     },
     156              :     /// Print details about a particular tenant, including all its shards' states.
     157              :     TenantDescribe {
     158              :         #[arg(long)]
     159            0 :         tenant_id: TenantId,
     160              :     },
     161              :     /// For a tenant which hasn't been onboarded to the storage controller yet, add it in secondary
     162              :     /// mode so that it can warm up content on a pageserver.
     163              :     TenantWarmup {
     164              :         #[arg(long)]
     165            0 :         tenant_id: TenantId,
     166              :     },
     167              :     TenantSetPreferredAz {
     168              :         #[arg(long)]
     169            0 :         tenant_id: TenantId,
     170              :         #[arg(long)]
     171              :         preferred_az: Option<String>,
     172              :     },
     173              :     /// Uncleanly drop a tenant from the storage controller: this doesn't delete anything from pageservers. Appropriate
     174              :     /// if you e.g. used `tenant-warmup` by mistake on a tenant ID that doesn't really exist, or is in some other region.
     175              :     TenantDrop {
     176              :         #[arg(long)]
     177            0 :         tenant_id: TenantId,
     178              :         #[arg(long)]
     179            0 :         unclean: bool,
     180              :     },
     181              :     NodeDrop {
     182              :         #[arg(long)]
     183            0 :         node_id: NodeId,
     184              :         #[arg(long)]
     185            0 :         unclean: bool,
     186              :     },
     187              :     TenantSetTimeBasedEviction {
     188              :         #[arg(long)]
     189            0 :         tenant_id: TenantId,
     190              :         #[arg(long)]
     191            0 :         period: humantime::Duration,
     192              :         #[arg(long)]
     193            0 :         threshold: humantime::Duration,
     194              :     },
     195              :     // Migrate away from a set of specified pageservers by moving the primary attachments to pageservers
     196              :     // outside of the specified set.
     197              :     BulkMigrate {
     198              :         // Set of pageserver node ids to drain.
     199              :         #[arg(long)]
     200            0 :         nodes: Vec<NodeId>,
     201              :         // Optional: migration concurrency (default is 8)
     202              :         #[arg(long)]
     203              :         concurrency: Option<usize>,
     204              :         // Optional: maximum number of shards to migrate
     205              :         #[arg(long)]
     206              :         max_shards: Option<usize>,
     207              :         // Optional: when set to true, nothing is migrated, but the plan is printed to stdout
     208              :         #[arg(long)]
     209              :         dry_run: Option<bool>,
     210              :     },
     211              :     /// Start draining the specified pageserver.
     212              :     /// The drain is complete when the schedulling policy returns to active.
     213              :     StartDrain {
     214              :         #[arg(long)]
     215            0 :         node_id: NodeId,
     216              :     },
     217              :     /// Cancel draining the specified pageserver and wait for `timeout`
     218              :     /// for the operation to be canceled. May be retried.
     219              :     CancelDrain {
     220              :         #[arg(long)]
     221            0 :         node_id: NodeId,
     222              :         #[arg(long)]
     223            0 :         timeout: humantime::Duration,
     224              :     },
     225              :     /// Start filling the specified pageserver.
     226              :     /// The drain is complete when the schedulling policy returns to active.
     227              :     StartFill {
     228              :         #[arg(long)]
     229            0 :         node_id: NodeId,
     230              :     },
     231              :     /// Cancel filling the specified pageserver and wait for `timeout`
     232              :     /// for the operation to be canceled. May be retried.
     233              :     CancelFill {
     234              :         #[arg(long)]
     235            0 :         node_id: NodeId,
     236              :         #[arg(long)]
     237            0 :         timeout: humantime::Duration,
     238              :     },
     239              :     /// List safekeepers known to the storage controller
     240              :     Safekeepers {},
     241              :     /// Set the scheduling policy of the specified safekeeper
     242              :     SafekeeperScheduling {
     243              :         #[arg(long)]
     244            0 :         node_id: NodeId,
     245              :         #[arg(long)]
     246            0 :         scheduling_policy: SkSchedulingPolicyArg,
     247              :     },
     248              :     /// Downloads any missing heatmap layers for all shard for a given timeline
     249              :     DownloadHeatmapLayers {
     250              :         /// Tenant ID or tenant shard ID. When an unsharded tenant ID is specified,
     251              :         /// the operation is performed on all shards. When a sharded tenant ID is
     252              :         /// specified, the operation is only performed on the specified shard.
     253              :         #[arg(long)]
     254            0 :         tenant_shard_id: TenantShardId,
     255              :         #[arg(long)]
     256            0 :         timeline_id: TimelineId,
     257              :         /// Optional: Maximum download concurrency (default is 16)
     258              :         #[arg(long)]
     259              :         concurrency: Option<usize>,
     260              :     },
     261              : }
     262              : 
     263              : #[derive(Parser)]
     264              : #[command(
     265              :     author,
     266              :     version,
     267              :     about,
     268              :     long_about = "CLI for Storage Controller Support/Debug"
     269              : )]
     270              : #[command(arg_required_else_help(true))]
     271              : struct Cli {
     272              :     #[arg(long)]
     273              :     /// URL to storage controller.  e.g. http://127.0.0.1:1234 when using `neon_local`
     274            0 :     api: Url,
     275              : 
     276              :     #[arg(long)]
     277              :     /// JWT token for authenticating with storage controller.  Depending on the API used, this
     278              :     /// should have either `pageserverapi` or `admin` scopes: for convenience, you should mint
     279              :     /// a token with both scopes to use with this tool.
     280              :     jwt: Option<String>,
     281              : 
     282              :     #[arg(long)]
     283              :     /// Trusted root CA certificate to use in https APIs.
     284              :     ssl_ca_file: Option<PathBuf>,
     285              : 
     286              :     #[command(subcommand)]
     287              :     command: Command,
     288              : }
     289              : 
     290              : #[derive(Debug, Clone)]
     291              : struct PlacementPolicyArg(PlacementPolicy);
     292              : 
     293              : impl FromStr for PlacementPolicyArg {
     294              :     type Err = anyhow::Error;
     295              : 
     296            0 :     fn from_str(s: &str) -> Result<Self, Self::Err> {
     297            0 :         match s {
     298            0 :             "detached" => Ok(Self(PlacementPolicy::Detached)),
     299            0 :             "secondary" => Ok(Self(PlacementPolicy::Secondary)),
     300            0 :             _ if s.starts_with("attached:") => {
     301            0 :                 let mut splitter = s.split(':');
     302            0 :                 let _prefix = splitter.next().unwrap();
     303            0 :                 match splitter.next().and_then(|s| s.parse::<usize>().ok()) {
     304            0 :                     Some(n) => Ok(Self(PlacementPolicy::Attached(n))),
     305            0 :                     None => Err(anyhow::anyhow!(
     306            0 :                         "Invalid format '{s}', a valid example is 'attached:1'"
     307            0 :                     )),
     308              :                 }
     309              :             }
     310            0 :             _ => Err(anyhow::anyhow!(
     311            0 :                 "Unknown placement policy '{s}', try detached,secondary,attached:<n>"
     312            0 :             )),
     313              :         }
     314            0 :     }
     315              : }
     316              : 
     317              : #[derive(Debug, Clone)]
     318              : struct SkSchedulingPolicyArg(SkSchedulingPolicy);
     319              : 
     320              : impl FromStr for SkSchedulingPolicyArg {
     321              :     type Err = anyhow::Error;
     322              : 
     323            0 :     fn from_str(s: &str) -> Result<Self, Self::Err> {
     324            0 :         SkSchedulingPolicy::from_str(s).map(Self)
     325            0 :     }
     326              : }
     327              : 
     328              : #[derive(Debug, Clone)]
     329              : struct ShardSchedulingPolicyArg(ShardSchedulingPolicy);
     330              : 
     331              : impl FromStr for ShardSchedulingPolicyArg {
     332              :     type Err = anyhow::Error;
     333              : 
     334            0 :     fn from_str(s: &str) -> Result<Self, Self::Err> {
     335            0 :         match s {
     336            0 :             "active" => Ok(Self(ShardSchedulingPolicy::Active)),
     337            0 :             "essential" => Ok(Self(ShardSchedulingPolicy::Essential)),
     338            0 :             "pause" => Ok(Self(ShardSchedulingPolicy::Pause)),
     339            0 :             "stop" => Ok(Self(ShardSchedulingPolicy::Stop)),
     340            0 :             _ => Err(anyhow::anyhow!(
     341            0 :                 "Unknown scheduling policy '{s}', try active,essential,pause,stop"
     342            0 :             )),
     343              :         }
     344            0 :     }
     345              : }
     346              : 
     347              : #[derive(Debug, Clone)]
     348              : struct NodeAvailabilityArg(NodeAvailabilityWrapper);
     349              : 
     350              : impl FromStr for NodeAvailabilityArg {
     351              :     type Err = anyhow::Error;
     352              : 
     353            0 :     fn from_str(s: &str) -> Result<Self, Self::Err> {
     354            0 :         match s {
     355            0 :             "active" => Ok(Self(NodeAvailabilityWrapper::Active)),
     356            0 :             "offline" => Ok(Self(NodeAvailabilityWrapper::Offline)),
     357            0 :             _ => Err(anyhow::anyhow!("Unknown availability state '{s}'")),
     358              :         }
     359            0 :     }
     360              : }
     361              : 
     362            0 : async fn wait_for_scheduling_policy<F>(
     363            0 :     client: Client,
     364            0 :     node_id: NodeId,
     365            0 :     timeout: Duration,
     366            0 :     f: F,
     367            0 : ) -> anyhow::Result<NodeSchedulingPolicy>
     368            0 : where
     369            0 :     F: Fn(NodeSchedulingPolicy) -> bool,
     370            0 : {
     371            0 :     let waiter = tokio::time::timeout(timeout, async move {
     372              :         loop {
     373            0 :             let node = client
     374            0 :                 .dispatch::<(), NodeDescribeResponse>(
     375            0 :                     Method::GET,
     376            0 :                     format!("control/v1/node/{node_id}"),
     377            0 :                     None,
     378            0 :                 )
     379            0 :                 .await?;
     380              : 
     381            0 :             if f(node.scheduling) {
     382            0 :                 return Ok::<NodeSchedulingPolicy, mgmt_api::Error>(node.scheduling);
     383            0 :             }
     384              :         }
     385            0 :     });
     386            0 : 
     387            0 :     Ok(waiter.await??)
     388            0 : }
     389              : 
     390              : #[tokio::main]
     391            0 : async fn main() -> anyhow::Result<()> {
     392            0 :     let cli = Cli::parse();
     393            0 : 
     394            0 :     let storcon_client = Client::new(cli.api.clone(), cli.jwt.clone());
     395            0 : 
     396            0 :     let ssl_ca_cert = match &cli.ssl_ca_file {
     397            0 :         Some(ssl_ca_file) => {
     398            0 :             let buf = tokio::fs::read(ssl_ca_file).await?;
     399            0 :             Some(reqwest::Certificate::from_pem(&buf)?)
     400            0 :         }
     401            0 :         None => None,
     402            0 :     };
     403            0 : 
     404            0 :     let mut trimmed = cli.api.to_string();
     405            0 :     trimmed.pop();
     406            0 :     let vps_client = mgmt_api::Client::new(trimmed, cli.jwt.as_deref(), ssl_ca_cert)?;
     407            0 : 
     408            0 :     match cli.command {
     409            0 :         Command::NodeRegister {
     410            0 :             node_id,
     411            0 :             listen_pg_addr,
     412            0 :             listen_pg_port,
     413            0 :             listen_http_addr,
     414            0 :             listen_http_port,
     415            0 :             listen_https_port,
     416            0 :             availability_zone_id,
     417            0 :         } => {
     418            0 :             storcon_client
     419            0 :                 .dispatch::<_, ()>(
     420            0 :                     Method::POST,
     421            0 :                     "control/v1/node".to_string(),
     422            0 :                     Some(NodeRegisterRequest {
     423            0 :                         node_id,
     424            0 :                         listen_pg_addr,
     425            0 :                         listen_pg_port,
     426            0 :                         listen_http_addr,
     427            0 :                         listen_http_port,
     428            0 :                         listen_https_port,
     429            0 :                         availability_zone_id: AvailabilityZone(availability_zone_id),
     430            0 :                     }),
     431            0 :                 )
     432            0 :                 .await?;
     433            0 :         }
     434            0 :         Command::TenantCreate { tenant_id } => {
     435            0 :             storcon_client
     436            0 :                 .dispatch::<_, ()>(
     437            0 :                     Method::POST,
     438            0 :                     "v1/tenant".to_string(),
     439            0 :                     Some(TenantCreateRequest {
     440            0 :                         new_tenant_id: TenantShardId::unsharded(tenant_id),
     441            0 :                         generation: None,
     442            0 :                         shard_parameters: ShardParameters::default(),
     443            0 :                         placement_policy: Some(PlacementPolicy::Attached(1)),
     444            0 :                         config: TenantConfig::default(),
     445            0 :                     }),
     446            0 :                 )
     447            0 :                 .await?;
     448            0 :         }
     449            0 :         Command::TenantDelete { tenant_id } => {
     450            0 :             let status = vps_client
     451            0 :                 .tenant_delete(TenantShardId::unsharded(tenant_id))
     452            0 :                 .await?;
     453            0 :             tracing::info!("Delete status: {}", status);
     454            0 :         }
     455            0 :         Command::Nodes {} => {
     456            0 :             let mut resp = storcon_client
     457            0 :                 .dispatch::<(), Vec<NodeDescribeResponse>>(
     458            0 :                     Method::GET,
     459            0 :                     "control/v1/node".to_string(),
     460            0 :                     None,
     461            0 :                 )
     462            0 :                 .await?;
     463            0 : 
     464            0 :             resp.sort_by(|a, b| a.listen_http_addr.cmp(&b.listen_http_addr));
     465            0 : 
     466            0 :             let mut table = comfy_table::Table::new();
     467            0 :             table.set_header(["Id", "Hostname", "AZ", "Scheduling", "Availability"]);
     468            0 :             for node in resp {
     469            0 :                 table.add_row([
     470            0 :                     format!("{}", node.id),
     471            0 :                     node.listen_http_addr,
     472            0 :                     node.availability_zone_id,
     473            0 :                     format!("{:?}", node.scheduling),
     474            0 :                     format!("{:?}", node.availability),
     475            0 :                 ]);
     476            0 :             }
     477            0 :             println!("{table}");
     478            0 :         }
     479            0 :         Command::NodeConfigure {
     480            0 :             node_id,
     481            0 :             availability,
     482            0 :             scheduling,
     483            0 :         } => {
     484            0 :             let req = NodeConfigureRequest {
     485            0 :                 node_id,
     486            0 :                 availability: availability.map(|a| a.0),
     487            0 :                 scheduling,
     488            0 :             };
     489            0 :             storcon_client
     490            0 :                 .dispatch::<_, ()>(
     491            0 :                     Method::PUT,
     492            0 :                     format!("control/v1/node/{node_id}/config"),
     493            0 :                     Some(req),
     494            0 :                 )
     495            0 :                 .await?;
     496            0 :         }
     497            0 :         Command::Tenants {
     498            0 :             node_id: Some(node_id),
     499            0 :         } => {
     500            0 :             let describe_response = storcon_client
     501            0 :                 .dispatch::<(), NodeShardResponse>(
     502            0 :                     Method::GET,
     503            0 :                     format!("control/v1/node/{node_id}/shards"),
     504            0 :                     None,
     505            0 :                 )
     506            0 :                 .await?;
     507            0 :             let shards = describe_response.shards;
     508            0 :             let mut table = comfy_table::Table::new();
     509            0 :             table.set_header([
     510            0 :                 "Shard",
     511            0 :                 "Intended Primary/Secondary",
     512            0 :                 "Observed Primary/Secondary",
     513            0 :             ]);
     514            0 :             for shard in shards {
     515            0 :                 table.add_row([
     516            0 :                     format!("{}", shard.tenant_shard_id),
     517            0 :                     match shard.is_intended_secondary {
     518            0 :                         None => "".to_string(),
     519            0 :                         Some(true) => "Secondary".to_string(),
     520            0 :                         Some(false) => "Primary".to_string(),
     521            0 :                     },
     522            0 :                     match shard.is_observed_secondary {
     523            0 :                         None => "".to_string(),
     524            0 :                         Some(true) => "Secondary".to_string(),
     525            0 :                         Some(false) => "Primary".to_string(),
     526            0 :                     },
     527            0 :                 ]);
     528            0 :             }
     529            0 :             println!("{table}");
     530            0 :         }
     531            0 :         Command::Tenants { node_id: None } => {
     532            0 :             // Set up output formatting
     533            0 :             let mut table = comfy_table::Table::new();
     534            0 :             table.set_header([
     535            0 :                 "TenantId",
     536            0 :                 "Preferred AZ",
     537            0 :                 "ShardCount",
     538            0 :                 "StripeSize",
     539            0 :                 "Placement",
     540            0 :                 "Scheduling",
     541            0 :             ]);
     542            0 : 
     543            0 :             // Pagination loop over listing API
     544            0 :             let mut start_after = None;
     545            0 :             const LIMIT: usize = 1000;
     546            0 :             loop {
     547            0 :                 let path = match start_after {
     548            0 :                     None => format!("control/v1/tenant?limit={LIMIT}"),
     549            0 :                     Some(start_after) => {
     550            0 :                         format!("control/v1/tenant?limit={LIMIT}&start_after={start_after}")
     551            0 :                     }
     552            0 :                 };
     553            0 : 
     554            0 :                 let resp = storcon_client
     555            0 :                     .dispatch::<(), Vec<TenantDescribeResponse>>(Method::GET, path, None)
     556            0 :                     .await?;
     557            0 : 
     558            0 :                 if resp.is_empty() {
     559            0 :                     // End of data reached
     560            0 :                     break;
     561            0 :                 }
     562            0 : 
     563            0 :                 // Give some visual feedback while we're building up the table (comfy_table doesn't have
     564            0 :                 // streaming output)
     565            0 :                 if resp.len() >= LIMIT {
     566            0 :                     eprint!(".");
     567            0 :                 }
     568            0 : 
     569            0 :                 start_after = Some(resp.last().unwrap().tenant_id);
     570            0 : 
     571            0 :                 for tenant in resp {
     572            0 :                     let shard_zero = tenant.shards.into_iter().next().unwrap();
     573            0 :                     table.add_row([
     574            0 :                         format!("{}", tenant.tenant_id),
     575            0 :                         shard_zero
     576            0 :                             .preferred_az_id
     577            0 :                             .as_ref()
     578            0 :                             .cloned()
     579            0 :                             .unwrap_or("".to_string()),
     580            0 :                         format!("{}", shard_zero.tenant_shard_id.shard_count.literal()),
     581            0 :                         format!("{:?}", tenant.stripe_size),
     582            0 :                         format!("{:?}", tenant.policy),
     583            0 :                         format!("{:?}", shard_zero.scheduling_policy),
     584            0 :                     ]);
     585            0 :                 }
     586            0 :             }
     587            0 : 
     588            0 :             // Terminate progress dots
     589            0 :             if table.row_count() > LIMIT {
     590            0 :                 eprint!("");
     591            0 :             }
     592            0 : 
     593            0 :             println!("{table}");
     594            0 :         }
     595            0 :         Command::TenantPolicy {
     596            0 :             tenant_id,
     597            0 :             placement,
     598            0 :             scheduling,
     599            0 :         } => {
     600            0 :             let req = TenantPolicyRequest {
     601            0 :                 scheduling: scheduling.map(|s| s.0),
     602            0 :                 placement: placement.map(|p| p.0),
     603            0 :             };
     604            0 :             storcon_client
     605            0 :                 .dispatch::<_, ()>(
     606            0 :                     Method::PUT,
     607            0 :                     format!("control/v1/tenant/{tenant_id}/policy"),
     608            0 :                     Some(req),
     609            0 :                 )
     610            0 :                 .await?;
     611            0 :         }
     612            0 :         Command::TenantShardSplit {
     613            0 :             tenant_id,
     614            0 :             shard_count,
     615            0 :             stripe_size,
     616            0 :         } => {
     617            0 :             let req = TenantShardSplitRequest {
     618            0 :                 new_shard_count: shard_count,
     619            0 :                 new_stripe_size: stripe_size.map(ShardStripeSize),
     620            0 :             };
     621            0 : 
     622            0 :             let response = storcon_client
     623            0 :                 .dispatch::<TenantShardSplitRequest, TenantShardSplitResponse>(
     624            0 :                     Method::PUT,
     625            0 :                     format!("control/v1/tenant/{tenant_id}/shard_split"),
     626            0 :                     Some(req),
     627            0 :                 )
     628            0 :                 .await?;
     629            0 :             println!(
     630            0 :                 "Split tenant {} into {} shards: {}",
     631            0 :                 tenant_id,
     632            0 :                 shard_count,
     633            0 :                 response
     634            0 :                     .new_shards
     635            0 :                     .iter()
     636            0 :                     .map(|s| format!("{:?}", s))
     637            0 :                     .collect::<Vec<_>>()
     638            0 :                     .join(",")
     639            0 :             );
     640            0 :         }
     641            0 :         Command::TenantShardMigrate {
     642            0 :             tenant_shard_id,
     643            0 :             node,
     644            0 :             prewarm,
     645            0 :             override_scheduler,
     646            0 :         } => {
     647            0 :             let migration_config = MigrationConfig {
     648            0 :                 prewarm,
     649            0 :                 override_scheduler,
     650            0 :                 ..Default::default()
     651            0 :             };
     652            0 : 
     653            0 :             let req = TenantShardMigrateRequest {
     654            0 :                 node_id: node,
     655            0 :                 origin_node_id: None,
     656            0 :                 migration_config,
     657            0 :             };
     658            0 : 
     659            0 :             match storcon_client
     660            0 :                 .dispatch::<TenantShardMigrateRequest, TenantShardMigrateResponse>(
     661            0 :                     Method::PUT,
     662            0 :                     format!("control/v1/tenant/{tenant_shard_id}/migrate"),
     663            0 :                     Some(req),
     664            0 :                 )
     665            0 :                 .await
     666            0 :             {
     667            0 :                 Err(mgmt_api::Error::ApiError(StatusCode::PRECONDITION_FAILED, msg)) => {
     668            0 :                     anyhow::bail!(
     669            0 :                         "Migration to {node} rejected, may require `--force` ({}) ",
     670            0 :                         msg
     671            0 :                     );
     672            0 :                 }
     673            0 :                 Err(e) => return Err(e.into()),
     674            0 :                 Ok(_) => {}
     675            0 :             }
     676            0 : 
     677            0 :             watch_tenant_shard(storcon_client, tenant_shard_id, Some(node)).await?;
     678            0 :         }
     679            0 :         Command::TenantShardWatch { tenant_shard_id } => {
     680            0 :             watch_tenant_shard(storcon_client, tenant_shard_id, None).await?;
     681            0 :         }
     682            0 :         Command::TenantShardMigrateSecondary {
     683            0 :             tenant_shard_id,
     684            0 :             node,
     685            0 :         } => {
     686            0 :             let req = TenantShardMigrateRequest {
     687            0 :                 node_id: node,
     688            0 :                 origin_node_id: None,
     689            0 :                 migration_config: MigrationConfig::default(),
     690            0 :             };
     691            0 : 
     692            0 :             storcon_client
     693            0 :                 .dispatch::<TenantShardMigrateRequest, TenantShardMigrateResponse>(
     694            0 :                     Method::PUT,
     695            0 :                     format!("control/v1/tenant/{tenant_shard_id}/migrate_secondary"),
     696            0 :                     Some(req),
     697            0 :                 )
     698            0 :                 .await?;
     699            0 :         }
     700            0 :         Command::TenantShardCancelReconcile { tenant_shard_id } => {
     701            0 :             storcon_client
     702            0 :                 .dispatch::<(), ()>(
     703            0 :                     Method::PUT,
     704            0 :                     format!("control/v1/tenant/{tenant_shard_id}/cancel_reconcile"),
     705            0 :                     None,
     706            0 :                 )
     707            0 :                 .await?;
     708            0 :         }
     709            0 :         Command::SetTenantConfig { tenant_id, config } => {
     710            0 :             let tenant_conf = serde_json::from_str(&config)?;
     711            0 : 
     712            0 :             vps_client
     713            0 :                 .set_tenant_config(&TenantConfigRequest {
     714            0 :                     tenant_id,
     715            0 :                     config: tenant_conf,
     716            0 :                 })
     717            0 :                 .await?;
     718            0 :         }
     719            0 :         Command::PatchTenantConfig { tenant_id, config } => {
     720            0 :             let tenant_conf = serde_json::from_str(&config)?;
     721            0 : 
     722            0 :             vps_client
     723            0 :                 .patch_tenant_config(&TenantConfigPatchRequest {
     724            0 :                     tenant_id,
     725            0 :                     config: tenant_conf,
     726            0 :                 })
     727            0 :                 .await?;
     728            0 :         }
     729            0 :         Command::TenantDescribe { tenant_id } => {
     730            0 :             let TenantDescribeResponse {
     731            0 :                 tenant_id,
     732            0 :                 shards,
     733            0 :                 stripe_size,
     734            0 :                 policy,
     735            0 :                 config,
     736            0 :             } = storcon_client
     737            0 :                 .dispatch::<(), TenantDescribeResponse>(
     738            0 :                     Method::GET,
     739            0 :                     format!("control/v1/tenant/{tenant_id}"),
     740            0 :                     None,
     741            0 :                 )
     742            0 :                 .await?;
     743            0 : 
     744            0 :             let nodes = storcon_client
     745            0 :                 .dispatch::<(), Vec<NodeDescribeResponse>>(
     746            0 :                     Method::GET,
     747            0 :                     "control/v1/node".to_string(),
     748            0 :                     None,
     749            0 :                 )
     750            0 :                 .await?;
     751            0 :             let nodes = nodes
     752            0 :                 .into_iter()
     753            0 :                 .map(|n| (n.id, n))
     754            0 :                 .collect::<HashMap<_, _>>();
     755            0 : 
     756            0 :             println!("Tenant {tenant_id}");
     757            0 :             let mut table = comfy_table::Table::new();
     758            0 :             table.add_row(["Policy", &format!("{:?}", policy)]);
     759            0 :             table.add_row(["Stripe size", &format!("{:?}", stripe_size)]);
     760            0 :             table.add_row(["Config", &serde_json::to_string_pretty(&config).unwrap()]);
     761            0 :             println!("{table}");
     762            0 :             println!("Shards:");
     763            0 :             let mut table = comfy_table::Table::new();
     764            0 :             table.set_header([
     765            0 :                 "Shard",
     766            0 :                 "Attached",
     767            0 :                 "Attached AZ",
     768            0 :                 "Secondary",
     769            0 :                 "Last error",
     770            0 :                 "status",
     771            0 :             ]);
     772            0 :             for shard in shards {
     773            0 :                 let secondary = shard
     774            0 :                     .node_secondary
     775            0 :                     .iter()
     776            0 :                     .map(|n| format!("{}", n))
     777            0 :                     .collect::<Vec<_>>()
     778            0 :                     .join(",");
     779            0 : 
     780            0 :                 let mut status_parts = Vec::new();
     781            0 :                 if shard.is_reconciling {
     782            0 :                     status_parts.push("reconciling");
     783            0 :                 }
     784            0 : 
     785            0 :                 if shard.is_pending_compute_notification {
     786            0 :                     status_parts.push("pending_compute");
     787            0 :                 }
     788            0 : 
     789            0 :                 if shard.is_splitting {
     790            0 :                     status_parts.push("splitting");
     791            0 :                 }
     792            0 :                 let status = status_parts.join(",");
     793            0 : 
     794            0 :                 let attached_node = shard
     795            0 :                     .node_attached
     796            0 :                     .as_ref()
     797            0 :                     .map(|id| nodes.get(id).expect("Shard references nonexistent node"));
     798            0 : 
     799            0 :                 table.add_row([
     800            0 :                     format!("{}", shard.tenant_shard_id),
     801            0 :                     attached_node
     802            0 :                         .map(|n| format!("{} ({})", n.listen_http_addr, n.id))
     803            0 :                         .unwrap_or(String::new()),
     804            0 :                     attached_node
     805            0 :                         .map(|n| n.availability_zone_id.clone())
     806            0 :                         .unwrap_or(String::new()),
     807            0 :                     secondary,
     808            0 :                     shard.last_error,
     809            0 :                     status,
     810            0 :                 ]);
     811            0 :             }
     812            0 :             println!("{table}");
     813            0 :         }
     814            0 :         Command::TenantSetPreferredAz {
     815            0 :             tenant_id,
     816            0 :             preferred_az,
     817            0 :         } => {
     818            0 :             // First learn about the tenant's shards
     819            0 :             let describe_response = storcon_client
     820            0 :                 .dispatch::<(), TenantDescribeResponse>(
     821            0 :                     Method::GET,
     822            0 :                     format!("control/v1/tenant/{tenant_id}"),
     823            0 :                     None,
     824            0 :                 )
     825            0 :                 .await?;
     826            0 : 
     827            0 :             // Learn about nodes to validate the AZ ID
     828            0 :             let nodes = storcon_client
     829            0 :                 .dispatch::<(), Vec<NodeDescribeResponse>>(
     830            0 :                     Method::GET,
     831            0 :                     "control/v1/node".to_string(),
     832            0 :                     None,
     833            0 :                 )
     834            0 :                 .await?;
     835            0 : 
     836            0 :             if let Some(preferred_az) = &preferred_az {
     837            0 :                 let azs = nodes
     838            0 :                     .into_iter()
     839            0 :                     .map(|n| (n.availability_zone_id))
     840            0 :                     .collect::<HashSet<_>>();
     841            0 :                 if !azs.contains(preferred_az) {
     842            0 :                     anyhow::bail!(
     843            0 :                         "AZ {} not found on any node: known AZs are: {:?}",
     844            0 :                         preferred_az,
     845            0 :                         azs
     846            0 :                     );
     847            0 :                 }
     848            0 :             } else {
     849            0 :                 // Make it obvious to the user that since they've omitted an AZ, we're clearing it
     850            0 :                 eprintln!("Clearing preferred AZ for tenant {}", tenant_id);
     851            0 :             }
     852            0 : 
     853            0 :             // Construct a request that modifies all the tenant's shards
     854            0 :             let req = ShardsPreferredAzsRequest {
     855            0 :                 preferred_az_ids: describe_response
     856            0 :                     .shards
     857            0 :                     .into_iter()
     858            0 :                     .map(|s| {
     859            0 :                         (
     860            0 :                             s.tenant_shard_id,
     861            0 :                             preferred_az.clone().map(AvailabilityZone),
     862            0 :                         )
     863            0 :                     })
     864            0 :                     .collect(),
     865            0 :             };
     866            0 :             storcon_client
     867            0 :                 .dispatch::<ShardsPreferredAzsRequest, ShardsPreferredAzsResponse>(
     868            0 :                     Method::PUT,
     869            0 :                     "control/v1/preferred_azs".to_string(),
     870            0 :                     Some(req),
     871            0 :                 )
     872            0 :                 .await?;
     873            0 :         }
     874            0 :         Command::TenantWarmup { tenant_id } => {
     875            0 :             let describe_response = storcon_client
     876            0 :                 .dispatch::<(), TenantDescribeResponse>(
     877            0 :                     Method::GET,
     878            0 :                     format!("control/v1/tenant/{tenant_id}"),
     879            0 :                     None,
     880            0 :                 )
     881            0 :                 .await;
     882            0 :             match describe_response {
     883            0 :                 Ok(describe) => {
     884            0 :                     if matches!(describe.policy, PlacementPolicy::Secondary) {
     885            0 :                         // Fine: it's already known to controller in secondary mode: calling
     886            0 :                         // again to put it into secondary mode won't cause problems.
     887            0 :                     } else {
     888            0 :                         anyhow::bail!("Tenant already present with policy {:?}", describe.policy);
     889            0 :                     }
     890            0 :                 }
     891            0 :                 Err(mgmt_api::Error::ApiError(StatusCode::NOT_FOUND, _)) => {
     892            0 :                     // Fine: this tenant isn't know to the storage controller yet.
     893            0 :                 }
     894            0 :                 Err(e) => {
     895            0 :                     // Unexpected API error
     896            0 :                     return Err(e.into());
     897            0 :                 }
     898            0 :             }
     899            0 : 
     900            0 :             vps_client
     901            0 :                 .location_config(
     902            0 :                     TenantShardId::unsharded(tenant_id),
     903            0 :                     pageserver_api::models::LocationConfig {
     904            0 :                         mode: pageserver_api::models::LocationConfigMode::Secondary,
     905            0 :                         generation: None,
     906            0 :                         secondary_conf: Some(LocationConfigSecondary { warm: true }),
     907            0 :                         shard_number: 0,
     908            0 :                         shard_count: 0,
     909            0 :                         shard_stripe_size: ShardParameters::DEFAULT_STRIPE_SIZE.0,
     910            0 :                         tenant_conf: TenantConfig::default(),
     911            0 :                     },
     912            0 :                     None,
     913            0 :                     true,
     914            0 :                 )
     915            0 :                 .await?;
     916            0 : 
     917            0 :             let describe_response = storcon_client
     918            0 :                 .dispatch::<(), TenantDescribeResponse>(
     919            0 :                     Method::GET,
     920            0 :                     format!("control/v1/tenant/{tenant_id}"),
     921            0 :                     None,
     922            0 :                 )
     923            0 :                 .await?;
     924            0 : 
     925            0 :             let secondary_ps_id = describe_response
     926            0 :                 .shards
     927            0 :                 .first()
     928            0 :                 .unwrap()
     929            0 :                 .node_secondary
     930            0 :                 .first()
     931            0 :                 .unwrap();
     932            0 : 
     933            0 :             println!("Tenant {tenant_id} warming up on pageserver {secondary_ps_id}");
     934            0 :             loop {
     935            0 :                 let (status, progress) = vps_client
     936            0 :                     .tenant_secondary_download(
     937            0 :                         TenantShardId::unsharded(tenant_id),
     938            0 :                         Some(Duration::from_secs(10)),
     939            0 :                     )
     940            0 :                     .await?;
     941            0 :                 println!(
     942            0 :                     "Progress: {}/{} layers, {}/{} bytes",
     943            0 :                     progress.layers_downloaded,
     944            0 :                     progress.layers_total,
     945            0 :                     progress.bytes_downloaded,
     946            0 :                     progress.bytes_total
     947            0 :                 );
     948            0 :                 match status {
     949            0 :                     StatusCode::OK => {
     950            0 :                         println!("Download complete");
     951            0 :                         break;
     952            0 :                     }
     953            0 :                     StatusCode::ACCEPTED => {
     954            0 :                         // Loop
     955            0 :                     }
     956            0 :                     _ => {
     957            0 :                         anyhow::bail!("Unexpected download status: {status}");
     958            0 :                     }
     959            0 :                 }
     960            0 :             }
     961            0 :         }
     962            0 :         Command::TenantDrop { tenant_id, unclean } => {
     963            0 :             if !unclean {
     964            0 :                 anyhow::bail!(
     965            0 :                     "This command is not a tenant deletion, and uncleanly drops all controller state for the tenant.  If you know what you're doing, add `--unclean` to proceed."
     966            0 :                 )
     967            0 :             }
     968            0 :             storcon_client
     969            0 :                 .dispatch::<(), ()>(
     970            0 :                     Method::POST,
     971            0 :                     format!("debug/v1/tenant/{tenant_id}/drop"),
     972            0 :                     None,
     973            0 :                 )
     974            0 :                 .await?;
     975            0 :         }
     976            0 :         Command::NodeDrop { node_id, unclean } => {
     977            0 :             if !unclean {
     978            0 :                 anyhow::bail!(
     979            0 :                     "This command is not a clean node decommission, and uncleanly drops all controller state for the node, without checking if any tenants still refer to it.  If you know what you're doing, add `--unclean` to proceed."
     980            0 :                 )
     981            0 :             }
     982            0 :             storcon_client
     983            0 :                 .dispatch::<(), ()>(Method::POST, format!("debug/v1/node/{node_id}/drop"), None)
     984            0 :                 .await?;
     985            0 :         }
     986            0 :         Command::NodeDelete { node_id } => {
     987            0 :             storcon_client
     988            0 :                 .dispatch::<(), ()>(Method::DELETE, format!("control/v1/node/{node_id}"), None)
     989            0 :                 .await?;
     990            0 :         }
     991            0 :         Command::TenantSetTimeBasedEviction {
     992            0 :             tenant_id,
     993            0 :             period,
     994            0 :             threshold,
     995            0 :         } => {
     996            0 :             vps_client
     997            0 :                 .set_tenant_config(&TenantConfigRequest {
     998            0 :                     tenant_id,
     999            0 :                     config: TenantConfig {
    1000            0 :                         eviction_policy: Some(EvictionPolicy::LayerAccessThreshold(
    1001            0 :                             EvictionPolicyLayerAccessThreshold {
    1002            0 :                                 period: period.into(),
    1003            0 :                                 threshold: threshold.into(),
    1004            0 :                             },
    1005            0 :                         )),
    1006            0 :                         heatmap_period: Some(Duration::from_secs(300)),
    1007            0 :                         ..Default::default()
    1008            0 :                     },
    1009            0 :                 })
    1010            0 :                 .await?;
    1011            0 :         }
    1012            0 :         Command::BulkMigrate {
    1013            0 :             nodes,
    1014            0 :             concurrency,
    1015            0 :             max_shards,
    1016            0 :             dry_run,
    1017            0 :         } => {
    1018            0 :             // Load the list of nodes, split them up into the drained and filled sets,
    1019            0 :             // and validate that draining is possible.
    1020            0 :             let node_descs = storcon_client
    1021            0 :                 .dispatch::<(), Vec<NodeDescribeResponse>>(
    1022            0 :                     Method::GET,
    1023            0 :                     "control/v1/node".to_string(),
    1024            0 :                     None,
    1025            0 :                 )
    1026            0 :                 .await?;
    1027            0 : 
    1028            0 :             let mut node_to_drain_descs = Vec::new();
    1029            0 :             let mut node_to_fill_descs = Vec::new();
    1030            0 : 
    1031            0 :             for desc in node_descs {
    1032            0 :                 let to_drain = nodes.iter().any(|id| *id == desc.id);
    1033            0 :                 if to_drain {
    1034            0 :                     node_to_drain_descs.push(desc);
    1035            0 :                 } else {
    1036            0 :                     node_to_fill_descs.push(desc);
    1037            0 :                 }
    1038            0 :             }
    1039            0 : 
    1040            0 :             if nodes.len() != node_to_drain_descs.len() {
    1041            0 :                 anyhow::bail!("Bulk migration requested away from node which doesn't exist.")
    1042            0 :             }
    1043            0 : 
    1044            0 :             node_to_fill_descs.retain(|desc| {
    1045            0 :                 matches!(desc.availability, NodeAvailabilityWrapper::Active)
    1046            0 :                     && matches!(
    1047            0 :                         desc.scheduling,
    1048            0 :                         NodeSchedulingPolicy::Active | NodeSchedulingPolicy::Filling
    1049            0 :                     )
    1050            0 :             });
    1051            0 : 
    1052            0 :             if node_to_fill_descs.is_empty() {
    1053            0 :                 anyhow::bail!("There are no nodes to migrate to")
    1054            0 :             }
    1055            0 : 
    1056            0 :             // Set the node scheduling policy to draining for the nodes which
    1057            0 :             // we plan to drain.
    1058            0 :             for node_desc in node_to_drain_descs.iter() {
    1059            0 :                 let req = NodeConfigureRequest {
    1060            0 :                     node_id: node_desc.id,
    1061            0 :                     availability: None,
    1062            0 :                     scheduling: Some(NodeSchedulingPolicy::Draining),
    1063            0 :                 };
    1064            0 : 
    1065            0 :                 storcon_client
    1066            0 :                     .dispatch::<_, ()>(
    1067            0 :                         Method::PUT,
    1068            0 :                         format!("control/v1/node/{}/config", node_desc.id),
    1069            0 :                         Some(req),
    1070            0 :                     )
    1071            0 :                     .await?;
    1072            0 :             }
    1073            0 : 
    1074            0 :             // Perform the migration: move each tenant shard scheduled on a node to
    1075            0 :             // be drained to a node which is being filled. A simple round robin
    1076            0 :             // strategy is used to pick the new node.
    1077            0 :             let tenants = storcon_client
    1078            0 :                 .dispatch::<(), Vec<TenantDescribeResponse>>(
    1079            0 :                     Method::GET,
    1080            0 :                     "control/v1/tenant".to_string(),
    1081            0 :                     None,
    1082            0 :                 )
    1083            0 :                 .await?;
    1084            0 : 
    1085            0 :             let mut selected_node_idx = 0;
    1086            0 : 
    1087            0 :             struct MigrationMove {
    1088            0 :                 tenant_shard_id: TenantShardId,
    1089            0 :                 from: NodeId,
    1090            0 :                 to: NodeId,
    1091            0 :             }
    1092            0 : 
    1093            0 :             let mut moves: Vec<MigrationMove> = Vec::new();
    1094            0 : 
    1095            0 :             let shards = tenants
    1096            0 :                 .into_iter()
    1097            0 :                 .flat_map(|tenant| tenant.shards.into_iter());
    1098            0 :             for shard in shards {
    1099            0 :                 if let Some(max_shards) = max_shards {
    1100            0 :                     if moves.len() >= max_shards {
    1101            0 :                         println!(
    1102            0 :                             "Stop planning shard moves since the requested maximum was reached"
    1103            0 :                         );
    1104            0 :                         break;
    1105            0 :                     }
    1106            0 :                 }
    1107            0 : 
    1108            0 :                 let should_migrate = {
    1109            0 :                     if let Some(attached_to) = shard.node_attached {
    1110            0 :                         node_to_drain_descs
    1111            0 :                             .iter()
    1112            0 :                             .map(|desc| desc.id)
    1113            0 :                             .any(|id| id == attached_to)
    1114            0 :                     } else {
    1115            0 :                         false
    1116            0 :                     }
    1117            0 :                 };
    1118            0 : 
    1119            0 :                 if !should_migrate {
    1120            0 :                     continue;
    1121            0 :                 }
    1122            0 : 
    1123            0 :                 moves.push(MigrationMove {
    1124            0 :                     tenant_shard_id: shard.tenant_shard_id,
    1125            0 :                     from: shard
    1126            0 :                         .node_attached
    1127            0 :                         .expect("We only migrate attached tenant shards"),
    1128            0 :                     to: node_to_fill_descs[selected_node_idx].id,
    1129            0 :                 });
    1130            0 :                 selected_node_idx = (selected_node_idx + 1) % node_to_fill_descs.len();
    1131            0 :             }
    1132            0 : 
    1133            0 :             let total_moves = moves.len();
    1134            0 : 
    1135            0 :             if dry_run == Some(true) {
    1136            0 :                 println!("Dryrun requested. Planned {total_moves} moves:");
    1137            0 :                 for mv in &moves {
    1138            0 :                     println!("{}: {} -> {}", mv.tenant_shard_id, mv.from, mv.to)
    1139            0 :                 }
    1140            0 : 
    1141            0 :                 return Ok(());
    1142            0 :             }
    1143            0 : 
    1144            0 :             const DEFAULT_MIGRATE_CONCURRENCY: usize = 8;
    1145            0 :             let mut stream = futures::stream::iter(moves)
    1146            0 :                 .map(|mv| {
    1147            0 :                     let client = Client::new(cli.api.clone(), cli.jwt.clone());
    1148            0 :                     async move {
    1149            0 :                         client
    1150            0 :                             .dispatch::<TenantShardMigrateRequest, TenantShardMigrateResponse>(
    1151            0 :                                 Method::PUT,
    1152            0 :                                 format!("control/v1/tenant/{}/migrate", mv.tenant_shard_id),
    1153            0 :                                 Some(TenantShardMigrateRequest {
    1154            0 :                                     node_id: mv.to,
    1155            0 :                                     origin_node_id: Some(mv.from),
    1156            0 :                                     migration_config: MigrationConfig::default(),
    1157            0 :                                 }),
    1158            0 :                             )
    1159            0 :                             .await
    1160            0 :                             .map_err(|e| (mv.tenant_shard_id, mv.from, mv.to, e))
    1161            0 :                     }
    1162            0 :                 })
    1163            0 :                 .buffered(concurrency.unwrap_or(DEFAULT_MIGRATE_CONCURRENCY));
    1164            0 : 
    1165            0 :             let mut success = 0;
    1166            0 :             let mut failure = 0;
    1167            0 : 
    1168            0 :             while let Some(res) = stream.next().await {
    1169            0 :                 match res {
    1170            0 :                     Ok(_) => {
    1171            0 :                         success += 1;
    1172            0 :                     }
    1173            0 :                     Err((tenant_shard_id, from, to, error)) => {
    1174            0 :                         failure += 1;
    1175            0 :                         println!(
    1176            0 :                             "Failed to migrate {} from node {} to node {}: {}",
    1177            0 :                             tenant_shard_id, from, to, error
    1178            0 :                         );
    1179            0 :                     }
    1180            0 :                 }
    1181            0 : 
    1182            0 :                 if (success + failure) % 20 == 0 {
    1183            0 :                     println!(
    1184            0 :                         "Processed {}/{} shards: {} succeeded, {} failed",
    1185            0 :                         success + failure,
    1186            0 :                         total_moves,
    1187            0 :                         success,
    1188            0 :                         failure
    1189            0 :                     );
    1190            0 :                 }
    1191            0 :             }
    1192            0 : 
    1193            0 :             println!(
    1194            0 :                 "Processed {}/{} shards: {} succeeded, {} failed",
    1195            0 :                 success + failure,
    1196            0 :                 total_moves,
    1197            0 :                 success,
    1198            0 :                 failure
    1199            0 :             );
    1200            0 :         }
    1201            0 :         Command::StartDrain { node_id } => {
    1202            0 :             storcon_client
    1203            0 :                 .dispatch::<(), ()>(
    1204            0 :                     Method::PUT,
    1205            0 :                     format!("control/v1/node/{node_id}/drain"),
    1206            0 :                     None,
    1207            0 :                 )
    1208            0 :                 .await?;
    1209            0 :             println!("Drain started for {node_id}");
    1210            0 :         }
    1211            0 :         Command::CancelDrain { node_id, timeout } => {
    1212            0 :             storcon_client
    1213            0 :                 .dispatch::<(), ()>(
    1214            0 :                     Method::DELETE,
    1215            0 :                     format!("control/v1/node/{node_id}/drain"),
    1216            0 :                     None,
    1217            0 :                 )
    1218            0 :                 .await?;
    1219            0 : 
    1220            0 :             println!("Waiting for node {node_id} to quiesce on scheduling policy ...");
    1221            0 : 
    1222            0 :             let final_policy =
    1223            0 :                 wait_for_scheduling_policy(storcon_client, node_id, *timeout, |sched| {
    1224            0 :                     use NodeSchedulingPolicy::*;
    1225            0 :                     matches!(sched, Active | PauseForRestart)
    1226            0 :                 })
    1227            0 :                 .await?;
    1228            0 : 
    1229            0 :             println!(
    1230            0 :                 "Drain was cancelled for node {node_id}. Schedulling policy is now {final_policy:?}"
    1231            0 :             );
    1232            0 :         }
    1233            0 :         Command::StartFill { node_id } => {
    1234            0 :             storcon_client
    1235            0 :                 .dispatch::<(), ()>(Method::PUT, format!("control/v1/node/{node_id}/fill"), None)
    1236            0 :                 .await?;
    1237            0 : 
    1238            0 :             println!("Fill started for {node_id}");
    1239            0 :         }
    1240            0 :         Command::CancelFill { node_id, timeout } => {
    1241            0 :             storcon_client
    1242            0 :                 .dispatch::<(), ()>(
    1243            0 :                     Method::DELETE,
    1244            0 :                     format!("control/v1/node/{node_id}/fill"),
    1245            0 :                     None,
    1246            0 :                 )
    1247            0 :                 .await?;
    1248            0 : 
    1249            0 :             println!("Waiting for node {node_id} to quiesce on scheduling policy ...");
    1250            0 : 
    1251            0 :             let final_policy =
    1252            0 :                 wait_for_scheduling_policy(storcon_client, node_id, *timeout, |sched| {
    1253            0 :                     use NodeSchedulingPolicy::*;
    1254            0 :                     matches!(sched, Active)
    1255            0 :                 })
    1256            0 :                 .await?;
    1257            0 : 
    1258            0 :             println!(
    1259            0 :                 "Fill was cancelled for node {node_id}. Schedulling policy is now {final_policy:?}"
    1260            0 :             );
    1261            0 :         }
    1262            0 :         Command::Safekeepers {} => {
    1263            0 :             let mut resp = storcon_client
    1264            0 :                 .dispatch::<(), Vec<SafekeeperDescribeResponse>>(
    1265            0 :                     Method::GET,
    1266            0 :                     "control/v1/safekeeper".to_string(),
    1267            0 :                     None,
    1268            0 :                 )
    1269            0 :                 .await?;
    1270            0 : 
    1271            0 :             resp.sort_by(|a, b| a.id.cmp(&b.id));
    1272            0 : 
    1273            0 :             let mut table = comfy_table::Table::new();
    1274            0 :             table.set_header([
    1275            0 :                 "Id",
    1276            0 :                 "Version",
    1277            0 :                 "Host",
    1278            0 :                 "Port",
    1279            0 :                 "Http Port",
    1280            0 :                 "AZ Id",
    1281            0 :                 "Scheduling",
    1282            0 :             ]);
    1283            0 :             for sk in resp {
    1284            0 :                 table.add_row([
    1285            0 :                     format!("{}", sk.id),
    1286            0 :                     format!("{}", sk.version),
    1287            0 :                     sk.host,
    1288            0 :                     format!("{}", sk.port),
    1289            0 :                     format!("{}", sk.http_port),
    1290            0 :                     sk.availability_zone_id.clone(),
    1291            0 :                     String::from(sk.scheduling_policy),
    1292            0 :                 ]);
    1293            0 :             }
    1294            0 :             println!("{table}");
    1295            0 :         }
    1296            0 :         Command::SafekeeperScheduling {
    1297            0 :             node_id,
    1298            0 :             scheduling_policy,
    1299            0 :         } => {
    1300            0 :             let scheduling_policy = scheduling_policy.0;
    1301            0 :             storcon_client
    1302            0 :                 .dispatch::<SafekeeperSchedulingPolicyRequest, ()>(
    1303            0 :                     Method::POST,
    1304            0 :                     format!("control/v1/safekeeper/{node_id}/scheduling_policy"),
    1305            0 :                     Some(SafekeeperSchedulingPolicyRequest { scheduling_policy }),
    1306            0 :                 )
    1307            0 :                 .await?;
    1308            0 :             println!(
    1309            0 :                 "Scheduling policy of {node_id} set to {}",
    1310            0 :                 String::from(scheduling_policy)
    1311            0 :             );
    1312            0 :         }
    1313            0 :         Command::DownloadHeatmapLayers {
    1314            0 :             tenant_shard_id,
    1315            0 :             timeline_id,
    1316            0 :             concurrency,
    1317            0 :         } => {
    1318            0 :             let mut path = format!(
    1319            0 :                 "/v1/tenant/{}/timeline/{}/download_heatmap_layers",
    1320            0 :                 tenant_shard_id, timeline_id,
    1321            0 :             );
    1322            0 : 
    1323            0 :             if let Some(c) = concurrency {
    1324            0 :                 path = format!("{path}?concurrency={c}");
    1325            0 :             }
    1326            0 : 
    1327            0 :             storcon_client
    1328            0 :                 .dispatch::<(), ()>(Method::POST, path, None)
    1329            0 :                 .await?;
    1330            0 :         }
    1331            0 :     }
    1332            0 : 
    1333            0 :     Ok(())
    1334            0 : }
    1335              : 
    1336              : static WATCH_INTERVAL: Duration = Duration::from_secs(5);
    1337              : 
    1338            0 : async fn watch_tenant_shard(
    1339            0 :     storcon_client: Client,
    1340            0 :     tenant_shard_id: TenantShardId,
    1341            0 :     until_migrated_to: Option<NodeId>,
    1342            0 : ) -> anyhow::Result<()> {
    1343            0 :     if let Some(until_migrated_to) = until_migrated_to {
    1344            0 :         println!(
    1345            0 :             "Waiting for tenant shard {} to be migrated to node {}",
    1346            0 :             tenant_shard_id, until_migrated_to
    1347            0 :         );
    1348            0 :     }
    1349              : 
    1350              :     loop {
    1351            0 :         let desc = storcon_client
    1352            0 :             .dispatch::<(), TenantDescribeResponse>(
    1353            0 :                 Method::GET,
    1354            0 :                 format!("control/v1/tenant/{}", tenant_shard_id.tenant_id),
    1355            0 :                 None,
    1356            0 :             )
    1357            0 :             .await?;
    1358              : 
    1359              :         // Output the current state of the tenant shard
    1360            0 :         let shard = desc
    1361            0 :             .shards
    1362            0 :             .iter()
    1363            0 :             .find(|s| s.tenant_shard_id == tenant_shard_id)
    1364            0 :             .ok_or(anyhow::anyhow!("Tenant shard not found"))?;
    1365            0 :         let summary = format!(
    1366            0 :             "attached: {} secondary: {} {}",
    1367            0 :             shard
    1368            0 :                 .node_attached
    1369            0 :                 .map(|n| format!("{}", n))
    1370            0 :                 .unwrap_or("none".to_string()),
    1371            0 :             shard
    1372            0 :                 .node_secondary
    1373            0 :                 .iter()
    1374            0 :                 .map(|n| n.to_string())
    1375            0 :                 .collect::<Vec<_>>()
    1376            0 :                 .join(","),
    1377            0 :             if shard.is_reconciling {
    1378            0 :                 "(reconciler active)"
    1379              :             } else {
    1380            0 :                 "(reconciler idle)"
    1381              :             }
    1382              :         );
    1383            0 :         println!("{}", summary);
    1384              : 
    1385              :         // Maybe drop out if we finished migration
    1386            0 :         if let Some(until_migrated_to) = until_migrated_to {
    1387            0 :             if shard.node_attached == Some(until_migrated_to) && !shard.is_reconciling {
    1388            0 :                 println!(
    1389            0 :                     "Tenant shard {} is now on node {}",
    1390            0 :                     tenant_shard_id, until_migrated_to
    1391            0 :                 );
    1392            0 :                 break;
    1393            0 :             }
    1394            0 :         }
    1395              : 
    1396            0 :         tokio::time::sleep(WATCH_INTERVAL).await;
    1397              :     }
    1398            0 :     Ok(())
    1399            0 : }
        

Generated by: LCOV version 2.1-beta