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

Generated by: LCOV version 2.1-beta