LCOV - code coverage report
Current view: top level - control_plane/src/bin - neon_local.rs (source / functions) Coverage Total Hit
Test: c8f8d331b83562868d9054d9e0e68f866772aeaa.info Lines: 0.0 % 1038 0
Test Date: 2025-07-26 17:20:05 Functions: 0.0 % 77 0

            Line data    Source code
       1              : //!
       2              : //! `neon_local` is an executable that can be used to create a local
       3              : //! Neon environment, for testing purposes. The local environment is
       4              : //! quite different from the cloud environment with Kubernetes, but it
       5              : //! easier to work with locally. The python tests in `test_runner`
       6              : //! rely on `neon_local` to set up the environment for each test.
       7              : //!
       8              : use std::borrow::Cow;
       9              : use std::collections::{BTreeSet, HashMap};
      10              : use std::fs::File;
      11              : use std::path::PathBuf;
      12              : use std::process::exit;
      13              : use std::str::FromStr;
      14              : use std::time::Duration;
      15              : 
      16              : use anyhow::{Context, Result, anyhow, bail};
      17              : use clap::Parser;
      18              : use compute_api::requests::ComputeClaimsScope;
      19              : use compute_api::spec::{ComputeMode, PageserverProtocol};
      20              : use control_plane::broker::StorageBroker;
      21              : use control_plane::endpoint::{ComputeControlPlane, EndpointTerminateMode};
      22              : use control_plane::endpoint_storage::{ENDPOINT_STORAGE_DEFAULT_ADDR, EndpointStorage};
      23              : use control_plane::local_env;
      24              : use control_plane::local_env::{
      25              :     EndpointStorageConf, InitForceMode, LocalEnv, NeonBroker, NeonLocalInitConf,
      26              :     NeonLocalInitPageserverConf, SafekeeperConf,
      27              : };
      28              : use control_plane::pageserver::PageServerNode;
      29              : use control_plane::safekeeper::SafekeeperNode;
      30              : use control_plane::storage_controller::{
      31              :     NeonStorageControllerStartArgs, NeonStorageControllerStopArgs, StorageController,
      32              : };
      33              : use nix::fcntl::{Flock, FlockArg};
      34              : use pageserver_api::config::{
      35              :     DEFAULT_GRPC_LISTEN_PORT as DEFAULT_PAGESERVER_GRPC_PORT,
      36              :     DEFAULT_HTTP_LISTEN_PORT as DEFAULT_PAGESERVER_HTTP_PORT,
      37              :     DEFAULT_PG_LISTEN_PORT as DEFAULT_PAGESERVER_PG_PORT,
      38              : };
      39              : use pageserver_api::controller_api::{
      40              :     NodeAvailabilityWrapper, PlacementPolicy, TenantCreateRequest,
      41              : };
      42              : use pageserver_api::models::{
      43              :     ShardParameters, TenantConfigRequest, TimelineCreateRequest, TimelineInfo,
      44              : };
      45              : use pageserver_api::shard::{DEFAULT_STRIPE_SIZE, ShardCount, ShardStripeSize, TenantShardId};
      46              : use postgres_backend::AuthType;
      47              : use postgres_connection::parse_host_port;
      48              : use safekeeper_api::membership::{SafekeeperGeneration, SafekeeperId};
      49              : use safekeeper_api::{
      50              :     DEFAULT_HTTP_LISTEN_PORT as DEFAULT_SAFEKEEPER_HTTP_PORT,
      51              :     DEFAULT_PG_LISTEN_PORT as DEFAULT_SAFEKEEPER_PG_PORT, PgMajorVersion, PgVersionId,
      52              : };
      53              : use storage_broker::DEFAULT_LISTEN_ADDR as DEFAULT_BROKER_ADDR;
      54              : use tokio::task::JoinSet;
      55              : use url::Host;
      56              : use utils::auth::{Claims, Scope};
      57              : use utils::id::{NodeId, TenantId, TenantTimelineId, TimelineId};
      58              : use utils::lsn::Lsn;
      59              : use utils::project_git_version;
      60              : 
      61              : // Default id of a safekeeper node, if not specified on the command line.
      62              : const DEFAULT_SAFEKEEPER_ID: NodeId = NodeId(1);
      63              : const DEFAULT_PAGESERVER_ID: NodeId = NodeId(1);
      64              : const DEFAULT_BRANCH_NAME: &str = "main";
      65              : project_git_version!(GIT_VERSION);
      66              : 
      67              : #[allow(dead_code)]
      68              : const DEFAULT_PG_VERSION: PgMajorVersion = PgMajorVersion::PG17;
      69              : const DEFAULT_PG_VERSION_NUM: &str = "17";
      70              : 
      71              : const DEFAULT_PAGESERVER_CONTROL_PLANE_API: &str = "http://127.0.0.1:1234/upcall/v1/";
      72              : 
      73              : #[derive(clap::Parser)]
      74              : #[command(version = GIT_VERSION, about, name = "Neon CLI")]
      75              : struct Cli {
      76              :     #[command(subcommand)]
      77              :     command: NeonLocalCmd,
      78              : }
      79              : 
      80              : #[derive(clap::Subcommand)]
      81              : enum NeonLocalCmd {
      82              :     Init(InitCmdArgs),
      83              : 
      84              :     #[command(subcommand)]
      85              :     Tenant(TenantCmd),
      86              :     #[command(subcommand)]
      87              :     Timeline(TimelineCmd),
      88              :     #[command(subcommand)]
      89              :     Pageserver(PageserverCmd),
      90              :     #[command(subcommand)]
      91              :     #[clap(alias = "storage_controller")]
      92              :     StorageController(StorageControllerCmd),
      93              :     #[command(subcommand)]
      94              :     #[clap(alias = "storage_broker")]
      95              :     StorageBroker(StorageBrokerCmd),
      96              :     #[command(subcommand)]
      97              :     Safekeeper(SafekeeperCmd),
      98              :     #[command(subcommand)]
      99              :     EndpointStorage(EndpointStorageCmd),
     100              :     #[command(subcommand)]
     101              :     Endpoint(EndpointCmd),
     102              :     #[command(subcommand)]
     103              :     Mappings(MappingsCmd),
     104              : 
     105              :     Start(StartCmdArgs),
     106              :     Stop(StopCmdArgs),
     107              : }
     108              : 
     109              : #[derive(clap::Args)]
     110              : #[clap(about = "Initialize a new Neon repository, preparing configs for services to start with")]
     111              : struct InitCmdArgs {
     112              :     #[clap(long, help("How many pageservers to create (default 1)"))]
     113              :     num_pageservers: Option<u16>,
     114              : 
     115              :     #[clap(long)]
     116              :     config: Option<PathBuf>,
     117              : 
     118              :     #[clap(long, help("Force initialization even if the repository is not empty"))]
     119              :     #[arg(value_parser)]
     120              :     #[clap(default_value = "must-not-exist")]
     121              :     force: InitForceMode,
     122              : }
     123              : 
     124              : #[derive(clap::Args)]
     125              : #[clap(about = "Start pageserver and safekeepers")]
     126              : struct StartCmdArgs {
     127              :     #[clap(long = "start-timeout", default_value = "10s")]
     128              :     timeout: humantime::Duration,
     129              : }
     130              : 
     131              : #[derive(clap::Args)]
     132              : #[clap(about = "Stop pageserver and safekeepers")]
     133              : struct StopCmdArgs {
     134              :     #[arg(value_enum)]
     135              :     #[clap(long, default_value_t = StopMode::Fast)]
     136              :     mode: StopMode,
     137              : }
     138              : 
     139              : #[derive(Clone, Copy, clap::ValueEnum)]
     140              : enum StopMode {
     141              :     Fast,
     142              :     Immediate,
     143              : }
     144              : 
     145              : #[derive(clap::Subcommand)]
     146              : #[clap(about = "Manage tenants")]
     147              : enum TenantCmd {
     148              :     List,
     149              :     Create(TenantCreateCmdArgs),
     150              :     SetDefault(TenantSetDefaultCmdArgs),
     151              :     Config(TenantConfigCmdArgs),
     152              :     Import(TenantImportCmdArgs),
     153              : }
     154              : 
     155              : #[derive(clap::Args)]
     156              : struct TenantCreateCmdArgs {
     157              :     #[clap(
     158              :         long = "tenant-id",
     159              :         help = "Tenant id. Represented as a hexadecimal string 32 symbols length"
     160              :     )]
     161              :     tenant_id: Option<TenantId>,
     162              : 
     163              :     #[clap(
     164              :         long,
     165              :         help = "Use a specific timeline id when creating a tenant and its initial timeline"
     166              :     )]
     167              :     timeline_id: Option<TimelineId>,
     168              : 
     169              :     #[clap(short = 'c')]
     170              :     config: Vec<String>,
     171              : 
     172              :     #[arg(default_value = DEFAULT_PG_VERSION_NUM)]
     173              :     #[clap(long, help = "Postgres version to use for the initial timeline")]
     174              :     pg_version: PgMajorVersion,
     175              : 
     176              :     #[clap(
     177              :         long,
     178              :         help = "Use this tenant in future CLI commands where tenant_id is needed, but not specified"
     179              :     )]
     180              :     set_default: bool,
     181              : 
     182              :     #[clap(long, help = "Number of shards in the new tenant")]
     183              :     #[arg(default_value_t = 0)]
     184              :     shard_count: u8,
     185              :     #[clap(long, help = "Sharding stripe size in pages")]
     186              :     shard_stripe_size: Option<u32>,
     187              : 
     188              :     #[clap(long, help = "Placement policy shards in this tenant")]
     189              :     #[arg(value_parser = parse_placement_policy)]
     190              :     placement_policy: Option<PlacementPolicy>,
     191              : }
     192              : 
     193            0 : fn parse_placement_policy(s: &str) -> anyhow::Result<PlacementPolicy> {
     194            0 :     Ok(serde_json::from_str::<PlacementPolicy>(s)?)
     195            0 : }
     196              : 
     197              : #[derive(clap::Args)]
     198              : #[clap(
     199              :     about = "Set a particular tenant as default in future CLI commands where tenant_id is needed, but not specified"
     200              : )]
     201              : struct TenantSetDefaultCmdArgs {
     202              :     #[clap(
     203              :         long = "tenant-id",
     204              :         help = "Tenant id. Represented as a hexadecimal string 32 symbols length"
     205              :     )]
     206              :     tenant_id: TenantId,
     207              : }
     208              : 
     209              : #[derive(clap::Args)]
     210              : struct TenantConfigCmdArgs {
     211              :     #[clap(
     212              :         long = "tenant-id",
     213              :         help = "Tenant id. Represented as a hexadecimal string 32 symbols length"
     214              :     )]
     215              :     tenant_id: Option<TenantId>,
     216              : 
     217              :     #[clap(short = 'c')]
     218              :     config: Vec<String>,
     219              : }
     220              : 
     221              : #[derive(clap::Args)]
     222              : #[clap(
     223              :     about = "Import a tenant that is present in remote storage, and create branches for its timelines"
     224              : )]
     225              : struct TenantImportCmdArgs {
     226              :     #[clap(
     227              :         long = "tenant-id",
     228              :         help = "Tenant id. Represented as a hexadecimal string 32 symbols length"
     229              :     )]
     230              :     tenant_id: TenantId,
     231              : }
     232              : 
     233              : #[derive(clap::Subcommand)]
     234              : #[clap(about = "Manage timelines")]
     235              : enum TimelineCmd {
     236              :     List(TimelineListCmdArgs),
     237              :     Branch(TimelineBranchCmdArgs),
     238              :     Create(TimelineCreateCmdArgs),
     239              :     Import(TimelineImportCmdArgs),
     240              : }
     241              : 
     242              : #[derive(clap::Args)]
     243              : #[clap(about = "List all timelines available to this pageserver")]
     244              : struct TimelineListCmdArgs {
     245              :     #[clap(
     246              :         long = "tenant-id",
     247              :         help = "Tenant id. Represented as a hexadecimal string 32 symbols length"
     248              :     )]
     249              :     tenant_shard_id: Option<TenantShardId>,
     250              : }
     251              : 
     252              : #[derive(clap::Args)]
     253              : #[clap(about = "Create a new timeline, branching off from another timeline")]
     254              : struct TimelineBranchCmdArgs {
     255              :     #[clap(
     256              :         long = "tenant-id",
     257              :         help = "Tenant id. Represented as a hexadecimal string 32 symbols length"
     258              :     )]
     259              :     tenant_id: Option<TenantId>,
     260              : 
     261              :     #[clap(long, help = "New timeline's ID")]
     262              :     timeline_id: Option<TimelineId>,
     263              : 
     264              :     #[clap(long, help = "Human-readable alias for the new timeline")]
     265              :     branch_name: String,
     266              : 
     267              :     #[clap(
     268              :         long,
     269              :         help = "Use last Lsn of another timeline (and its data) as base when creating the new timeline. The timeline gets resolved by its branch name."
     270              :     )]
     271              :     ancestor_branch_name: Option<String>,
     272              : 
     273              :     #[clap(
     274              :         long,
     275              :         help = "When using another timeline as base, use a specific Lsn in it instead of the latest one"
     276              :     )]
     277              :     ancestor_start_lsn: Option<Lsn>,
     278              : }
     279              : 
     280              : #[derive(clap::Args)]
     281              : #[clap(about = "Create a new blank timeline")]
     282              : struct TimelineCreateCmdArgs {
     283              :     #[clap(
     284              :         long = "tenant-id",
     285              :         help = "Tenant id. Represented as a hexadecimal string 32 symbols length"
     286              :     )]
     287              :     tenant_id: Option<TenantId>,
     288              : 
     289              :     #[clap(long, help = "New timeline's ID")]
     290              :     timeline_id: Option<TimelineId>,
     291              : 
     292              :     #[clap(long, help = "Human-readable alias for the new timeline")]
     293              :     branch_name: String,
     294              : 
     295              :     #[arg(default_value = DEFAULT_PG_VERSION_NUM)]
     296              :     #[clap(long, help = "Postgres version")]
     297              :     pg_version: PgMajorVersion,
     298              : }
     299              : 
     300              : #[derive(clap::Args)]
     301              : #[clap(about = "Import timeline from a basebackup directory")]
     302              : struct TimelineImportCmdArgs {
     303              :     #[clap(
     304              :         long = "tenant-id",
     305              :         help = "Tenant id. Represented as a hexadecimal string 32 symbols length"
     306              :     )]
     307              :     tenant_id: Option<TenantId>,
     308              : 
     309              :     #[clap(long, help = "New timeline's ID")]
     310              :     timeline_id: TimelineId,
     311              : 
     312              :     #[clap(long, help = "Human-readable alias for the new timeline")]
     313              :     branch_name: String,
     314              : 
     315              :     #[clap(long, help = "Basebackup tarfile to import")]
     316              :     base_tarfile: PathBuf,
     317              : 
     318              :     #[clap(long, help = "Lsn the basebackup starts at")]
     319              :     base_lsn: Lsn,
     320              : 
     321              :     #[clap(long, help = "Wal to add after base")]
     322              :     wal_tarfile: Option<PathBuf>,
     323              : 
     324              :     #[clap(long, help = "Lsn the basebackup ends at")]
     325              :     end_lsn: Option<Lsn>,
     326              : 
     327              :     #[arg(default_value = DEFAULT_PG_VERSION_NUM)]
     328              :     #[clap(long, help = "Postgres version of the backup being imported")]
     329              :     pg_version: PgMajorVersion,
     330              : }
     331              : 
     332              : #[derive(clap::Subcommand)]
     333              : #[clap(about = "Manage pageservers")]
     334              : enum PageserverCmd {
     335              :     Status(PageserverStatusCmdArgs),
     336              :     Start(PageserverStartCmdArgs),
     337              :     Stop(PageserverStopCmdArgs),
     338              :     Restart(PageserverRestartCmdArgs),
     339              : }
     340              : 
     341              : #[derive(clap::Args)]
     342              : #[clap(about = "Show status of a local pageserver")]
     343              : struct PageserverStatusCmdArgs {
     344              :     #[clap(long = "id", help = "pageserver id")]
     345              :     pageserver_id: Option<NodeId>,
     346              : }
     347              : 
     348              : #[derive(clap::Args)]
     349              : #[clap(about = "Start local pageserver")]
     350              : struct PageserverStartCmdArgs {
     351              :     #[clap(long = "id", help = "pageserver id")]
     352              :     pageserver_id: Option<NodeId>,
     353              : 
     354              :     #[clap(short = 't', long, help = "timeout until we fail the command")]
     355              :     #[arg(default_value = "10s")]
     356              :     start_timeout: humantime::Duration,
     357              : }
     358              : 
     359              : #[derive(clap::Args)]
     360              : #[clap(about = "Stop local pageserver")]
     361              : struct PageserverStopCmdArgs {
     362              :     #[clap(long = "id", help = "pageserver id")]
     363              :     pageserver_id: Option<NodeId>,
     364              : 
     365              :     #[clap(
     366              :         short = 'm',
     367              :         help = "If 'immediate', don't flush repository data at shutdown"
     368              :     )]
     369              :     #[arg(value_enum, default_value = "fast")]
     370              :     stop_mode: StopMode,
     371              : }
     372              : 
     373              : #[derive(clap::Args)]
     374              : #[clap(about = "Restart local pageserver")]
     375              : struct PageserverRestartCmdArgs {
     376              :     #[clap(long = "id", help = "pageserver id")]
     377              :     pageserver_id: Option<NodeId>,
     378              : 
     379              :     #[clap(short = 't', long, help = "timeout until we fail the command")]
     380              :     #[arg(default_value = "10s")]
     381              :     start_timeout: humantime::Duration,
     382              : }
     383              : 
     384              : #[derive(clap::Subcommand)]
     385              : #[clap(about = "Manage storage controller")]
     386              : enum StorageControllerCmd {
     387              :     Start(StorageControllerStartCmdArgs),
     388              :     Stop(StorageControllerStopCmdArgs),
     389              : }
     390              : 
     391              : #[derive(clap::Args)]
     392              : #[clap(about = "Start storage controller")]
     393              : struct StorageControllerStartCmdArgs {
     394              :     #[clap(short = 't', long, help = "timeout until we fail the command")]
     395              :     #[arg(default_value = "10s")]
     396              :     start_timeout: humantime::Duration,
     397              : 
     398              :     #[clap(
     399              :         long,
     400              :         help = "Identifier used to distinguish storage controller instances"
     401              :     )]
     402              :     #[arg(default_value_t = 1)]
     403              :     instance_id: u8,
     404              : 
     405              :     #[clap(
     406              :         long,
     407              :         help = "Base port for the storage controller instance idenfified by instance-id (defaults to pageserver cplane api)"
     408              :     )]
     409              :     base_port: Option<u16>,
     410              : 
     411              :     #[clap(
     412              :         long,
     413              :         help = "Whether the storage controller should handle pageserver-reported local disk loss events."
     414              :     )]
     415              :     handle_ps_local_disk_loss: Option<bool>,
     416              : }
     417              : 
     418              : #[derive(clap::Args)]
     419              : #[clap(about = "Stop storage controller")]
     420              : struct StorageControllerStopCmdArgs {
     421              :     #[clap(
     422              :         short = 'm',
     423              :         help = "If 'immediate', don't flush repository data at shutdown"
     424              :     )]
     425              :     #[arg(value_enum, default_value = "fast")]
     426              :     stop_mode: StopMode,
     427              : 
     428              :     #[clap(
     429              :         long,
     430              :         help = "Identifier used to distinguish storage controller instances"
     431              :     )]
     432              :     #[arg(default_value_t = 1)]
     433              :     instance_id: u8,
     434              : }
     435              : 
     436              : #[derive(clap::Subcommand)]
     437              : #[clap(about = "Manage storage broker")]
     438              : enum StorageBrokerCmd {
     439              :     Start(StorageBrokerStartCmdArgs),
     440              :     Stop(StorageBrokerStopCmdArgs),
     441              : }
     442              : 
     443              : #[derive(clap::Args)]
     444              : #[clap(about = "Start broker")]
     445              : struct StorageBrokerStartCmdArgs {
     446              :     #[clap(short = 't', long, help = "timeout until we fail the command")]
     447              :     #[arg(default_value = "10s")]
     448              :     start_timeout: humantime::Duration,
     449              : }
     450              : 
     451              : #[derive(clap::Args)]
     452              : #[clap(about = "stop broker")]
     453              : struct StorageBrokerStopCmdArgs {
     454              :     #[clap(
     455              :         short = 'm',
     456              :         help = "If 'immediate', don't flush repository data at shutdown"
     457              :     )]
     458              :     #[arg(value_enum, default_value = "fast")]
     459              :     stop_mode: StopMode,
     460              : }
     461              : 
     462              : #[derive(clap::Subcommand)]
     463              : #[clap(about = "Manage safekeepers")]
     464              : enum SafekeeperCmd {
     465              :     Start(SafekeeperStartCmdArgs),
     466              :     Stop(SafekeeperStopCmdArgs),
     467              :     Restart(SafekeeperRestartCmdArgs),
     468              : }
     469              : 
     470              : #[derive(clap::Subcommand)]
     471              : #[clap(about = "Manage object storage")]
     472              : enum EndpointStorageCmd {
     473              :     Start(EndpointStorageStartCmd),
     474              :     Stop(EndpointStorageStopCmd),
     475              : }
     476              : 
     477              : #[derive(clap::Args)]
     478              : #[clap(about = "Start object storage")]
     479              : struct EndpointStorageStartCmd {
     480              :     #[clap(short = 't', long, help = "timeout until we fail the command")]
     481              :     #[arg(default_value = "10s")]
     482              :     start_timeout: humantime::Duration,
     483              : }
     484              : 
     485              : #[derive(clap::Args)]
     486              : #[clap(about = "Stop object storage")]
     487              : struct EndpointStorageStopCmd {
     488              :     #[arg(value_enum, default_value = "fast")]
     489              :     #[clap(
     490              :         short = 'm',
     491              :         help = "If 'immediate', don't flush repository data at shutdown"
     492              :     )]
     493              :     stop_mode: StopMode,
     494              : }
     495              : 
     496              : #[derive(clap::Args)]
     497              : #[clap(about = "Start local safekeeper")]
     498              : struct SafekeeperStartCmdArgs {
     499              :     #[clap(help = "safekeeper id")]
     500              :     #[arg(default_value_t = NodeId(1))]
     501              :     id: NodeId,
     502              : 
     503              :     #[clap(
     504              :         short = 'e',
     505              :         long = "safekeeper-extra-opt",
     506              :         help = "Additional safekeeper invocation options, e.g. -e=--http-auth-public-key-path=foo"
     507              :     )]
     508              :     extra_opt: Vec<String>,
     509              : 
     510              :     #[clap(short = 't', long, help = "timeout until we fail the command")]
     511              :     #[arg(default_value = "10s")]
     512              :     start_timeout: humantime::Duration,
     513              : }
     514              : 
     515              : #[derive(clap::Args)]
     516              : #[clap(about = "Stop local safekeeper")]
     517              : struct SafekeeperStopCmdArgs {
     518              :     #[clap(help = "safekeeper id")]
     519              :     #[arg(default_value_t = NodeId(1))]
     520              :     id: NodeId,
     521              : 
     522              :     #[arg(value_enum, default_value = "fast")]
     523              :     #[clap(
     524              :         short = 'm',
     525              :         help = "If 'immediate', don't flush repository data at shutdown"
     526              :     )]
     527              :     stop_mode: StopMode,
     528              : }
     529              : 
     530              : #[derive(clap::Args)]
     531              : #[clap(about = "Restart local safekeeper")]
     532              : struct SafekeeperRestartCmdArgs {
     533              :     #[clap(help = "safekeeper id")]
     534              :     #[arg(default_value_t = NodeId(1))]
     535              :     id: NodeId,
     536              : 
     537              :     #[arg(value_enum, default_value = "fast")]
     538              :     #[clap(
     539              :         short = 'm',
     540              :         help = "If 'immediate', don't flush repository data at shutdown"
     541              :     )]
     542              :     stop_mode: StopMode,
     543              : 
     544              :     #[clap(
     545              :         short = 'e',
     546              :         long = "safekeeper-extra-opt",
     547              :         help = "Additional safekeeper invocation options, e.g. -e=--http-auth-public-key-path=foo"
     548              :     )]
     549              :     extra_opt: Vec<String>,
     550              : 
     551              :     #[clap(short = 't', long, help = "timeout until we fail the command")]
     552              :     #[arg(default_value = "10s")]
     553              :     start_timeout: humantime::Duration,
     554              : }
     555              : 
     556              : #[derive(clap::Subcommand)]
     557              : #[clap(about = "Manage Postgres instances")]
     558              : enum EndpointCmd {
     559              :     List(EndpointListCmdArgs),
     560              :     Create(EndpointCreateCmdArgs),
     561              :     Start(EndpointStartCmdArgs),
     562              :     Reconfigure(EndpointReconfigureCmdArgs),
     563              :     RefreshConfiguration(EndpointRefreshConfigurationArgs),
     564              :     Stop(EndpointStopCmdArgs),
     565              :     UpdatePageservers(EndpointUpdatePageserversCmdArgs),
     566              :     GenerateJwt(EndpointGenerateJwtCmdArgs),
     567              : }
     568              : 
     569              : #[derive(clap::Args)]
     570              : #[clap(about = "List endpoints")]
     571              : struct EndpointListCmdArgs {
     572              :     #[clap(
     573              :         long = "tenant-id",
     574              :         help = "Tenant id. Represented as a hexadecimal string 32 symbols length"
     575              :     )]
     576              :     tenant_shard_id: Option<TenantShardId>,
     577              : }
     578              : 
     579              : #[derive(clap::Args)]
     580              : #[clap(about = "Create a compute endpoint")]
     581              : struct EndpointCreateCmdArgs {
     582              :     #[clap(
     583              :         long = "tenant-id",
     584              :         help = "Tenant id. Represented as a hexadecimal string 32 symbols length"
     585              :     )]
     586              :     tenant_id: Option<TenantId>,
     587              : 
     588              :     #[clap(help = "Postgres endpoint id")]
     589              :     endpoint_id: Option<String>,
     590              :     #[clap(long, help = "Name of the branch the endpoint will run on")]
     591              :     branch_name: Option<String>,
     592              :     #[clap(
     593              :         long,
     594              :         help = "Specify Lsn on the timeline to start from. By default, end of the timeline would be used"
     595              :     )]
     596              :     lsn: Option<Lsn>,
     597              :     #[clap(long)]
     598              :     pg_port: Option<u16>,
     599              :     #[clap(long, alias = "http-port")]
     600              :     external_http_port: Option<u16>,
     601              :     #[clap(long)]
     602              :     internal_http_port: Option<u16>,
     603              :     #[clap(long = "pageserver-id")]
     604              :     endpoint_pageserver_id: Option<NodeId>,
     605              : 
     606              :     #[clap(
     607              :         long,
     608              :         help = "Don't do basebackup, create endpoint directory with only config files",
     609              :         action = clap::ArgAction::Set,
     610              :         default_value_t = false
     611              :     )]
     612              :     config_only: bool,
     613              : 
     614              :     #[arg(default_value = DEFAULT_PG_VERSION_NUM)]
     615              :     #[clap(long, help = "Postgres version")]
     616              :     pg_version: PgMajorVersion,
     617              : 
     618              :     /// Use gRPC to communicate with Pageservers, by generating grpc:// connstrings.
     619              :     ///
     620              :     /// Specified on creation such that it's retained across reconfiguration and restarts.
     621              :     ///
     622              :     /// NB: not yet supported by computes.
     623              :     #[clap(long)]
     624              :     grpc: bool,
     625              : 
     626              :     #[clap(
     627              :         long,
     628              :         help = "If set, the node will be a hot replica on the specified timeline",
     629              :         action = clap::ArgAction::Set,
     630              :         default_value_t = false
     631              :     )]
     632              :     hot_standby: bool,
     633              : 
     634              :     #[clap(long, help = "If set, will set up the catalog for neon_superuser")]
     635              :     update_catalog: bool,
     636              : 
     637              :     #[clap(
     638              :         long,
     639              :         help = "Allow multiple primary endpoints running on the same branch. Shouldn't be used normally, but useful for tests."
     640              :     )]
     641              :     allow_multiple: bool,
     642              : 
     643              :     /// Only allow changing it on creation
     644              :     #[clap(long, help = "Name of the privileged role for the endpoint")]
     645              :     privileged_role_name: Option<String>,
     646              : }
     647              : 
     648              : #[derive(clap::Args)]
     649              : #[clap(about = "Start postgres. If the endpoint doesn't exist yet, it is created.")]
     650              : struct EndpointStartCmdArgs {
     651              :     #[clap(help = "Postgres endpoint id")]
     652              :     endpoint_id: String,
     653              :     #[clap(long = "pageserver-id")]
     654              :     endpoint_pageserver_id: Option<NodeId>,
     655              : 
     656              :     #[clap(
     657              :         long,
     658              :         help = "Safekeepers membership generation to prefix neon.safekeepers with. Normally neon_local sets it on its own, but this option allows to override. Non zero value forces endpoint to use membership configurations."
     659              :     )]
     660              :     safekeepers_generation: Option<u32>,
     661              :     #[clap(
     662              :         long,
     663              :         help = "List of safekeepers endpoint will talk to. Normally neon_local chooses them on its own, but this option allows to override."
     664              :     )]
     665              :     safekeepers: Option<String>,
     666              : 
     667              :     #[clap(
     668              :         long,
     669              :         help = "Configure the remote extensions storage proxy gateway URL to request for extensions.",
     670              :         alias = "remote-ext-config"
     671              :     )]
     672              :     remote_ext_base_url: Option<String>,
     673              : 
     674              :     #[clap(
     675              :         long,
     676              :         help = "If set, will create test user `user` and `neondb` database. Requires `update-catalog = true`"
     677              :     )]
     678              :     create_test_user: bool,
     679              : 
     680              :     #[clap(
     681              :         long,
     682              :         help = "Allow multiple primary endpoints running on the same branch. Shouldn't be used normally, but useful for tests."
     683              :     )]
     684              :     allow_multiple: bool,
     685              : 
     686              :     #[clap(short = 't', long, value_parser= humantime::parse_duration, help = "timeout until we fail the command")]
     687              :     #[arg(default_value = "90s")]
     688              :     start_timeout: Duration,
     689              : 
     690              :     #[clap(
     691              :         long,
     692              :         help = "Download LFC cache from endpoint storage on endpoint startup",
     693              :         default_value = "false"
     694              :     )]
     695              :     autoprewarm: bool,
     696              : 
     697              :     #[clap(long, help = "Upload LFC cache to endpoint storage periodically")]
     698              :     offload_lfc_interval_seconds: Option<std::num::NonZeroU64>,
     699              : 
     700              :     #[clap(
     701              :         long,
     702              :         help = "Run in development mode, skipping VM-specific operations like process termination",
     703              :         action = clap::ArgAction::SetTrue
     704              :     )]
     705              :     dev: bool,
     706              : }
     707              : 
     708              : #[derive(clap::Args)]
     709              : #[clap(about = "Reconfigure an endpoint")]
     710              : struct EndpointReconfigureCmdArgs {
     711              :     #[clap(
     712              :         long = "tenant-id",
     713              :         help = "Tenant id. Represented as a hexadecimal string 32 symbols length"
     714              :     )]
     715              :     tenant_id: Option<TenantId>,
     716              : 
     717              :     #[clap(help = "Postgres endpoint id")]
     718              :     endpoint_id: String,
     719              :     #[clap(long = "pageserver-id")]
     720              :     endpoint_pageserver_id: Option<NodeId>,
     721              : 
     722              :     #[clap(long)]
     723              :     safekeepers: Option<String>,
     724              : }
     725              : 
     726              : #[derive(clap::Args)]
     727              : #[clap(about = "Refresh the endpoint's configuration by forcing it reload it's spec")]
     728              : struct EndpointRefreshConfigurationArgs {
     729              :     #[clap(help = "Postgres endpoint id")]
     730              :     endpoint_id: String,
     731              : }
     732              : 
     733              : #[derive(clap::Args)]
     734              : #[clap(about = "Stop an endpoint")]
     735              : struct EndpointStopCmdArgs {
     736              :     #[clap(help = "Postgres endpoint id")]
     737              :     endpoint_id: String,
     738              : 
     739              :     #[clap(
     740              :         long,
     741              :         help = "Also delete data directory (now optional, should be default in future)"
     742              :     )]
     743              :     destroy: bool,
     744              : 
     745              :     #[clap(long, help = "Postgres shutdown mode")]
     746              :     #[clap(default_value = "fast")]
     747              :     mode: EndpointTerminateMode,
     748              : }
     749              : 
     750              : #[derive(clap::Args)]
     751              : #[clap(about = "Update the pageservers in the spec file of the compute endpoint")]
     752              : struct EndpointUpdatePageserversCmdArgs {
     753              :     #[clap(help = "Postgres endpoint id")]
     754              :     endpoint_id: String,
     755              : 
     756              :     #[clap(short = 'p', long, help = "Specified pageserver id")]
     757              :     pageserver_id: Option<NodeId>,
     758              : }
     759              : 
     760              : #[derive(clap::Args)]
     761              : #[clap(about = "Generate a JWT for an endpoint")]
     762              : struct EndpointGenerateJwtCmdArgs {
     763              :     #[clap(help = "Postgres endpoint id")]
     764              :     endpoint_id: String,
     765              : 
     766              :     #[clap(short = 's', long, help = "Scope to generate the JWT with", value_parser = ComputeClaimsScope::from_str)]
     767              :     scope: Option<ComputeClaimsScope>,
     768              : }
     769              : 
     770              : #[derive(clap::Subcommand)]
     771              : #[clap(about = "Manage neon_local branch name mappings")]
     772              : enum MappingsCmd {
     773              :     Map(MappingsMapCmdArgs),
     774              : }
     775              : 
     776              : #[derive(clap::Args)]
     777              : #[clap(about = "Create new mapping which cannot exist already")]
     778              : struct MappingsMapCmdArgs {
     779              :     #[clap(
     780              :         long,
     781              :         help = "Tenant id. Represented as a hexadecimal string 32 symbols length"
     782              :     )]
     783              :     tenant_id: TenantId,
     784              :     #[clap(
     785              :         long,
     786              :         help = "Timeline id. Represented as a hexadecimal string 32 symbols length"
     787              :     )]
     788              :     timeline_id: TimelineId,
     789              :     #[clap(long, help = "Branch name to give to the timeline")]
     790              :     branch_name: String,
     791              : }
     792              : 
     793              : ///
     794              : /// Timelines tree element used as a value in the HashMap.
     795              : ///
     796              : struct TimelineTreeEl {
     797              :     /// `TimelineInfo` received from the `pageserver` via the `timeline_list` http API call.
     798              :     pub info: TimelineInfo,
     799              :     /// Name, recovered from neon config mappings
     800              :     pub name: Option<String>,
     801              :     /// Holds all direct children of this timeline referenced using `timeline_id`.
     802              :     pub children: BTreeSet<TimelineId>,
     803              : }
     804              : 
     805              : /// A flock-based guard over the neon_local repository directory
     806              : struct RepoLock {
     807              :     _file: Flock<File>,
     808              : }
     809              : 
     810              : impl RepoLock {
     811            0 :     fn new() -> Result<Self> {
     812            0 :         let repo_dir = File::open(local_env::base_path())?;
     813            0 :         match Flock::lock(repo_dir, FlockArg::LockExclusive) {
     814            0 :             Ok(f) => Ok(Self { _file: f }),
     815            0 :             Err((_, e)) => Err(e).context("flock error"),
     816              :         }
     817            0 :     }
     818              : }
     819              : 
     820              : // Main entry point for the 'neon_local' CLI utility
     821              : //
     822              : // This utility helps to manage neon installation. That includes following:
     823              : //   * Management of local postgres installations running on top of the
     824              : //     pageserver.
     825              : //   * Providing CLI api to the pageserver
     826              : //   * TODO: export/import to/from usual postgres
     827            0 : fn main() -> Result<()> {
     828            0 :     let cli = Cli::parse();
     829              : 
     830              :     // Check for 'neon init' command first.
     831            0 :     let (subcommand_result, _lock) = if let NeonLocalCmd::Init(args) = cli.command {
     832            0 :         (handle_init(&args).map(|env| Some(Cow::Owned(env))), None)
     833              :     } else {
     834              :         // This tool uses a collection of simple files to store its state, and consequently
     835              :         // it is not generally safe to run multiple commands concurrently.  Rather than expect
     836              :         // all callers to know this, use a lock file to protect against concurrent execution.
     837            0 :         let _repo_lock = RepoLock::new().unwrap();
     838              : 
     839              :         // all other commands need an existing config
     840            0 :         let env = LocalEnv::load_config(&local_env::base_path()).context("Error loading config")?;
     841            0 :         let original_env = env.clone();
     842            0 :         let env = Box::leak(Box::new(env));
     843            0 :         let rt = tokio::runtime::Builder::new_current_thread()
     844            0 :             .enable_all()
     845            0 :             .build()
     846            0 :             .unwrap();
     847              : 
     848            0 :         let subcommand_result = match cli.command {
     849            0 :             NeonLocalCmd::Init(_) => unreachable!("init was handled earlier already"),
     850            0 :             NeonLocalCmd::Start(args) => rt.block_on(handle_start_all(&args, env)),
     851            0 :             NeonLocalCmd::Stop(args) => rt.block_on(handle_stop_all(&args, env)),
     852            0 :             NeonLocalCmd::Tenant(subcmd) => rt.block_on(handle_tenant(&subcmd, env)),
     853            0 :             NeonLocalCmd::Timeline(subcmd) => rt.block_on(handle_timeline(&subcmd, env)),
     854            0 :             NeonLocalCmd::Pageserver(subcmd) => rt.block_on(handle_pageserver(&subcmd, env)),
     855            0 :             NeonLocalCmd::StorageController(subcmd) => {
     856            0 :                 rt.block_on(handle_storage_controller(&subcmd, env))
     857              :             }
     858            0 :             NeonLocalCmd::StorageBroker(subcmd) => rt.block_on(handle_storage_broker(&subcmd, env)),
     859            0 :             NeonLocalCmd::Safekeeper(subcmd) => rt.block_on(handle_safekeeper(&subcmd, env)),
     860            0 :             NeonLocalCmd::EndpointStorage(subcmd) => {
     861            0 :                 rt.block_on(handle_endpoint_storage(&subcmd, env))
     862              :             }
     863            0 :             NeonLocalCmd::Endpoint(subcmd) => rt.block_on(handle_endpoint(&subcmd, env)),
     864            0 :             NeonLocalCmd::Mappings(subcmd) => handle_mappings(&subcmd, env),
     865              :         };
     866              : 
     867            0 :         let subcommand_result = if &original_env != env {
     868            0 :             subcommand_result.map(|()| Some(Cow::Borrowed(env)))
     869              :         } else {
     870            0 :             subcommand_result.map(|()| None)
     871              :         };
     872            0 :         (subcommand_result, Some(_repo_lock))
     873              :     };
     874              : 
     875            0 :     match subcommand_result {
     876            0 :         Ok(Some(updated_env)) => updated_env.persist_config()?,
     877            0 :         Ok(None) => (),
     878            0 :         Err(e) => {
     879            0 :             eprintln!("command failed: {e:?}");
     880            0 :             exit(1);
     881              :         }
     882              :     }
     883            0 :     Ok(())
     884            0 : }
     885              : 
     886              : ///
     887              : /// Prints timelines list as a tree-like structure.
     888              : ///
     889            0 : fn print_timelines_tree(
     890            0 :     timelines: Vec<TimelineInfo>,
     891            0 :     mut timeline_name_mappings: HashMap<TenantTimelineId, String>,
     892            0 : ) -> Result<()> {
     893            0 :     let mut timelines_hash = timelines
     894            0 :         .iter()
     895            0 :         .map(|t| {
     896            0 :             (
     897            0 :                 t.timeline_id,
     898            0 :                 TimelineTreeEl {
     899            0 :                     info: t.clone(),
     900            0 :                     children: BTreeSet::new(),
     901            0 :                     name: timeline_name_mappings
     902            0 :                         .remove(&TenantTimelineId::new(t.tenant_id.tenant_id, t.timeline_id)),
     903            0 :                 },
     904            0 :             )
     905            0 :         })
     906            0 :         .collect::<HashMap<_, _>>();
     907              : 
     908              :     // Memorize all direct children of each timeline.
     909            0 :     for timeline in timelines.iter() {
     910            0 :         if let Some(ancestor_timeline_id) = timeline.ancestor_timeline_id {
     911            0 :             timelines_hash
     912            0 :                 .get_mut(&ancestor_timeline_id)
     913            0 :                 .context("missing timeline info in the HashMap")?
     914              :                 .children
     915            0 :                 .insert(timeline.timeline_id);
     916            0 :         }
     917              :     }
     918              : 
     919            0 :     for timeline in timelines_hash.values() {
     920              :         // Start with root local timelines (no ancestors) first.
     921            0 :         if timeline.info.ancestor_timeline_id.is_none() {
     922            0 :             print_timeline(0, &Vec::from([true]), timeline, &timelines_hash)?;
     923            0 :         }
     924              :     }
     925              : 
     926            0 :     Ok(())
     927            0 : }
     928              : 
     929              : ///
     930              : /// Recursively prints timeline info with all its children.
     931              : ///
     932            0 : fn print_timeline(
     933            0 :     nesting_level: usize,
     934            0 :     is_last: &[bool],
     935            0 :     timeline: &TimelineTreeEl,
     936            0 :     timelines: &HashMap<TimelineId, TimelineTreeEl>,
     937            0 : ) -> Result<()> {
     938            0 :     if nesting_level > 0 {
     939            0 :         let ancestor_lsn = match timeline.info.ancestor_lsn {
     940            0 :             Some(lsn) => lsn.to_string(),
     941            0 :             None => "Unknown Lsn".to_string(),
     942              :         };
     943              : 
     944            0 :         let mut br_sym = "┣━";
     945              : 
     946              :         // Draw each nesting padding with proper style
     947              :         // depending on whether its timeline ended or not.
     948            0 :         if nesting_level > 1 {
     949            0 :             for l in &is_last[1..is_last.len() - 1] {
     950            0 :                 if *l {
     951            0 :                     print!("   ");
     952            0 :                 } else {
     953            0 :                     print!("┃  ");
     954            0 :                 }
     955              :             }
     956            0 :         }
     957              : 
     958              :         // We are the last in this sub-timeline
     959            0 :         if *is_last.last().unwrap() {
     960            0 :             br_sym = "┗━";
     961            0 :         }
     962              : 
     963            0 :         print!("{br_sym} @{ancestor_lsn}: ");
     964            0 :     }
     965              : 
     966              :     // Finally print a timeline id and name with new line
     967            0 :     println!(
     968            0 :         "{} [{}]",
     969            0 :         timeline.name.as_deref().unwrap_or("_no_name_"),
     970              :         timeline.info.timeline_id
     971              :     );
     972              : 
     973            0 :     let len = timeline.children.len();
     974            0 :     let mut i: usize = 0;
     975            0 :     let mut is_last_new = Vec::from(is_last);
     976            0 :     is_last_new.push(false);
     977              : 
     978            0 :     for child in &timeline.children {
     979            0 :         i += 1;
     980              : 
     981              :         // Mark that the last padding is the end of the timeline
     982            0 :         if i == len {
     983            0 :             if let Some(last) = is_last_new.last_mut() {
     984            0 :                 *last = true;
     985            0 :             }
     986            0 :         }
     987              : 
     988            0 :         print_timeline(
     989            0 :             nesting_level + 1,
     990            0 :             &is_last_new,
     991            0 :             timelines
     992            0 :                 .get(child)
     993            0 :                 .context("missing timeline info in the HashMap")?,
     994            0 :             timelines,
     995            0 :         )?;
     996              :     }
     997              : 
     998            0 :     Ok(())
     999            0 : }
    1000              : 
    1001              : /// Helper function to get tenant id from an optional --tenant_id option or from the config file
    1002            0 : fn get_tenant_id(
    1003            0 :     tenant_id_arg: Option<TenantId>,
    1004            0 :     env: &local_env::LocalEnv,
    1005            0 : ) -> anyhow::Result<TenantId> {
    1006            0 :     if let Some(tenant_id_from_arguments) = tenant_id_arg {
    1007            0 :         Ok(tenant_id_from_arguments)
    1008            0 :     } else if let Some(default_id) = env.default_tenant_id {
    1009            0 :         Ok(default_id)
    1010              :     } else {
    1011            0 :         anyhow::bail!("No tenant id. Use --tenant-id, or set a default tenant");
    1012              :     }
    1013            0 : }
    1014              : 
    1015              : /// Helper function to get tenant-shard ID from an optional --tenant_id option or from the config file,
    1016              : /// for commands that accept a shard suffix
    1017            0 : fn get_tenant_shard_id(
    1018            0 :     tenant_shard_id_arg: Option<TenantShardId>,
    1019            0 :     env: &local_env::LocalEnv,
    1020            0 : ) -> anyhow::Result<TenantShardId> {
    1021            0 :     if let Some(tenant_id_from_arguments) = tenant_shard_id_arg {
    1022            0 :         Ok(tenant_id_from_arguments)
    1023            0 :     } else if let Some(default_id) = env.default_tenant_id {
    1024            0 :         Ok(TenantShardId::unsharded(default_id))
    1025              :     } else {
    1026            0 :         anyhow::bail!("No tenant shard id. Use --tenant-id, or set a default tenant");
    1027              :     }
    1028            0 : }
    1029              : 
    1030            0 : fn handle_init(args: &InitCmdArgs) -> anyhow::Result<LocalEnv> {
    1031              :     // Create the in-memory `LocalEnv` that we'd normally load from disk in `load_config`.
    1032            0 :     let init_conf: NeonLocalInitConf = if let Some(config_path) = &args.config {
    1033              :         // User (likely the Python test suite) provided a description of the environment.
    1034            0 :         if args.num_pageservers.is_some() {
    1035            0 :             bail!(
    1036            0 :                 "Cannot specify both --num-pageservers and --config, use key `pageservers` in the --config file instead"
    1037              :             );
    1038            0 :         }
    1039              :         // load and parse the file
    1040            0 :         let contents = std::fs::read_to_string(config_path).with_context(|| {
    1041            0 :             format!(
    1042            0 :                 "Could not read configuration file '{}'",
    1043            0 :                 config_path.display()
    1044              :             )
    1045            0 :         })?;
    1046            0 :         toml_edit::de::from_str(&contents)?
    1047              :     } else {
    1048              :         // User (likely interactive) did not provide a description of the environment, give them the default
    1049              :         NeonLocalInitConf {
    1050            0 :             control_plane_api: Some(DEFAULT_PAGESERVER_CONTROL_PLANE_API.parse().unwrap()),
    1051            0 :             broker: NeonBroker {
    1052            0 :                 listen_addr: Some(DEFAULT_BROKER_ADDR.parse().unwrap()),
    1053            0 :                 listen_https_addr: None,
    1054            0 :             },
    1055            0 :             safekeepers: vec![SafekeeperConf {
    1056            0 :                 id: DEFAULT_SAFEKEEPER_ID,
    1057            0 :                 pg_port: DEFAULT_SAFEKEEPER_PG_PORT,
    1058            0 :                 http_port: DEFAULT_SAFEKEEPER_HTTP_PORT,
    1059            0 :                 ..Default::default()
    1060            0 :             }],
    1061            0 :             pageservers: (0..args.num_pageservers.unwrap_or(1))
    1062            0 :                 .map(|i| {
    1063            0 :                     let pageserver_id = NodeId(DEFAULT_PAGESERVER_ID.0 + i as u64);
    1064            0 :                     let pg_port = DEFAULT_PAGESERVER_PG_PORT + i;
    1065            0 :                     let http_port = DEFAULT_PAGESERVER_HTTP_PORT + i;
    1066            0 :                     let grpc_port = DEFAULT_PAGESERVER_GRPC_PORT + i;
    1067            0 :                     NeonLocalInitPageserverConf {
    1068            0 :                         id: pageserver_id,
    1069            0 :                         listen_pg_addr: format!("127.0.0.1:{pg_port}"),
    1070            0 :                         listen_http_addr: format!("127.0.0.1:{http_port}"),
    1071            0 :                         listen_https_addr: None,
    1072            0 :                         listen_grpc_addr: Some(format!("127.0.0.1:{grpc_port}")),
    1073            0 :                         pg_auth_type: AuthType::Trust,
    1074            0 :                         http_auth_type: AuthType::Trust,
    1075            0 :                         grpc_auth_type: AuthType::Trust,
    1076            0 :                         other: Default::default(),
    1077            0 :                         // Typical developer machines use disks with slow fsync, and we don't care
    1078            0 :                         // about data integrity: disable disk syncs.
    1079            0 :                         no_sync: true,
    1080            0 :                     }
    1081            0 :                 })
    1082            0 :                 .collect(),
    1083            0 :             endpoint_storage: EndpointStorageConf {
    1084            0 :                 listen_addr: ENDPOINT_STORAGE_DEFAULT_ADDR,
    1085            0 :             },
    1086            0 :             pg_distrib_dir: None,
    1087            0 :             neon_distrib_dir: None,
    1088            0 :             default_tenant_id: TenantId::from_array(std::array::from_fn(|_| 0)),
    1089            0 :             storage_controller: None,
    1090            0 :             control_plane_hooks_api: None,
    1091              :             generate_local_ssl_certs: false,
    1092              :         }
    1093              :     };
    1094              : 
    1095            0 :     LocalEnv::init(init_conf, &args.force)
    1096            0 :         .context("materialize initial neon_local environment on disk")?;
    1097            0 :     Ok(LocalEnv::load_config(&local_env::base_path())
    1098            0 :         .expect("freshly written config should be loadable"))
    1099            0 : }
    1100              : 
    1101              : /// The default pageserver is the one where CLI tenant/timeline operations are sent by default.
    1102              : /// For typical interactive use, one would just run with a single pageserver.  Scenarios with
    1103              : /// tenant/timeline placement across multiple pageservers are managed by python test code rather
    1104              : /// than this CLI.
    1105            0 : fn get_default_pageserver(env: &local_env::LocalEnv) -> PageServerNode {
    1106            0 :     let ps_conf = env
    1107            0 :         .pageservers
    1108            0 :         .first()
    1109            0 :         .expect("Config is validated to contain at least one pageserver");
    1110            0 :     PageServerNode::from_env(env, ps_conf)
    1111            0 : }
    1112              : 
    1113            0 : async fn handle_tenant(subcmd: &TenantCmd, env: &mut local_env::LocalEnv) -> anyhow::Result<()> {
    1114            0 :     let pageserver = get_default_pageserver(env);
    1115            0 :     match subcmd {
    1116              :         TenantCmd::List => {
    1117            0 :             for t in pageserver.tenant_list().await? {
    1118            0 :                 println!("{} {:?}", t.id, t.state);
    1119            0 :             }
    1120              :         }
    1121            0 :         TenantCmd::Import(args) => {
    1122            0 :             let tenant_id = args.tenant_id;
    1123              : 
    1124            0 :             let storage_controller = StorageController::from_env(env);
    1125            0 :             let create_response = storage_controller.tenant_import(tenant_id).await?;
    1126              : 
    1127            0 :             let shard_zero = create_response
    1128            0 :                 .shards
    1129            0 :                 .first()
    1130            0 :                 .expect("Import response omitted shards");
    1131              : 
    1132            0 :             let attached_pageserver_id = shard_zero.node_id;
    1133            0 :             let pageserver =
    1134            0 :                 PageServerNode::from_env(env, env.get_pageserver_conf(attached_pageserver_id)?);
    1135              : 
    1136            0 :             println!(
    1137            0 :                 "Imported tenant {tenant_id}, attached to pageserver {attached_pageserver_id}"
    1138              :             );
    1139              : 
    1140            0 :             let timelines = pageserver
    1141            0 :                 .http_client
    1142            0 :                 .list_timelines(shard_zero.shard_id)
    1143            0 :                 .await?;
    1144              : 
    1145              :             // Pick a 'main' timeline that has no ancestors, the rest will get arbitrary names
    1146            0 :             let main_timeline = timelines
    1147            0 :                 .iter()
    1148            0 :                 .find(|t| t.ancestor_timeline_id.is_none())
    1149            0 :                 .expect("No timelines found")
    1150              :                 .timeline_id;
    1151              : 
    1152            0 :             let mut branch_i = 0;
    1153            0 :             for timeline in timelines.iter() {
    1154            0 :                 let branch_name = if timeline.timeline_id == main_timeline {
    1155            0 :                     "main".to_string()
    1156              :                 } else {
    1157            0 :                     branch_i += 1;
    1158            0 :                     format!("branch_{branch_i}")
    1159              :                 };
    1160              : 
    1161            0 :                 println!(
    1162            0 :                     "Importing timeline {tenant_id}/{} as branch {branch_name}",
    1163              :                     timeline.timeline_id
    1164              :                 );
    1165              : 
    1166            0 :                 env.register_branch_mapping(branch_name, tenant_id, timeline.timeline_id)?;
    1167              :             }
    1168              :         }
    1169            0 :         TenantCmd::Create(args) => {
    1170            0 :             let tenant_conf: HashMap<_, _> =
    1171            0 :                 args.config.iter().flat_map(|c| c.split_once(':')).collect();
    1172              : 
    1173            0 :             let tenant_conf = PageServerNode::parse_config(tenant_conf)?;
    1174              : 
    1175              :             // If tenant ID was not specified, generate one
    1176            0 :             let tenant_id = args.tenant_id.unwrap_or_else(TenantId::generate);
    1177              : 
    1178              :             // We must register the tenant with the storage controller, so
    1179              :             // that when the pageserver restarts, it will be re-attached.
    1180            0 :             let storage_controller = StorageController::from_env(env);
    1181            0 :             storage_controller
    1182            0 :                 .tenant_create(TenantCreateRequest {
    1183            0 :                     // Note that ::unsharded here isn't actually because the tenant is unsharded, its because the
    1184            0 :                     // storage controller expects a shard-naive tenant_id in this attribute, and the TenantCreateRequest
    1185            0 :                     // type is used both in the storage controller (for creating tenants) and in the pageserver (for
    1186            0 :                     // creating shards)
    1187            0 :                     new_tenant_id: TenantShardId::unsharded(tenant_id),
    1188            0 :                     generation: None,
    1189            0 :                     shard_parameters: ShardParameters {
    1190            0 :                         count: ShardCount::new(args.shard_count),
    1191            0 :                         stripe_size: args
    1192            0 :                             .shard_stripe_size
    1193            0 :                             .map(ShardStripeSize)
    1194            0 :                             .unwrap_or(DEFAULT_STRIPE_SIZE),
    1195            0 :                     },
    1196            0 :                     placement_policy: args.placement_policy.clone(),
    1197            0 :                     config: tenant_conf,
    1198            0 :                 })
    1199            0 :                 .await?;
    1200            0 :             println!("tenant {tenant_id} successfully created on the pageserver");
    1201              : 
    1202              :             // Create an initial timeline for the new tenant
    1203            0 :             let new_timeline_id = args.timeline_id.unwrap_or(TimelineId::generate());
    1204              : 
    1205              :             // FIXME: passing None for ancestor_start_lsn is not kosher in a sharded world: we can't have
    1206              :             // different shards picking different start lsns.  Maybe we have to teach storage controller
    1207              :             // to let shard 0 branch first and then propagate the chosen LSN to other shards.
    1208            0 :             storage_controller
    1209            0 :                 .tenant_timeline_create(
    1210            0 :                     tenant_id,
    1211            0 :                     TimelineCreateRequest {
    1212            0 :                         new_timeline_id,
    1213            0 :                         mode: pageserver_api::models::TimelineCreateRequestMode::Bootstrap {
    1214            0 :                             existing_initdb_timeline_id: None,
    1215            0 :                             pg_version: Some(args.pg_version),
    1216            0 :                         },
    1217            0 :                     },
    1218            0 :                 )
    1219            0 :                 .await?;
    1220              : 
    1221            0 :             env.register_branch_mapping(
    1222            0 :                 DEFAULT_BRANCH_NAME.to_string(),
    1223            0 :                 tenant_id,
    1224            0 :                 new_timeline_id,
    1225            0 :             )?;
    1226              : 
    1227            0 :             println!("Created an initial timeline '{new_timeline_id}' for tenant: {tenant_id}",);
    1228              : 
    1229            0 :             if args.set_default {
    1230            0 :                 println!("Setting tenant {tenant_id} as a default one");
    1231            0 :                 env.default_tenant_id = Some(tenant_id);
    1232            0 :             }
    1233              :         }
    1234            0 :         TenantCmd::SetDefault(args) => {
    1235            0 :             println!("Setting tenant {} as a default one", args.tenant_id);
    1236            0 :             env.default_tenant_id = Some(args.tenant_id);
    1237            0 :         }
    1238            0 :         TenantCmd::Config(args) => {
    1239            0 :             let tenant_id = get_tenant_id(args.tenant_id, env)?;
    1240            0 :             let tenant_conf: HashMap<_, _> =
    1241            0 :                 args.config.iter().flat_map(|c| c.split_once(':')).collect();
    1242            0 :             let config = PageServerNode::parse_config(tenant_conf)?;
    1243              : 
    1244            0 :             let req = TenantConfigRequest { tenant_id, config };
    1245              : 
    1246            0 :             let storage_controller = StorageController::from_env(env);
    1247            0 :             storage_controller
    1248            0 :                 .set_tenant_config(&req)
    1249            0 :                 .await
    1250            0 :                 .with_context(|| format!("Tenant config failed for tenant with id {tenant_id}"))?;
    1251            0 :             println!("tenant {tenant_id} successfully configured via storcon");
    1252              :         }
    1253              :     }
    1254            0 :     Ok(())
    1255            0 : }
    1256              : 
    1257            0 : async fn handle_timeline(cmd: &TimelineCmd, env: &mut local_env::LocalEnv) -> Result<()> {
    1258            0 :     let pageserver = get_default_pageserver(env);
    1259              : 
    1260            0 :     match cmd {
    1261            0 :         TimelineCmd::List(args) => {
    1262              :             // TODO(sharding): this command shouldn't have to specify a shard ID: we should ask the storage controller
    1263              :             // where shard 0 is attached, and query there.
    1264            0 :             let tenant_shard_id = get_tenant_shard_id(args.tenant_shard_id, env)?;
    1265            0 :             let timelines = pageserver.timeline_list(&tenant_shard_id).await?;
    1266            0 :             print_timelines_tree(timelines, env.timeline_name_mappings())?;
    1267              :         }
    1268            0 :         TimelineCmd::Create(args) => {
    1269            0 :             let tenant_id = get_tenant_id(args.tenant_id, env)?;
    1270            0 :             let new_branch_name = &args.branch_name;
    1271            0 :             let new_timeline_id_opt = args.timeline_id;
    1272            0 :             let new_timeline_id = new_timeline_id_opt.unwrap_or(TimelineId::generate());
    1273              : 
    1274            0 :             let storage_controller = StorageController::from_env(env);
    1275            0 :             let create_req = TimelineCreateRequest {
    1276            0 :                 new_timeline_id,
    1277            0 :                 mode: pageserver_api::models::TimelineCreateRequestMode::Bootstrap {
    1278            0 :                     existing_initdb_timeline_id: None,
    1279            0 :                     pg_version: Some(args.pg_version),
    1280            0 :                 },
    1281            0 :             };
    1282            0 :             let timeline_info = storage_controller
    1283            0 :                 .tenant_timeline_create(tenant_id, create_req)
    1284            0 :                 .await?;
    1285              : 
    1286            0 :             let last_record_lsn = timeline_info.last_record_lsn;
    1287            0 :             env.register_branch_mapping(new_branch_name.to_string(), tenant_id, new_timeline_id)?;
    1288              : 
    1289            0 :             println!(
    1290            0 :                 "Created timeline '{}' at Lsn {last_record_lsn} for tenant: {tenant_id}",
    1291              :                 timeline_info.timeline_id
    1292              :             );
    1293              :         }
    1294              :         // TODO: rename to import-basebackup-plus-wal
    1295            0 :         TimelineCmd::Import(args) => {
    1296            0 :             let tenant_id = get_tenant_id(args.tenant_id, env)?;
    1297            0 :             let timeline_id = args.timeline_id;
    1298            0 :             let branch_name = &args.branch_name;
    1299              : 
    1300              :             // Parse base inputs
    1301            0 :             let base = (args.base_lsn, args.base_tarfile.clone());
    1302              : 
    1303              :             // Parse pg_wal inputs
    1304            0 :             let wal_tarfile = args.wal_tarfile.clone();
    1305            0 :             let end_lsn = args.end_lsn;
    1306              :             // TODO validate both or none are provided
    1307            0 :             let pg_wal = end_lsn.zip(wal_tarfile);
    1308              : 
    1309            0 :             println!("Importing timeline into pageserver ...");
    1310            0 :             pageserver
    1311            0 :                 .timeline_import(tenant_id, timeline_id, base, pg_wal, args.pg_version)
    1312            0 :                 .await?;
    1313            0 :             if env.storage_controller.timelines_onto_safekeepers {
    1314            0 :                 println!("Creating timeline on safekeeper ...");
    1315            0 :                 let timeline_info = pageserver
    1316            0 :                     .timeline_info(
    1317            0 :                         TenantShardId::unsharded(tenant_id),
    1318            0 :                         timeline_id,
    1319            0 :                         pageserver_client::mgmt_api::ForceAwaitLogicalSize::No,
    1320            0 :                     )
    1321            0 :                     .await?;
    1322            0 :                 let default_sk = SafekeeperNode::from_env(env, env.safekeepers.first().unwrap());
    1323            0 :                 let default_host = default_sk
    1324            0 :                     .conf
    1325            0 :                     .listen_addr
    1326            0 :                     .clone()
    1327            0 :                     .unwrap_or_else(|| "localhost".to_string());
    1328            0 :                 let mconf = safekeeper_api::membership::Configuration {
    1329            0 :                     generation: SafekeeperGeneration::new(1),
    1330            0 :                     members: safekeeper_api::membership::MemberSet {
    1331            0 :                         m: vec![SafekeeperId {
    1332            0 :                             host: default_host,
    1333            0 :                             id: default_sk.conf.id,
    1334            0 :                             pg_port: default_sk.conf.pg_port,
    1335            0 :                         }],
    1336            0 :                     },
    1337            0 :                     new_members: None,
    1338            0 :                 };
    1339            0 :                 let pg_version = PgVersionId::from(args.pg_version);
    1340            0 :                 let req = safekeeper_api::models::TimelineCreateRequest {
    1341            0 :                     tenant_id,
    1342            0 :                     timeline_id,
    1343            0 :                     mconf,
    1344            0 :                     pg_version,
    1345            0 :                     system_id: None,
    1346            0 :                     wal_seg_size: None,
    1347            0 :                     start_lsn: timeline_info.last_record_lsn,
    1348            0 :                     commit_lsn: None,
    1349            0 :                 };
    1350            0 :                 default_sk.create_timeline(&req).await?;
    1351            0 :             }
    1352            0 :             env.register_branch_mapping(branch_name.to_string(), tenant_id, timeline_id)?;
    1353            0 :             println!("Done");
    1354              :         }
    1355            0 :         TimelineCmd::Branch(args) => {
    1356            0 :             let tenant_id = get_tenant_id(args.tenant_id, env)?;
    1357            0 :             let new_timeline_id = args.timeline_id.unwrap_or(TimelineId::generate());
    1358            0 :             let new_branch_name = &args.branch_name;
    1359            0 :             let ancestor_branch_name = args
    1360            0 :                 .ancestor_branch_name
    1361            0 :                 .clone()
    1362            0 :                 .unwrap_or(DEFAULT_BRANCH_NAME.to_owned());
    1363            0 :             let ancestor_timeline_id = env
    1364            0 :                 .get_branch_timeline_id(&ancestor_branch_name, tenant_id)
    1365            0 :                 .ok_or_else(|| {
    1366            0 :                     anyhow!("Found no timeline id for branch name '{ancestor_branch_name}'")
    1367            0 :                 })?;
    1368              : 
    1369            0 :             let start_lsn = args.ancestor_start_lsn;
    1370            0 :             let storage_controller = StorageController::from_env(env);
    1371            0 :             let create_req = TimelineCreateRequest {
    1372            0 :                 new_timeline_id,
    1373            0 :                 mode: pageserver_api::models::TimelineCreateRequestMode::Branch {
    1374            0 :                     ancestor_timeline_id,
    1375            0 :                     ancestor_start_lsn: start_lsn,
    1376            0 :                     read_only: false,
    1377            0 :                     pg_version: None,
    1378            0 :                 },
    1379            0 :             };
    1380            0 :             let timeline_info = storage_controller
    1381            0 :                 .tenant_timeline_create(tenant_id, create_req)
    1382            0 :                 .await?;
    1383              : 
    1384            0 :             let last_record_lsn = timeline_info.last_record_lsn;
    1385              : 
    1386            0 :             env.register_branch_mapping(new_branch_name.to_string(), tenant_id, new_timeline_id)?;
    1387              : 
    1388            0 :             println!(
    1389            0 :                 "Created timeline '{}' at Lsn {last_record_lsn} for tenant: {tenant_id}. Ancestor timeline: '{ancestor_branch_name}'",
    1390              :                 timeline_info.timeline_id
    1391              :             );
    1392              :         }
    1393              :     }
    1394              : 
    1395            0 :     Ok(())
    1396            0 : }
    1397              : 
    1398            0 : async fn handle_endpoint(subcmd: &EndpointCmd, env: &local_env::LocalEnv) -> Result<()> {
    1399            0 :     let mut cplane = ComputeControlPlane::load(env.clone())?;
    1400              : 
    1401            0 :     match subcmd {
    1402            0 :         EndpointCmd::List(args) => {
    1403              :             // TODO(sharding): this command shouldn't have to specify a shard ID: we should ask the storage controller
    1404              :             // where shard 0 is attached, and query there.
    1405            0 :             let tenant_shard_id = get_tenant_shard_id(args.tenant_shard_id, env)?;
    1406              : 
    1407            0 :             let timeline_name_mappings = env.timeline_name_mappings();
    1408              : 
    1409            0 :             let mut table = comfy_table::Table::new();
    1410              : 
    1411            0 :             table.load_preset(comfy_table::presets::NOTHING);
    1412              : 
    1413            0 :             table.set_header([
    1414            0 :                 "ENDPOINT",
    1415            0 :                 "ADDRESS",
    1416            0 :                 "TIMELINE",
    1417            0 :                 "BRANCH NAME",
    1418            0 :                 "LSN",
    1419            0 :                 "STATUS",
    1420            0 :             ]);
    1421              : 
    1422            0 :             for (endpoint_id, endpoint) in cplane
    1423            0 :                 .endpoints
    1424            0 :                 .iter()
    1425            0 :                 .filter(|(_, endpoint)| endpoint.tenant_id == tenant_shard_id.tenant_id)
    1426              :             {
    1427            0 :                 let lsn_str = match endpoint.mode {
    1428            0 :                     ComputeMode::Static(lsn) => {
    1429              :                         // -> read-only endpoint
    1430              :                         // Use the node's LSN.
    1431            0 :                         lsn.to_string()
    1432              :                     }
    1433              :                     _ => {
    1434              :                         // As the LSN here refers to the one that the compute is started with,
    1435              :                         // we display nothing as it is a primary/hot standby compute.
    1436            0 :                         "---".to_string()
    1437              :                     }
    1438              :                 };
    1439              : 
    1440            0 :                 let branch_name = timeline_name_mappings
    1441            0 :                     .get(&TenantTimelineId::new(
    1442            0 :                         tenant_shard_id.tenant_id,
    1443            0 :                         endpoint.timeline_id,
    1444            0 :                     ))
    1445            0 :                     .map(|name| name.as_str())
    1446            0 :                     .unwrap_or("?");
    1447              : 
    1448            0 :                 table.add_row([
    1449            0 :                     endpoint_id.as_str(),
    1450            0 :                     &endpoint.pg_address.to_string(),
    1451            0 :                     &endpoint.timeline_id.to_string(),
    1452            0 :                     branch_name,
    1453            0 :                     lsn_str.as_str(),
    1454            0 :                     &format!("{}", endpoint.status()),
    1455            0 :                 ]);
    1456              :             }
    1457              : 
    1458            0 :             println!("{table}");
    1459              :         }
    1460            0 :         EndpointCmd::Create(args) => {
    1461            0 :             let tenant_id = get_tenant_id(args.tenant_id, env)?;
    1462            0 :             let branch_name = args
    1463            0 :                 .branch_name
    1464            0 :                 .clone()
    1465            0 :                 .unwrap_or(DEFAULT_BRANCH_NAME.to_owned());
    1466            0 :             let endpoint_id = args
    1467            0 :                 .endpoint_id
    1468            0 :                 .clone()
    1469            0 :                 .unwrap_or_else(|| format!("ep-{branch_name}"));
    1470              : 
    1471            0 :             let timeline_id = env
    1472            0 :                 .get_branch_timeline_id(&branch_name, tenant_id)
    1473            0 :                 .ok_or_else(|| anyhow!("Found no timeline id for branch name '{branch_name}'"))?;
    1474              : 
    1475            0 :             let mode = match (args.lsn, args.hot_standby) {
    1476            0 :                 (Some(lsn), false) => ComputeMode::Static(lsn),
    1477            0 :                 (None, true) => ComputeMode::Replica,
    1478            0 :                 (None, false) => ComputeMode::Primary,
    1479            0 :                 (Some(_), true) => anyhow::bail!("cannot specify both lsn and hot-standby"),
    1480              :             };
    1481              : 
    1482            0 :             match (mode, args.hot_standby) {
    1483              :                 (ComputeMode::Static(_), true) => {
    1484            0 :                     bail!(
    1485            0 :                         "Cannot start a node in hot standby mode when it is already configured as a static replica"
    1486              :                     )
    1487              :                 }
    1488              :                 (ComputeMode::Primary, true) => {
    1489            0 :                     bail!(
    1490            0 :                         "Cannot start a node as a hot standby replica, it is already configured as primary node"
    1491              :                     )
    1492              :                 }
    1493            0 :                 _ => {}
    1494              :             }
    1495              : 
    1496            0 :             if !args.allow_multiple {
    1497            0 :                 cplane.check_conflicting_endpoints(mode, tenant_id, timeline_id)?;
    1498            0 :             }
    1499              : 
    1500            0 :             cplane.new_endpoint(
    1501            0 :                 &endpoint_id,
    1502            0 :                 tenant_id,
    1503            0 :                 timeline_id,
    1504            0 :                 args.pg_port,
    1505            0 :                 args.external_http_port,
    1506            0 :                 args.internal_http_port,
    1507            0 :                 args.pg_version,
    1508            0 :                 mode,
    1509            0 :                 args.grpc,
    1510            0 :                 !args.update_catalog,
    1511              :                 false,
    1512            0 :                 args.privileged_role_name.clone(),
    1513            0 :             )?;
    1514              :         }
    1515            0 :         EndpointCmd::Start(args) => {
    1516            0 :             let endpoint_id = &args.endpoint_id;
    1517            0 :             let pageserver_id = args.endpoint_pageserver_id;
    1518            0 :             let remote_ext_base_url = &args.remote_ext_base_url;
    1519              : 
    1520            0 :             let default_generation = env
    1521            0 :                 .storage_controller
    1522            0 :                 .timelines_onto_safekeepers
    1523            0 :                 .then_some(1);
    1524            0 :             let safekeepers_generation = args
    1525            0 :                 .safekeepers_generation
    1526            0 :                 .or(default_generation)
    1527            0 :                 .map(SafekeeperGeneration::new);
    1528              :             // If --safekeepers argument is given, use only the listed
    1529              :             // safekeeper nodes; otherwise all from the env.
    1530            0 :             let safekeepers = if let Some(safekeepers) = parse_safekeepers(&args.safekeepers)? {
    1531            0 :                 safekeepers
    1532              :             } else {
    1533            0 :                 env.safekeepers.iter().map(|sk| sk.id).collect()
    1534              :             };
    1535              : 
    1536            0 :             let endpoint = cplane
    1537            0 :                 .endpoints
    1538            0 :                 .get(endpoint_id.as_str())
    1539            0 :                 .ok_or_else(|| anyhow!("endpoint {endpoint_id} not found"))?;
    1540              : 
    1541            0 :             if !args.allow_multiple {
    1542            0 :                 cplane.check_conflicting_endpoints(
    1543            0 :                     endpoint.mode,
    1544            0 :                     endpoint.tenant_id,
    1545            0 :                     endpoint.timeline_id,
    1546            0 :                 )?;
    1547            0 :             }
    1548              : 
    1549            0 :             let (pageservers, stripe_size) = if let Some(pageserver_id) = pageserver_id {
    1550            0 :                 let conf = env.get_pageserver_conf(pageserver_id).unwrap();
    1551              :                 // Use gRPC if requested.
    1552            0 :                 let pageserver = if endpoint.grpc {
    1553            0 :                     let grpc_addr = conf.listen_grpc_addr.as_ref().expect("bad config");
    1554            0 :                     let (host, port) = parse_host_port(grpc_addr)?;
    1555            0 :                     let port = port.unwrap_or(DEFAULT_PAGESERVER_GRPC_PORT);
    1556            0 :                     (PageserverProtocol::Grpc, host, port)
    1557              :                 } else {
    1558            0 :                     let (host, port) = parse_host_port(&conf.listen_pg_addr)?;
    1559            0 :                     let port = port.unwrap_or(5432);
    1560            0 :                     (PageserverProtocol::Libpq, host, port)
    1561              :                 };
    1562              :                 // If caller is telling us what pageserver to use, this is not a tenant which is
    1563              :                 // fully managed by storage controller, therefore not sharded.
    1564            0 :                 (vec![pageserver], DEFAULT_STRIPE_SIZE)
    1565              :             } else {
    1566              :                 // Look up the currently attached location of the tenant, and its striping metadata,
    1567              :                 // to pass these on to postgres.
    1568            0 :                 let storage_controller = StorageController::from_env(env);
    1569            0 :                 let locate_result = storage_controller.tenant_locate(endpoint.tenant_id).await?;
    1570            0 :                 let pageservers = futures::future::try_join_all(
    1571            0 :                     locate_result.shards.into_iter().map(|shard| async move {
    1572            0 :                         if let ComputeMode::Static(lsn) = endpoint.mode {
    1573              :                             // Initialize LSN leases for static computes.
    1574            0 :                             let conf = env.get_pageserver_conf(shard.node_id).unwrap();
    1575            0 :                             let pageserver = PageServerNode::from_env(env, conf);
    1576              : 
    1577            0 :                             pageserver
    1578            0 :                                 .http_client
    1579            0 :                                 .timeline_init_lsn_lease(shard.shard_id, endpoint.timeline_id, lsn)
    1580            0 :                                 .await?;
    1581            0 :                         }
    1582              : 
    1583            0 :                         let pageserver = if endpoint.grpc {
    1584              :                             (
    1585            0 :                                 PageserverProtocol::Grpc,
    1586            0 :                                 Host::parse(&shard.listen_grpc_addr.expect("no gRPC address"))?,
    1587            0 :                                 shard.listen_grpc_port.expect("no gRPC port"),
    1588              :                             )
    1589              :                         } else {
    1590              :                             (
    1591            0 :                                 PageserverProtocol::Libpq,
    1592            0 :                                 Host::parse(&shard.listen_pg_addr)?,
    1593            0 :                                 shard.listen_pg_port,
    1594              :                             )
    1595              :                         };
    1596            0 :                         anyhow::Ok(pageserver)
    1597            0 :                     }),
    1598              :                 )
    1599            0 :                 .await?;
    1600            0 :                 let stripe_size = locate_result.shard_params.stripe_size;
    1601              : 
    1602            0 :                 (pageservers, stripe_size)
    1603              :             };
    1604            0 :             assert!(!pageservers.is_empty());
    1605              : 
    1606            0 :             let ps_conf = env.get_pageserver_conf(DEFAULT_PAGESERVER_ID)?;
    1607            0 :             let auth_token = if matches!(ps_conf.pg_auth_type, AuthType::NeonJWT) {
    1608            0 :                 let claims = Claims::new(Some(endpoint.tenant_id), Scope::Tenant);
    1609              : 
    1610            0 :                 Some(env.generate_auth_token(&claims)?)
    1611              :             } else {
    1612            0 :                 None
    1613              :             };
    1614              : 
    1615            0 :             let exp = (std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH)?
    1616            0 :                 + Duration::from_secs(86400))
    1617            0 :             .as_secs();
    1618            0 :             let claims = endpoint_storage::claims::EndpointStorageClaims {
    1619            0 :                 tenant_id: endpoint.tenant_id,
    1620            0 :                 timeline_id: endpoint.timeline_id,
    1621            0 :                 endpoint_id: endpoint_id.to_string(),
    1622            0 :                 exp,
    1623            0 :             };
    1624              : 
    1625            0 :             let endpoint_storage_token = env.generate_auth_token(&claims)?;
    1626            0 :             let endpoint_storage_addr = env.endpoint_storage.listen_addr.to_string();
    1627              : 
    1628            0 :             let args = control_plane::endpoint::EndpointStartArgs {
    1629            0 :                 auth_token,
    1630            0 :                 endpoint_storage_token,
    1631            0 :                 endpoint_storage_addr,
    1632            0 :                 safekeepers_generation,
    1633            0 :                 safekeepers,
    1634            0 :                 pageservers,
    1635            0 :                 remote_ext_base_url: remote_ext_base_url.clone(),
    1636            0 :                 shard_stripe_size: stripe_size.0 as usize,
    1637            0 :                 create_test_user: args.create_test_user,
    1638            0 :                 start_timeout: args.start_timeout,
    1639            0 :                 autoprewarm: args.autoprewarm,
    1640            0 :                 offload_lfc_interval_seconds: args.offload_lfc_interval_seconds,
    1641            0 :                 dev: args.dev,
    1642            0 :             };
    1643              : 
    1644            0 :             println!("Starting existing endpoint {endpoint_id}...");
    1645            0 :             endpoint.start(args).await?;
    1646              :         }
    1647            0 :         EndpointCmd::UpdatePageservers(args) => {
    1648            0 :             let endpoint_id = &args.endpoint_id;
    1649            0 :             let endpoint = cplane
    1650            0 :                 .endpoints
    1651            0 :                 .get(endpoint_id.as_str())
    1652            0 :                 .with_context(|| format!("postgres endpoint {endpoint_id} is not found"))?;
    1653            0 :             let pageservers = match args.pageserver_id {
    1654            0 :                 Some(pageserver_id) => {
    1655            0 :                     let pageserver =
    1656            0 :                         PageServerNode::from_env(env, env.get_pageserver_conf(pageserver_id)?);
    1657              : 
    1658            0 :                     vec![(
    1659            0 :                         PageserverProtocol::Libpq,
    1660            0 :                         pageserver.pg_connection_config.host().clone(),
    1661            0 :                         pageserver.pg_connection_config.port(),
    1662            0 :                     )]
    1663              :                 }
    1664              :                 None => {
    1665            0 :                     let storage_controller = StorageController::from_env(env);
    1666            0 :                     storage_controller
    1667            0 :                         .tenant_locate(endpoint.tenant_id)
    1668            0 :                         .await?
    1669              :                         .shards
    1670            0 :                         .into_iter()
    1671            0 :                         .map(|shard| {
    1672            0 :                             (
    1673            0 :                                 PageserverProtocol::Libpq,
    1674            0 :                                 Host::parse(&shard.listen_pg_addr)
    1675            0 :                                     .expect("Storage controller reported malformed host"),
    1676            0 :                                 shard.listen_pg_port,
    1677            0 :                             )
    1678            0 :                         })
    1679            0 :                         .collect::<Vec<_>>()
    1680              :                 }
    1681              :             };
    1682              : 
    1683            0 :             endpoint.update_pageservers_in_config(pageservers).await?;
    1684              :         }
    1685            0 :         EndpointCmd::Reconfigure(args) => {
    1686            0 :             let endpoint_id = &args.endpoint_id;
    1687            0 :             let endpoint = cplane
    1688            0 :                 .endpoints
    1689            0 :                 .get(endpoint_id.as_str())
    1690            0 :                 .with_context(|| format!("postgres endpoint {endpoint_id} is not found"))?;
    1691            0 :             let pageservers = if let Some(ps_id) = args.endpoint_pageserver_id {
    1692            0 :                 let conf = env.get_pageserver_conf(ps_id)?;
    1693              :                 // Use gRPC if requested.
    1694            0 :                 let pageserver = if endpoint.grpc {
    1695            0 :                     let grpc_addr = conf.listen_grpc_addr.as_ref().expect("bad config");
    1696            0 :                     let (host, port) = parse_host_port(grpc_addr)?;
    1697            0 :                     let port = port.unwrap_or(DEFAULT_PAGESERVER_GRPC_PORT);
    1698            0 :                     (PageserverProtocol::Grpc, host, port)
    1699              :                 } else {
    1700            0 :                     let (host, port) = parse_host_port(&conf.listen_pg_addr)?;
    1701            0 :                     let port = port.unwrap_or(5432);
    1702            0 :                     (PageserverProtocol::Libpq, host, port)
    1703              :                 };
    1704            0 :                 vec![pageserver]
    1705              :             } else {
    1706            0 :                 let storage_controller = StorageController::from_env(env);
    1707            0 :                 storage_controller
    1708            0 :                     .tenant_locate(endpoint.tenant_id)
    1709            0 :                     .await?
    1710              :                     .shards
    1711            0 :                     .into_iter()
    1712            0 :                     .map(|shard| {
    1713              :                         // Use gRPC if requested.
    1714            0 :                         if endpoint.grpc {
    1715            0 :                             (
    1716            0 :                                 PageserverProtocol::Grpc,
    1717            0 :                                 Host::parse(&shard.listen_grpc_addr.expect("no gRPC address"))
    1718            0 :                                     .expect("bad hostname"),
    1719            0 :                                 shard.listen_grpc_port.expect("no gRPC port"),
    1720            0 :                             )
    1721              :                         } else {
    1722            0 :                             (
    1723            0 :                                 PageserverProtocol::Libpq,
    1724            0 :                                 Host::parse(&shard.listen_pg_addr).expect("bad hostname"),
    1725            0 :                                 shard.listen_pg_port,
    1726            0 :                             )
    1727              :                         }
    1728            0 :                     })
    1729            0 :                     .collect::<Vec<_>>()
    1730              :             };
    1731              :             // If --safekeepers argument is given, use only the listed
    1732              :             // safekeeper nodes; otherwise all from the env.
    1733            0 :             let safekeepers = parse_safekeepers(&args.safekeepers)?;
    1734            0 :             endpoint
    1735            0 :                 .reconfigure(Some(pageservers), None, safekeepers, None)
    1736            0 :                 .await?;
    1737              :         }
    1738            0 :         EndpointCmd::RefreshConfiguration(args) => {
    1739            0 :             let endpoint_id = &args.endpoint_id;
    1740            0 :             let endpoint = cplane
    1741            0 :                 .endpoints
    1742            0 :                 .get(endpoint_id.as_str())
    1743            0 :                 .with_context(|| format!("postgres endpoint {endpoint_id} is not found"))?;
    1744            0 :             endpoint.refresh_configuration().await?;
    1745              :         }
    1746            0 :         EndpointCmd::Stop(args) => {
    1747            0 :             let endpoint_id = &args.endpoint_id;
    1748            0 :             let endpoint = cplane
    1749            0 :                 .endpoints
    1750            0 :                 .get(endpoint_id)
    1751            0 :                 .with_context(|| format!("postgres endpoint {endpoint_id} is not found"))?;
    1752            0 :             match endpoint.stop(args.mode, args.destroy).await?.lsn {
    1753            0 :                 Some(lsn) => println!("{lsn}"),
    1754            0 :                 None => println!("null"),
    1755              :             }
    1756              :         }
    1757            0 :         EndpointCmd::GenerateJwt(args) => {
    1758            0 :             let endpoint = {
    1759            0 :                 let endpoint_id = &args.endpoint_id;
    1760              : 
    1761            0 :                 cplane
    1762            0 :                     .endpoints
    1763            0 :                     .get(endpoint_id)
    1764            0 :                     .with_context(|| format!("postgres endpoint {endpoint_id} is not found"))?
    1765              :             };
    1766              : 
    1767            0 :             let jwt = endpoint.generate_jwt(args.scope)?;
    1768              : 
    1769            0 :             print!("{jwt}");
    1770              :         }
    1771              :     }
    1772              : 
    1773            0 :     Ok(())
    1774            0 : }
    1775              : 
    1776              : /// Parse --safekeepers as list of safekeeper ids.
    1777            0 : fn parse_safekeepers(safekeepers_str: &Option<String>) -> Result<Option<Vec<NodeId>>> {
    1778            0 :     if let Some(safekeepers_str) = safekeepers_str {
    1779            0 :         let mut safekeepers: Vec<NodeId> = Vec::new();
    1780            0 :         for sk_id in safekeepers_str.split(',').map(str::trim) {
    1781            0 :             let sk_id = NodeId(
    1782            0 :                 u64::from_str(sk_id)
    1783            0 :                     .map_err(|_| anyhow!("invalid node ID \"{sk_id}\" in --safekeepers list"))?,
    1784              :             );
    1785            0 :             safekeepers.push(sk_id);
    1786              :         }
    1787            0 :         Ok(Some(safekeepers))
    1788              :     } else {
    1789            0 :         Ok(None)
    1790              :     }
    1791            0 : }
    1792              : 
    1793            0 : fn handle_mappings(subcmd: &MappingsCmd, env: &mut local_env::LocalEnv) -> Result<()> {
    1794            0 :     match subcmd {
    1795            0 :         MappingsCmd::Map(args) => {
    1796            0 :             env.register_branch_mapping(
    1797            0 :                 args.branch_name.to_owned(),
    1798            0 :                 args.tenant_id,
    1799            0 :                 args.timeline_id,
    1800            0 :             )?;
    1801              : 
    1802            0 :             Ok(())
    1803              :         }
    1804              :     }
    1805            0 : }
    1806              : 
    1807            0 : fn get_pageserver(
    1808            0 :     env: &local_env::LocalEnv,
    1809            0 :     pageserver_id_arg: Option<NodeId>,
    1810            0 : ) -> Result<PageServerNode> {
    1811            0 :     let node_id = pageserver_id_arg.unwrap_or(DEFAULT_PAGESERVER_ID);
    1812              : 
    1813            0 :     Ok(PageServerNode::from_env(
    1814            0 :         env,
    1815            0 :         env.get_pageserver_conf(node_id)?,
    1816              :     ))
    1817            0 : }
    1818              : 
    1819            0 : async fn handle_pageserver(subcmd: &PageserverCmd, env: &local_env::LocalEnv) -> Result<()> {
    1820            0 :     match subcmd {
    1821            0 :         PageserverCmd::Start(args) => {
    1822            0 :             if let Err(e) = get_pageserver(env, args.pageserver_id)?
    1823            0 :                 .start(&args.start_timeout)
    1824            0 :                 .await
    1825              :             {
    1826            0 :                 eprintln!("pageserver start failed: {e}");
    1827            0 :                 exit(1);
    1828            0 :             }
    1829              :         }
    1830              : 
    1831            0 :         PageserverCmd::Stop(args) => {
    1832            0 :             let immediate = match args.stop_mode {
    1833            0 :                 StopMode::Fast => false,
    1834            0 :                 StopMode::Immediate => true,
    1835              :             };
    1836            0 :             if let Err(e) = get_pageserver(env, args.pageserver_id)?.stop(immediate) {
    1837            0 :                 eprintln!("pageserver stop failed: {e}");
    1838            0 :                 exit(1);
    1839            0 :             }
    1840              :         }
    1841              : 
    1842            0 :         PageserverCmd::Restart(args) => {
    1843            0 :             let pageserver = get_pageserver(env, args.pageserver_id)?;
    1844              :             //TODO what shutdown strategy should we use here?
    1845            0 :             if let Err(e) = pageserver.stop(false) {
    1846            0 :                 eprintln!("pageserver stop failed: {e}");
    1847            0 :                 exit(1);
    1848            0 :             }
    1849              : 
    1850            0 :             if let Err(e) = pageserver.start(&args.start_timeout).await {
    1851            0 :                 eprintln!("pageserver start failed: {e}");
    1852            0 :                 exit(1);
    1853            0 :             }
    1854              :         }
    1855              : 
    1856            0 :         PageserverCmd::Status(args) => {
    1857            0 :             match get_pageserver(env, args.pageserver_id)?
    1858            0 :                 .check_status()
    1859            0 :                 .await
    1860              :             {
    1861            0 :                 Ok(_) => println!("Page server is up and running"),
    1862            0 :                 Err(err) => {
    1863            0 :                     eprintln!("Page server is not available: {err}");
    1864            0 :                     exit(1);
    1865              :                 }
    1866              :             }
    1867              :         }
    1868              :     }
    1869            0 :     Ok(())
    1870            0 : }
    1871              : 
    1872            0 : async fn handle_storage_controller(
    1873            0 :     subcmd: &StorageControllerCmd,
    1874            0 :     env: &local_env::LocalEnv,
    1875            0 : ) -> Result<()> {
    1876            0 :     let svc = StorageController::from_env(env);
    1877            0 :     match subcmd {
    1878            0 :         StorageControllerCmd::Start(args) => {
    1879            0 :             let start_args = NeonStorageControllerStartArgs {
    1880            0 :                 instance_id: args.instance_id,
    1881            0 :                 base_port: args.base_port,
    1882            0 :                 start_timeout: args.start_timeout,
    1883            0 :                 handle_ps_local_disk_loss: args.handle_ps_local_disk_loss,
    1884            0 :             };
    1885              : 
    1886            0 :             if let Err(e) = svc.start(start_args).await {
    1887            0 :                 eprintln!("start failed: {e}");
    1888            0 :                 exit(1);
    1889            0 :             }
    1890              :         }
    1891              : 
    1892            0 :         StorageControllerCmd::Stop(args) => {
    1893            0 :             let stop_args = NeonStorageControllerStopArgs {
    1894            0 :                 instance_id: args.instance_id,
    1895            0 :                 immediate: match args.stop_mode {
    1896            0 :                     StopMode::Fast => false,
    1897            0 :                     StopMode::Immediate => true,
    1898              :                 },
    1899              :             };
    1900            0 :             if let Err(e) = svc.stop(stop_args).await {
    1901            0 :                 eprintln!("stop failed: {e}");
    1902            0 :                 exit(1);
    1903            0 :             }
    1904              :         }
    1905              :     }
    1906            0 :     Ok(())
    1907            0 : }
    1908              : 
    1909            0 : fn get_safekeeper(env: &local_env::LocalEnv, id: NodeId) -> Result<SafekeeperNode> {
    1910            0 :     if let Some(node) = env.safekeepers.iter().find(|node| node.id == id) {
    1911            0 :         Ok(SafekeeperNode::from_env(env, node))
    1912              :     } else {
    1913            0 :         bail!("could not find safekeeper {id}")
    1914              :     }
    1915            0 : }
    1916              : 
    1917            0 : async fn handle_safekeeper(subcmd: &SafekeeperCmd, env: &local_env::LocalEnv) -> Result<()> {
    1918            0 :     match subcmd {
    1919            0 :         SafekeeperCmd::Start(args) => {
    1920            0 :             let safekeeper = get_safekeeper(env, args.id)?;
    1921              : 
    1922            0 :             if let Err(e) = safekeeper.start(&args.extra_opt, &args.start_timeout).await {
    1923            0 :                 eprintln!("safekeeper start failed: {e}");
    1924            0 :                 exit(1);
    1925            0 :             }
    1926              :         }
    1927              : 
    1928            0 :         SafekeeperCmd::Stop(args) => {
    1929            0 :             let safekeeper = get_safekeeper(env, args.id)?;
    1930            0 :             let immediate = match args.stop_mode {
    1931            0 :                 StopMode::Fast => false,
    1932            0 :                 StopMode::Immediate => true,
    1933              :             };
    1934            0 :             if let Err(e) = safekeeper.stop(immediate) {
    1935            0 :                 eprintln!("safekeeper stop failed: {e}");
    1936            0 :                 exit(1);
    1937            0 :             }
    1938              :         }
    1939              : 
    1940            0 :         SafekeeperCmd::Restart(args) => {
    1941            0 :             let safekeeper = get_safekeeper(env, args.id)?;
    1942            0 :             let immediate = match args.stop_mode {
    1943            0 :                 StopMode::Fast => false,
    1944            0 :                 StopMode::Immediate => true,
    1945              :             };
    1946              : 
    1947            0 :             if let Err(e) = safekeeper.stop(immediate) {
    1948            0 :                 eprintln!("safekeeper stop failed: {e}");
    1949            0 :                 exit(1);
    1950            0 :             }
    1951              : 
    1952            0 :             if let Err(e) = safekeeper.start(&args.extra_opt, &args.start_timeout).await {
    1953            0 :                 eprintln!("safekeeper start failed: {e}");
    1954            0 :                 exit(1);
    1955            0 :             }
    1956              :         }
    1957              :     }
    1958            0 :     Ok(())
    1959            0 : }
    1960              : 
    1961            0 : async fn handle_endpoint_storage(
    1962            0 :     subcmd: &EndpointStorageCmd,
    1963            0 :     env: &local_env::LocalEnv,
    1964            0 : ) -> Result<()> {
    1965              :     use EndpointStorageCmd::*;
    1966            0 :     let storage = EndpointStorage::from_env(env);
    1967              : 
    1968              :     // In tests like test_forward_compatibility or test_graceful_cluster_restart
    1969              :     // old neon binaries (without endpoint_storage) are present
    1970            0 :     if !storage.bin.exists() {
    1971            0 :         eprintln!(
    1972            0 :             "{} binary not found. Ignore if this is a compatibility test",
    1973              :             storage.bin
    1974              :         );
    1975            0 :         return Ok(());
    1976            0 :     }
    1977              : 
    1978            0 :     match subcmd {
    1979            0 :         Start(EndpointStorageStartCmd { start_timeout }) => {
    1980            0 :             if let Err(e) = storage.start(start_timeout).await {
    1981            0 :                 eprintln!("endpoint_storage start failed: {e}");
    1982            0 :                 exit(1);
    1983            0 :             }
    1984              :         }
    1985            0 :         Stop(EndpointStorageStopCmd { stop_mode }) => {
    1986            0 :             let immediate = match stop_mode {
    1987            0 :                 StopMode::Fast => false,
    1988            0 :                 StopMode::Immediate => true,
    1989              :             };
    1990            0 :             if let Err(e) = storage.stop(immediate) {
    1991            0 :                 eprintln!("proxy stop failed: {e}");
    1992            0 :                 exit(1);
    1993            0 :             }
    1994              :         }
    1995              :     };
    1996            0 :     Ok(())
    1997            0 : }
    1998              : 
    1999            0 : async fn handle_storage_broker(subcmd: &StorageBrokerCmd, env: &local_env::LocalEnv) -> Result<()> {
    2000            0 :     match subcmd {
    2001            0 :         StorageBrokerCmd::Start(args) => {
    2002            0 :             let storage_broker = StorageBroker::from_env(env);
    2003            0 :             if let Err(e) = storage_broker.start(&args.start_timeout).await {
    2004            0 :                 eprintln!("broker start failed: {e}");
    2005            0 :                 exit(1);
    2006            0 :             }
    2007              :         }
    2008              : 
    2009            0 :         StorageBrokerCmd::Stop(_args) => {
    2010              :             // FIXME: stop_mode unused
    2011            0 :             let storage_broker = StorageBroker::from_env(env);
    2012            0 :             if let Err(e) = storage_broker.stop() {
    2013            0 :                 eprintln!("broker stop failed: {e}");
    2014            0 :                 exit(1);
    2015            0 :             }
    2016              :         }
    2017              :     }
    2018            0 :     Ok(())
    2019            0 : }
    2020              : 
    2021            0 : async fn handle_start_all(
    2022            0 :     args: &StartCmdArgs,
    2023            0 :     env: &'static local_env::LocalEnv,
    2024            0 : ) -> anyhow::Result<()> {
    2025              :     // FIXME: this was called "retry_timeout", is it right?
    2026            0 :     let Err(errors) = handle_start_all_impl(env, args.timeout).await else {
    2027            0 :         neon_start_status_check(env, args.timeout.as_ref())
    2028            0 :             .await
    2029            0 :             .context("status check after successful startup of all services")?;
    2030            0 :         return Ok(());
    2031              :     };
    2032              : 
    2033            0 :     eprintln!("startup failed because one or more services could not be started");
    2034              : 
    2035            0 :     for e in errors {
    2036            0 :         eprintln!("{e}");
    2037            0 :         let debug_repr = format!("{e:?}");
    2038            0 :         for line in debug_repr.lines() {
    2039            0 :             eprintln!("  {line}");
    2040            0 :         }
    2041              :     }
    2042              : 
    2043            0 :     try_stop_all(env, true).await;
    2044              : 
    2045            0 :     exit(2);
    2046            0 : }
    2047              : 
    2048              : /// Returns Ok() if and only if all services could be started successfully.
    2049              : /// Otherwise, returns the list of errors that occurred during startup.
    2050            0 : async fn handle_start_all_impl(
    2051            0 :     env: &'static local_env::LocalEnv,
    2052            0 :     retry_timeout: humantime::Duration,
    2053            0 : ) -> Result<(), Vec<anyhow::Error>> {
    2054              :     // Endpoints are not started automatically
    2055              : 
    2056            0 :     let mut js = JoinSet::new();
    2057              : 
    2058              :     // force infalliblity through closure
    2059              :     #[allow(clippy::redundant_closure_call)]
    2060            0 :     (|| {
    2061            0 :         js.spawn(async move {
    2062            0 :             let storage_broker = StorageBroker::from_env(env);
    2063            0 :             storage_broker
    2064            0 :                 .start(&retry_timeout)
    2065            0 :                 .await
    2066            0 :                 .map_err(|e| e.context("start storage_broker"))
    2067            0 :         });
    2068              : 
    2069            0 :         js.spawn(async move {
    2070            0 :             let storage_controller = StorageController::from_env(env);
    2071            0 :             storage_controller
    2072            0 :                 .start(NeonStorageControllerStartArgs::with_default_instance_id(
    2073            0 :                     retry_timeout,
    2074            0 :                 ))
    2075            0 :                 .await
    2076            0 :                 .map_err(|e| e.context("start storage_controller"))
    2077            0 :         });
    2078              : 
    2079            0 :         for ps_conf in &env.pageservers {
    2080            0 :             js.spawn(async move {
    2081            0 :                 let pageserver = PageServerNode::from_env(env, ps_conf);
    2082            0 :                 pageserver
    2083            0 :                     .start(&retry_timeout)
    2084            0 :                     .await
    2085            0 :                     .map_err(|e| e.context(format!("start pageserver {}", ps_conf.id)))
    2086            0 :             });
    2087              :         }
    2088              : 
    2089            0 :         for node in env.safekeepers.iter() {
    2090            0 :             js.spawn(async move {
    2091            0 :                 let safekeeper = SafekeeperNode::from_env(env, node);
    2092            0 :                 safekeeper
    2093            0 :                     .start(&[], &retry_timeout)
    2094            0 :                     .await
    2095            0 :                     .map_err(|e| e.context(format!("start safekeeper {}", safekeeper.id)))
    2096            0 :             });
    2097              :         }
    2098              : 
    2099            0 :         js.spawn(async move {
    2100            0 :             EndpointStorage::from_env(env)
    2101            0 :                 .start(&retry_timeout)
    2102            0 :                 .await
    2103            0 :                 .map_err(|e| e.context("start endpoint_storage"))
    2104            0 :         });
    2105              :     })();
    2106              : 
    2107            0 :     let mut errors = Vec::new();
    2108            0 :     while let Some(result) = js.join_next().await {
    2109            0 :         let result = result.expect("we don't panic or cancel the tasks");
    2110            0 :         if let Err(e) = result {
    2111            0 :             errors.push(e);
    2112            0 :         }
    2113              :     }
    2114              : 
    2115            0 :     if !errors.is_empty() {
    2116            0 :         return Err(errors);
    2117            0 :     }
    2118              : 
    2119            0 :     Ok(())
    2120            0 : }
    2121              : 
    2122            0 : async fn neon_start_status_check(
    2123            0 :     env: &local_env::LocalEnv,
    2124            0 :     retry_timeout: &Duration,
    2125            0 : ) -> anyhow::Result<()> {
    2126              :     const RETRY_INTERVAL: Duration = Duration::from_millis(100);
    2127              :     const NOTICE_AFTER_RETRIES: Duration = Duration::from_secs(5);
    2128              : 
    2129            0 :     let storcon = StorageController::from_env(env);
    2130              : 
    2131            0 :     let retries = retry_timeout.as_millis() / RETRY_INTERVAL.as_millis();
    2132            0 :     let notice_after_retries = retry_timeout.as_millis() / NOTICE_AFTER_RETRIES.as_millis();
    2133              : 
    2134            0 :     println!("\nRunning neon status check");
    2135              : 
    2136            0 :     for retry in 0..retries {
    2137            0 :         if retry == notice_after_retries {
    2138            0 :             println!("\nNeon status check has not passed yet, continuing to wait")
    2139            0 :         }
    2140              : 
    2141            0 :         let mut passed = true;
    2142            0 :         let mut nodes = storcon.node_list().await?;
    2143            0 :         let mut pageservers = env.pageservers.clone();
    2144              : 
    2145            0 :         if nodes.len() != pageservers.len() {
    2146            0 :             continue;
    2147            0 :         }
    2148              : 
    2149            0 :         nodes.sort_by_key(|ps| ps.id);
    2150            0 :         pageservers.sort_by_key(|ps| ps.id);
    2151              : 
    2152            0 :         for (idx, pageserver) in pageservers.iter().enumerate() {
    2153            0 :             let node = &nodes[idx];
    2154            0 :             if node.id != pageserver.id {
    2155            0 :                 passed = false;
    2156            0 :                 break;
    2157            0 :             }
    2158              : 
    2159            0 :             if !matches!(node.availability, NodeAvailabilityWrapper::Active) {
    2160            0 :                 passed = false;
    2161            0 :                 break;
    2162            0 :             }
    2163              :         }
    2164              : 
    2165            0 :         if passed {
    2166            0 :             println!("\nNeon started and passed status check");
    2167            0 :             return Ok(());
    2168            0 :         }
    2169              : 
    2170            0 :         tokio::time::sleep(RETRY_INTERVAL).await;
    2171              :     }
    2172              : 
    2173            0 :     anyhow::bail!("\nNeon passed status check")
    2174            0 : }
    2175              : 
    2176            0 : async fn handle_stop_all(args: &StopCmdArgs, env: &local_env::LocalEnv) -> Result<()> {
    2177            0 :     let immediate = match args.mode {
    2178            0 :         StopMode::Fast => false,
    2179            0 :         StopMode::Immediate => true,
    2180              :     };
    2181              : 
    2182            0 :     try_stop_all(env, immediate).await;
    2183              : 
    2184            0 :     Ok(())
    2185            0 : }
    2186              : 
    2187            0 : async fn try_stop_all(env: &local_env::LocalEnv, immediate: bool) {
    2188            0 :     let mode = if immediate {
    2189            0 :         EndpointTerminateMode::Immediate
    2190              :     } else {
    2191            0 :         EndpointTerminateMode::Fast
    2192              :     };
    2193              :     // Stop all endpoints
    2194            0 :     match ComputeControlPlane::load(env.clone()) {
    2195            0 :         Ok(cplane) => {
    2196            0 :             for (_k, node) in cplane.endpoints {
    2197            0 :                 if let Err(e) = node.stop(mode, false).await {
    2198            0 :                     eprintln!("postgres stop failed: {e:#}");
    2199            0 :                 }
    2200              :             }
    2201              :         }
    2202            0 :         Err(e) => {
    2203            0 :             eprintln!("postgres stop failed, could not restore control plane data from env: {e:#}")
    2204              :         }
    2205              :     }
    2206              : 
    2207            0 :     let storage = EndpointStorage::from_env(env);
    2208            0 :     if let Err(e) = storage.stop(immediate) {
    2209            0 :         eprintln!("endpoint_storage stop failed: {e:#}");
    2210            0 :     }
    2211              : 
    2212            0 :     for ps_conf in &env.pageservers {
    2213            0 :         let pageserver = PageServerNode::from_env(env, ps_conf);
    2214            0 :         if let Err(e) = pageserver.stop(immediate) {
    2215            0 :             eprintln!("pageserver {} stop failed: {:#}", ps_conf.id, e);
    2216            0 :         }
    2217              :     }
    2218              : 
    2219            0 :     for node in env.safekeepers.iter() {
    2220            0 :         let safekeeper = SafekeeperNode::from_env(env, node);
    2221            0 :         if let Err(e) = safekeeper.stop(immediate) {
    2222            0 :             eprintln!("safekeeper {} stop failed: {:#}", safekeeper.id, e);
    2223            0 :         }
    2224              :     }
    2225              : 
    2226            0 :     let storage_broker = StorageBroker::from_env(env);
    2227            0 :     if let Err(e) = storage_broker.stop() {
    2228            0 :         eprintln!("neon broker stop failed: {e:#}");
    2229            0 :     }
    2230              : 
    2231              :     // Stop all storage controller instances. In the most common case there's only one,
    2232              :     // but iterate though the base data directory in order to discover the instances.
    2233            0 :     let storcon_instances = env
    2234            0 :         .storage_controller_instances()
    2235            0 :         .await
    2236            0 :         .expect("Must inspect data dir");
    2237            0 :     for (instance_id, _instance_dir_path) in storcon_instances {
    2238            0 :         let storage_controller = StorageController::from_env(env);
    2239            0 :         let stop_args = NeonStorageControllerStopArgs {
    2240            0 :             instance_id,
    2241            0 :             immediate,
    2242            0 :         };
    2243              : 
    2244            0 :         if let Err(e) = storage_controller.stop(stop_args).await {
    2245            0 :             eprintln!("Storage controller instance {instance_id} stop failed: {e:#}");
    2246            0 :         }
    2247              :     }
    2248            0 : }
        

Generated by: LCOV version 2.1-beta