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

Generated by: LCOV version 2.1-beta