hacky way to make tests pass

Closes: https://github.com/neondatabase/neon/issues/8969

Cargo.lock

@@ -82,12 +82,27 @@ dependencies = [
  "anstyle",
  "anstyle-parse",
  "anstyle-query",
- "anstyle-wincon",
+ "anstyle-wincon 1.0.1",
  "colorchoice",
  "is-terminal",
  "utf8parse",
 ]
 
+[[package]]
+name = "anstream"
+version = "0.6.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "64e15c1ab1f89faffbf04a634d5e1962e9074f2741eef6d97f3c4e322426d526"
+dependencies = [
+ "anstyle",
+ "anstyle-parse",
+ "anstyle-query",
+ "anstyle-wincon 3.0.4",
+ "colorchoice",
+ "is_terminal_polyfill",
+ "utf8parse",
+]
+
 [[package]]
 name = "anstyle"
 version = "1.0.8"
@@ -122,6 +137,16 @@ dependencies = [
  "windows-sys 0.48.0",
 ]
 
+[[package]]
+name = "anstyle-wincon"
+version = "3.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5bf74e1b6e971609db8ca7a9ce79fd5768ab6ae46441c572e46cf596f59e57f8"
+dependencies = [
+ "anstyle",
+ "windows-sys 0.52.0",
+]
+
 [[package]]
 name = "anyhow"
 version = "1.0.71"
@@ -674,6 +699,7 @@ checksum = "3b829e4e32b91e643de6eafe82b1d90675f5874230191a4ffbc1b336dec4d6bf"
 dependencies = [
  "async-trait",
  "axum-core 0.3.4",
+ "base64 0.21.1",
  "bitflags 1.3.2",
  "bytes",
  "futures-util",
@@ -688,7 +714,13 @@ dependencies = [
  "pin-project-lite",
  "rustversion",
  "serde",
+ "serde_json",
+ "serde_path_to_error",
+ "serde_urlencoded",
+ "sha1",
  "sync_wrapper 0.1.2",
+ "tokio",
+ "tokio-tungstenite",
  "tower",
  "tower-layer",
  "tower-service",
@@ -702,14 +734,11 @@ checksum = "3a6c9af12842a67734c9a2e355436e5d03b22383ed60cf13cd0c18fbfe3dcbcf"
 dependencies = [
  "async-trait",
  "axum-core 0.4.5",
- "base64 0.21.1",
  "bytes",
  "futures-util",
  "http 1.1.0",
  "http-body 1.0.0",
  "http-body-util",
- "hyper 1.4.1",
- "hyper-util",
  "itoa",
  "matchit 0.7.0",
  "memchr",
@@ -718,17 +747,10 @@ dependencies = [
  "pin-project-lite",
  "rustversion",
  "serde",
- "serde_json",
- "serde_path_to_error",
- "serde_urlencoded",
- "sha1",
  "sync_wrapper 1.0.1",
- "tokio",
- "tokio-tungstenite",
  "tower",
  "tower-layer",
  "tower-service",
- "tracing",
 ]
 
 [[package]]
@@ -766,7 +788,6 @@ dependencies = [
  "sync_wrapper 1.0.1",
  "tower-layer",
  "tower-service",
- "tracing",
 ]
 
 [[package]]
@@ -971,7 +992,7 @@ dependencies = [
  "clang-sys",
  "itertools 0.12.1",
  "log",
- "prettyplease 0.2.17",
+ "prettyplease",
  "proc-macro2",
  "quote",
  "regex",
@@ -1189,7 +1210,7 @@ version = "4.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4f423e341edefb78c9caba2d9c7f7687d0e72e89df3ce3394554754393ac3990"
 dependencies = [
- "anstream",
+ "anstream 0.3.2",
  "anstyle",
  "bitflags 1.3.2",
  "clap_lex",
@@ -1265,7 +1286,6 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "bytes",
- "camino",
  "cfg-if",
  "chrono",
  "clap",
@@ -1929,6 +1949,15 @@ dependencies = [
  "syn 2.0.52",
 ]
 
+[[package]]
+name = "env_filter"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4f2c92ceda6ceec50f43169f9ee8424fe2db276791afde7b2cd8bc084cb376ab"
+dependencies = [
+ "log",
+]
+
 [[package]]
 name = "env_logger"
 version = "0.10.2"
@@ -1942,6 +1971,18 @@ dependencies = [
  "termcolor",
 ]
 
+[[package]]
+name = "env_logger"
+version = "0.11.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6c012a26a7f605efc424dd53697843a72be7dc86ad2d01f7814337794a12231d"
+dependencies = [
+ "anstream 0.6.15",
+ "anstyle",
+ "env_filter",
+ "log",
+]
+
 [[package]]
 name = "equivalent"
 version = "1.0.1"
@@ -2454,15 +2495,6 @@ dependencies = [
  "digest",
 ]
 
-[[package]]
-name = "home"
-version = "0.5.9"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e3d1354bf6b7235cb4a0576c2619fd4ed18183f689b12b006a0ee7329eeff9a5"
-dependencies = [
- "windows-sys 0.52.0",
-]
-
 [[package]]
 name = "hostname"
 version = "0.4.0"
@@ -2657,14 +2689,15 @@ dependencies = [
 
 [[package]]
 name = "hyper-timeout"
-version = "0.4.1"
+version = "0.5.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bbb958482e8c7be4bc3cf272a766a2b0bf1a6755e7a6ae777f017a31d11b13b1"
+checksum = "3203a961e5c83b6f5498933e78b6b263e208c197b63e9c6c53cc82ffd3f63793"
 dependencies = [
- "hyper 0.14.30",
+ "hyper 1.4.1",
+ "hyper-util",
  "pin-project-lite",
  "tokio",
- "tokio-io-timeout",
+ "tower-service",
 ]
 
 [[package]]
@@ -2824,6 +2857,12 @@ dependencies = [
  "windows-sys 0.52.0",
 ]
 
+[[package]]
+name = "is_terminal_polyfill"
+version = "1.70.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf"
+
 [[package]]
 name = "itertools"
 version = "0.10.5"
@@ -3267,6 +3306,16 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "nu-ansi-term"
+version = "0.46.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84"
+dependencies = [
+ "overload",
+ "winapi",
+]
+
 [[package]]
 name = "num"
 version = "0.4.1"
@@ -3470,7 +3519,7 @@ dependencies = [
  "opentelemetry-http",
  "opentelemetry-proto",
  "opentelemetry_sdk",
- "prost 0.13.3",
+ "prost",
  "reqwest 0.12.4",
  "thiserror",
 ]
@@ -3483,8 +3532,8 @@ checksum = "30ee9f20bff9c984511a02f082dc8ede839e4a9bf15cc2487c8d6fea5ad850d9"
 dependencies = [
  "opentelemetry",
  "opentelemetry_sdk",
- "prost 0.13.3",
- "tonic 0.12.3",
+ "prost",
+ "tonic",
 ]
 
 [[package]]
@@ -3550,6 +3599,12 @@ version = "0.5.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4030760ffd992bef45b0ae3f10ce1aba99e33464c90d14dd7c039884963ddc7a"
 
+[[package]]
+name = "overload"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39"
+
 [[package]]
 name = "p256"
 version = "0.11.1"
@@ -4132,7 +4187,7 @@ dependencies = [
  "bindgen",
  "bytes",
  "crc32c",
- "env_logger",
+ "env_logger 0.10.2",
  "log",
  "memoffset 0.9.0",
  "once_cell",
@@ -4178,16 +4233,6 @@ dependencies = [
  "tokio",
 ]
 
-[[package]]
-name = "prettyplease"
-version = "0.1.25"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6c8646e95016a7a6c4adea95bafa8a16baab64b583356217f2c85db4a39d9a86"
-dependencies = [
- "proc-macro2",
- "syn 1.0.109",
-]
-
 [[package]]
 name = "prettyplease"
 version = "0.2.17"
@@ -4258,16 +4303,6 @@ dependencies = [
  "thiserror",
 ]
 
-[[package]]
-name = "prost"
-version = "0.11.9"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0b82eaa1d779e9a4bc1c3217db8ffbeabaae1dca241bf70183242128d48681cd"
-dependencies = [
- "bytes",
- "prost-derive 0.11.9",
-]
-
 [[package]]
 name = "prost"
 version = "0.13.3"
@@ -4275,42 +4310,28 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7b0487d90e047de87f984913713b85c601c05609aad5b0df4b4573fbf69aa13f"
 dependencies = [
  "bytes",
- "prost-derive 0.13.3",
+ "prost-derive",
 ]
 
 [[package]]
 name = "prost-build"
-version = "0.11.9"
+version = "0.13.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "119533552c9a7ffacc21e099c24a0ac8bb19c2a2a3f363de84cd9b844feab270"
+checksum = "0c1318b19085f08681016926435853bbf7858f9c082d0999b80550ff5d9abe15"
 dependencies = [
  "bytes",
- "heck 0.4.1",
- "itertools 0.10.5",
- "lazy_static",
+ "heck 0.5.0",
+ "itertools 0.12.1",
  "log",
  "multimap",
+ "once_cell",
  "petgraph",
- "prettyplease 0.1.25",
- "prost 0.11.9",
+ "prettyplease",
+ "prost",
  "prost-types",
  "regex",
- "syn 1.0.109",
+ "syn 2.0.52",
  "tempfile",
- "which",
-]
-
-[[package]]
-name = "prost-derive"
-version = "0.11.9"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e5d2d8d10f3c6ded6da8b05b5fb3b8a5082514344d56c9f871412d29b4e075b4"
-dependencies = [
- "anyhow",
- "itertools 0.10.5",
- "proc-macro2",
- "quote",
- "syn 1.0.109",
 ]
 
 [[package]]
@@ -4328,11 +4349,11 @@ dependencies = [
 
 [[package]]
 name = "prost-types"
-version = "0.11.9"
+version = "0.13.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "213622a1460818959ac1181aaeb2dc9c7f63df720db7d788b3e24eacd1983e13"
+checksum = "4759aa0d3a6232fb8dbdb97b61de2c20047c68aca932c7ed76da9d788508d670"
 dependencies = [
- "prost 0.11.9",
+ "prost",
 ]
 
 [[package]]
@@ -4359,7 +4380,7 @@ dependencies = [
  "consumption_metrics",
  "dashmap",
  "ecdsa 0.16.9",
- "env_logger",
+ "env_logger 0.10.2",
  "fallible-iterator",
  "framed-websockets",
  "futures",
@@ -5094,6 +5115,21 @@ dependencies = [
  "zeroize",
 ]
 
+[[package]]
+name = "rustls"
+version = "0.23.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ebbbdb961df0ad3f2652da8f3fdc4b36122f568f968f45ad3316f26c025c677b"
+dependencies = [
+ "log",
+ "once_cell",
+ "ring",
+ "rustls-pki-types",
+ "rustls-webpki 0.102.2",
+ "subtle",
+ "zeroize",
+]
+
 [[package]]
 name = "rustls-native-certs"
 version = "0.6.2"
@@ -5119,6 +5155,19 @@ dependencies = [
  "security-framework",
 ]
 
+[[package]]
+name = "rustls-native-certs"
+version = "0.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fcaf18a4f2be7326cd874a5fa579fae794320a0f388d365dca7e480e55f83f8a"
+dependencies = [
+ "openssl-probe",
+ "rustls-pemfile 2.1.1",
+ "rustls-pki-types",
+ "schannel",
+ "security-framework",
+]
+
 [[package]]
 name = "rustls-pemfile"
 version = "1.0.2"
@@ -5194,6 +5243,7 @@ dependencies = [
  "fail",
  "futures",
  "hex",
+ "http 1.1.0",
  "humantime",
  "hyper 0.14.30",
  "metrics",
@@ -5750,19 +5800,22 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "async-stream",
+ "bytes",
  "clap",
  "const_format",
  "futures",
  "futures-core",
  "futures-util",
+ "http-body-util",
  "humantime",
- "hyper 0.14.30",
+ "hyper 1.4.1",
+ "hyper-util",
  "metrics",
  "once_cell",
  "parking_lot 0.12.1",
- "prost 0.11.9",
+ "prost",
  "tokio",
- "tonic 0.9.2",
+ "tonic",
  "tonic-build",
  "tracing",
  "utils",
@@ -5802,6 +5855,7 @@ dependencies = [
  "serde_json",
  "strum",
  "strum_macros",
+ "test-log",
  "thiserror",
  "tokio",
  "tokio-util",
@@ -6054,6 +6108,28 @@ dependencies = [
  "syn 2.0.52",
 ]
 
+[[package]]
+name = "test-log"
+version = "0.2.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3dffced63c2b5c7be278154d76b479f9f9920ed34e7574201407f0b14e2bbb93"
+dependencies = [
+ "env_logger 0.11.2",
+ "test-log-macros",
+ "tracing-subscriber",
+]
+
+[[package]]
+name = "test-log-macros"
+version = "0.2.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5999e24eaa32083191ba4e425deb75cdf25efefabe5aaccb7446dd0d4122a3f5"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.52",
+]
+
 [[package]]
 name = "thiserror"
 version = "1.0.57"
@@ -6306,6 +6382,17 @@ dependencies = [
  "tokio",
 ]
 
+[[package]]
+name = "tokio-rustls"
+version = "0.26.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0c7bc40d0e5a97695bb96e27995cd3a08538541b0a846f65bba7a359f36700d4"
+dependencies = [
+ "rustls 0.23.7",
+ "rustls-pki-types",
+ "tokio",
+]
+
 [[package]]
 name = "tokio-stream"
 version = "0.1.16"
@@ -6334,9 +6421,9 @@ dependencies = [
 
 [[package]]
 name = "tokio-tungstenite"
-version = "0.21.0"
+version = "0.20.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c83b561d025642014097b66e6c1bb422783339e0909e4429cde4749d1990bc38"
+checksum = "2b2dbec703c26b00d74844519606ef15d09a7d6857860f84ad223dec002ddea2"
 dependencies = [
  "futures-util",
  "log",
@@ -6397,29 +6484,30 @@ dependencies = [
 
 [[package]]
 name = "tonic"
-version = "0.9.2"
+version = "0.12.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3082666a3a6433f7f511c7192923fa1fe07c69332d3c6a2e6bb040b569199d5a"
+checksum = "877c5b330756d856ffcc4553ab34a5684481ade925ecc54bcd1bf02b1d0d4d52"
 dependencies = [
  "async-stream",
  "async-trait",
- "axum 0.6.20",
- "base64 0.21.1",
+ "axum 0.7.5",
+ "base64 0.22.1",
  "bytes",
- "futures-core",
- "futures-util",
- "h2 0.3.26",
- "http 0.2.9",
- "http-body 0.4.5",
- "hyper 0.14.30",
+ "h2 0.4.4",
+ "http 1.1.0",
+ "http-body 1.0.0",
+ "http-body-util",
+ "hyper 1.4.1",
  "hyper-timeout",
+ "hyper-util",
  "percent-encoding",
  "pin-project",
- "prost 0.11.9",
- "rustls-native-certs 0.6.2",
- "rustls-pemfile 1.0.2",
+ "prost",
+ "rustls-native-certs 0.8.0",
+ "rustls-pemfile 2.1.1",
+ "socket2",
  "tokio",
- "tokio-rustls 0.24.0",
+ "tokio-rustls 0.26.0",
  "tokio-stream",
  "tower",
  "tower-layer",
@@ -6428,37 +6516,17 @@ dependencies = [
 ]
 
 [[package]]
-name = "tonic"
+name = "tonic-build"
 version = "0.12.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "877c5b330756d856ffcc4553ab34a5684481ade925ecc54bcd1bf02b1d0d4d52"
+checksum = "9557ce109ea773b399c9b9e5dca39294110b74f1f342cb347a80d1fce8c26a11"
 dependencies = [
- "async-trait",
- "base64 0.22.1",
- "bytes",
- "http 1.1.0",
- "http-body 1.0.0",
- "http-body-util",
- "percent-encoding",
- "pin-project",
- "prost 0.13.3",
- "tokio-stream",
- "tower-layer",
- "tower-service",
- "tracing",
-]
-
-[[package]]
-name = "tonic-build"
-version = "0.9.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a6fdaae4c2c638bb70fe42803a26fbd6fc6ac8c72f5c59f67ecc2a2dcabf4b07"
-dependencies = [
- "prettyplease 0.1.25",
+ "prettyplease",
  "proc-macro2",
  "prost-build",
+ "prost-types",
  "quote",
- "syn 1.0.109",
+ "syn 2.0.52",
 ]
 
 [[package]]
@@ -6593,6 +6661,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b"
 dependencies = [
  "matchers",
+ "nu-ansi-term",
  "once_cell",
  "regex",
  "serde",
@@ -6629,14 +6698,14 @@ checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b"
 
 [[package]]
 name = "tungstenite"
-version = "0.21.0"
+version = "0.20.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9ef1a641ea34f399a848dea702823bbecfb4c486f911735368f1f137cb8257e1"
+checksum = "9e3dac10fd62eaf6617d3a904ae222845979aec67c615d1c842b4002c7666fb9"
 dependencies = [
  "byteorder",
  "bytes",
  "data-encoding",
- "http 1.1.0",
+ "http 0.2.9",
  "httparse",
  "log",
  "rand 0.8.5",
@@ -6864,7 +6933,7 @@ name = "vm_monitor"
 version = "0.1.0"
 dependencies = [
  "anyhow",
- "axum 0.7.5",
+ "axum 0.6.20",
  "cgroups-rs",
  "clap",
  "futures",
@@ -6897,7 +6966,7 @@ dependencies = [
  "anyhow",
  "camino-tempfile",
  "clap",
- "env_logger",
+ "env_logger 0.10.2",
  "log",
  "postgres",
  "postgres_ffi",
@@ -7095,18 +7164,6 @@ dependencies = [
  "rustls-pki-types",
 ]
 
-[[package]]
-name = "which"
-version = "4.4.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7"
-dependencies = [
- "either",
- "home",
- "once_cell",
- "rustix",
-]
-
 [[package]]
 name = "whoami"
 version = "1.5.1"
@@ -7335,9 +7392,14 @@ version = "0.1.0"
 dependencies = [
  "ahash",
  "anyhow",
+ "aws-config",
+ "aws-runtime",
+ "aws-sigv4",
+ "aws-smithy-async",
+ "aws-smithy-http",
+ "aws-smithy-types",
  "base64 0.21.1",
  "base64ct",
- "bitflags 2.4.1",
  "bytes",
  "camino",
  "cc",
@@ -7365,7 +7427,6 @@ dependencies = [
  "hyper 1.4.1",
  "hyper-util",
  "indexmap 1.9.3",
- "itertools 0.10.5",
  "itertools 0.12.1",
  "lazy_static",
  "libc",
@@ -7377,8 +7438,9 @@ dependencies = [
  "num-traits",
  "once_cell",
  "parquet",
+ "prettyplease",
  "proc-macro2",
- "prost 0.11.9",
+ "prost",
  "quote",
  "rand 0.8.5",
  "regex",
@@ -7401,13 +7463,15 @@ dependencies = [
  "time",
  "time-macros",
  "tokio",
- "tokio-rustls 0.24.0",
+ "tokio-stream",
  "tokio-util",
  "toml_edit",
+ "tonic",
  "tower",
  "tracing",
  "tracing-core",
  "url",
+ "uuid",
  "zeroize",
  "zstd",
  "zstd-safe",

Cargo.toml

@@ -53,7 +53,7 @@ azure_storage_blobs = { version = "0.19", default-features = false, features = [
 flate2 = "1.0.26"
 async-stream = "0.3"
 async-trait = "0.1"
-aws-config = { version = "1.5", default-features = false, features=["rustls", "sso"] }
+aws-config = { version = "1.5", default-features = false, features=["rustls"] }
 aws-sdk-s3 = "1.52"
 aws-sdk-iam = "1.46.0"
 aws-smithy-async = { version = "1.2.1", default-features = false, features=["rt-tokio"] }
@@ -61,7 +61,7 @@ aws-smithy-types = "1.2"
 aws-credential-types = "1.2.0"
 aws-sigv4 = { version = "1.2", features = ["sign-http"] }
 aws-types = "1.3"
-axum = { version = "0.7.5", features = ["ws"] }
+axum = { version = "0.6.20", features = ["ws"] }
 base64 = "0.13.0"
 bincode = "1.3"
 bindgen = "0.70"
@@ -99,10 +99,10 @@ http-types = { version = "2", default-features = false }
 http-body-util = "0.1.2"
 humantime = "2.1"
 humantime-serde = "1.1.1"
-hyper0 = { package = "hyper", version = "0.14" }
-hyper = "1.4"
+hyper = "0.14"
+hyper_1 = { package = "hyper", version = "1.4" }
 hyper-util = "0.1"
-tokio-tungstenite = "0.21.0"
+tokio-tungstenite = "0.20.0"
 indexmap = "2"
 indoc = "2"
 ipnet = "2.9.0"
@@ -130,7 +130,7 @@ pbkdf2 = { version = "0.12.1", features = ["simple", "std"] }
 pin-project-lite = "0.2"
 procfs = "0.16"
 prometheus = {version = "0.13", default-features=false, features = ["process"]} # removes protobuf dependency
-prost = "0.11"
+prost = "0.13"
 rand = "0.8"
 redis = { version = "0.25.2", features = ["tokio-rustls-comp", "keep-alive"] }
 regex = "1.10.2"
@@ -178,7 +178,7 @@ tokio-tar = "0.3"
 tokio-util = { version = "0.7.10", features = ["io", "rt"] }
 toml = "0.8"
 toml_edit = "0.22"
-tonic = {version = "0.9", features = ["tls", "tls-roots"]}
+tonic = {version = "0.12.3", features = ["tls", "tls-roots"]}
 tower-service = "0.3.2"
 tracing = "0.1"
 tracing-error = "0.2"
@@ -246,7 +246,7 @@ criterion = "0.5.1"
 rcgen = "0.12"
 rstest = "0.18"
 camino-tempfile = "1.0.2"
-tonic-build = "0.9"
+tonic-build = "0.12"
 
 [patch.crates-io]
 

README.md

@@ -58,7 +58,7 @@ curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
 1. Install XCode and dependencies
 ```
 xcode-select --install
-brew install protobuf openssl flex bison icu4c pkg-config m4
+brew install protobuf openssl flex bison icu4c pkg-config
 
 # add openssl to PATH, required for ed25519 keys generation in neon_local
 echo 'export PATH="$(brew --prefix openssl)/bin:$PATH"' >> ~/.zshrc

compute/Dockerfile.compute-node

@@ -880,6 +880,9 @@ RUN case "${PG_VERSION}" in "v17") \
     mkdir pg_session_jwt-src && cd pg_session_jwt-src && tar xzf ../pg_session_jwt.tar.gz --strip-components=1 -C . && \
     sed -i 's/pgrx = "=0.11.3"/pgrx = { version = "=0.11.3", features = [ "unsafe-postgres" ] }/g' Cargo.toml && \
     cargo pgrx install --release
+    # it's needed to enable extension because it uses untrusted C language
+    # sed -i 's/superuser = false/superuser = true/g' /usr/local/pgsql/share/extension/pg_session_jwt.control && \
+    # echo "trusted = true" >> /usr/local/pgsql/share/extension/pg_session_jwt.control
 
 #########################################################################################
 #
@@ -1075,20 +1078,6 @@ RUN set -e \
     && make -j $(nproc) dist_man_MANS= \
     && make install dist_man_MANS=
 
-#########################################################################################
-#
-# Compile the Neon-specific `local_proxy` binary
-#
-#########################################################################################
-FROM $REPOSITORY/$IMAGE:$TAG AS local_proxy
-ARG BUILD_TAG
-ENV BUILD_TAG=$BUILD_TAG
-
-USER nonroot
-# Copy entire project to get Cargo.* files with proper dependencies for the whole project
-COPY --chown=nonroot . .
-RUN mold -run cargo build --locked --profile release-line-debug-size-lto --bin local_proxy
-
 #########################################################################################
 #
 # Layers "postgres-exporter" and "sql-exporter"
@@ -1227,10 +1216,6 @@ COPY --from=compute-tools --chown=postgres /home/nonroot/target/release-line-deb
 COPY --from=pgbouncer         /usr/local/pgbouncer/bin/pgbouncer /usr/local/bin/pgbouncer
 COPY --chmod=0666 --chown=postgres compute/etc/pgbouncer.ini /etc/pgbouncer.ini
 
-# local_proxy and its config
-COPY --from=local_proxy --chown=postgres /home/nonroot/target/release-line-debug-size-lto/local_proxy /usr/local/bin/local_proxy
-RUN mkdir -p /etc/local_proxy && chown postgres:postgres /etc/local_proxy
-
 # Metrics exporter binaries and  configuration files
 COPY --from=postgres-exporter /bin/postgres_exporter /bin/postgres_exporter
 COPY --from=sql-exporter      /bin/sql_exporter      /bin/sql_exporter

compute/vm-image-spec.yaml

@@ -19,10 +19,6 @@ commands:
     user: postgres
     sysvInitAction: respawn
     shell: '/usr/local/bin/pgbouncer /etc/pgbouncer.ini'
-  - name: local_proxy
-    user: postgres
-    sysvInitAction: respawn
-    shell: '/usr/local/bin/local_proxy --config-path /etc/local_proxy/config.json --pid-path /etc/local_proxy/pid --http 0.0.0.0:10432'
   - name: postgres-exporter
     user: nobody
     sysvInitAction: respawn

compute_tools/Cargo.toml

@@ -11,13 +11,12 @@ testing = []
 
 [dependencies]
 anyhow.workspace = true
-camino.workspace = true
 chrono.workspace = true
 cfg-if.workspace = true
 clap.workspace = true
 flate2.workspace = true
 futures.workspace = true
-hyper0 = { workspace = true, features = ["full"] }
+hyper = { workspace = true, features = ["full"] }
 nix.workspace = true
 notify.workspace = true
 num_cpus.workspace = true

compute_tools/src/compute.rs

@@ -34,7 +34,6 @@ use nix::sys::signal::{kill, Signal};
 use remote_storage::{DownloadError, RemotePath};
 
 use crate::checker::create_availability_check_data;
-use crate::local_proxy;
 use crate::logger::inlinify;
 use crate::pg_helpers::*;
 use crate::spec::*;
@@ -887,11 +886,6 @@ impl ComputeNode {
         // 'Close' connection
         drop(client);
 
-        if let Some(ref local_proxy) = spec.local_proxy_config {
-            info!("configuring local_proxy");
-            local_proxy::configure(local_proxy).context("apply_config local_proxy")?;
-        }
-
         // Run migrations separately to not hold up cold starts
         thread::spawn(move || {
             let mut connstr = connstr.clone();
@@ -942,19 +936,6 @@ impl ComputeNode {
             });
         }
 
-        if let Some(ref local_proxy) = spec.local_proxy_config {
-            info!("configuring local_proxy");
-
-            // Spawn a thread to do the configuration,
-            // so that we don't block the main thread that starts Postgres.
-            let local_proxy = local_proxy.clone();
-            let _handle = Some(thread::spawn(move || {
-                if let Err(err) = local_proxy::configure(&local_proxy) {
-                    error!("error while configuring local_proxy: {err:?}");
-                }
-            }));
-        }
-
         // Write new config
         let pgdata_path = Path::new(&self.pgdata);
         let postgresql_conf_path = pgdata_path.join("postgresql.conf");
@@ -1042,19 +1023,6 @@ impl ComputeNode {
             });
         }
 
-        if let Some(local_proxy) = &pspec.spec.local_proxy_config {
-            info!("configuring local_proxy");
-
-            // Spawn a thread to do the configuration,
-            // so that we don't block the main thread that starts Postgres.
-            let local_proxy = local_proxy.clone();
-            let _handle = thread::spawn(move || {
-                if let Err(err) = local_proxy::configure(&local_proxy) {
-                    error!("error while configuring local_proxy: {err:?}");
-                }
-            });
-        }
-
         info!(
             "start_compute spec.remote_extensions {:?}",
             pspec.spec.remote_extensions

compute_tools/src/lib.rs

@@ -2,9 +2,6 @@
 //! configuration.
 #![deny(unsafe_code)]
 #![deny(clippy::undocumented_unsafe_blocks)]
-
-extern crate hyper0 as hyper;
-
 pub mod checker;
 pub mod config;
 pub mod configurator;
@@ -15,7 +12,6 @@ pub mod catalog;
 pub mod compute;
 pub mod disk_quota;
 pub mod extension_server;
-pub mod local_proxy;
 pub mod lsn_lease;
 mod migration;
 pub mod monitor;

compute_tools/src/local_proxy.rs

@@ -1,56 +0,0 @@
-//! Local Proxy is a feature of our BaaS Neon Authorize project.
-//!
-//! Local Proxy validates JWTs and manages the pg_session_jwt extension.
-//! It also maintains a connection pool to postgres.
-
-use anyhow::{Context, Result};
-use camino::Utf8Path;
-use compute_api::spec::LocalProxySpec;
-use nix::sys::signal::Signal;
-use utils::pid_file::{self, PidFileRead};
-
-pub fn configure(local_proxy: &LocalProxySpec) -> Result<()> {
-    write_local_proxy_conf("/etc/local_proxy/config.json".as_ref(), local_proxy)?;
-    notify_local_proxy("/etc/local_proxy/pid".as_ref())?;
-
-    Ok(())
-}
-
-/// Create or completely rewrite configuration file specified by `path`
-fn write_local_proxy_conf(path: &Utf8Path, local_proxy: &LocalProxySpec) -> Result<()> {
-    let config =
-        serde_json::to_string_pretty(local_proxy).context("serializing LocalProxySpec to json")?;
-    std::fs::write(path, config).with_context(|| format!("writing {path}"))?;
-
-    Ok(())
-}
-
-/// Notify local proxy about a new config file.
-fn notify_local_proxy(path: &Utf8Path) -> Result<()> {
-    match pid_file::read(path)? {
-        // if the file doesn't exist, or isn't locked, local_proxy isn't running
-        // and will naturally pick up our config later
-        PidFileRead::NotExist | PidFileRead::NotHeldByAnyProcess(_) => {}
-        PidFileRead::LockedByOtherProcess(pid) => {
-            // From the pid_file docs:
-            //
-            // > 1. The other process might exit at any time, turning the given PID stale.
-            // > 2. There is a small window in which `claim_for_current_process` has already
-            // >    locked the file but not yet updates its contents. [`read`] will return
-            // >    this variant here, but with the old file contents, i.e., a stale PID.
-            // >
-            // > The kernel is free to recycle PID once it has been `wait(2)`ed upon by
-            // > its creator. Thus, acting upon a stale PID, e.g., by issuing a `kill`
-            // > system call on it, bears the risk of killing an unrelated process.
-            // > This is an inherent limitation of using pidfiles.
-            // > The only race-free solution is to have a supervisor-process with a lifetime
-            // > that exceeds that of all of its child-processes (e.g., `runit`, `supervisord`).
-            //
-            // This is an ok risk as we only send a SIGHUP which likely won't actually
-            // kill the process, only reload config.
-            nix::sys::signal::kill(pid, Signal::SIGHUP).context("sending signal to local_proxy")?;
-        }
-    }
-
-    Ok(())
-}

control_plane/Cargo.toml

@@ -14,7 +14,7 @@ humantime.workspace = true
 nix.workspace = true
 once_cell.workspace = true
 humantime-serde.workspace = true
-hyper0.workspace = true
+hyper.workspace = true
 regex.workspace = true
 reqwest = { workspace = true, features = ["blocking", "json"] }
 scopeguard.workspace = true

control_plane/src/endpoint.rs

@@ -599,7 +599,6 @@ impl Endpoint {
             remote_extensions,
             pgbouncer_settings: None,
             shard_stripe_size: Some(shard_stripe_size),
-            local_proxy_config: None,
         };
         let spec_path = self.endpoint_path().join("spec.json");
         std::fs::write(spec_path, serde_json::to_string_pretty(&spec)?)?;

control_plane/src/storage_controller.rs

@@ -3,7 +3,7 @@ use crate::{
     local_env::{LocalEnv, NeonStorageControllerConf},
 };
 use camino::{Utf8Path, Utf8PathBuf};
-use hyper0::Uri;
+use hyper::Uri;
 use nix::unistd::Pid;
 use pageserver_api::{
     controller_api::{

docs/docker.md

@@ -5,7 +5,7 @@
 Currently we build two main images:
 
 - [neondatabase/neon](https://hub.docker.com/repository/docker/neondatabase/neon) — image with pre-built `pageserver`, `safekeeper` and `proxy` binaries and all the required runtime dependencies. Built from [/Dockerfile](/Dockerfile).
-- [neondatabase/compute-node-v16](https://hub.docker.com/repository/docker/neondatabase/compute-node-v16) — compute node image with pre-built Postgres binaries from [neondatabase/postgres](https://github.com/neondatabase/postgres). Similar images exist for v15 and v14. Built from [/compute-node/Dockerfile](/compute/Dockerfile.compute-node).
+- [neondatabase/compute-node-v16](https://hub.docker.com/repository/docker/neondatabase/compute-node-v16) — compute node image with pre-built Postgres binaries from [neondatabase/postgres](https://github.com/neondatabase/postgres). Similar images exist for v15 and v14.
 
 And additional intermediate image:
 

libs/compute_api/src/spec.rs

@@ -106,10 +106,6 @@ pub struct ComputeSpec {
     // Stripe size for pageserver sharding, in pages
     #[serde(default)]
     pub shard_stripe_size: Option<usize>,
-
-    /// Local Proxy configuration used for JWT authentication
-    #[serde(default)]
-    pub local_proxy_config: Option<LocalProxySpec>,
 }
 
 /// Feature flag to signal `compute_ctl` to enable certain experimental functionality.
@@ -282,13 +278,11 @@ pub struct GenericOption {
 /// declare a `trait` on it.
 pub type GenericOptions = Option<Vec<GenericOption>>;
 
-/// Configured the local_proxy application with the relevant JWKS and roles it should
+/// Configured the local-proxy application with the relevant JWKS and roles it should
 /// use for authorizing connect requests using JWT.
 #[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct LocalProxySpec {
-    #[serde(default)]
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub jwks: Option<Vec<JwksSettings>>,
+    pub jwks: Vec<JwksSettings>,
 }
 
 #[derive(Clone, Debug, Deserialize, Serialize)]

libs/pageserver_api/src/config.rs

@@ -296,14 +296,7 @@ pub mod defaults {
 
     pub const DEFAULT_INGEST_BATCH_SIZE: u64 = 100;
 
-    /// Soft limit for the maximum size of a vectored read.
-    ///
-    /// This is determined by the largest NeonWalRecord that can exist (minus dbdir and reldir keys
-    /// which are bounded by the blob io limits only). As of this writing, that is a `NeonWalRecord::ClogSetCommitted` record,
-    /// with 32k xids. That's the max number of XIDS on a single CLOG page. The size of such a record
-    /// is `sizeof(Transactionid) * 32768 + (some fixed overhead from 'timestamp`, the Vec length and whatever extra serde serialization adds)`.
-    /// That is, slightly above 128 kB.
-    pub const DEFAULT_MAX_VECTORED_READ_BYTES: usize = 130 * 1024; // 130 KiB
+    pub const DEFAULT_MAX_VECTORED_READ_BYTES: usize = 128 * 1024; // 128 KiB
 
     pub const DEFAULT_IMAGE_COMPRESSION: ImageCompressionAlgorithm =
         ImageCompressionAlgorithm::Zstd { level: Some(1) };

libs/pageserver_api/src/key.rs

@@ -748,16 +748,6 @@ impl Key {
         self.field1 == 0x00 && self.field4 != 0 && self.field6 != 0xffffffff
     }
 
-    #[inline(always)]
-    pub fn is_rel_dir_key(&self) -> bool {
-        self.field1 == 0x00
-            && self.field2 != 0
-            && self.field3 != 0
-            && self.field4 == 0
-            && self.field5 == 0
-            && self.field6 == 1
-    }
-
     /// Guaranteed to return `Ok()` if [`Self::is_rel_block_key`] returns `true` for `key`.
     #[inline(always)]
     pub fn to_rel_block(self) -> anyhow::Result<(RelTag, BlockNumber)> {

libs/remote_storage/Cargo.toml

@@ -16,7 +16,7 @@ aws-sdk-s3.workspace = true
 bytes.workspace = true
 camino = { workspace = true, features = ["serde1"] }
 humantime-serde.workspace = true
-hyper0 = { workspace = true, features = ["stream"] }
+hyper = { workspace = true, features = ["stream"] }
 futures.workspace = true
 serde.workspace = true
 serde_json.workspace = true

libs/remote_storage/src/azure_blob.rs

@@ -14,7 +14,7 @@ use std::time::SystemTime;
 
 use super::REMOTE_STORAGE_PREFIX_SEPARATOR;
 use anyhow::Result;
-use azure_core::request_options::{IfMatchCondition, MaxResults, Metadata, Range};
+use azure_core::request_options::{MaxResults, Metadata, Range};
 use azure_core::{Continuable, RetryOptions};
 use azure_identity::DefaultAzureCredential;
 use azure_storage::StorageCredentials;
@@ -33,10 +33,10 @@ use tracing::debug;
 use utils::backoff;
 
 use crate::metrics::{start_measuring_requests, AttemptOutcome, RequestKind};
+use crate::ListingObject;
 use crate::{
-    config::AzureConfig, error::Cancelled, ConcurrencyLimiter, Download, DownloadError,
-    DownloadOpts, Listing, ListingMode, ListingObject, RemotePath, RemoteStorage, StorageMetadata,
-    TimeTravelError, TimeoutOrCancel,
+    config::AzureConfig, error::Cancelled, ConcurrencyLimiter, Download, DownloadError, Listing,
+    ListingMode, RemotePath, RemoteStorage, StorageMetadata, TimeTravelError, TimeoutOrCancel,
 };
 
 pub struct AzureBlobStorage {
@@ -259,7 +259,6 @@ fn to_download_error(error: azure_core::Error) -> DownloadError {
     if let Some(http_err) = error.as_http_error() {
         match http_err.status() {
             StatusCode::NotFound => DownloadError::NotFound,
-            StatusCode::NotModified => DownloadError::Unmodified,
             StatusCode::BadRequest => DownloadError::BadInput(anyhow::Error::new(error)),
             _ => DownloadError::Other(anyhow::Error::new(error)),
         }
@@ -485,16 +484,11 @@ impl RemoteStorage for AzureBlobStorage {
     async fn download(
         &self,
         from: &RemotePath,
-        opts: &DownloadOpts,
         cancel: &CancellationToken,
     ) -> Result<Download, DownloadError> {
         let blob_client = self.client.blob_client(self.relative_path_to_name(from));
 
-        let mut builder = blob_client.get();
-
-        if let Some(ref etag) = opts.etag {
-            builder = builder.if_match(IfMatchCondition::NotMatch(etag.to_string()))
-        }
+        let builder = blob_client.get();
 
         self.download_for_builder(builder, cancel).await
     }

libs/remote_storage/src/error.rs

@@ -5,8 +5,6 @@ pub enum DownloadError {
     BadInput(anyhow::Error),
     /// The file was not found in the remote storage.
     NotFound,
-    /// The caller provided an ETag, and the file was not modified.
-    Unmodified,
     /// A cancellation token aborted the download, typically during
     /// tenant detach or process shutdown.
     Cancelled,
@@ -26,7 +24,6 @@ impl std::fmt::Display for DownloadError {
                 write!(f, "Failed to download a remote file due to user input: {e}")
             }
             DownloadError::NotFound => write!(f, "No file found for the remote object id given"),
-            DownloadError::Unmodified => write!(f, "File was not modified"),
             DownloadError::Cancelled => write!(f, "Cancelled, shutting down"),
             DownloadError::Timeout => write!(f, "timeout"),
             DownloadError::Other(e) => write!(f, "Failed to download a remote file: {e:?}"),
@@ -41,7 +38,7 @@ impl DownloadError {
     pub fn is_permanent(&self) -> bool {
         use DownloadError::*;
         match self {
-            BadInput(_) | NotFound | Unmodified | Cancelled => true,
+            BadInput(_) | NotFound | Cancelled => true,
             Timeout | Other(_) => false,
         }
     }

libs/remote_storage/src/lib.rs

@@ -161,14 +161,6 @@ pub struct Listing {
     pub keys: Vec<ListingObject>,
 }
 
-/// Options for downloads. The default value is a plain GET.
-#[derive(Default)]
-pub struct DownloadOpts {
-    /// If given, returns [`DownloadError::Unmodified`] if the object still has
-    /// the same ETag (using If-None-Match).
-    pub etag: Option<Etag>,
-}
-
 /// Storage (potentially remote) API to manage its state.
 /// This storage tries to be unaware of any layered repository context,
 /// providing basic CRUD operations for storage files.
@@ -253,7 +245,6 @@ pub trait RemoteStorage: Send + Sync + 'static {
     async fn download(
         &self,
         from: &RemotePath,
-        opts: &DownloadOpts,
         cancel: &CancellationToken,
     ) -> Result<Download, DownloadError>;
 
@@ -410,18 +401,16 @@ impl<Other: RemoteStorage> GenericRemoteStorage<Arc<Other>> {
         }
     }
 
-    /// See [`RemoteStorage::download`]
     pub async fn download(
         &self,
         from: &RemotePath,
-        opts: &DownloadOpts,
         cancel: &CancellationToken,
     ) -> Result<Download, DownloadError> {
         match self {
-            Self::LocalFs(s) => s.download(from, opts, cancel).await,
-            Self::AwsS3(s) => s.download(from, opts, cancel).await,
-            Self::AzureBlob(s) => s.download(from, opts, cancel).await,
-            Self::Unreliable(s) => s.download(from, opts, cancel).await,
+            Self::LocalFs(s) => s.download(from, cancel).await,
+            Self::AwsS3(s) => s.download(from, cancel).await,
+            Self::AzureBlob(s) => s.download(from, cancel).await,
+            Self::Unreliable(s) => s.download(from, cancel).await,
         }
     }
 
@@ -583,7 +572,7 @@ impl GenericRemoteStorage {
     ) -> Result<Download, DownloadError> {
         match byte_range {
             Some((start, end)) => self.download_byte_range(from, start, end, cancel).await,
-            None => self.download(from, &DownloadOpts::default(), cancel).await,
+            None => self.download(from, cancel).await,
         }
     }
 

libs/remote_storage/src/local_fs.rs

@@ -23,8 +23,8 @@ use tokio_util::{io::ReaderStream, sync::CancellationToken};
 use utils::crashsafe::path_with_suffix_extension;
 
 use crate::{
-    Download, DownloadError, DownloadOpts, Listing, ListingMode, ListingObject, RemotePath,
-    TimeTravelError, TimeoutOrCancel, REMOTE_STORAGE_PREFIX_SEPARATOR,
+    Download, DownloadError, Listing, ListingMode, ListingObject, RemotePath, TimeTravelError,
+    TimeoutOrCancel, REMOTE_STORAGE_PREFIX_SEPARATOR,
 };
 
 use super::{RemoteStorage, StorageMetadata};
@@ -494,17 +494,11 @@ impl RemoteStorage for LocalFs {
     async fn download(
         &self,
         from: &RemotePath,
-        opts: &DownloadOpts,
         cancel: &CancellationToken,
     ) -> Result<Download, DownloadError> {
         let target_path = from.with_base(&self.storage_root);
 
         let file_metadata = file_metadata(&target_path).await?;
-        let etag = mock_etag(&file_metadata);
-
-        if opts.etag.as_ref() == Some(&etag) {
-            return Err(DownloadError::Unmodified);
-        }
 
         let source = ReaderStream::new(
             fs::OpenOptions::new()
@@ -525,6 +519,7 @@ impl RemoteStorage for LocalFs {
         let cancel_or_timeout = crate::support::cancel_or_timeout(self.timeout, cancel.clone());
         let source = crate::support::DownloadStream::new(cancel_or_timeout, source);
 
+        let etag = mock_etag(&file_metadata);
         Ok(Download {
             metadata,
             last_modified: file_metadata
@@ -697,7 +692,7 @@ mod fs_tests {
     ) -> anyhow::Result<String> {
         let cancel = CancellationToken::new();
         let download = storage
-            .download(remote_storage_path, &DownloadOpts::default(), &cancel)
+            .download(remote_storage_path, &cancel)
             .await
             .map_err(|e| anyhow::anyhow!("Download failed: {e}"))?;
         ensure!(
@@ -778,8 +773,8 @@ mod fs_tests {
             "We should upload and download the same contents"
         );
 
-        let non_existing_path = RemotePath::new(Utf8Path::new("somewhere/else"))?;
-        match storage.download(&non_existing_path, &DownloadOpts::default(), &cancel).await {
+        let non_existing_path = "somewhere/else";
+        match storage.download(&RemotePath::new(Utf8Path::new(non_existing_path))?, &cancel).await {
             Err(DownloadError::NotFound) => {} // Should get NotFound for non existing keys
             other => panic!("Should get a NotFound error when downloading non-existing storage files, but got: {other:?}"),
         }
@@ -1106,13 +1101,7 @@ mod fs_tests {
             storage.upload(body, len, &path, None, &cancel).await?;
         }
 
-        let read = aggregate(
-            storage
-                .download(&path, &DownloadOpts::default(), &cancel)
-                .await?
-                .download_stream,
-        )
-        .await?;
+        let read = aggregate(storage.download(&path, &cancel).await?.download_stream).await?;
         assert_eq!(body, read);
 
         let shorter = Bytes::from_static(b"shorter body");
@@ -1123,13 +1112,7 @@ mod fs_tests {
             storage.upload(body, len, &path, None, &cancel).await?;
         }
 
-        let read = aggregate(
-            storage
-                .download(&path, &DownloadOpts::default(), &cancel)
-                .await?
-                .download_stream,
-        )
-        .await?;
+        let read = aggregate(storage.download(&path, &cancel).await?.download_stream).await?;
         assert_eq!(shorter, read);
         Ok(())
     }
@@ -1162,13 +1145,7 @@ mod fs_tests {
             storage.upload(body, len, &path, None, &cancel).await?;
         }
 
-        let read = aggregate(
-            storage
-                .download(&path, &DownloadOpts::default(), &cancel)
-                .await?
-                .download_stream,
-        )
-        .await?;
+        let read = aggregate(storage.download(&path, &cancel).await?.download_stream).await?;
         assert_eq!(body, read);
 
         Ok(())

libs/remote_storage/src/s3_bucket.rs

@@ -28,13 +28,12 @@ use aws_sdk_s3::{
     Client,
 };
 use aws_smithy_async::rt::sleep::TokioSleep;
-use http_types::StatusCode;
 
 use aws_smithy_types::{body::SdkBody, DateTime};
 use aws_smithy_types::{byte_stream::ByteStream, date_time::ConversionError};
 use bytes::Bytes;
 use futures::stream::Stream;
-use hyper0::Body;
+use hyper::Body;
 use scopeguard::ScopeGuard;
 use tokio_util::sync::CancellationToken;
 use utils::backoff;
@@ -45,8 +44,8 @@ use crate::{
     error::Cancelled,
     metrics::{start_counting_cancelled_wait, start_measuring_requests},
     support::PermitCarrying,
-    ConcurrencyLimiter, Download, DownloadError, DownloadOpts, Listing, ListingMode, ListingObject,
-    RemotePath, RemoteStorage, TimeTravelError, TimeoutOrCancel, MAX_KEYS_PER_DELETE,
+    ConcurrencyLimiter, Download, DownloadError, Listing, ListingMode, ListingObject, RemotePath,
+    RemoteStorage, TimeTravelError, TimeoutOrCancel, MAX_KEYS_PER_DELETE,
     REMOTE_STORAGE_PREFIX_SEPARATOR,
 };
 
@@ -68,7 +67,6 @@ pub struct S3Bucket {
 struct GetObjectRequest {
     bucket: String,
     key: String,
-    etag: Option<String>,
     range: Option<String>,
 }
 impl S3Bucket {
@@ -250,18 +248,13 @@ impl S3Bucket {
 
         let started_at = start_measuring_requests(kind);
 
-        let mut builder = self
+        let get_object = self
             .client
             .get_object()
             .bucket(request.bucket)
             .key(request.key)
-            .set_range(request.range);
-
-        if let Some(etag) = request.etag {
-            builder = builder.if_none_match(etag);
-        }
-
-        let get_object = builder.send();
+            .set_range(request.range)
+            .send();
 
         let get_object = tokio::select! {
             res = get_object => res,
@@ -284,20 +277,6 @@ impl S3Bucket {
                 );
                 return Err(DownloadError::NotFound);
             }
-            Err(SdkError::ServiceError(e))
-                // aws_smithy_runtime_api::http::response::StatusCode isn't
-                // re-exported by any aws crates, so just check the numeric
-                // status against http_types::StatusCode instead of pulling it.
-                if e.raw().status().as_u16() == StatusCode::NotModified =>
-            {
-                // Count an unmodified file as a success.
-                crate::metrics::BUCKET_METRICS.req_seconds.observe_elapsed(
-                    kind,
-                    AttemptOutcome::Ok,
-                    started_at,
-                );
-                return Err(DownloadError::Unmodified);
-            }
             Err(e) => {
                 crate::metrics::BUCKET_METRICS.req_seconds.observe_elapsed(
                     kind,
@@ -794,7 +773,6 @@ impl RemoteStorage for S3Bucket {
     async fn download(
         &self,
         from: &RemotePath,
-        opts: &DownloadOpts,
         cancel: &CancellationToken,
     ) -> Result<Download, DownloadError> {
         // if prefix is not none then download file `prefix/from`
@@ -803,7 +781,6 @@ impl RemoteStorage for S3Bucket {
             GetObjectRequest {
                 bucket: self.bucket_name.clone(),
                 key: self.relative_path_to_s3_object(from),
-                etag: opts.etag.as_ref().map(|e| e.to_string()),
                 range: None,
             },
             cancel,
@@ -830,7 +807,6 @@ impl RemoteStorage for S3Bucket {
             GetObjectRequest {
                 bucket: self.bucket_name.clone(),
                 key: self.relative_path_to_s3_object(from),
-                etag: None,
                 range,
             },
             cancel,

libs/remote_storage/src/simulate_failures.rs

@@ -12,8 +12,8 @@ use std::{collections::hash_map::Entry, sync::Arc};
 use tokio_util::sync::CancellationToken;
 
 use crate::{
-    Download, DownloadError, DownloadOpts, GenericRemoteStorage, Listing, ListingMode, RemotePath,
-    RemoteStorage, StorageMetadata, TimeTravelError,
+    Download, DownloadError, GenericRemoteStorage, Listing, ListingMode, RemotePath, RemoteStorage,
+    StorageMetadata, TimeTravelError,
 };
 
 pub struct UnreliableWrapper {
@@ -167,12 +167,11 @@ impl RemoteStorage for UnreliableWrapper {
     async fn download(
         &self,
         from: &RemotePath,
-        opts: &DownloadOpts,
         cancel: &CancellationToken,
     ) -> Result<Download, DownloadError> {
         self.attempt(RemoteOp::Download(from.clone()))
             .map_err(DownloadError::Other)?;
-        self.inner.download(from, opts, cancel).await
+        self.inner.download(from, cancel).await
     }
 
     async fn download_byte_range(

libs/remote_storage/tests/common/tests.rs

@@ -1,7 +1,8 @@
 use anyhow::Context;
 use camino::Utf8Path;
 use futures::StreamExt;
-use remote_storage::{DownloadError, DownloadOpts, ListingMode, ListingObject, RemotePath};
+use remote_storage::ListingMode;
+use remote_storage::RemotePath;
 use std::sync::Arc;
 use std::{collections::HashSet, num::NonZeroU32};
 use test_context::test_context;
@@ -283,10 +284,7 @@ async fn upload_download_works(ctx: &mut MaybeEnabledStorage) -> anyhow::Result<
     ctx.client.upload(data, len, &path, None, &cancel).await?;
 
     // Normal download request
-    let dl = ctx
-        .client
-        .download(&path, &DownloadOpts::default(), &cancel)
-        .await?;
+    let dl = ctx.client.download(&path, &cancel).await?;
     let buf = download_to_vec(dl).await?;
     assert_eq!(&buf, &orig);
 
@@ -339,54 +337,6 @@ async fn upload_download_works(ctx: &mut MaybeEnabledStorage) -> anyhow::Result<
     Ok(())
 }
 
-/// Tests that conditional downloads work properly, by returning
-/// DownloadError::Unmodified when the object ETag matches the given ETag.
-#[test_context(MaybeEnabledStorage)]
-#[tokio::test]
-async fn download_conditional(ctx: &mut MaybeEnabledStorage) -> anyhow::Result<()> {
-    let MaybeEnabledStorage::Enabled(ctx) = ctx else {
-        return Ok(());
-    };
-    let cancel = CancellationToken::new();
-
-    // Create a file.
-    let path = RemotePath::new(Utf8Path::new(format!("{}/file", ctx.base_prefix).as_str()))?;
-    let data = bytes::Bytes::from_static("foo".as_bytes());
-    let (stream, len) = wrap_stream(data);
-    ctx.client.upload(stream, len, &path, None, &cancel).await?;
-
-    // Download it to obtain its etag.
-    let mut opts = DownloadOpts::default();
-    let download = ctx.client.download(&path, &opts, &cancel).await?;
-
-    // Download with the etag yields DownloadError::Unmodified.
-    opts.etag = Some(download.etag);
-    let result = ctx.client.download(&path, &opts, &cancel).await;
-    assert!(
-        matches!(result, Err(DownloadError::Unmodified)),
-        "expected DownloadError::Unmodified, got {result:?}"
-    );
-
-    // Replace the file contents.
-    let data = bytes::Bytes::from_static("bar".as_bytes());
-    let (stream, len) = wrap_stream(data);
-    ctx.client.upload(stream, len, &path, None, &cancel).await?;
-
-    // A download with the old etag should yield the new file.
-    let download = ctx.client.download(&path, &opts, &cancel).await?;
-    assert_ne!(download.etag, opts.etag.unwrap(), "ETag did not change");
-
-    // A download with the new etag should yield Unmodified again.
-    opts.etag = Some(download.etag);
-    let result = ctx.client.download(&path, &opts, &cancel).await;
-    assert!(
-        matches!(result, Err(DownloadError::Unmodified)),
-        "expected DownloadError::Unmodified, got {result:?}"
-    );
-
-    Ok(())
-}
-
 #[test_context(MaybeEnabledStorage)]
 #[tokio::test]
 async fn copy_works(ctx: &mut MaybeEnabledStorage) -> anyhow::Result<()> {
@@ -414,10 +364,7 @@ async fn copy_works(ctx: &mut MaybeEnabledStorage) -> anyhow::Result<()> {
     // Normal download request
     ctx.client.copy_object(&path, &path_dest, &cancel).await?;
 
-    let dl = ctx
-        .client
-        .download(&path_dest, &DownloadOpts::default(), &cancel)
-        .await?;
+    let dl = ctx.client.download(&path_dest, &cancel).await?;
     let buf = download_to_vec(dl).await?;
     assert_eq!(&buf, &orig);
 
@@ -429,56 +376,3 @@ async fn copy_works(ctx: &mut MaybeEnabledStorage) -> anyhow::Result<()> {
 
     Ok(())
 }
-
-/// Tests that head_object works properly.
-#[test_context(MaybeEnabledStorage)]
-#[tokio::test]
-async fn head_object(ctx: &mut MaybeEnabledStorage) -> anyhow::Result<()> {
-    let MaybeEnabledStorage::Enabled(ctx) = ctx else {
-        return Ok(());
-    };
-    let cancel = CancellationToken::new();
-
-    let path = RemotePath::new(Utf8Path::new(format!("{}/file", ctx.base_prefix).as_str()))?;
-
-    // Errors on missing file.
-    let result = ctx.client.head_object(&path, &cancel).await;
-    assert!(
-        matches!(result, Err(DownloadError::NotFound)),
-        "expected NotFound, got {result:?}"
-    );
-
-    // Create the file.
-    let data = bytes::Bytes::from_static("foo".as_bytes());
-    let (stream, len) = wrap_stream(data);
-    ctx.client.upload(stream, len, &path, None, &cancel).await?;
-
-    // Fetch the head metadata.
-    let object = ctx.client.head_object(&path, &cancel).await?;
-    assert_eq!(
-        object,
-        ListingObject {
-            key: path.clone(),
-            last_modified: object.last_modified, // ignore
-            size: 3
-        }
-    );
-
-    // Wait for a couple of seconds, and then update the file to check the last
-    // modified timestamp.
-    tokio::time::sleep(std::time::Duration::from_secs(2)).await;
-
-    let data = bytes::Bytes::from_static("bar".as_bytes());
-    let (stream, len) = wrap_stream(data);
-    ctx.client.upload(stream, len, &path, None, &cancel).await?;
-    let new = ctx.client.head_object(&path, &cancel).await?;
-
-    assert!(
-        !new.last_modified
-            .duration_since(object.last_modified)?
-            .is_zero(),
-        "last_modified did not advance"
-    );
-
-    Ok(())
-}

libs/remote_storage/tests/test_real_s3.rs

@@ -12,8 +12,8 @@ use anyhow::Context;
 use camino::Utf8Path;
 use futures_util::StreamExt;
 use remote_storage::{
-    DownloadError, DownloadOpts, GenericRemoteStorage, ListingMode, RemotePath,
-    RemoteStorageConfig, RemoteStorageKind, S3Config,
+    DownloadError, GenericRemoteStorage, ListingMode, RemotePath, RemoteStorageConfig,
+    RemoteStorageKind, S3Config,
 };
 use test_context::test_context;
 use test_context::AsyncTestContext;
@@ -121,8 +121,7 @@ async fn s3_time_travel_recovery_works(ctx: &mut MaybeEnabledStorage) -> anyhow:
 
     // A little check to ensure that our clock is not too far off from the S3 clock
     {
-        let opts = DownloadOpts::default();
-        let dl = retry(|| ctx.client.download(&path2, &opts, &cancel)).await?;
+        let dl = retry(|| ctx.client.download(&path2, &cancel)).await?;
         let last_modified = dl.last_modified;
         let half_wt = WAIT_TIME.mul_f32(0.5);
         let t0_hwt = t0 + half_wt;
@@ -160,12 +159,7 @@ async fn s3_time_travel_recovery_works(ctx: &mut MaybeEnabledStorage) -> anyhow:
     let t2_files_recovered = list_files(&ctx.client, &cancel).await?;
     println!("after recovery to t2: {t2_files_recovered:?}");
     assert_eq!(t2_files, t2_files_recovered);
-    let path2_recovered_t2 = download_to_vec(
-        ctx.client
-            .download(&path2, &DownloadOpts::default(), &cancel)
-            .await?,
-    )
-    .await?;
+    let path2_recovered_t2 = download_to_vec(ctx.client.download(&path2, &cancel).await?).await?;
     assert_eq!(path2_recovered_t2, new_data.as_bytes());
 
     // after recovery to t1: path1 is back, path2 has the old content
@@ -176,12 +170,7 @@ async fn s3_time_travel_recovery_works(ctx: &mut MaybeEnabledStorage) -> anyhow:
     let t1_files_recovered = list_files(&ctx.client, &cancel).await?;
     println!("after recovery to t1: {t1_files_recovered:?}");
     assert_eq!(t1_files, t1_files_recovered);
-    let path2_recovered_t1 = download_to_vec(
-        ctx.client
-            .download(&path2, &DownloadOpts::default(), &cancel)
-            .await?,
-    )
-    .await?;
+    let path2_recovered_t1 = download_to_vec(ctx.client.download(&path2, &cancel).await?).await?;
     assert_eq!(path2_recovered_t1, old_data.as_bytes());
 
     // after recovery to t0: everything is gone except for path1
@@ -427,7 +416,7 @@ async fn download_is_timeouted(ctx: &mut MaybeEnabledStorage) {
     let started_at = std::time::Instant::now();
     let mut stream = ctx
         .client
-        .download(&path, &DownloadOpts::default(), &cancel)
+        .download(&path, &cancel)
         .await
         .expect("download succeeds")
         .download_stream;
@@ -502,7 +491,7 @@ async fn download_is_cancelled(ctx: &mut MaybeEnabledStorage) {
     {
         let stream = ctx
             .client
-            .download(&path, &DownloadOpts::default(), &cancel)
+            .download(&path, &cancel)
             .await
             .expect("download succeeds")
             .download_stream;

libs/tracing-utils/Cargo.toml

@@ -5,7 +5,7 @@ edition.workspace = true
 license.workspace = true
 
 [dependencies]
-hyper0.workspace = true
+hyper.workspace = true
 opentelemetry = { workspace = true, features = ["trace"] }
 opentelemetry_sdk = { workspace = true, features = ["rt-tokio"] }
 opentelemetry-otlp = { workspace = true, default-features = false, features = ["http-proto", "trace", "http", "reqwest-client"] }

libs/tracing-utils/src/http.rs

@@ -1,7 +1,7 @@
 //! Tracing wrapper for Hyper HTTP server
 
-use hyper0::HeaderMap;
-use hyper0::{Body, Request, Response};
+use hyper::HeaderMap;
+use hyper::{Body, Request, Response};
 use std::future::Future;
 use tracing::Instrument;
 use tracing_opentelemetry::OpenTelemetrySpanExt;

libs/utils/Cargo.toml

@@ -22,7 +22,7 @@ chrono.workspace = true
 git-version.workspace = true
 hex = { workspace = true, features = ["serde"] }
 humantime.workspace = true
-hyper0 = { workspace = true, features = ["full"] }
+hyper = { workspace = true, features = ["full"] }
 fail.workspace = true
 futures = { workspace = true}
 jsonwebtoken.workspace = true

libs/utils/src/lib.rs

@@ -2,8 +2,6 @@
 //! between other crates in this repository.
 #![deny(clippy::undocumented_unsafe_blocks)]
 
-extern crate hyper0 as hyper;
-
 pub mod backoff;
 
 /// `Lsn` type implements common tasks on Log Sequence Numbers

libs/vm_monitor/src/lib.rs

@@ -7,13 +7,11 @@ use axum::{
     extract::{ws::WebSocket, State, WebSocketUpgrade},
     response::Response,
 };
-use axum::{routing::get, Router};
+use axum::{routing::get, Router, Server};
 use clap::Parser;
 use futures::Future;
-use std::net::SocketAddr;
 use std::{fmt::Debug, time::Duration};
 use sysinfo::{RefreshKind, System, SystemExt};
-use tokio::net::TcpListener;
 use tokio::{sync::broadcast, task::JoinHandle};
 use tokio_util::sync::CancellationToken;
 use tracing::{error, info};
@@ -134,14 +132,14 @@ pub async fn start(args: &'static Args, token: CancellationToken) -> anyhow::Res
             args,
         });
 
-    let addr_str = args.addr();
-    let addr: SocketAddr = addr_str.parse().expect("parsing address should not fail");
-
-    let listener = TcpListener::bind(&addr)
-        .await
+    let addr = args.addr();
+    let bound = Server::try_bind(&addr.parse().expect("parsing address should not fail"))
         .with_context(|| format!("failed to bind to {addr}"))?;
-    info!(addr_str, "server bound");
-    axum::serve(listener, app.into_make_service())
+
+    info!(addr, "server bound");
+
+    bound
+        .serve(app.into_make_service())
         .await
         .context("server exited")?;
 

pageserver/Cargo.toml

@@ -30,7 +30,7 @@ futures.workspace = true
 hex.workspace = true
 humantime.workspace = true
 humantime-serde.workspace = true
-hyper0.workspace = true
+hyper.workspace = true
 itertools.workspace = true
 md5.workspace = true
 nix.workspace = true

pageserver/src/bin/pageserver.rs

@@ -575,7 +575,7 @@ fn start_pageserver(
             .build()
             .map_err(|err| anyhow!(err))?;
         let service = utils::http::RouterService::new(router).unwrap();
-        let server = hyper0::Server::from_tcp(http_listener)?
+        let server = hyper::Server::from_tcp(http_listener)?
             .serve(service)
             .with_graceful_shutdown({
                 let cancel = cancel.clone();

pageserver/src/lib.rs

@@ -13,8 +13,6 @@ pub mod http;
 pub mod import_datadir;
 pub mod l0_flush;
 
-extern crate hyper0 as hyper;
-
 use futures::{stream::FuturesUnordered, StreamExt};
 pub use pageserver_api::keyspace;
 use tokio_util::sync::CancellationToken;

pageserver/src/tenant/remote_timeline_client/download.rs

@@ -27,7 +27,7 @@ use crate::tenant::Generation;
 use crate::virtual_file::owned_buffers_io::io_buf_ext::IoBufExt;
 use crate::virtual_file::{on_fatal_io_error, MaybeFatalIo, VirtualFile};
 use crate::TEMP_FILE_SUFFIX;
-use remote_storage::{DownloadError, DownloadOpts, GenericRemoteStorage, ListingMode, RemotePath};
+use remote_storage::{DownloadError, GenericRemoteStorage, ListingMode, RemotePath};
 use utils::crashsafe::path_with_suffix_extension;
 use utils::id::{TenantId, TimelineId};
 use utils::pausable_failpoint;
@@ -153,9 +153,7 @@ async fn download_object<'a>(
                     .with_context(|| format!("create a destination file for layer '{dst_path}'"))
                     .map_err(DownloadError::Other)?;
 
-                let download = storage
-                    .download(src_path, &DownloadOpts::default(), cancel)
-                    .await?;
+                let download = storage.download(src_path, cancel).await?;
 
                 pausable_failpoint!("before-downloading-layer-stream-pausable");
 
@@ -206,9 +204,7 @@ async fn download_object<'a>(
                     .with_context(|| format!("create a destination file for layer '{dst_path}'"))
                     .map_err(DownloadError::Other)?;
 
-                let mut download = storage
-                    .download(src_path, &DownloadOpts::default(), cancel)
-                    .await?;
+                let mut download = storage.download(src_path, cancel).await?;
 
                 pausable_failpoint!("before-downloading-layer-stream-pausable");
 
@@ -348,9 +344,7 @@ async fn do_download_index_part(
 
     let index_part_bytes = download_retry_forever(
         || async {
-            let download = storage
-                .download(&remote_path, &DownloadOpts::default(), cancel)
-                .await?;
+            let download = storage.download(&remote_path, cancel).await?;
 
             let mut bytes = Vec::new();
 
@@ -532,15 +526,10 @@ pub(crate) async fn download_initdb_tar_zst(
                 .with_context(|| format!("tempfile creation {temp_path}"))
                 .map_err(DownloadError::Other)?;
 
-            let download = match storage
-                .download(&remote_path, &DownloadOpts::default(), cancel)
-                .await
-            {
+            let download = match storage.download(&remote_path, cancel).await {
                 Ok(dl) => dl,
                 Err(DownloadError::NotFound) => {
-                    storage
-                        .download(&remote_preserved_path, &DownloadOpts::default(), cancel)
-                        .await?
+                    storage.download(&remote_preserved_path, cancel).await?
                 }
                 Err(other) => Err(other)?,
             };

pageserver/src/tenant/secondary/downloader.rs

@@ -49,7 +49,7 @@ use futures::Future;
 use metrics::UIntGauge;
 use pageserver_api::models::SecondaryProgress;
 use pageserver_api::shard::TenantShardId;
-use remote_storage::{DownloadError, DownloadOpts, Etag, GenericRemoteStorage};
+use remote_storage::{DownloadError, Etag, GenericRemoteStorage};
 
 use tokio_util::sync::CancellationToken;
 use tracing::{info_span, instrument, warn, Instrument};
@@ -944,34 +944,36 @@ impl<'a> TenantDownloader<'a> {
     ) -> Result<HeatMapDownload, UpdateError> {
         debug_assert_current_span_has_tenant_id();
         let tenant_shard_id = self.secondary_state.get_tenant_shard_id();
+        // TODO: pull up etag check into the request, to do a conditional GET rather than
+        // issuing a GET and then maybe ignoring the response body
+        // (https://github.com/neondatabase/neon/issues/6199)
         tracing::debug!("Downloading heatmap for secondary tenant",);
 
         let heatmap_path = remote_heatmap_path(tenant_shard_id);
         let cancel = &self.secondary_state.cancel;
-        let opts = DownloadOpts {
-            etag: prev_etag.cloned(),
-        };
 
         backoff::retry(
             || async {
-                let download = match self
+                let download = self
                     .remote_storage
-                    .download(&heatmap_path, &opts, cancel)
+                    .download(&heatmap_path, cancel)
                     .await
-                {
-                    Ok(download) => download,
-                    Err(DownloadError::Unmodified) => return Ok(HeatMapDownload::Unmodified),
-                    Err(err) => return Err(err.into()),
-                };
+                    .map_err(UpdateError::from)?;
 
-                let mut heatmap_bytes = Vec::new();
-                let mut body = tokio_util::io::StreamReader::new(download.download_stream);
-                let _size = tokio::io::copy_buf(&mut body, &mut heatmap_bytes).await?;
-                Ok(HeatMapDownload::Modified(HeatMapModified {
-                    etag: download.etag,
-                    last_modified: download.last_modified,
-                    bytes: heatmap_bytes,
-                }))
+                SECONDARY_MODE.download_heatmap.inc();
+
+                if Some(&download.etag) == prev_etag {
+                    Ok(HeatMapDownload::Unmodified)
+                } else {
+                    let mut heatmap_bytes = Vec::new();
+                    let mut body = tokio_util::io::StreamReader::new(download.download_stream);
+                    let _size = tokio::io::copy_buf(&mut body, &mut heatmap_bytes).await?;
+                    Ok(HeatMapDownload::Modified(HeatMapModified {
+                        etag: download.etag,
+                        last_modified: download.last_modified,
+                        bytes: heatmap_bytes,
+                    }))
+                }
             },
             |e| matches!(e, UpdateError::NoData | UpdateError::Cancelled),
             FAILED_DOWNLOAD_WARN_THRESHOLD,
@@ -982,7 +984,6 @@ impl<'a> TenantDownloader<'a> {
         .await
         .ok_or_else(|| UpdateError::Cancelled)
         .and_then(|x| x)
-        .inspect(|_| SECONDARY_MODE.download_heatmap.inc())
     }
 
     /// Download heatmap layers that are not present on local disk, or update their

pageserver/src/tenant/storage_layer/delta_layer.rs

@@ -53,7 +53,6 @@ use camino::{Utf8Path, Utf8PathBuf};
 use futures::StreamExt;
 use itertools::Itertools;
 use pageserver_api::config::MaxVectoredReadBytes;
-use pageserver_api::key::DBDIR_KEY;
 use pageserver_api::keyspace::KeySpace;
 use pageserver_api::models::ImageCompressionAlgorithm;
 use pageserver_api::shard::TenantShardId;
@@ -964,25 +963,14 @@ impl DeltaLayerInner {
                 .blobs_at
                 .as_slice()
                 .iter()
-                .filter_map(|(_, blob_meta)| {
-                    if blob_meta.key.is_rel_dir_key() || blob_meta.key == DBDIR_KEY {
-                        // The size of values for these keys is unbounded and can
-                        // grow very large in pathological cases.
-                        None
-                    } else {
-                        Some(format!("{}@{}", blob_meta.key, blob_meta.lsn))
-                    }
-                })
+                .map(|(_, blob_meta)| format!("{}@{}", blob_meta.key, blob_meta.lsn))
                 .join(", ");
-
-            if !offenders.is_empty() {
-                tracing::warn!(
-                    "Oversized vectored read ({} > {}) for keys {}",
-                    largest_read_size,
-                    read_size_soft_max,
-                    offenders
-                );
-            }
+            tracing::warn!(
+                "Oversized vectored read ({} > {}) for keys {}",
+                largest_read_size,
+                read_size_soft_max,
+                offenders
+            );
         }
 
         largest_read_size

pageserver/src/tenant/storage_layer/image_layer.rs

@@ -49,7 +49,6 @@ use camino::{Utf8Path, Utf8PathBuf};
 use hex;
 use itertools::Itertools;
 use pageserver_api::config::MaxVectoredReadBytes;
-use pageserver_api::key::DBDIR_KEY;
 use pageserver_api::keyspace::KeySpace;
 use pageserver_api::shard::{ShardIdentity, TenantShardId};
 use rand::{distributions::Alphanumeric, Rng};
@@ -588,25 +587,14 @@ impl ImageLayerInner {
                     .blobs_at
                     .as_slice()
                     .iter()
-                    .filter_map(|(_, blob_meta)| {
-                        if blob_meta.key.is_rel_dir_key() || blob_meta.key == DBDIR_KEY {
-                            // The size of values for these keys is unbounded and can
-                            // grow very large in pathological cases.
-                            None
-                        } else {
-                            Some(format!("{}@{}", blob_meta.key, blob_meta.lsn))
-                        }
-                    })
+                    .map(|(_, blob_meta)| format!("{}@{}", blob_meta.key, blob_meta.lsn))
                     .join(", ");
-
-                if !offenders.is_empty() {
-                    tracing::warn!(
-                        "Oversized vectored read ({} > {}) for keys {}",
-                        buf_size,
-                        max_vectored_read_bytes,
-                        offenders
-                    );
-                }
+                tracing::warn!(
+                    "Oversized vectored read ({} > {}) for keys {}",
+                    buf_size,
+                    max_vectored_read_bytes,
+                    offenders
+                );
             }
 
             let buf = BytesMut::with_capacity(buf_size);

pgxn/neon/walsender_hooks.c

@@ -191,14 +191,13 @@ NeonOnDemandXLogReaderRoutines(XLogReaderRoutine *xlr)
 
 	if (!wal_reader)
 	{
-		XLogRecPtr	basebackupLsn = GetRedoStartLsn();
+		XLogRecPtr	epochStartLsn = pg_atomic_read_u64(&GetWalpropShmemState()->propEpochStartLsn);
 
-		/* should never happen */
-		if (basebackupLsn == 0)
+		if (epochStartLsn == 0)
 		{
-			elog(ERROR, "unable to start walsender when basebackupLsn is 0");
+			elog(ERROR, "Unable to start walsender when propEpochStartLsn is 0!");
 		}
-		wal_reader = NeonWALReaderAllocate(wal_segment_size, basebackupLsn, "[walsender] ");
+		wal_reader = NeonWALReaderAllocate(wal_segment_size, epochStartLsn, "[walsender] ");
 	}
 	xlr->page_read = NeonWALPageRead;
 	xlr->segment_open = NeonWALReadSegmentOpen;

proxy/Cargo.toml

@@ -38,7 +38,7 @@ hostname.workspace = true
 http.workspace = true
 humantime.workspace = true
 humantime-serde.workspace = true
-hyper0.workspace = true
+hyper.workspace = true
 hyper1 = { package = "hyper", version = "1.2", features = ["server"] }
 hyper-util = { version = "0.1", features = ["server", "http1", "http2", "tokio"] }
 http-body-util = { version = "0.1" }

proxy/src/bin/local_proxy.rs

@@ -77,10 +77,10 @@ struct LocalProxyCliArgs {
     #[clap(long, default_value = "127.0.0.1:5432")]
     compute: SocketAddr,
     /// Path of the local proxy config file
-    #[clap(long, default_value = "./local_proxy.json")]
+    #[clap(long, default_value = "./localproxy.json")]
     config_path: Utf8PathBuf,
     /// Path of the local proxy PID file
-    #[clap(long, default_value = "./local_proxy.pid")]
+    #[clap(long, default_value = "./localproxy.pid")]
     pid_path: Utf8PathBuf,
 }
 
@@ -109,7 +109,7 @@ struct SqlOverHttpArgs {
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
-    let _logging_guard = proxy::logging::init_local_proxy()?;
+    let _logging_guard = proxy::logging::init().await?;
     let _panic_hook_guard = utils::logging::replace_panic_hook_with_tracing_panic_hook();
     let _sentry_guard = init_sentry(Some(GIT_VERSION.into()), &[]);
 
@@ -138,7 +138,7 @@ async fn main() -> anyhow::Result<()> {
     // in order to trigger the appropriate SIGHUP on config change.
     //
     // This also claims a "lock" that makes sure only one instance
-    // of local_proxy runs at a time.
+    // of local-proxy runs at a time.
     let _process_guard = loop {
         match pid_file::claim_for_current_process(&args.pid_path) {
             Ok(guard) => break guard,
@@ -164,6 +164,12 @@ async fn main() -> anyhow::Result<()> {
         16,
     ));
 
+    // write the process ID to a file so that compute-ctl can find our process later
+    // in order to trigger the appropriate SIGHUP on config change.
+    let pid = std::process::id();
+    info!("process running in PID {pid}");
+    std::fs::write(args.pid_path, format!("{pid}\n")).context("writing PID to file")?;
+
     let mut maintenance_tasks = JoinSet::new();
 
     let refresh_config_notify = Arc::new(Notify::new());
@@ -176,9 +182,9 @@ async fn main() -> anyhow::Result<()> {
 
     // trigger the first config load **after** setting up the signal hook
     // to avoid the race condition where:
-    // 1. No config file registered when local_proxy starts up
+    // 1. No config file registered when local-proxy starts up
     // 2. The config file is written but the signal hook is not yet received
-    // 3. local_proxy completes startup but has no config loaded, despite there being a registerd config.
+    // 3. local-proxy completes startup but has no config loaded, despite there being a registerd config.
     refresh_config_notify.notify_one();
     tokio::spawn(refresh_config_loop(args.config_path, refresh_config_notify));
 
@@ -305,7 +311,7 @@ async fn refresh_config_inner(path: &Utf8Path) -> anyhow::Result<()> {
 
     let mut jwks_set = vec![];
 
-    for jwks in data.jwks.into_iter().flatten() {
+    for jwks in data.jwks {
         let mut jwks_url = url::Url::from_str(&jwks.jwks_url).context("parsing JWKS url")?;
 
         ensure!(

proxy/src/compute.rs

@@ -20,7 +20,7 @@ use tokio_postgres::tls::MakeTlsConnect;
 use tokio_postgres_rustls::MakeRustlsConnect;
 use tracing::{error, info, warn};
 
-pub const COULD_NOT_CONNECT: &str = "Couldn't connect to compute node";
+const COULD_NOT_CONNECT: &str = "Couldn't connect to compute node";
 
 #[derive(Debug, Error)]
 pub(crate) enum ConnectionError {

proxy/src/lib.rs

@@ -90,8 +90,6 @@ use tokio::task::JoinError;
 use tokio_util::sync::CancellationToken;
 use tracing::warn;
 
-extern crate hyper0 as hyper;
-
 pub mod auth;
 pub mod cache;
 pub mod cancellation;

proxy/src/logging.rs

@@ -1,13 +1,6 @@
-use tracing::Subscriber;
 use tracing_subscriber::{
     filter::{EnvFilter, LevelFilter},
-    fmt::{
-        format::{Format, Full},
-        time::SystemTime,
-        FormatEvent, FormatFields,
-    },
     prelude::*,
-    registry::LookupSpan,
 };
 
 /// Initialize logging and OpenTelemetry tracing and exporter.
@@ -40,45 +33,6 @@ pub async fn init() -> anyhow::Result<LoggingGuard> {
     Ok(LoggingGuard)
 }
 
-/// Initialize logging for local_proxy with log prefix and no opentelemetry.
-///
-/// Logging can be configured using `RUST_LOG` environment variable.
-pub fn init_local_proxy() -> anyhow::Result<LoggingGuard> {
-    let env_filter = EnvFilter::builder()
-        .with_default_directive(LevelFilter::INFO.into())
-        .from_env_lossy();
-
-    let fmt_layer = tracing_subscriber::fmt::layer()
-        .with_ansi(false)
-        .with_writer(std::io::stderr)
-        .event_format(LocalProxyFormatter(Format::default().with_target(false)));
-
-    tracing_subscriber::registry()
-        .with(env_filter)
-        .with(fmt_layer)
-        .try_init()?;
-
-    Ok(LoggingGuard)
-}
-
-pub struct LocalProxyFormatter(Format<Full, SystemTime>);
-
-impl<S, N> FormatEvent<S, N> for LocalProxyFormatter
-where
-    S: Subscriber + for<'a> LookupSpan<'a>,
-    N: for<'a> FormatFields<'a> + 'static,
-{
-    fn format_event(
-        &self,
-        ctx: &tracing_subscriber::fmt::FmtContext<'_, S, N>,
-        mut writer: tracing_subscriber::fmt::format::Writer<'_>,
-        event: &tracing::Event<'_>,
-    ) -> std::fmt::Result {
-        writer.write_str("[local_proxy] ")?;
-        self.0.format_event(ctx, writer, event)
-    }
-}
-
 pub struct LoggingGuard;
 
 impl Drop for LoggingGuard {

proxy/src/proxy/connect_compute.rs

@@ -1,6 +1,5 @@
 use crate::{
     auth::backend::ComputeCredentialKeys,
-    compute::COULD_NOT_CONNECT,
     compute::{self, PostgresConnection},
     config::RetryConfig,
     console::{self, errors::WakeComputeError, locks::ApiLocks, CachedNodeInfo, NodeInfo},
@@ -16,7 +15,7 @@ use crate::{
 use async_trait::async_trait;
 use pq_proto::StartupMessageParams;
 use tokio::time;
-use tracing::{debug, info, warn};
+use tracing::{error, info, warn};
 
 use super::retry::ShouldRetryWakeCompute;
 
@@ -117,6 +116,7 @@ where
 
     node_info.set_keys(user_info.get_keys());
     node_info.allow_self_signed_compute = allow_self_signed_compute;
+    // let mut node_info = credentials.get_node_info(ctx, user_info).await?;
     mechanism.update_connect_config(&mut node_info.config);
     let retry_type = RetryType::ConnectToCompute;
 
@@ -139,10 +139,10 @@ where
         Err(e) => e,
     };
 
-    debug!(error = ?err, COULD_NOT_CONNECT);
+    error!(error = ?err, "could not connect to compute node");
 
     let node_info = if !node_info.cached() || !err.should_retry_wake_compute() {
-        // If we just recieved this from cplane and didn't get it from cache, we shouldn't retry.
+        // If we just recieved this from cplane and dodn't get it from cache, we shouldn't retry.
         // Do not need to retrieve a new node_info, just return the old one.
         if should_retry(&err, num_retries, connect_to_compute_retry_config) {
             Metrics::get().proxy.retries_metric.observe(
@@ -191,7 +191,7 @@ where
             }
             Err(e) => {
                 if !should_retry(&e, num_retries, connect_to_compute_retry_config) {
-                    // Don't log an error here, caller will print the error
+                    error!(error = ?e, num_retries, retriable = false, "couldn't connect to compute node");
                     Metrics::get().proxy.retries_metric.observe(
                         RetriesMetricGroup {
                             outcome: ConnectOutcome::Failed,
@@ -202,7 +202,7 @@ where
                     return Err(e.into());
                 }
 
-                warn!(error = ?e, num_retries, retriable = true, COULD_NOT_CONNECT);
+                warn!(error = ?e, num_retries, retriable = true, "couldn't connect to compute node");
             }
         };
 

safekeeper/Cargo.toml

@@ -23,7 +23,8 @@ crc32c.workspace = true
 fail.workspace = true
 hex.workspace = true
 humantime.workspace = true
-hyper0.workspace = true
+http.workspace = true
+hyper.workspace = true
 futures.workspace = true
 once_cell.workspace = true
 parking_lot.workspace = true

safekeeper/src/debug_dump.rs

@@ -253,13 +253,6 @@ pub async fn build(args: Args) -> Result<Response> {
         });
     }
 
-    // Tokio forbids to drop runtime in async context, so this is a stupid way
-    // to drop it in non async context.
-    tokio::task::spawn_blocking(move || {
-        let _r = runtime;
-    })
-    .await?;
-
     Ok(Response {
         start_time,
         finish_time: Utc::now(),

safekeeper/src/lib.rs

@@ -1,7 +1,4 @@
 #![deny(clippy::undocumented_unsafe_blocks)]
-
-extern crate hyper0 as hyper;
-
 use camino::Utf8PathBuf;
 use once_cell::sync::Lazy;
 use remote_storage::RemoteStorageConfig;

safekeeper/src/remove_wal.rs

@@ -2,29 +2,21 @@ use utils::lsn::Lsn;
 
 use crate::timeline_manager::StateSnapshot;
 
-/// Get oldest LSN we still need to keep.
-///
-/// We hold WAL till it is consumed by
-/// 1) pageserver (remote_consistent_lsn)
-/// 2) s3 offloading.
-/// 3) Additionally we must store WAL since last local commit_lsn because
-///    that's where we start looking for last WAL record on start.
-///
-/// If some peer safekeeper misses data it will fetch it from the remote
-/// storage. While it is safe to use inmem values for determining horizon, we
-/// use persistent to make possible normal states less surprising. All segments
-/// covering LSNs before horizon_lsn can be removed.
+/// Get oldest LSN we still need to keep. We hold WAL till it is consumed
+/// by all of 1) pageserver (remote_consistent_lsn) 2) peers 3) s3
+/// offloading.
+/// While it is safe to use inmem values for determining horizon,
+/// we use persistent to make possible normal states less surprising.
+/// All segments covering LSNs before horizon_lsn can be removed.
 pub(crate) fn calc_horizon_lsn(state: &StateSnapshot, extra_horizon_lsn: Option<Lsn>) -> Lsn {
     use std::cmp::min;
 
-    let mut horizon_lsn = state.cfile_remote_consistent_lsn;
+    let mut horizon_lsn = min(
+        state.cfile_remote_consistent_lsn,
+        state.cfile_peer_horizon_lsn,
+    );
     // we don't want to remove WAL that is not yet offloaded to s3
     horizon_lsn = min(horizon_lsn, state.cfile_backup_lsn);
-    // Min by local commit_lsn to be able to begin reading WAL from somewhere on
-    // sk start. Technically we don't allow local commit_lsn to be higher than
-    // flush_lsn, but let's be double safe by including it as well.
-    horizon_lsn = min(horizon_lsn, state.cfile_commit_lsn);
-    horizon_lsn = min(horizon_lsn, state.flush_lsn);
     if let Some(extra_horizon_lsn) = extra_horizon_lsn {
         horizon_lsn = min(horizon_lsn, extra_horizon_lsn);
     }

safekeeper/src/timeline_manager.rs

@@ -47,7 +47,7 @@ pub(crate) struct StateSnapshot {
     pub(crate) remote_consistent_lsn: Lsn,
 
     // persistent control file values
-    pub(crate) cfile_commit_lsn: Lsn,
+    pub(crate) cfile_peer_horizon_lsn: Lsn,
     pub(crate) cfile_remote_consistent_lsn: Lsn,
     pub(crate) cfile_backup_lsn: Lsn,
 
@@ -70,7 +70,7 @@ impl StateSnapshot {
             commit_lsn: state.inmem.commit_lsn,
             backup_lsn: state.inmem.backup_lsn,
             remote_consistent_lsn: state.inmem.remote_consistent_lsn,
-            cfile_commit_lsn: state.commit_lsn,
+            cfile_peer_horizon_lsn: state.peer_horizon_lsn,
             cfile_remote_consistent_lsn: state.remote_consistent_lsn,
             cfile_backup_lsn: state.backup_lsn,
             flush_lsn: read_guard.sk.flush_lsn(),

safekeeper/tests/walproposer_sim/safekeeper.rs

@@ -13,7 +13,7 @@ use desim::{
     node_os::NodeOs,
     proto::{AnyMessage, NetEvent, NodeEvent},
 };
-use hyper0::Uri;
+use http::Uri;
 use safekeeper::{
     safekeeper::{ProposerAcceptorMessage, SafeKeeper, ServerInfo, UNKNOWN_SERVER_VERSION},
     state::{TimelinePersistentState, TimelineState},

storage_broker/Cargo.toml

@@ -10,13 +10,16 @@ bench = []
 [dependencies]
 anyhow.workspace = true
 async-stream.workspace = true
+bytes.workspace = true
 clap = { workspace = true, features = ["derive"] }
 const_format.workspace = true
 futures.workspace = true
 futures-core.workspace = true
 futures-util.workspace = true
 humantime.workspace = true
-hyper0 = { workspace = true, features = ["full"] }
+hyper_1 = { workspace = true, features = ["full"] }
+http-body-util.workspace = true
+hyper-util = "0.1"
 once_cell.workspace = true
 parking_lot.workspace = true
 prost.workspace = true

storage_broker/src/bin/storage_broker.rs

@@ -10,16 +10,16 @@
 //!
 //! Only safekeeper message is supported, but it is not hard to add something
 //! else with generics.
-
-extern crate hyper0 as hyper;
-
 use clap::{command, Parser};
 use futures_core::Stream;
 use futures_util::StreamExt;
+use http_body_util::Full;
 use hyper::header::CONTENT_TYPE;
-use hyper::server::conn::AddrStream;
-use hyper::service::{make_service_fn, service_fn};
-use hyper::{Body, Method, StatusCode};
+use hyper::service::service_fn;
+use hyper::{Method, StatusCode};
+use hyper_1 as hyper;
+use hyper_1::body::Incoming;
+use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
 use parking_lot::RwLock;
 use std::collections::HashMap;
 use std::convert::Infallible;
@@ -27,9 +27,11 @@ use std::net::SocketAddr;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::time::Duration;
+use tokio::net::TcpListener;
 use tokio::sync::broadcast;
 use tokio::sync::broadcast::error::RecvError;
 use tokio::time;
+use tonic::body::{self, empty_body, BoxBody};
 use tonic::codegen::Service;
 use tonic::transport::server::Connected;
 use tonic::Code;
@@ -48,9 +50,7 @@ use storage_broker::proto::{
     FilterTenantTimelineId, MessageType, SafekeeperDiscoveryRequest, SafekeeperDiscoveryResponse,
     SafekeeperTimelineInfo, SubscribeByFilterRequest, SubscribeSafekeeperInfoRequest, TypedMessage,
 };
-use storage_broker::{
-    parse_proto_ttid, EitherBody, DEFAULT_KEEPALIVE_INTERVAL, DEFAULT_LISTEN_ADDR,
-};
+use storage_broker::{parse_proto_ttid, DEFAULT_KEEPALIVE_INTERVAL, DEFAULT_LISTEN_ADDR};
 use utils::id::TenantTimelineId;
 use utils::logging::{self, LogFormat};
 use utils::sentry_init::init_sentry;
@@ -602,8 +602,8 @@ impl BrokerService for Broker {
 
 // We serve only metrics and healthcheck through http1.
 async fn http1_handler(
-    req: hyper::Request<hyper::body::Body>,
-) -> Result<hyper::Response<Body>, Infallible> {
+    req: hyper::Request<Incoming>,
+) -> Result<hyper::Response<BoxBody>, Infallible> {
     let resp = match (req.method(), req.uri().path()) {
         (&Method::GET, "/metrics") => {
             let mut buffer = vec![];
@@ -614,16 +614,16 @@ async fn http1_handler(
             hyper::Response::builder()
                 .status(StatusCode::OK)
                 .header(CONTENT_TYPE, encoder.format_type())
-                .body(Body::from(buffer))
+                .body(body::boxed(Full::new(bytes::Bytes::from(buffer))))
                 .unwrap()
         }
         (&Method::GET, "/status") => hyper::Response::builder()
             .status(StatusCode::OK)
-            .body(Body::empty())
+            .body(empty_body())
             .unwrap(),
         _ => hyper::Response::builder()
             .status(StatusCode::NOT_FOUND)
-            .body(Body::empty())
+            .body(empty_body())
             .unwrap(),
     };
     Ok(resp)
@@ -665,52 +665,74 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     };
     let storage_broker_server = BrokerServiceServer::new(storage_broker_impl);
 
-    info!("listening on {}", &args.listen_addr);
-
     // grpc is served along with http1 for metrics on a single port, hence we
     // don't use tonic's Server.
-    hyper::Server::bind(&args.listen_addr)
-        .http2_keep_alive_interval(Some(args.http2_keepalive_interval))
-        .serve(make_service_fn(move |conn: &AddrStream| {
-            let storage_broker_server_cloned = storage_broker_server.clone();
-            let connect_info = conn.connect_info();
-            async move {
-                Ok::<_, Infallible>(service_fn(move |mut req| {
-                    // That's what tonic's MakeSvc.call does to pass conninfo to
-                    // the request handler (and where its request.remote_addr()
-                    // expects it to find).
-                    req.extensions_mut().insert(connect_info.clone());
-
-                    // Technically this second clone is not needed, but consume
-                    // by async block is apparently unavoidable. BTW, error
-                    // message is enigmatic, see
-                    // https://github.com/rust-lang/rust/issues/68119
-                    //
-                    // We could get away without async block at all, but then we
-                    // need to resort to futures::Either to merge the result,
-                    // which doesn't caress an eye as well.
-                    let mut storage_broker_server_svc = storage_broker_server_cloned.clone();
-                    async move {
-                        if req.headers().get("content-type").map(|x| x.as_bytes())
-                            == Some(b"application/grpc")
-                        {
-                            let res_resp = storage_broker_server_svc.call(req).await;
-                            // Grpc and http1 handlers have slightly different
-                            // Response types: it is UnsyncBoxBody for the
-                            // former one (not sure why) and plain hyper::Body
-                            // for the latter. Both implement HttpBody though,
-                            // and EitherBody is used to merge them.
-                            res_resp.map(|resp| resp.map(EitherBody::Left))
-                        } else {
-                            let res_resp = http1_handler(req).await;
-                            res_resp.map(|resp| resp.map(EitherBody::Right))
-                        }
-                    }
-                }))
+    let tcp_listener = TcpListener::bind(&args.listen_addr).await?;
+    info!("listening on {}", &args.listen_addr);
+    loop {
+        let (stream, addr) = match tcp_listener.accept().await {
+            Ok(v) => v,
+            Err(e) => {
+                info!("couldn't accept connection: {e}");
+                continue;
             }
-        }))
-        .await?;
-    Ok(())
+        };
+
+        let mut builder = hyper_util::server::conn::auto::Builder::new(TokioExecutor::new());
+        builder.http1().timer(TokioTimer::new());
+        builder
+            .http2()
+            .timer(TokioTimer::new())
+            .keep_alive_interval(Some(args.http2_keepalive_interval));
+
+        let storage_broker_server_cloned = storage_broker_server.clone();
+        let connect_info = stream.connect_info();
+        let service_fn_ = async move {
+            service_fn(move |mut req| {
+                // That's what tonic's MakeSvc.call does to pass conninfo to
+                // the request handler (and where its request.remote_addr()
+                // expects it to find).
+                req.extensions_mut().insert(connect_info.clone());
+
+                // Technically this second clone is not needed, but consume
+                // by async block is apparently unavoidable. BTW, error
+                // message is enigmatic, see
+                // https://github.com/rust-lang/rust/issues/68119
+                //
+                // We could get away without async block at all, but then we
+                // need to resort to futures::Either to merge the result,
+                // which doesn't caress an eye as well.
+                let mut storage_broker_server_svc = storage_broker_server_cloned.clone();
+                async move {
+                    if req.headers().get("content-type").map(|x| x.as_bytes())
+                        == Some(b"application/grpc")
+                    {
+                        let res_resp = storage_broker_server_svc.call(req).await;
+                        // Grpc and http1 handlers have slightly different
+                        // Response types: it is UnsyncBoxBody for the
+                        // former one (not sure why) and plain hyper::Body
+                        // for the latter. Both implement HttpBody though,
+                        // and `Either` is used to merge them.
+                        res_resp.map(|resp| resp.map(http_body_util::Either::Left))
+                    } else {
+                        let res_resp = http1_handler(req).await;
+                        res_resp.map(|resp| resp.map(http_body_util::Either::Right))
+                    }
+                }
+            })
+        }
+        .await;
+
+        tokio::task::spawn(async move {
+            let res = builder
+                .serve_connection(TokioIo::new(stream), service_fn_)
+                .await;
+
+            if let Err(e) = res {
+                info!("error serving connection from {addr}: {e}");
+            }
+        });
+    }
 }
 
 #[cfg(test)]

storage_broker/src/lib.rs

@@ -1,8 +1,4 @@
-extern crate hyper0 as hyper;
-
-use hyper::body::HttpBody;
-use std::pin::Pin;
-use std::task::{Context, Poll};
+use hyper_1 as hyper;
 use std::time::Duration;
 use tonic::codegen::StdError;
 use tonic::transport::{ClientTlsConfig, Endpoint};
@@ -96,56 +92,3 @@ pub fn parse_proto_ttid(proto_ttid: &ProtoTenantTimelineId) -> Result<TenantTime
         timeline_id,
     })
 }
-
-// These several usages don't justify anyhow dependency, though it would work as
-// well.
-type AnyError = Box<dyn std::error::Error + Send + Sync + 'static>;
-
-// Provides impl HttpBody for two different types implementing it. Inspired by
-// https://github.com/hyperium/tonic/blob/master/examples/src/hyper_warp/server.rs
-pub enum EitherBody<A, B> {
-    Left(A),
-    Right(B),
-}
-
-impl<A, B> HttpBody for EitherBody<A, B>
-where
-    A: HttpBody + Send + Unpin,
-    B: HttpBody<Data = A::Data> + Send + Unpin,
-    A::Error: Into<AnyError>,
-    B::Error: Into<AnyError>,
-{
-    type Data = A::Data;
-    type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
-
-    fn is_end_stream(&self) -> bool {
-        match self {
-            EitherBody::Left(b) => b.is_end_stream(),
-            EitherBody::Right(b) => b.is_end_stream(),
-        }
-    }
-
-    fn poll_data(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-    ) -> Poll<Option<Result<Self::Data, Self::Error>>> {
-        match self.get_mut() {
-            EitherBody::Left(b) => Pin::new(b).poll_data(cx).map(map_option_err),
-            EitherBody::Right(b) => Pin::new(b).poll_data(cx).map(map_option_err),
-        }
-    }
-
-    fn poll_trailers(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-    ) -> Poll<Result<Option<hyper::HeaderMap>, Self::Error>> {
-        match self.get_mut() {
-            EitherBody::Left(b) => Pin::new(b).poll_trailers(cx).map_err(Into::into),
-            EitherBody::Right(b) => Pin::new(b).poll_trailers(cx).map_err(Into::into),
-        }
-    }
-}
-
-fn map_option_err<T, U: Into<AnyError>>(err: Option<Result<T, U>>) -> Option<Result<T, AnyError>> {
-    err.map(|e| e.map_err(Into::into))
-}

storage_controller/Cargo.toml

@@ -21,7 +21,7 @@ clap.workspace = true
 fail.workspace = true
 futures.workspace = true
 hex.workspace = true
-hyper0.workspace = true
+hyper.workspace = true
 humantime.workspace = true
 itertools.workspace = true
 lasso.workspace = true
@@ -56,3 +56,6 @@ utils = { path = "../libs/utils/" }
 metrics = { path = "../libs/metrics/" }
 control_plane = { path = "../control_plane" }
 workspace_hack = { version = "0.1", path = "../workspace_hack" }
+
+[dev-dependencies]
+test-log = "*"

storage_controller/src/lib.rs

@@ -1,8 +1,6 @@
 use serde::Serialize;
 use utils::seqwait::MonotonicCounter;
 
-extern crate hyper0 as hyper;
-
 mod auth;
 mod background_node_operations;
 mod compute_hook;

storage_controller/src/main.rs

@@ -1,6 +1,6 @@
 use anyhow::{anyhow, Context};
 use clap::Parser;
-use hyper0::Uri;
+use hyper::Uri;
 use metrics::launch_timestamp::LaunchTimestamp;
 use metrics::BuildInfo;
 use std::path::PathBuf;
@@ -324,7 +324,7 @@ async fn async_main() -> anyhow::Result<()> {
 
     // Start HTTP server
     let server_shutdown = CancellationToken::new();
-    let server = hyper0::Server::from_tcp(http_listener)?
+    let server = hyper::Server::from_tcp(http_listener)?
         .serve(router_service)
         .with_graceful_shutdown({
             let server_shutdown = server_shutdown.clone();

storage_controller/src/scheduler.rs

@@ -206,6 +206,10 @@ pub(crate) struct NodeSecondarySchedulingScore {
     /// The number of shards belonging to the tenant currently being
     /// scheduled that are attached to this node.
     affinity_score: AffinityScore,
+    /// Size of [`ScheduleContext::attached_nodes`] for the current node.
+    /// This normally tracks the number of attached shards belonging to the
+    /// tenant being scheduled that are already on this node.
+    secondary_shards_in_context: usize,
     /// Utilisation score that combines shard count and disk utilisation
     utilization_score: u64,
     /// Total number of shards attached to this node. When nodes have identical utilisation, this
@@ -231,6 +235,7 @@ impl NodeSchedulingScore for NodeSecondarySchedulingScore {
 
         Some(Self {
             az_match: SecondaryAzMatch(AzMatch::new(&node.az, preferred_az.as_ref())),
+            secondary_shards_in_context: context.secondary_nodes.get(node_id).copied().unwrap_or(0),
             affinity_score: context
                 .nodes
                 .get(node_id)
@@ -327,6 +332,9 @@ pub(crate) struct ScheduleContext {
     /// Specifically how many _attached_ locations are on each node
     pub(crate) attached_nodes: HashMap<NodeId, usize>,
 
+    /// Specifically how many _secondary_ locations are on each node
+    pub(crate) secondary_nodes: HashMap<NodeId, usize>,
+
     pub(crate) mode: ScheduleMode,
 }
 
@@ -345,6 +353,11 @@ impl ScheduleContext {
         *entry += 1;
     }
 
+    pub(crate) fn push_secondary(&mut self, node_id: NodeId) {
+        let entry = self.secondary_nodes.entry(node_id).or_default();
+        *entry += 1;
+    }
+
     pub(crate) fn get_node_affinity(&self, node_id: NodeId) -> AffinityScore {
         self.nodes
             .get(&node_id)
@@ -786,7 +799,14 @@ pub(crate) mod test_utils {
 
 #[cfg(test)]
 mod tests {
-    use pageserver_api::{controller_api::NodeAvailability, models::utilization::test_utilization};
+    use pageserver_api::{
+        controller_api::NodeAvailability, models::utilization::test_utilization,
+        shard::ShardIdentity, shard::TenantShardId,
+    };
+    use utils::{
+        id::TenantId,
+        shard::{ShardCount, ShardNumber},
+    };
 
     use super::*;
 
@@ -1074,4 +1094,171 @@ mod tests {
             intent.clear(&mut scheduler);
         }
     }
+
+    #[test]
+    fn repro_foo() {
+        let az_tag = AvailabilityZone("az-a".to_string());
+
+        let nodes = test_utils::make_test_nodes(
+            5,
+            &[
+                az_tag.clone(),
+                az_tag.clone(),
+                az_tag.clone(),
+                az_tag.clone(),
+                az_tag.clone(),
+            ],
+        );
+        let mut scheduler = Scheduler::new(nodes.values());
+
+        // Need to keep these alive because they contribute to shard counts via RAII
+        let mut scheduled_shards = Vec::new();
+
+        let mut context = ScheduleContext::default();
+
+        fn schedule_shard(
+            tenant_shard_id: TenantShardId,
+            expect_attached: NodeId,
+            expect_secondary: NodeId,
+            scheduled_shards: &mut Vec<TenantShard>,
+            scheduler: &mut Scheduler,
+            context: &mut ScheduleContext,
+        ) {
+            let shard_identity = ShardIdentity::new(
+                tenant_shard_id.shard_number,
+                tenant_shard_id.shard_count,
+                pageserver_api::shard::ShardStripeSize(1),
+            )
+            .unwrap();
+            let mut shard = TenantShard::new(
+                tenant_shard_id,
+                shard_identity,
+                pageserver_api::controller_api::PlacementPolicy::Attached(1),
+            );
+
+            shard.schedule(scheduler, context).unwrap();
+
+            assert_eq!(shard.intent.get_attached().unwrap(), expect_attached);
+            assert_eq!(
+                shard.intent.get_secondary().first().unwrap(),
+                &expect_secondary
+            );
+
+            scheduled_shards.push(shard);
+        }
+
+        let tenant_id = TenantId::generate();
+
+        schedule_shard(
+            TenantShardId {
+                tenant_id,
+                shard_number: ShardNumber(0),
+                shard_count: ShardCount(8),
+            },
+            NodeId(1),
+            NodeId(2),
+            &mut scheduled_shards,
+            &mut scheduler,
+            &mut context,
+        );
+
+        schedule_shard(
+            TenantShardId {
+                tenant_id,
+                shard_number: ShardNumber(1),
+                shard_count: ShardCount(8),
+            },
+            NodeId(3),
+            NodeId(4),
+            &mut scheduled_shards,
+            &mut scheduler,
+            &mut context,
+        );
+
+        schedule_shard(
+            TenantShardId {
+                tenant_id,
+                shard_number: ShardNumber(2),
+                shard_count: ShardCount(8),
+            },
+            NodeId(5),
+            NodeId(1),
+            &mut scheduled_shards,
+            &mut scheduler,
+            &mut context,
+        );
+
+        schedule_shard(
+            TenantShardId {
+                tenant_id,
+                shard_number: ShardNumber(3),
+                shard_count: ShardCount(8),
+            },
+            NodeId(2),
+            NodeId(3),
+            &mut scheduled_shards,
+            &mut scheduler,
+            &mut context,
+        );
+
+        schedule_shard(
+            TenantShardId {
+                tenant_id,
+                shard_number: ShardNumber(4),
+                shard_count: ShardCount(8),
+            },
+            NodeId(4),
+            NodeId(5),
+            &mut scheduled_shards,
+            &mut scheduler,
+            &mut context,
+        );
+
+        schedule_shard(
+            TenantShardId {
+                tenant_id,
+                shard_number: ShardNumber(5),
+                shard_count: ShardCount(8),
+            },
+            NodeId(1),
+            NodeId(2),
+            &mut scheduled_shards,
+            &mut scheduler,
+            &mut context,
+        );
+
+        schedule_shard(
+            TenantShardId {
+                tenant_id,
+                shard_number: ShardNumber(6),
+                shard_count: ShardCount(8),
+            },
+            NodeId(3),
+            NodeId(4),
+            &mut scheduled_shards,
+            &mut scheduler,
+            &mut context,
+        );
+
+        schedule_shard(
+            TenantShardId {
+                tenant_id,
+                shard_number: ShardNumber(7),
+                shard_count: ShardCount(8),
+            },
+            NodeId(5),
+            NodeId(1),
+            &mut scheduled_shards,
+            &mut scheduler,
+            &mut context,
+        );
+
+        for shard in &scheduled_shards {
+            assert_eq!(shard.optimize_attachment(&nodes, &context), None);
+        }
+
+        for mut shard in scheduled_shards {
+            shard.intent.clear(&mut scheduler);
+        }
+    }
 }

storage_controller/src/service.rs

@@ -526,21 +526,6 @@ pub(crate) enum ReconcileResultRequest {
     Stop,
 }
 
-#[derive(Clone)]
-struct MutationLocation {
-    node: Node,
-    generation: Generation,
-}
-
-#[derive(Clone)]
-struct ShardMutationLocations {
-    latest: MutationLocation,
-    other: Vec<MutationLocation>,
-}
-
-#[derive(Default, Clone)]
-struct TenantMutationLocations(BTreeMap<TenantShardId, ShardMutationLocations>);
-
 impl Service {
     pub fn get_config(&self) -> &Config {
         &self.config
@@ -3002,83 +2987,38 @@ impl Service {
         failpoint_support::sleep_millis_async!("tenant-create-timeline-shared-lock");
 
         self.tenant_remote_mutation(tenant_id, move |mut targets| async move {
-            if targets.0.is_empty() {
+            if targets.is_empty() {
                 return Err(ApiError::NotFound(
                     anyhow::anyhow!("Tenant not found").into(),
                 ));
             };
-
-            let (shard_zero_tid, shard_zero_locations) =
-                targets.0.pop_first().expect("Must have at least one shard");
-            assert!(shard_zero_tid.is_shard_zero());
+            let shard_zero = targets.remove(0);
 
             async fn create_one(
                 tenant_shard_id: TenantShardId,
-                locations: ShardMutationLocations,
+                node: Node,
                 jwt: Option<String>,
                 create_req: TimelineCreateRequest,
             ) -> Result<TimelineInfo, ApiError> {
-                let latest = locations.latest.node;
-
                 tracing::info!(
-                    "Creating timeline on shard {}/{}, attached to node {latest} in generation {:?}",
+                    "Creating timeline on shard {}/{}, attached to node {node}",
                     tenant_shard_id,
                     create_req.new_timeline_id,
-                    locations.latest.generation
                 );
+                let client = PageserverClient::new(node.get_id(), node.base_url(), jwt.as_deref());
 
-                let client =
-                    PageserverClient::new(latest.get_id(), latest.base_url(), jwt.as_deref());
-
-                let timeline_info = client
+                client
                     .timeline_create(tenant_shard_id, &create_req)
                     .await
-                    .map_err(|e| passthrough_api_error(&latest, e))?;
-
-                // We propagate timeline creations to all attached locations such that a compute
-                // for the new timeline is able to start regardless of the current state of the
-                // tenant shard reconciliation.
-                for location in locations.other {
-                    tracing::info!(
-                        "Creating timeline on shard {}/{}, stale attached to node {} in generation {:?}",
-                        tenant_shard_id,
-                        create_req.new_timeline_id,
-                        location.node,
-                        location.generation
-                    );
-
-                    let client = PageserverClient::new(
-                        location.node.get_id(),
-                        location.node.base_url(),
-                        jwt.as_deref(),
-                    );
-
-                    let res = client
-                        .timeline_create(tenant_shard_id, &create_req)
-                        .await;
-
-                    if let Err(e) = res {
-                        match e {
-                            mgmt_api::Error::ApiError(StatusCode::NOT_FOUND, _) => {
-                                // Tenant might have been detached from the stale location,
-                                // so ignore 404s.
-                            },
-                            _ => {
-                                return Err(passthrough_api_error(&location.node, e));
-                            }
-                        }
-                    }
-                }
-
-                Ok(timeline_info)
+                    .map_err(|e| passthrough_api_error(&node, e))
             }
 
             // Because the caller might not provide an explicit LSN, we must do the creation first on a single shard, and then
             // use whatever LSN that shard picked when creating on subsequent shards.  We arbitrarily use shard zero as the shard
             // that will get the first creation request, and propagate the LSN to all the >0 shards.
             let timeline_info = create_one(
-                shard_zero_tid,
-                shard_zero_locations,
+                shard_zero.0,
+                shard_zero.1,
                 self.config.jwt_token.clone(),
                 create_req.clone(),
             )
@@ -3091,24 +3031,14 @@ impl Service {
             }
 
             // Create timeline on remaining shards with number >0
-            if !targets.0.is_empty() {
+            if !targets.is_empty() {
                 // If we had multiple shards, issue requests for the remainder now.
                 let jwt = &self.config.jwt_token;
                 self.tenant_for_shards(
-                    targets
-                        .0
-                        .iter()
-                        .map(|t| (*t.0, t.1.latest.node.clone()))
-                        .collect(),
-                    |tenant_shard_id: TenantShardId, _node: Node| {
+                    targets.iter().map(|t| (t.0, t.1.clone())).collect(),
+                    |tenant_shard_id: TenantShardId, node: Node| {
                         let create_req = create_req.clone();
-                        let mutation_locations = targets.0.remove(&tenant_shard_id).unwrap();
-                        Box::pin(create_one(
-                            tenant_shard_id,
-                            mutation_locations,
-                            jwt.clone(),
-                            create_req,
-                        ))
+                        Box::pin(create_one(tenant_shard_id, node, jwt.clone(), create_req))
                     },
                 )
                 .await?;
@@ -3138,7 +3068,7 @@ impl Service {
         .await;
 
         self.tenant_remote_mutation(tenant_id, move |targets| async move {
-            if targets.0.is_empty() {
+            if targets.is_empty() {
                 return Err(ApiError::NotFound(
                     anyhow::anyhow!("Tenant not found").into(),
                 ));
@@ -3169,9 +3099,8 @@ impl Service {
 
             // no shard needs to go first/last; the operation should be idempotent
             // TODO: it would be great to ensure that all shards return the same error
-            let locations = targets.0.iter().map(|t| (*t.0, t.1.latest.node.clone())).collect();
             let results = self
-                .tenant_for_shards(locations, |tenant_shard_id, node| {
+                .tenant_for_shards(targets, |tenant_shard_id, node| {
                     futures::FutureExt::boxed(config_one(
                         tenant_shard_id,
                         timeline_id,
@@ -3202,7 +3131,7 @@ impl Service {
         .await;
 
         self.tenant_remote_mutation(tenant_id, move |targets| async move {
-            if targets.0.is_empty() {
+            if targets.is_empty() {
                 return Err(ApiError::NotFound(
                     anyhow::anyhow!("Tenant not found").into(),
                 ));
@@ -3250,9 +3179,8 @@ impl Service {
             }
 
             // no shard needs to go first/last; the operation should be idempotent
-            let locations = targets.0.iter().map(|t| (*t.0, t.1.latest.node.clone())).collect();
             let mut results = self
-                .tenant_for_shards(locations, |tenant_shard_id, node| {
+                .tenant_for_shards(targets, |tenant_shard_id, node| {
                     futures::FutureExt::boxed(detach_one(
                         tenant_shard_id,
                         timeline_id,
@@ -3299,7 +3227,7 @@ impl Service {
         .await;
 
         self.tenant_remote_mutation(tenant_id, move |targets| async move {
-            if targets.0.is_empty() {
+            if targets.is_empty() {
                 return Err(ApiError::NotFound(
                     anyhow::anyhow!("Tenant not found").into(),
                 ));
@@ -3321,12 +3249,7 @@ impl Service {
             }
 
             // no shard needs to go first/last; the operation should be idempotent
-            let locations = targets
-                .0
-                .iter()
-                .map(|t| (*t.0, t.1.latest.node.clone()))
-                .collect();
-            self.tenant_for_shards(locations, |tenant_shard_id, node| {
+            self.tenant_for_shards(targets, |tenant_shard_id, node| {
                 futures::FutureExt::boxed(do_one(
                     tenant_shard_id,
                     timeline_id,
@@ -3421,11 +3344,11 @@ impl Service {
         op: O,
     ) -> Result<R, ApiError>
     where
-        O: FnOnce(TenantMutationLocations) -> F,
+        O: FnOnce(Vec<(TenantShardId, Node)>) -> F,
         F: std::future::Future<Output = R>,
     {
-        let mutation_locations = {
-            let mut locations = TenantMutationLocations::default();
+        let target_gens = {
+            let mut targets = Vec::new();
 
             // Load the currently attached pageservers for the latest generation of each shard.  This can
             // run concurrently with reconciliations, and it is not guaranteed that the node we find here
@@ -3476,50 +3399,14 @@ impl Service {
                     .ok_or(ApiError::Conflict(format!(
                         "Raced with removal of node {node_id}"
                     )))?;
-                let generation = generation.expect("Checked above");
-
-                let tenant = locked.tenants.get(&tenant_shard_id);
-
-                // TODO(vlad): Abstract the logic that finds stale attached locations
-                // from observed state into a [`Service`] method.
-                let other_locations = match tenant {
-                    Some(tenant) => {
-                        let mut other = tenant.attached_locations();
-                        let latest_location_index =
-                            other.iter().position(|&l| l == (node.get_id(), generation));
-                        if let Some(idx) = latest_location_index {
-                            other.remove(idx);
-                        }
-
-                        other
-                    }
-                    None => Vec::default(),
-                };
-
-                let location = ShardMutationLocations {
-                    latest: MutationLocation {
-                        node: node.clone(),
-                        generation,
-                    },
-                    other: other_locations
-                        .into_iter()
-                        .filter_map(|(node_id, generation)| {
-                            let node = locked.nodes.get(&node_id)?;
-
-                            Some(MutationLocation {
-                                node: node.clone(),
-                                generation,
-                            })
-                        })
-                        .collect(),
-                };
-                locations.0.insert(tenant_shard_id, location);
+                targets.push((tenant_shard_id, node.clone(), generation));
             }
 
-            locations
+            targets
         };
 
-        let result = op(mutation_locations.clone()).await;
+        let targets = target_gens.iter().map(|t| (t.0, t.1.clone())).collect();
+        let result = op(targets).await;
 
         // Post-check: are all the generations of all the shards the same as they were initially?  This proves that
         // our remote operation executed on the latest generation and is therefore persistent.
@@ -3535,10 +3422,9 @@ impl Service {
                      }| (tenant_shard_id, generation),
                 )
                 .collect::<Vec<_>>()
-                != mutation_locations
-                    .0
+                != target_gens
                     .into_iter()
-                    .map(|i| (i.0, Some(i.1.latest.generation)))
+                    .map(|i| (i.0, i.2))
                     .collect::<Vec<_>>()
             {
                 // We raced with something that incremented the generation, and therefore cannot be
@@ -3568,14 +3454,12 @@ impl Service {
         .await;
 
         self.tenant_remote_mutation(tenant_id, move |mut targets| async move {
-            if targets.0.is_empty() {
+            if targets.is_empty() {
                 return Err(ApiError::NotFound(
                     anyhow::anyhow!("Tenant not found").into(),
                 ));
             }
-
-            let (shard_zero_tid, shard_zero_locations) = targets.0.pop_first().expect("Must have at least one shard");
-            assert!(shard_zero_tid.is_shard_zero());
+            let shard_zero = targets.remove(0);
 
             async fn delete_one(
                 tenant_shard_id: TenantShardId,
@@ -3598,9 +3482,8 @@ impl Service {
                     })
             }
 
-            let locations = targets.0.iter().map(|t| (*t.0, t.1.latest.node.clone())).collect();
             let statuses = self
-                .tenant_for_shards(locations, |tenant_shard_id: TenantShardId, node: Node| {
+                .tenant_for_shards(targets, |tenant_shard_id: TenantShardId, node: Node| {
                     Box::pin(delete_one(
                         tenant_shard_id,
                         timeline_id,
@@ -3618,9 +3501,9 @@ impl Service {
             // Delete shard zero last: this is not strictly necessary, but since a caller's GET on a timeline will be routed
             // to shard zero, it gives a more obvious behavior that a GET returns 404 once the deletion is done.
             let shard_zero_status = delete_one(
-                shard_zero_tid,
+                shard_zero.0,
                 timeline_id,
-                shard_zero_locations.latest.node,
+                shard_zero.1,
                 self.config.jwt_token.clone(),
             )
             .await?;
@@ -5975,10 +5858,7 @@ impl Service {
 
             // Accumulate the schedule context for all the shards in a tenant: we must have
             // the total view of all shards before we can try to optimize any of them.
-            schedule_context.avoid(&shard.intent.all_pageservers());
-            if let Some(attached) = shard.intent.get_attached() {
-                schedule_context.push_attached(*attached);
-            }
+            shard.populate_context(&mut schedule_context);
             tenant_shards.push(shard);
 
             // Once we have seen the last shard in the tenant, proceed to search across all shards

storage_controller/src/tenant_shard.rs

@@ -17,7 +17,6 @@ use crate::{
     service::ReconcileResultRequest,
 };
 use futures::future::{self, Either};
-use itertools::Itertools;
 use pageserver_api::controller_api::{
     AvailabilityZone, NodeSchedulingPolicy, PlacementPolicy, ShardSchedulingPolicy,
 };
@@ -567,10 +566,7 @@ impl TenantShard {
     ) -> Result<(), ScheduleError> {
         let r = self.do_schedule(scheduler, context);
 
-        context.avoid(&self.intent.all_pageservers());
-        if let Some(attached) = self.intent.get_attached() {
-            context.push_attached(*attached);
-        }
+        self.populate_context(context);
 
         r
     }
@@ -677,6 +673,19 @@ impl TenantShard {
         Ok(())
     }
 
+    /// When building the ScheduleContext of a tenant, call this on each shard to
+    /// add its contribution to the context.
+    pub(crate) fn populate_context(&self, context: &mut ScheduleContext) {
+        context.avoid(&self.intent.all_pageservers());
+
+        if let Some(attached) = self.intent.get_attached() {
+            context.push_attached(*attached);
+        }
+        for secondary in self.intent.get_secondary() {
+            context.push_secondary(*secondary);
+        }
+    }
+
     /// Reschedule this tenant shard to one of its secondary locations. Returns a scheduling error
     /// if the swap is not possible and leaves the intent state in its original state.
     ///
@@ -824,10 +833,13 @@ impl TenantShard {
                 continue;
             };
 
+            // TODO: make this AZ aware: secondary should be chosen "As if I am an attachment, but
+            // in a different AZ to my actual preferred AZ"
+
             // Let the scheduler suggest a node, where it would put us if we were scheduling afresh
             // This implicitly limits the choice to nodes that are available, and prefers nodes
             // with lower utilization.
-            let Ok(candidate_node) = scheduler.schedule_shard::<SecondaryShardTag>(
+            let Ok(candidate_node) = scheduler.schedule_shard::<AttachedShardTag>(
                 &self.intent.all_pageservers(),
                 &self.preferred_az_id,
                 schedule_context,
@@ -1411,32 +1423,6 @@ impl TenantShard {
     pub(crate) fn set_preferred_az(&mut self, preferred_az_id: AvailabilityZone) {
         self.preferred_az_id = Some(preferred_az_id);
     }
-
-    /// Returns all the nodes to which this tenant shard is attached according to the
-    /// observed state and the generations. Return vector is sorted from latest generation
-    /// to earliest.
-    pub(crate) fn attached_locations(&self) -> Vec<(NodeId, Generation)> {
-        self.observed
-            .locations
-            .iter()
-            .filter_map(|(node_id, observed)| {
-                use LocationConfigMode::{AttachedMulti, AttachedSingle, AttachedStale};
-
-                let conf = observed.conf.as_ref()?;
-
-                match (conf.generation, conf.mode) {
-                    (Some(gen), AttachedMulti | AttachedSingle | AttachedStale) => {
-                        Some((*node_id, gen))
-                    }
-                    _ => None,
-                }
-            })
-            .sorted_by(|(_lhs_node_id, lhs_gen), (_rhs_node_id, rhs_gen)| {
-                lhs_gen.cmp(rhs_gen).reverse()
-            })
-            .map(|(node_id, gen)| (node_id, Generation::new(gen)))
-            .collect()
-    }
 }
 
 #[cfg(test)]
@@ -1659,10 +1645,8 @@ pub(crate) mod tests {
         shard_b.intent.push_secondary(&mut scheduler, NodeId(3));
 
         let mut schedule_context = ScheduleContext::default();
-        schedule_context.avoid(&shard_a.intent.all_pageservers());
-        schedule_context.push_attached(shard_a.intent.get_attached().unwrap());
-        schedule_context.avoid(&shard_b.intent.all_pageservers());
-        schedule_context.push_attached(shard_b.intent.get_attached().unwrap());
+        shard_a.populate_context(&mut schedule_context);
+        shard_b.populate_context(&mut schedule_context);
 
         let optimization_a = shard_a.optimize_attachment(&nodes, &schedule_context);
 
@@ -1726,10 +1710,8 @@ pub(crate) mod tests {
         shard_b.intent.push_secondary(&mut scheduler, NodeId(3));
 
         let mut schedule_context = ScheduleContext::default();
-        schedule_context.avoid(&shard_a.intent.all_pageservers());
-        schedule_context.push_attached(shard_a.intent.get_attached().unwrap());
-        schedule_context.avoid(&shard_b.intent.all_pageservers());
-        schedule_context.push_attached(shard_b.intent.get_attached().unwrap());
+        shard_a.populate_context(&mut schedule_context);
+        shard_b.populate_context(&mut schedule_context);
 
         let optimization_a = shard_a.optimize_secondary(&mut scheduler, &schedule_context);
 
@@ -1771,10 +1753,7 @@ pub(crate) mod tests {
             let mut any_changed = false;
 
             for shard in shards.iter() {
-                schedule_context.avoid(&shard.intent.all_pageservers());
-                if let Some(attached) = shard.intent.get_attached() {
-                    schedule_context.push_attached(*attached);
-                }
+                shard.populate_context(&mut schedule_context);
             }
 
             for shard in shards.iter_mut() {

storage_scrubber/Cargo.toml

@@ -5,7 +5,6 @@ edition.workspace = true
 license.workspace = true
 
 [dependencies]
-aws-config.workspace = true
 aws-sdk-s3.workspace = true
 either.workspace = true
 anyhow.workspace = true
@@ -32,6 +31,7 @@ storage_controller_client.workspace = true
 tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
 chrono = { workspace = true, default-features = false, features = ["clock", "serde"] }
 reqwest = { workspace = true, default-features = false, features = ["rustls-tls", "json"] }
+aws-config = { workspace = true, default-features = false, features = ["rustls", "sso"] }
 
 pageserver = { path = "../pageserver" }
 pageserver_api = { path = "../libs/pageserver_api" }

storage_scrubber/src/lib.rs

@@ -28,9 +28,8 @@ use pageserver::tenant::remote_timeline_client::{remote_tenant_path, remote_time
 use pageserver::tenant::TENANTS_SEGMENT_NAME;
 use pageserver_api::shard::TenantShardId;
 use remote_storage::{
-    DownloadOpts, GenericRemoteStorage, Listing, ListingMode, RemotePath, RemoteStorageConfig,
-    RemoteStorageKind, S3Config, DEFAULT_MAX_KEYS_PER_LIST_RESPONSE,
-    DEFAULT_REMOTE_STORAGE_S3_CONCURRENCY_LIMIT,
+    GenericRemoteStorage, Listing, ListingMode, RemotePath, RemoteStorageConfig, RemoteStorageKind,
+    S3Config, DEFAULT_MAX_KEYS_PER_LIST_RESPONSE, DEFAULT_REMOTE_STORAGE_S3_CONCURRENCY_LIMIT,
 };
 use reqwest::Url;
 use serde::{Deserialize, Serialize};
@@ -489,10 +488,7 @@ async fn download_object_with_retries(
     let cancel = CancellationToken::new();
     for trial in 0..MAX_RETRIES {
         let mut buf = Vec::new();
-        let download = match remote_client
-            .download(key, &DownloadOpts::default(), &cancel)
-            .await
-        {
+        let download = match remote_client.download(key, &cancel).await {
             Ok(response) => response,
             Err(e) => {
                 error!("Failed to download object for key {key}: {e}");

test_runner/conftest.py

@@ -5,7 +5,6 @@ pytest_plugins = (
     "fixtures.compute_reconfigure",
     "fixtures.storage_controller_proxy",
     "fixtures.neon_fixtures",
-    "fixtures.neon_tenant",
     "fixtures.benchmark_fixture",
     "fixtures.pg_stats",
     "fixtures.compare_fixtures",

test_runner/fixtures/neon_cli.py

@@ -1,662 +0,0 @@
-from __future__ import annotations
-
-import abc
-import json
-import os
-import re
-import subprocess
-import tempfile
-import textwrap
-from itertools import chain, product
-from pathlib import Path
-from typing import (
-    Any,
-    Dict,
-    List,
-    Optional,
-    Tuple,
-    TypeVar,
-    cast,
-)
-
-import toml
-
-from fixtures.common_types import Lsn, TenantId, TimelineId
-from fixtures.log_helper import log
-from fixtures.pageserver.common_types import IndexPartDump
-from fixtures.pg_version import PgVersion
-from fixtures.utils import AuxFileStore
-
-T = TypeVar("T")
-
-
-class AbstractNeonCli(abc.ABC):
-    """
-    A typed wrapper around an arbitrary Neon CLI tool.
-    Supports a way to run arbitrary command directly via CLI.
-    Do not use directly, use specific subclasses instead.
-    """
-
-    def __init__(self, extra_env: Optional[Dict[str, str]], binpath: Path):
-        self.extra_env = extra_env
-        self.binpath = binpath
-
-    COMMAND: str = cast(str, None)  # To be overwritten by the derived class.
-
-    def raw_cli(
-        self,
-        arguments: List[str],
-        extra_env_vars: Optional[Dict[str, str]] = None,
-        check_return_code=True,
-        timeout=None,
-    ) -> "subprocess.CompletedProcess[str]":
-        """
-        Run the command with the specified arguments.
-
-        Arguments must be in list form, e.g. ['endpoint', 'create']
-
-        Return both stdout and stderr, which can be accessed as
-
-        >>> result = env.neon_cli.raw_cli(...)
-        >>> assert result.stderr == ""
-        >>> log.info(result.stdout)
-
-        If `check_return_code`, on non-zero exit code logs failure and raises.
-        """
-
-        assert isinstance(arguments, list)
-        assert isinstance(self.COMMAND, str)
-
-        command_path = str(self.binpath / self.COMMAND)
-
-        args = [command_path] + arguments
-        log.info('Running command "{}"'.format(" ".join(args)))
-
-        env_vars = os.environ.copy()
-
-        # extra env
-        for extra_env_key, extra_env_value in (self.extra_env or {}).items():
-            env_vars[extra_env_key] = extra_env_value
-        for extra_env_key, extra_env_value in (extra_env_vars or {}).items():
-            env_vars[extra_env_key] = extra_env_value
-
-        # Pass through coverage settings
-        var = "LLVM_PROFILE_FILE"
-        val = os.environ.get(var)
-        if val:
-            env_vars[var] = val
-
-        # Intercept CalledProcessError and print more info
-        try:
-            res = subprocess.run(
-                args,
-                env=env_vars,
-                check=False,
-                universal_newlines=True,
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
-                timeout=timeout,
-            )
-        except subprocess.TimeoutExpired as e:
-            if e.stderr:
-                stderr = e.stderr.decode(errors="replace")
-            else:
-                stderr = ""
-
-            if e.stdout:
-                stdout = e.stdout.decode(errors="replace")
-            else:
-                stdout = ""
-
-            log.warn(f"CLI timeout: stderr={stderr}, stdout={stdout}")
-            raise
-
-        indent = "  "
-        if not res.returncode:
-            stripped = res.stdout.strip()
-            lines = stripped.splitlines()
-            if len(lines) < 2:
-                log.debug(f"Run {res.args} success: {stripped}")
-            else:
-                log.debug("Run %s success:\n%s" % (res.args, textwrap.indent(stripped, indent)))
-        elif check_return_code:
-            # this way command output will be in recorded and shown in CI in failure message
-            indent = indent * 2
-            msg = textwrap.dedent(
-                """\
-            Run %s failed:
-              stdout:
-            %s
-              stderr:
-            %s
-            """
-            )
-            msg = msg % (
-                res.args,
-                textwrap.indent(res.stdout.strip(), indent),
-                textwrap.indent(res.stderr.strip(), indent),
-            )
-            log.info(msg)
-            raise RuntimeError(msg) from subprocess.CalledProcessError(
-                res.returncode, res.args, res.stdout, res.stderr
-            )
-        return res
-
-
-class NeonLocalCli(AbstractNeonCli):
-    """A typed wrapper around the `neon_local` CLI tool.
-    Supports main commands via typed methods and a way to run arbitrary command directly via CLI.
-
-    Note: The methods in this class are supposed to be faithful wrappers of the underlying
-    'neon_local' commands. If you're tempted to add any logic here, please consider putting it
-    in the caller instead!
-
-    There are a few exceptions where these wrapper methods intentionally differ from the
-    underlying commands, however:
-    - Many 'neon_local' commands take an optional 'tenant_id' argument and use the default from
-      the config file if it's omitted. The corresponding wrappers require an explicit 'tenant_id'
-      argument. The idea is that we don't want to rely on the config file's default in tests,
-      because NeonEnv has its own 'initial_tenant'. They are currently always the same, but we
-      want to rely on the Neonenv's default instead of the config file default in tests.
-
-    - Similarly, --pg_version argument is always required in the wrappers, even when it's
-      optional in the 'neon_local' command. The default in 'neon_local' is a specific
-      hardcoded version, but in tests, we never want to accidentally rely on that;, we
-      always want to use the version from the test fixtures.
-
-    - Wrappers for commands that create a new tenant or timeline ID require the new tenant
-      or timeline ID to be passed by the caller, while the 'neon_local' commands will
-      generate a random ID if it's not specified. This is because we don't want to have to
-      parse the ID from the 'neon_local' output. Making it required ensures that the
-      caller has to generate it.
-    """
-
-    COMMAND = "neon_local"
-
-    def __init__(
-        self,
-        extra_env: Optional[Dict[str, str]],
-        binpath: Path,
-        repo_dir: Path,
-        pg_distrib_dir: Path,
-    ):
-        if extra_env is None:
-            env_vars = {}
-        else:
-            env_vars = extra_env.copy()
-        env_vars["NEON_REPO_DIR"] = str(repo_dir)
-        env_vars["POSTGRES_DISTRIB_DIR"] = str(pg_distrib_dir)
-
-        super().__init__(env_vars, binpath)
-
-    def raw_cli(self, *args, **kwargs) -> subprocess.CompletedProcess[str]:
-        return super().raw_cli(*args, **kwargs)
-
-    def tenant_create(
-        self,
-        tenant_id: TenantId,
-        timeline_id: TimelineId,
-        pg_version: PgVersion,
-        conf: Optional[Dict[str, Any]] = None,
-        shard_count: Optional[int] = None,
-        shard_stripe_size: Optional[int] = None,
-        placement_policy: Optional[str] = None,
-        set_default: bool = False,
-        aux_file_policy: Optional[AuxFileStore] = None,
-    ):
-        """
-        Creates a new tenant, returns its id and its initial timeline's id.
-        """
-        args = [
-            "tenant",
-            "create",
-            "--tenant-id",
-            str(tenant_id),
-            "--timeline-id",
-            str(timeline_id),
-            "--pg-version",
-            pg_version,
-        ]
-        if conf is not None:
-            args.extend(
-                chain.from_iterable(
-                    product(["-c"], (f"{key}:{value}" for key, value in conf.items()))
-                )
-            )
-
-        if aux_file_policy is AuxFileStore.V2:
-            args.extend(["-c", "switch_aux_file_policy:v2"])
-        elif aux_file_policy is AuxFileStore.V1:
-            args.extend(["-c", "switch_aux_file_policy:v1"])
-        elif aux_file_policy is AuxFileStore.CrossValidation:
-            args.extend(["-c", "switch_aux_file_policy:cross-validation"])
-
-        if set_default:
-            args.append("--set-default")
-
-        if shard_count is not None:
-            args.extend(["--shard-count", str(shard_count)])
-
-        if shard_stripe_size is not None:
-            args.extend(["--shard-stripe-size", str(shard_stripe_size)])
-
-        if placement_policy is not None:
-            args.extend(["--placement-policy", str(placement_policy)])
-
-        res = self.raw_cli(args)
-        res.check_returncode()
-
-    def tenant_import(self, tenant_id: TenantId):
-        args = ["tenant", "import", "--tenant-id", str(tenant_id)]
-        res = self.raw_cli(args)
-        res.check_returncode()
-
-    def tenant_set_default(self, tenant_id: TenantId):
-        """
-        Update default tenant for future operations that require tenant_id.
-        """
-        res = self.raw_cli(["tenant", "set-default", "--tenant-id", str(tenant_id)])
-        res.check_returncode()
-
-    def tenant_config(self, tenant_id: TenantId, conf: Dict[str, str]):
-        """
-        Update tenant config.
-        """
-
-        args = ["tenant", "config", "--tenant-id", str(tenant_id)]
-        if conf is not None:
-            args.extend(
-                chain.from_iterable(
-                    product(["-c"], (f"{key}:{value}" for key, value in conf.items()))
-                )
-            )
-
-        res = self.raw_cli(args)
-        res.check_returncode()
-
-    def tenant_list(self) -> "subprocess.CompletedProcess[str]":
-        res = self.raw_cli(["tenant", "list"])
-        res.check_returncode()
-        return res
-
-    def timeline_create(
-        self,
-        new_branch_name: str,
-        tenant_id: TenantId,
-        timeline_id: TimelineId,
-        pg_version: PgVersion,
-    ) -> TimelineId:
-        if timeline_id is None:
-            timeline_id = TimelineId.generate()
-
-        cmd = [
-            "timeline",
-            "create",
-            "--branch-name",
-            new_branch_name,
-            "--tenant-id",
-            str(tenant_id),
-            "--timeline-id",
-            str(timeline_id),
-            "--pg-version",
-            pg_version,
-        ]
-
-        res = self.raw_cli(cmd)
-        res.check_returncode()
-
-        return timeline_id
-
-    def timeline_branch(
-        self,
-        tenant_id: TenantId,
-        timeline_id: TimelineId,
-        new_branch_name,
-        ancestor_branch_name: Optional[str] = None,
-        ancestor_start_lsn: Optional[Lsn] = None,
-    ):
-        cmd = [
-            "timeline",
-            "branch",
-            "--branch-name",
-            new_branch_name,
-            "--timeline-id",
-            str(timeline_id),
-            "--tenant-id",
-            str(tenant_id),
-        ]
-        if ancestor_branch_name is not None:
-            cmd.extend(["--ancestor-branch-name", ancestor_branch_name])
-        if ancestor_start_lsn is not None:
-            cmd.extend(["--ancestor-start-lsn", str(ancestor_start_lsn)])
-
-        res = self.raw_cli(cmd)
-        res.check_returncode()
-
-    def timeline_import(
-        self,
-        tenant_id: TenantId,
-        timeline_id: TimelineId,
-        new_branch_name: str,
-        base_lsn: Lsn,
-        base_tarfile: Path,
-        pg_version: PgVersion,
-        end_lsn: Optional[Lsn] = None,
-        wal_tarfile: Optional[Path] = None,
-    ):
-        cmd = [
-            "timeline",
-            "import",
-            "--tenant-id",
-            str(tenant_id),
-            "--timeline-id",
-            str(timeline_id),
-            "--pg-version",
-            pg_version,
-            "--branch-name",
-            new_branch_name,
-            "--base-lsn",
-            str(base_lsn),
-            "--base-tarfile",
-            str(base_tarfile),
-        ]
-        if end_lsn is not None:
-            cmd.extend(["--end-lsn", str(end_lsn)])
-        if wal_tarfile is not None:
-            cmd.extend(["--wal-tarfile", str(wal_tarfile)])
-
-        res = self.raw_cli(cmd)
-        res.check_returncode()
-
-    def timeline_list(self, tenant_id: TenantId) -> List[Tuple[str, TimelineId]]:
-        """
-        Returns a list of (branch_name, timeline_id) tuples out of parsed `neon timeline list` CLI output.
-        """
-
-        # main [b49f7954224a0ad25cc0013ea107b54b]
-        # ┣━ @0/16B5A50: test_cli_branch_list_main [20f98c79111b9015d84452258b7d5540]
-        TIMELINE_DATA_EXTRACTOR: re.Pattern = re.compile(  # type: ignore[type-arg]
-            r"\s?(?P<branch_name>[^\s]+)\s\[(?P<timeline_id>[^\]]+)\]", re.MULTILINE
-        )
-        res = self.raw_cli(["timeline", "list", "--tenant-id", str(tenant_id)])
-        timelines_cli = sorted(
-            map(
-                lambda branch_and_id: (branch_and_id[0], TimelineId(branch_and_id[1])),
-                TIMELINE_DATA_EXTRACTOR.findall(res.stdout),
-            )
-        )
-        return timelines_cli
-
-    def init(
-        self,
-        init_config: Dict[str, Any],
-        force: Optional[str] = None,
-    ) -> "subprocess.CompletedProcess[str]":
-        with tempfile.NamedTemporaryFile(mode="w+") as init_config_tmpfile:
-            init_config_tmpfile.write(toml.dumps(init_config))
-            init_config_tmpfile.flush()
-
-            cmd = [
-                "init",
-                f"--config={init_config_tmpfile.name}",
-            ]
-
-            if force is not None:
-                cmd.extend(["--force", force])
-
-            res = self.raw_cli(cmd)
-            res.check_returncode()
-        return res
-
-    def storage_controller_start(
-        self,
-        timeout_in_seconds: Optional[int] = None,
-        instance_id: Optional[int] = None,
-        base_port: Optional[int] = None,
-    ):
-        cmd = ["storage_controller", "start"]
-        if timeout_in_seconds is not None:
-            cmd.append(f"--start-timeout={timeout_in_seconds}s")
-        if instance_id is not None:
-            cmd.append(f"--instance-id={instance_id}")
-        if base_port is not None:
-            cmd.append(f"--base-port={base_port}")
-        return self.raw_cli(cmd)
-
-    def storage_controller_stop(self, immediate: bool, instance_id: Optional[int] = None):
-        cmd = ["storage_controller", "stop"]
-        if immediate:
-            cmd.extend(["-m", "immediate"])
-        if instance_id is not None:
-            cmd.append(f"--instance-id={instance_id}")
-        return self.raw_cli(cmd)
-
-    def pageserver_start(
-        self,
-        id: int,
-        extra_env_vars: Optional[Dict[str, str]] = None,
-        timeout_in_seconds: Optional[int] = None,
-    ) -> "subprocess.CompletedProcess[str]":
-        start_args = ["pageserver", "start", f"--id={id}"]
-        if timeout_in_seconds is not None:
-            start_args.append(f"--start-timeout={timeout_in_seconds}s")
-        return self.raw_cli(start_args, extra_env_vars=extra_env_vars)
-
-    def pageserver_stop(self, id: int, immediate=False) -> "subprocess.CompletedProcess[str]":
-        cmd = ["pageserver", "stop", f"--id={id}"]
-        if immediate:
-            cmd.extend(["-m", "immediate"])
-
-        log.info(f"Stopping pageserver with {cmd}")
-        return self.raw_cli(cmd)
-
-    def safekeeper_start(
-        self,
-        id: int,
-        extra_opts: Optional[List[str]] = None,
-        extra_env_vars: Optional[Dict[str, str]] = None,
-        timeout_in_seconds: Optional[int] = None,
-    ) -> "subprocess.CompletedProcess[str]":
-        if extra_opts is not None:
-            extra_opts = [f"-e={opt}" for opt in extra_opts]
-        else:
-            extra_opts = []
-        if timeout_in_seconds is not None:
-            extra_opts.append(f"--start-timeout={timeout_in_seconds}s")
-        return self.raw_cli(
-            ["safekeeper", "start", str(id), *extra_opts], extra_env_vars=extra_env_vars
-        )
-
-    def safekeeper_stop(
-        self, id: Optional[int] = None, immediate=False
-    ) -> "subprocess.CompletedProcess[str]":
-        args = ["safekeeper", "stop"]
-        if id is not None:
-            args.append(str(id))
-        if immediate:
-            args.extend(["-m", "immediate"])
-        return self.raw_cli(args)
-
-    def storage_broker_start(
-        self, timeout_in_seconds: Optional[int] = None
-    ) -> "subprocess.CompletedProcess[str]":
-        cmd = ["storage_broker", "start"]
-        if timeout_in_seconds is not None:
-            cmd.append(f"--start-timeout={timeout_in_seconds}s")
-        return self.raw_cli(cmd)
-
-    def storage_broker_stop(self) -> "subprocess.CompletedProcess[str]":
-        cmd = ["storage_broker", "stop"]
-        return self.raw_cli(cmd)
-
-    def endpoint_create(
-        self,
-        branch_name: str,
-        pg_port: int,
-        http_port: int,
-        tenant_id: TenantId,
-        pg_version: PgVersion,
-        endpoint_id: Optional[str] = None,
-        hot_standby: bool = False,
-        lsn: Optional[Lsn] = None,
-        pageserver_id: Optional[int] = None,
-        allow_multiple=False,
-    ) -> "subprocess.CompletedProcess[str]":
-        args = [
-            "endpoint",
-            "create",
-            "--tenant-id",
-            str(tenant_id),
-            "--branch-name",
-            branch_name,
-            "--pg-version",
-            pg_version,
-        ]
-        if lsn is not None:
-            args.extend(["--lsn", str(lsn)])
-        if pg_port is not None:
-            args.extend(["--pg-port", str(pg_port)])
-        if http_port is not None:
-            args.extend(["--http-port", str(http_port)])
-        if endpoint_id is not None:
-            args.append(endpoint_id)
-        if hot_standby:
-            args.extend(["--hot-standby", "true"])
-        if pageserver_id is not None:
-            args.extend(["--pageserver-id", str(pageserver_id)])
-        if allow_multiple:
-            args.extend(["--allow-multiple"])
-
-        res = self.raw_cli(args)
-        res.check_returncode()
-        return res
-
-    def endpoint_start(
-        self,
-        endpoint_id: str,
-        safekeepers: Optional[List[int]] = None,
-        remote_ext_config: Optional[str] = None,
-        pageserver_id: Optional[int] = None,
-        allow_multiple=False,
-        basebackup_request_tries: Optional[int] = None,
-    ) -> "subprocess.CompletedProcess[str]":
-        args = [
-            "endpoint",
-            "start",
-        ]
-        extra_env_vars = {}
-        if basebackup_request_tries is not None:
-            extra_env_vars["NEON_COMPUTE_TESTING_BASEBACKUP_TRIES"] = str(basebackup_request_tries)
-        if remote_ext_config is not None:
-            args.extend(["--remote-ext-config", remote_ext_config])
-
-        if safekeepers is not None:
-            args.extend(["--safekeepers", (",".join(map(str, safekeepers)))])
-        if endpoint_id is not None:
-            args.append(endpoint_id)
-        if pageserver_id is not None:
-            args.extend(["--pageserver-id", str(pageserver_id)])
-        if allow_multiple:
-            args.extend(["--allow-multiple"])
-
-        res = self.raw_cli(args, extra_env_vars)
-        res.check_returncode()
-        return res
-
-    def endpoint_reconfigure(
-        self,
-        endpoint_id: str,
-        tenant_id: Optional[TenantId] = None,
-        pageserver_id: Optional[int] = None,
-        safekeepers: Optional[List[int]] = None,
-        check_return_code=True,
-    ) -> "subprocess.CompletedProcess[str]":
-        args = ["endpoint", "reconfigure", endpoint_id]
-        if tenant_id is not None:
-            args.extend(["--tenant-id", str(tenant_id)])
-        if pageserver_id is not None:
-            args.extend(["--pageserver-id", str(pageserver_id)])
-        if safekeepers is not None:
-            args.extend(["--safekeepers", (",".join(map(str, safekeepers)))])
-        return self.raw_cli(args, check_return_code=check_return_code)
-
-    def endpoint_stop(
-        self,
-        endpoint_id: str,
-        destroy=False,
-        check_return_code=True,
-        mode: Optional[str] = None,
-    ) -> "subprocess.CompletedProcess[str]":
-        args = [
-            "endpoint",
-            "stop",
-        ]
-        if destroy:
-            args.append("--destroy")
-        if mode is not None:
-            args.append(f"--mode={mode}")
-        if endpoint_id is not None:
-            args.append(endpoint_id)
-
-        return self.raw_cli(args, check_return_code=check_return_code)
-
-    def mappings_map_branch(
-        self, name: str, tenant_id: TenantId, timeline_id: TimelineId
-    ) -> "subprocess.CompletedProcess[str]":
-        """
-        Map tenant id and timeline id to a neon_local branch name. They do not have to exist.
-        Usually needed when creating branches via PageserverHttpClient and not neon_local.
-
-        After creating a name mapping, you can use EndpointFactory.create_start
-        with this registered branch name.
-        """
-        args = [
-            "mappings",
-            "map",
-            "--branch-name",
-            name,
-            "--tenant-id",
-            str(tenant_id),
-            "--timeline-id",
-            str(timeline_id),
-        ]
-
-        return self.raw_cli(args, check_return_code=True)
-
-    def start(self, check_return_code=True) -> "subprocess.CompletedProcess[str]":
-        return self.raw_cli(["start"], check_return_code=check_return_code)
-
-    def stop(self, check_return_code=True) -> "subprocess.CompletedProcess[str]":
-        return self.raw_cli(["stop"], check_return_code=check_return_code)
-
-
-class WalCraft(AbstractNeonCli):
-    """
-    A typed wrapper around the `wal_craft` CLI tool.
-    Supports main commands via typed methods and a way to run arbitrary command directly via CLI.
-    """
-
-    COMMAND = "wal_craft"
-
-    def postgres_config(self) -> List[str]:
-        res = self.raw_cli(["print-postgres-config"])
-        res.check_returncode()
-        return res.stdout.split("\n")
-
-    def in_existing(self, type: str, connection: str) -> None:
-        res = self.raw_cli(["in-existing", type, connection])
-        res.check_returncode()
-
-
-class Pagectl(AbstractNeonCli):
-    """
-    A typed wrapper around the `pagectl` utility CLI tool.
-    """
-
-    COMMAND = "pagectl"
-
-    def dump_index_part(self, path: Path) -> IndexPartDump:
-        res = self.raw_cli(["index-part", "dump", str(path)])
-        res.check_returncode()
-        parsed = json.loads(res.stdout)
-        return IndexPartDump.from_json(parsed)

test_runner/fixtures/neon_fixtures.py

@@ -9,6 +9,8 @@ import os
 import re
 import shutil
 import subprocess
+import tempfile
+import textwrap
 import threading
 import time
 import uuid
@@ -19,6 +21,7 @@ from datetime import datetime
 from enum import Enum
 from fcntl import LOCK_EX, LOCK_UN, flock
 from functools import cached_property
+from itertools import chain, product
 from pathlib import Path
 from types import TracebackType
 from typing import (
@@ -61,12 +64,11 @@ from fixtures.common_types import Lsn, NodeId, TenantId, TenantShardId, Timeline
 from fixtures.endpoint.http import EndpointHttpClient
 from fixtures.log_helper import log
 from fixtures.metrics import Metrics, MetricsGetter, parse_metrics
-from fixtures.neon_cli import NeonLocalCli, Pagectl
 from fixtures.pageserver.allowed_errors import (
     DEFAULT_PAGESERVER_ALLOWED_ERRORS,
     DEFAULT_STORAGE_CONTROLLER_ALLOWED_ERRORS,
 )
-from fixtures.pageserver.common_types import LayerName, parse_layer_file_name
+from fixtures.pageserver.common_types import IndexPartDump, LayerName, parse_layer_file_name
 from fixtures.pageserver.http import PageserverHttpClient
 from fixtures.pageserver.utils import (
     wait_for_last_record_lsn,
@@ -454,10 +456,9 @@ class NeonEnvBuilder:
 
         self.pageserver_io_buffer_alignment = pageserver_io_buffer_alignment
 
-        # FIXME
-        # assert test_name.startswith(
-        #    "test_"
-        # ), "Unexpectedly instantiated from outside a test function"
+        assert test_name.startswith(
+            "test_"
+        ), "Unexpectedly instantiated from outside a test function"
         self.test_name = test_name
 
     def init_configs(self, default_remote_storage_if_missing: bool = True) -> NeonEnv:
@@ -490,7 +491,7 @@ class NeonEnvBuilder:
         log.debug(
             f"Services started, creating initial tenant {env.initial_tenant} and its initial timeline"
         )
-        initial_tenant, initial_timeline = env.create_tenant(
+        initial_tenant, initial_timeline = env.neon_cli.create_tenant(
             tenant_id=env.initial_tenant,
             conf=initial_tenant_conf,
             timeline_id=env.initial_timeline,
@@ -951,16 +952,10 @@ class NeonEnv:
 
     initial_tenant - tenant ID of the initial tenant created in the repository
 
-    neon_cli - can be used to run the 'neon_local' CLI tool
+    neon_cli - can be used to run the 'neon' CLI tool
 
-    create_tenant() - initializes a new tenant and an initial empty timeline on it,
-        returns the tenant and timeline id
-
-    create_branch() - branch a new timeline from an existing one, returns
-        the new timeline id
-
-    create_timeline() - initializes a new timeline by running initdb, returns
-        the new timeline id
+    create_tenant() - initializes a new tenant in the page server, returns
+        the tenant id
     """
 
     BASE_PAGESERVER_ID = 1
@@ -971,6 +966,8 @@ class NeonEnv:
         self.rust_log_override = config.rust_log_override
         self.port_distributor = config.port_distributor
         self.s3_mock_server = config.mock_s3_server
+        self.neon_cli = NeonCli(env=self)
+        self.pagectl = Pagectl(env=self)
         self.endpoints = EndpointFactory(self)
         self.safekeepers: List[Safekeeper] = []
         self.pageservers: List[NeonPageserver] = []
@@ -990,21 +987,6 @@ class NeonEnv:
         self.initial_tenant = config.initial_tenant
         self.initial_timeline = config.initial_timeline
 
-        neon_local_env_vars = {}
-        if self.rust_log_override is not None:
-            neon_local_env_vars["RUST_LOG"] = self.rust_log_override
-        self.neon_cli = NeonLocalCli(
-            extra_env=neon_local_env_vars,
-            binpath=self.neon_local_binpath,
-            repo_dir=self.repo_dir,
-            pg_distrib_dir=self.pg_distrib_dir,
-        )
-
-        pagectl_env_vars = {}
-        if self.rust_log_override is not None:
-            pagectl_env_vars["RUST_LOG"] = self.rust_log_override
-        self.pagectl = Pagectl(extra_env=pagectl_env_vars, binpath=self.neon_binpath)
-
         # The URL for the pageserver to use as its control_plane_api config
         if config.storage_controller_port_override is not None:
             log.info(
@@ -1328,74 +1310,6 @@ class NeonEnv:
         self.endpoint_counter += 1
         return "ep-" + str(self.endpoint_counter)
 
-    def create_tenant(
-        self,
-        tenant_id: Optional[TenantId] = None,
-        timeline_id: Optional[TimelineId] = None,
-        conf: Optional[Dict[str, Any]] = None,
-        shard_count: Optional[int] = None,
-        shard_stripe_size: Optional[int] = None,
-        placement_policy: Optional[str] = None,
-        set_default: bool = False,
-        aux_file_policy: Optional[AuxFileStore] = None,
-    ) -> Tuple[TenantId, TimelineId]:
-        """
-        Creates a new tenant, returns its id and its initial timeline's id.
-        """
-        tenant_id = tenant_id or TenantId.generate()
-        timeline_id = timeline_id or TimelineId.generate()
-
-        self.neon_cli.tenant_create(
-            tenant_id=tenant_id,
-            timeline_id=timeline_id,
-            pg_version=self.pg_version,
-            conf=conf,
-            shard_count=shard_count,
-            shard_stripe_size=shard_stripe_size,
-            placement_policy=placement_policy,
-            set_default=set_default,
-            aux_file_policy=aux_file_policy,
-        )
-
-        return tenant_id, timeline_id
-
-    def config_tenant(self, tenant_id: Optional[TenantId], conf: Dict[str, str]):
-        """
-        Update tenant config.
-        """
-        tenant_id = tenant_id or self.initial_tenant
-        self.neon_cli.tenant_config(tenant_id, conf)
-
-    def create_branch(
-        self,
-        new_branch_name: str = DEFAULT_BRANCH_NAME,
-        tenant_id: Optional[TenantId] = None,
-        ancestor_branch_name: Optional[str] = None,
-        ancestor_start_lsn: Optional[Lsn] = None,
-        new_timeline_id: Optional[TimelineId] = None,
-    ) -> TimelineId:
-        new_timeline_id = new_timeline_id or TimelineId.generate()
-        tenant_id = tenant_id or self.initial_tenant
-
-        self.neon_cli.timeline_branch(
-            tenant_id, new_timeline_id, new_branch_name, ancestor_branch_name, ancestor_start_lsn
-        )
-
-        return new_timeline_id
-
-    def create_timeline(
-        self,
-        new_branch_name: str,
-        tenant_id: Optional[TenantId] = None,
-        timeline_id: Optional[TimelineId] = None,
-    ) -> TimelineId:
-        timeline_id = timeline_id or TimelineId.generate()
-        tenant_id = tenant_id or self.initial_tenant
-
-        self.neon_cli.timeline_create(new_branch_name, tenant_id, timeline_id, self.pg_version)
-
-        return timeline_id
-
 
 @pytest.fixture(scope="function")
 def neon_simple_env(
@@ -1511,6 +1425,597 @@ class PageserverPort:
     http: int
 
 
+class AbstractNeonCli(abc.ABC):
+    """
+    A typed wrapper around an arbitrary Neon CLI tool.
+    Supports a way to run arbitrary command directly via CLI.
+    Do not use directly, use specific subclasses instead.
+    """
+
+    def __init__(self, env: NeonEnv):
+        self.env = env
+
+    COMMAND: str = cast(str, None)  # To be overwritten by the derived class.
+
+    def raw_cli(
+        self,
+        arguments: List[str],
+        extra_env_vars: Optional[Dict[str, str]] = None,
+        check_return_code=True,
+        timeout=None,
+        local_binpath=False,
+    ) -> "subprocess.CompletedProcess[str]":
+        """
+        Run the command with the specified arguments.
+
+        Arguments must be in list form, e.g. ['pg', 'create']
+
+        Return both stdout and stderr, which can be accessed as
+
+        >>> result = env.neon_cli.raw_cli(...)
+        >>> assert result.stderr == ""
+        >>> log.info(result.stdout)
+
+        If `check_return_code`, on non-zero exit code logs failure and raises.
+
+        If `local_binpath` is true, then we are invoking a test utility
+        """
+
+        assert isinstance(arguments, list)
+        assert isinstance(self.COMMAND, str)
+
+        if local_binpath:
+            # Test utility
+            bin_neon = str(self.env.neon_local_binpath / self.COMMAND)
+        else:
+            # Normal binary
+            bin_neon = str(self.env.neon_binpath / self.COMMAND)
+
+        args = [bin_neon] + arguments
+        log.info('Running command "{}"'.format(" ".join(args)))
+
+        env_vars = os.environ.copy()
+        env_vars["NEON_REPO_DIR"] = str(self.env.repo_dir)
+        env_vars["POSTGRES_DISTRIB_DIR"] = str(self.env.pg_distrib_dir)
+        if self.env.rust_log_override is not None:
+            env_vars["RUST_LOG"] = self.env.rust_log_override
+        for extra_env_key, extra_env_value in (extra_env_vars or {}).items():
+            env_vars[extra_env_key] = extra_env_value
+
+        # Pass coverage settings
+        var = "LLVM_PROFILE_FILE"
+        val = os.environ.get(var)
+        if val:
+            env_vars[var] = val
+
+        # Intercept CalledProcessError and print more info
+        try:
+            res = subprocess.run(
+                args,
+                env=env_vars,
+                check=False,
+                universal_newlines=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                timeout=timeout,
+            )
+        except subprocess.TimeoutExpired as e:
+            if e.stderr:
+                stderr = e.stderr.decode(errors="replace")
+            else:
+                stderr = ""
+
+            if e.stdout:
+                stdout = e.stdout.decode(errors="replace")
+            else:
+                stdout = ""
+
+            log.warn(f"CLI timeout: stderr={stderr}, stdout={stdout}")
+            raise
+
+        indent = "  "
+        if not res.returncode:
+            stripped = res.stdout.strip()
+            lines = stripped.splitlines()
+            if len(lines) < 2:
+                log.debug(f"Run {res.args} success: {stripped}")
+            else:
+                log.debug("Run %s success:\n%s" % (res.args, textwrap.indent(stripped, indent)))
+        elif check_return_code:
+            # this way command output will be in recorded and shown in CI in failure message
+            indent = indent * 2
+            msg = textwrap.dedent(
+                """\
+            Run %s failed:
+              stdout:
+            %s
+              stderr:
+            %s
+            """
+            )
+            msg = msg % (
+                res.args,
+                textwrap.indent(res.stdout.strip(), indent),
+                textwrap.indent(res.stderr.strip(), indent),
+            )
+            log.info(msg)
+            raise RuntimeError(msg) from subprocess.CalledProcessError(
+                res.returncode, res.args, res.stdout, res.stderr
+            )
+        return res
+
+
+class NeonCli(AbstractNeonCli):
+    """
+    A typed wrapper around the `neon` CLI tool.
+    Supports main commands via typed methods and a way to run arbitrary command directly via CLI.
+    """
+
+    COMMAND = "neon_local"
+
+    def raw_cli(self, *args, **kwargs) -> subprocess.CompletedProcess[str]:
+        kwargs["local_binpath"] = True
+        return super().raw_cli(*args, **kwargs)
+
+    def create_tenant(
+        self,
+        tenant_id: Optional[TenantId] = None,
+        timeline_id: Optional[TimelineId] = None,
+        conf: Optional[Dict[str, Any]] = None,
+        shard_count: Optional[int] = None,
+        shard_stripe_size: Optional[int] = None,
+        placement_policy: Optional[str] = None,
+        set_default: bool = False,
+        aux_file_policy: Optional[AuxFileStore] = None,
+    ) -> Tuple[TenantId, TimelineId]:
+        """
+        Creates a new tenant, returns its id and its initial timeline's id.
+        """
+        tenant_id = tenant_id or TenantId.generate()
+        timeline_id = timeline_id or TimelineId.generate()
+
+        args = [
+            "tenant",
+            "create",
+            "--tenant-id",
+            str(tenant_id),
+            "--timeline-id",
+            str(timeline_id),
+            "--pg-version",
+            self.env.pg_version,
+        ]
+        if conf is not None:
+            args.extend(
+                chain.from_iterable(
+                    product(["-c"], (f"{key}:{value}" for key, value in conf.items()))
+                )
+            )
+
+        if aux_file_policy is AuxFileStore.V2:
+            args.extend(["-c", "switch_aux_file_policy:v2"])
+        elif aux_file_policy is AuxFileStore.V1:
+            args.extend(["-c", "switch_aux_file_policy:v1"])
+        elif aux_file_policy is AuxFileStore.CrossValidation:
+            args.extend(["-c", "switch_aux_file_policy:cross-validation"])
+
+        if set_default:
+            args.append("--set-default")
+
+        if shard_count is not None:
+            args.extend(["--shard-count", str(shard_count)])
+
+        if shard_stripe_size is not None:
+            args.extend(["--shard-stripe-size", str(shard_stripe_size)])
+
+        if placement_policy is not None:
+            args.extend(["--placement-policy", str(placement_policy)])
+
+        res = self.raw_cli(args)
+        res.check_returncode()
+        return tenant_id, timeline_id
+
+    def import_tenant(self, tenant_id: TenantId):
+        args = ["tenant", "import", "--tenant-id", str(tenant_id)]
+        res = self.raw_cli(args)
+        res.check_returncode()
+
+    def set_default(self, tenant_id: TenantId):
+        """
+        Update default tenant for future operations that require tenant_id.
+        """
+        res = self.raw_cli(["tenant", "set-default", "--tenant-id", str(tenant_id)])
+        res.check_returncode()
+
+    def config_tenant(self, tenant_id: TenantId, conf: Dict[str, str]):
+        """
+        Update tenant config.
+        """
+
+        args = ["tenant", "config", "--tenant-id", str(tenant_id)]
+        if conf is not None:
+            args.extend(
+                chain.from_iterable(
+                    product(["-c"], (f"{key}:{value}" for key, value in conf.items()))
+                )
+            )
+
+        res = self.raw_cli(args)
+        res.check_returncode()
+
+    def list_tenants(self) -> "subprocess.CompletedProcess[str]":
+        res = self.raw_cli(["tenant", "list"])
+        res.check_returncode()
+        return res
+
+    def create_timeline(
+        self,
+        new_branch_name: str,
+        tenant_id: Optional[TenantId] = None,
+        timeline_id: Optional[TimelineId] = None,
+    ) -> TimelineId:
+        if timeline_id is None:
+            timeline_id = TimelineId.generate()
+
+        cmd = [
+            "timeline",
+            "create",
+            "--branch-name",
+            new_branch_name,
+            "--tenant-id",
+            str(tenant_id or self.env.initial_tenant),
+            "--timeline-id",
+            str(timeline_id),
+            "--pg-version",
+            self.env.pg_version,
+        ]
+
+        res = self.raw_cli(cmd)
+        res.check_returncode()
+
+        return timeline_id
+
+    def create_branch(
+        self,
+        new_branch_name: str = DEFAULT_BRANCH_NAME,
+        ancestor_branch_name: Optional[str] = None,
+        tenant_id: Optional[TenantId] = None,
+        ancestor_start_lsn: Optional[Lsn] = None,
+        new_timeline_id: Optional[TimelineId] = None,
+    ) -> TimelineId:
+        if new_timeline_id is None:
+            new_timeline_id = TimelineId.generate()
+        cmd = [
+            "timeline",
+            "branch",
+            "--branch-name",
+            new_branch_name,
+            "--timeline-id",
+            str(new_timeline_id),
+            "--tenant-id",
+            str(tenant_id or self.env.initial_tenant),
+        ]
+        if ancestor_branch_name is not None:
+            cmd.extend(["--ancestor-branch-name", ancestor_branch_name])
+        if ancestor_start_lsn is not None:
+            cmd.extend(["--ancestor-start-lsn", str(ancestor_start_lsn)])
+
+        res = self.raw_cli(cmd)
+        res.check_returncode()
+
+        return TimelineId(str(new_timeline_id))
+
+    def list_timelines(self, tenant_id: Optional[TenantId] = None) -> List[Tuple[str, TimelineId]]:
+        """
+        Returns a list of (branch_name, timeline_id) tuples out of parsed `neon timeline list` CLI output.
+        """
+
+        # main [b49f7954224a0ad25cc0013ea107b54b]
+        # ┣━ @0/16B5A50: test_cli_branch_list_main [20f98c79111b9015d84452258b7d5540]
+        TIMELINE_DATA_EXTRACTOR: re.Pattern = re.compile(  # type: ignore[type-arg]
+            r"\s?(?P<branch_name>[^\s]+)\s\[(?P<timeline_id>[^\]]+)\]", re.MULTILINE
+        )
+        res = self.raw_cli(
+            ["timeline", "list", "--tenant-id", str(tenant_id or self.env.initial_tenant)]
+        )
+        timelines_cli = sorted(
+            map(
+                lambda branch_and_id: (branch_and_id[0], TimelineId(branch_and_id[1])),
+                TIMELINE_DATA_EXTRACTOR.findall(res.stdout),
+            )
+        )
+        return timelines_cli
+
+    def init(
+        self,
+        init_config: Dict[str, Any],
+        force: Optional[str] = None,
+    ) -> "subprocess.CompletedProcess[str]":
+        with tempfile.NamedTemporaryFile(mode="w+") as init_config_tmpfile:
+            init_config_tmpfile.write(toml.dumps(init_config))
+            init_config_tmpfile.flush()
+
+            cmd = [
+                "init",
+                f"--config={init_config_tmpfile.name}",
+            ]
+
+            if force is not None:
+                cmd.extend(["--force", force])
+
+            res = self.raw_cli(cmd)
+            res.check_returncode()
+        return res
+
+    def storage_controller_start(
+        self,
+        timeout_in_seconds: Optional[int] = None,
+        instance_id: Optional[int] = None,
+        base_port: Optional[int] = None,
+    ):
+        cmd = ["storage_controller", "start"]
+        if timeout_in_seconds is not None:
+            cmd.append(f"--start-timeout={timeout_in_seconds}s")
+        if instance_id is not None:
+            cmd.append(f"--instance-id={instance_id}")
+        if base_port is not None:
+            cmd.append(f"--base-port={base_port}")
+        return self.raw_cli(cmd)
+
+    def storage_controller_stop(self, immediate: bool, instance_id: Optional[int] = None):
+        cmd = ["storage_controller", "stop"]
+        if immediate:
+            cmd.extend(["-m", "immediate"])
+        if instance_id is not None:
+            cmd.append(f"--instance-id={instance_id}")
+        return self.raw_cli(cmd)
+
+    def pageserver_start(
+        self,
+        id: int,
+        extra_env_vars: Optional[Dict[str, str]] = None,
+        timeout_in_seconds: Optional[int] = None,
+    ) -> "subprocess.CompletedProcess[str]":
+        start_args = ["pageserver", "start", f"--id={id}"]
+        if timeout_in_seconds is not None:
+            start_args.append(f"--start-timeout={timeout_in_seconds}s")
+        storage = self.env.pageserver_remote_storage
+
+        if isinstance(storage, S3Storage):
+            s3_env_vars = storage.access_env_vars()
+            extra_env_vars = (extra_env_vars or {}) | s3_env_vars
+
+        return self.raw_cli(start_args, extra_env_vars=extra_env_vars)
+
+    def pageserver_stop(self, id: int, immediate=False) -> "subprocess.CompletedProcess[str]":
+        cmd = ["pageserver", "stop", f"--id={id}"]
+        if immediate:
+            cmd.extend(["-m", "immediate"])
+
+        log.info(f"Stopping pageserver with {cmd}")
+        return self.raw_cli(cmd)
+
+    def safekeeper_start(
+        self,
+        id: int,
+        extra_opts: Optional[List[str]] = None,
+        timeout_in_seconds: Optional[int] = None,
+    ) -> "subprocess.CompletedProcess[str]":
+        s3_env_vars = None
+        if isinstance(self.env.safekeepers_remote_storage, S3Storage):
+            s3_env_vars = self.env.safekeepers_remote_storage.access_env_vars()
+
+        if extra_opts is not None:
+            extra_opts = [f"-e={opt}" for opt in extra_opts]
+        else:
+            extra_opts = []
+        if timeout_in_seconds is not None:
+            extra_opts.append(f"--start-timeout={timeout_in_seconds}s")
+        return self.raw_cli(
+            ["safekeeper", "start", str(id), *extra_opts], extra_env_vars=s3_env_vars
+        )
+
+    def safekeeper_stop(
+        self, id: Optional[int] = None, immediate=False
+    ) -> "subprocess.CompletedProcess[str]":
+        args = ["safekeeper", "stop"]
+        if id is not None:
+            args.append(str(id))
+        if immediate:
+            args.extend(["-m", "immediate"])
+        return self.raw_cli(args)
+
+    def broker_start(
+        self, timeout_in_seconds: Optional[int] = None
+    ) -> "subprocess.CompletedProcess[str]":
+        cmd = ["storage_broker", "start"]
+        if timeout_in_seconds is not None:
+            cmd.append(f"--start-timeout={timeout_in_seconds}s")
+        return self.raw_cli(cmd)
+
+    def broker_stop(self) -> "subprocess.CompletedProcess[str]":
+        cmd = ["storage_broker", "stop"]
+        return self.raw_cli(cmd)
+
+    def endpoint_create(
+        self,
+        branch_name: str,
+        pg_port: int,
+        http_port: int,
+        endpoint_id: Optional[str] = None,
+        tenant_id: Optional[TenantId] = None,
+        hot_standby: bool = False,
+        lsn: Optional[Lsn] = None,
+        pageserver_id: Optional[int] = None,
+        allow_multiple=False,
+    ) -> "subprocess.CompletedProcess[str]":
+        args = [
+            "endpoint",
+            "create",
+            "--tenant-id",
+            str(tenant_id or self.env.initial_tenant),
+            "--branch-name",
+            branch_name,
+            "--pg-version",
+            self.env.pg_version,
+        ]
+        if lsn is not None:
+            args.extend(["--lsn", str(lsn)])
+        if pg_port is not None:
+            args.extend(["--pg-port", str(pg_port)])
+        if http_port is not None:
+            args.extend(["--http-port", str(http_port)])
+        if endpoint_id is not None:
+            args.append(endpoint_id)
+        if hot_standby:
+            args.extend(["--hot-standby", "true"])
+        if pageserver_id is not None:
+            args.extend(["--pageserver-id", str(pageserver_id)])
+        if allow_multiple:
+            args.extend(["--allow-multiple"])
+
+        res = self.raw_cli(args)
+        res.check_returncode()
+        return res
+
+    def endpoint_start(
+        self,
+        endpoint_id: str,
+        safekeepers: Optional[List[int]] = None,
+        remote_ext_config: Optional[str] = None,
+        pageserver_id: Optional[int] = None,
+        allow_multiple=False,
+        basebackup_request_tries: Optional[int] = None,
+    ) -> "subprocess.CompletedProcess[str]":
+        args = [
+            "endpoint",
+            "start",
+        ]
+        extra_env_vars = {}
+        if basebackup_request_tries is not None:
+            extra_env_vars["NEON_COMPUTE_TESTING_BASEBACKUP_TRIES"] = str(basebackup_request_tries)
+        if remote_ext_config is not None:
+            args.extend(["--remote-ext-config", remote_ext_config])
+
+        if safekeepers is not None:
+            args.extend(["--safekeepers", (",".join(map(str, safekeepers)))])
+        if endpoint_id is not None:
+            args.append(endpoint_id)
+        if pageserver_id is not None:
+            args.extend(["--pageserver-id", str(pageserver_id)])
+        if allow_multiple:
+            args.extend(["--allow-multiple"])
+
+        res = self.raw_cli(args, extra_env_vars)
+        res.check_returncode()
+        return res
+
+    def endpoint_reconfigure(
+        self,
+        endpoint_id: str,
+        tenant_id: Optional[TenantId] = None,
+        pageserver_id: Optional[int] = None,
+        safekeepers: Optional[List[int]] = None,
+        check_return_code=True,
+    ) -> "subprocess.CompletedProcess[str]":
+        args = ["endpoint", "reconfigure", endpoint_id]
+        if tenant_id is not None:
+            args.extend(["--tenant-id", str(tenant_id)])
+        if pageserver_id is not None:
+            args.extend(["--pageserver-id", str(pageserver_id)])
+        if safekeepers is not None:
+            args.extend(["--safekeepers", (",".join(map(str, safekeepers)))])
+        return self.raw_cli(args, check_return_code=check_return_code)
+
+    def endpoint_stop(
+        self,
+        endpoint_id: str,
+        destroy=False,
+        check_return_code=True,
+        mode: Optional[str] = None,
+    ) -> "subprocess.CompletedProcess[str]":
+        args = [
+            "endpoint",
+            "stop",
+        ]
+        if destroy:
+            args.append("--destroy")
+        if mode is not None:
+            args.append(f"--mode={mode}")
+        if endpoint_id is not None:
+            args.append(endpoint_id)
+
+        return self.raw_cli(args, check_return_code=check_return_code)
+
+    def map_branch(
+        self, name: str, tenant_id: TenantId, timeline_id: TimelineId
+    ) -> "subprocess.CompletedProcess[str]":
+        """
+        Map tenant id and timeline id to a neon_local branch name. They do not have to exist.
+        Usually needed when creating branches via PageserverHttpClient and not neon_local.
+
+        After creating a name mapping, you can use EndpointFactory.create_start
+        with this registered branch name.
+        """
+        args = [
+            "mappings",
+            "map",
+            "--branch-name",
+            name,
+            "--tenant-id",
+            str(tenant_id),
+            "--timeline-id",
+            str(timeline_id),
+        ]
+
+        return self.raw_cli(args, check_return_code=True)
+
+    def start(self, check_return_code=True) -> "subprocess.CompletedProcess[str]":
+        return self.raw_cli(["start"], check_return_code=check_return_code)
+
+    def stop(self, check_return_code=True) -> "subprocess.CompletedProcess[str]":
+        return self.raw_cli(["stop"], check_return_code=check_return_code)
+
+
+class WalCraft(AbstractNeonCli):
+    """
+    A typed wrapper around the `wal_craft` CLI tool.
+    Supports main commands via typed methods and a way to run arbitrary command directly via CLI.
+    """
+
+    COMMAND = "wal_craft"
+
+    def postgres_config(self) -> List[str]:
+        res = self.raw_cli(["print-postgres-config"])
+        res.check_returncode()
+        return res.stdout.split("\n")
+
+    def in_existing(self, type: str, connection: str) -> None:
+        res = self.raw_cli(["in-existing", type, connection])
+        res.check_returncode()
+
+
+class ComputeCtl(AbstractNeonCli):
+    """
+    A typed wrapper around the `compute_ctl` CLI tool.
+    """
+
+    COMMAND = "compute_ctl"
+
+
+class Pagectl(AbstractNeonCli):
+    """
+    A typed wrapper around the `pagectl` utility CLI tool.
+    """
+
+    COMMAND = "pagectl"
+
+    def dump_index_part(self, path: Path) -> IndexPartDump:
+        res = self.raw_cli(["index-part", "dump", str(path)])
+        res.check_returncode()
+        parsed = json.loads(res.stdout)
+        return IndexPartDump.from_json(parsed)
+
+
 class LogUtils:
     """
     A mixin class which provides utilities for inspecting the logs of a service.
@@ -2428,10 +2933,6 @@ class NeonPageserver(PgProtocol, LogUtils):
         """
         assert self.running is False
 
-        storage = self.env.pageserver_remote_storage
-        if isinstance(storage, S3Storage):
-            s3_env_vars = storage.access_env_vars()
-            extra_env_vars = (extra_env_vars or {}) | s3_env_vars
         self.env.neon_cli.pageserver_start(
             self.id, extra_env_vars=extra_env_vars, timeout_in_seconds=timeout_in_seconds
         )
@@ -3452,7 +3953,6 @@ class Endpoint(PgProtocol, LogUtils):
             hot_standby=hot_standby,
             pg_port=self.pg_port,
             http_port=self.http_port,
-            pg_version=self.env.pg_version,
             pageserver_id=pageserver_id,
             allow_multiple=allow_multiple,
         )
@@ -3895,16 +4395,8 @@ class Safekeeper(LogUtils):
             extra_opts = self.extra_opts
 
         assert self.running is False
-
-        s3_env_vars = None
-        if isinstance(self.env.safekeepers_remote_storage, S3Storage):
-            s3_env_vars = self.env.safekeepers_remote_storage.access_env_vars()
-
         self.env.neon_cli.safekeeper_start(
-            self.id,
-            extra_opts=extra_opts,
-            timeout_in_seconds=timeout_in_seconds,
-            extra_env_vars=s3_env_vars,
+            self.id, extra_opts=extra_opts, timeout_in_seconds=timeout_in_seconds
         )
         self.running = True
         # wait for wal acceptor start by checking its status
@@ -4050,7 +4542,7 @@ class Safekeeper(LogUtils):
         1) wait for remote_consistent_lsn and wal_backup_lsn on safekeeper to reach it.
         2) checkpoint timeline on safekeeper, which should remove WAL before this LSN; optionally wait for that.
         """
-        client = self.http_client()
+        cli = self.http_client()
 
         target_segment_file = lsn.segment_name()
 
@@ -4062,7 +4554,7 @@ class Safekeeper(LogUtils):
             assert all(target_segment_file <= s for s in segments)
 
         def are_lsns_advanced():
-            stat = client.timeline_status(tenant_id, timeline_id)
+            stat = cli.timeline_status(tenant_id, timeline_id)
             log.info(
                 f"waiting for remote_consistent_lsn and backup_lsn on sk {self.id} to reach {lsn}, currently remote_consistent_lsn={stat.remote_consistent_lsn}, backup_lsn={stat.backup_lsn}"
             )
@@ -4071,7 +4563,7 @@ class Safekeeper(LogUtils):
         # xxx: max wait is long because we might be waiting for reconnection from
         # pageserver to this safekeeper
         wait_until(30, 1, are_lsns_advanced)
-        client.checkpoint(tenant_id, timeline_id)
+        cli.checkpoint(tenant_id, timeline_id)
         if wait_wal_removal:
             wait_until(30, 1, are_segments_removed)
 
@@ -4099,13 +4591,13 @@ class NeonBroker(LogUtils):
         timeout_in_seconds: Optional[int] = None,
     ):
         assert not self.running
-        self.env.neon_cli.storage_broker_start(timeout_in_seconds)
+        self.env.neon_cli.broker_start(timeout_in_seconds)
         self.running = True
         return self
 
     def stop(self):
         if self.running:
-            self.env.neon_cli.storage_broker_stop()
+            self.env.neon_cli.broker_stop()
             self.running = False
         return self
 
@@ -4734,10 +5226,10 @@ def flush_ep_to_pageserver(
     commit_lsn: Lsn = Lsn(0)
     # In principle in the absense of failures polling single sk would be enough.
     for sk in env.safekeepers:
-        client = sk.http_client()
+        cli = sk.http_client()
         # wait until compute connections are gone
-        wait_walreceivers_absent(client, tenant, timeline)
-        commit_lsn = max(client.get_commit_lsn(tenant, timeline), commit_lsn)
+        wait_walreceivers_absent(cli, tenant, timeline)
+        commit_lsn = max(cli.get_commit_lsn(tenant, timeline), commit_lsn)
 
     # Note: depending on WAL filtering implementation, probably most shards
     # won't be able to reach commit_lsn (unless gaps are also ack'ed), so this
@@ -4790,12 +5282,7 @@ def fork_at_current_lsn(
     the WAL up to that LSN to arrive in the pageserver before creating the branch.
     """
     current_lsn = endpoint.safe_psql("SELECT pg_current_wal_lsn()")[0][0]
-    return env.create_branch(
-        new_branch_name=new_branch_name,
-        tenant_id=tenant_id,
-        ancestor_branch_name=ancestor_branch_name,
-        ancestor_start_lsn=current_lsn,
-    )
+    return env.neon_cli.create_branch(new_branch_name, ancestor_branch_name, tenant_id, current_lsn)
 
 
 def import_timeline_from_vanilla_postgres(
@@ -4814,9 +5301,9 @@ def import_timeline_from_vanilla_postgres(
     """
 
     # Take backup of the existing PostgreSQL server with pg_basebackup
-    basebackup_dir = test_output_dir / "basebackup"
-    base_tar = basebackup_dir / "base.tar"
-    wal_tar = basebackup_dir / "pg_wal.tar"
+    basebackup_dir = os.path.join(test_output_dir, "basebackup")
+    base_tar = os.path.join(basebackup_dir, "base.tar")
+    wal_tar = os.path.join(basebackup_dir, "pg_wal.tar")
     os.mkdir(basebackup_dir)
     pg_bin.run(
         [
@@ -4826,28 +5313,40 @@ def import_timeline_from_vanilla_postgres(
             "-d",
             vanilla_pg_connstr,
             "-D",
-            str(basebackup_dir),
+            basebackup_dir,
         ]
     )
 
     # Extract start_lsn and end_lsn form the backup manifest file
     with open(os.path.join(basebackup_dir, "backup_manifest")) as f:
         manifest = json.load(f)
-        start_lsn = Lsn(manifest["WAL-Ranges"][0]["Start-LSN"])
-        end_lsn = Lsn(manifest["WAL-Ranges"][0]["End-LSN"])
+        start_lsn = manifest["WAL-Ranges"][0]["Start-LSN"]
+        end_lsn = manifest["WAL-Ranges"][0]["End-LSN"]
 
     # Import the backup tarballs into the pageserver
-    env.neon_cli.timeline_import(
-        tenant_id=tenant_id,
-        timeline_id=timeline_id,
-        new_branch_name=branch_name,
-        base_lsn=start_lsn,
-        base_tarfile=base_tar,
-        end_lsn=end_lsn,
-        wal_tarfile=wal_tar,
-        pg_version=env.pg_version,
+    env.neon_cli.raw_cli(
+        [
+            "timeline",
+            "import",
+            "--tenant-id",
+            str(tenant_id),
+            "--timeline-id",
+            str(timeline_id),
+            "--branch-name",
+            branch_name,
+            "--base-lsn",
+            start_lsn,
+            "--base-tarfile",
+            base_tar,
+            "--end-lsn",
+            end_lsn,
+            "--wal-tarfile",
+            wal_tar,
+            "--pg-version",
+            env.pg_version,
+        ]
     )
-    wait_for_last_record_lsn(env.pageserver.http_client(), tenant_id, timeline_id, end_lsn)
+    wait_for_last_record_lsn(env.pageserver.http_client(), tenant_id, timeline_id, Lsn(end_lsn))
 
 
 def last_flush_lsn_upload(

test_runner/fixtures/neon_tenant.py

@@ -1,352 +0,0 @@
-from __future__ import annotations
-
-import random
-import string
-import threading
-import uuid
-from dataclasses import dataclass
-from pathlib import Path
-from types import TracebackType
-from typing import (
-    Any,
-    Dict,
-    Iterator,
-    List,
-    Optional,
-    Type,
-    cast,
-)
-
-import pytest
-from _pytest.config import Config
-from _pytest.fixtures import FixtureRequest
-
-from fixtures.common_types import Lsn, TenantId, TimelineId
-from fixtures.log_helper import log
-from fixtures.neon_fixtures import Endpoint, NeonEnv, NeonEnvBuilder
-from fixtures.pg_version import PgVersion
-from fixtures.port_distributor import PortDistributor
-from fixtures.remote_storage import (
-    MockS3Server,
-)
-from fixtures.utils import AuxFileStore
-
-DEFAULT_BRANCH_NAME: str = "main"
-
-
-@pytest.fixture(scope="function")
-def neon_tenant(shared_env: NeonEnv) -> Iterator[NeonTestTenant]:
-    tenant = NeonTestTenant(shared_env)
-    tenant.create()
-    yield tenant
-    # TODO:  clean up the tenant
-
-
-@dataclass(frozen=True)
-class NeonEnvDiscriminants:
-    """The options that define which environments can be shared"""
-
-    neon_binpath: Path
-    pageserver_virtual_file_io_engine: str
-    pageserver_aux_file_policy: Optional[AuxFileStore]
-    pageserver_default_tenant_config_compaction_algorithm: Optional[Dict[str, Any]]
-    pageserver_io_buffer_alignment: Optional[int]
-
-
-class NeonSharedEnvs:
-    def __init__(
-        self,
-        port_distributor: PortDistributor,
-        run_id: uuid.UUID,
-        mock_s3_server: MockS3Server,
-        pg_distrib_dir: Path,
-        top_output_dir: Path,
-        preserve_database_files: bool,
-    ):
-        self.port_distributor = port_distributor
-        self.run_id = run_id
-        self.mock_s3_server = mock_s3_server
-        self.pg_distrib_dir = pg_distrib_dir
-        self.top_output_dir = top_output_dir
-        self.preserve_database_files = preserve_database_files
-
-        self.lock = threading.Lock()
-        self.envs: Dict[NeonEnvDiscriminants, NeonEnv] = {}
-
-        self.builders: List[NeonEnvBuilder] = []
-
-    def get_repo_dir(self, disc: NeonEnvDiscriminants) -> Path:
-        # FIXME use discriminants
-        randstr = "".join(
-            random.SystemRandom().choice(string.ascii_uppercase + string.digits) for _ in range(10)
-        )
-        s = Path(f"shared-{randstr}")
-        return self.top_output_dir / s
-
-    def get_or_create(self, disc: NeonEnvDiscriminants) -> NeonEnv:
-        with self.lock:
-            env = self.envs.get(disc)
-            if env is None:
-                builder = NeonEnvBuilder(
-                    repo_dir=self.get_repo_dir(disc),
-                    port_distributor=self.port_distributor,
-                    run_id=self.run_id,
-                    mock_s3_server=self.mock_s3_server,
-                    pg_distrib_dir=self.pg_distrib_dir,
-                    preserve_database_files=self.preserve_database_files,
-                    pg_version=PgVersion("17"),  # FIXME: this should go unused. Pass None?
-                    test_name="shared",  # FIXME
-                    test_output_dir=Path("shared"),  # FIXME
-                    test_overlay_dir=None,
-                    top_output_dir=self.top_output_dir,
-                    neon_binpath=disc.neon_binpath,
-                    pageserver_virtual_file_io_engine=disc.pageserver_virtual_file_io_engine,
-                    pageserver_aux_file_policy=disc.pageserver_aux_file_policy,
-                    pageserver_default_tenant_config_compaction_algorithm=disc.pageserver_default_tenant_config_compaction_algorithm,
-                    # FIXME: only support defaults for these currently
-                    # pageserver_remote_storage
-                    # pageserver_config_override
-                    # num_safekeepers
-                    # num_pageservers
-                    # safekeepers_id_start
-                    # safekeepers_enable_fsync
-                    # auth_enabled
-                    # rust_log_override
-                    # default_branch_name
-                    initial_tenant=None,  # FIXME should go unused
-                    initial_timeline=None,  # FIXME should go unused
-                    # safekeeper_extra_opts: Optional[list[str]] = None,
-                    # storage_controller_port_override: Optional[int] = None,
-                    # pageserver_io_buffer_alignment: Optional[int] = None,
-                )
-                env = builder.init_start()
-
-                self.envs[disc] = env
-        return env
-
-    def __enter__(self) -> "NeonSharedEnvs":
-        return self
-
-    def __exit__(
-        self,
-        exc_type: Optional[Type[BaseException]],
-        exc: Optional[BaseException],
-        tb: Optional[TracebackType],
-    ):
-        for env in self.envs.values():
-            env.stop(immediate=True)
-
-
-@pytest.fixture(scope="session")
-def shared_environments(
-    # session fixtures
-    port_distributor: PortDistributor,
-    run_id: uuid.UUID,
-    mock_s3_server: MockS3Server,
-    pg_distrib_dir: Path,
-    top_output_dir: Path,
-    pytestconfig: Config,
-) -> Iterator[NeonSharedEnvs]:
-    with NeonSharedEnvs(
-        port_distributor=port_distributor,
-        run_id=run_id,
-        mock_s3_server=mock_s3_server,
-        pg_distrib_dir=pg_distrib_dir,
-        top_output_dir=top_output_dir,
-        # rust_log_override=rust_log_override, # FIXME
-        preserve_database_files=cast(bool, pytestconfig.getoption("--preserve-database-files")),
-    ) as envs:
-        yield envs
-
-
-@pytest.fixture(scope="function")
-def shared_env(
-    request: FixtureRequest,
-    # session fixture holding all the envs
-    shared_environments: NeonSharedEnvs,
-    # other session fixtures
-    port_distributor: PortDistributor,
-    mock_s3_server: MockS3Server,
-    run_id: uuid.UUID,
-    top_output_dir: Path,
-    pg_distrib_dir: Path,
-    # these define the env to use
-    neon_binpath: Path,
-    pageserver_virtual_file_io_engine: str,
-    pageserver_aux_file_policy: Optional[AuxFileStore],
-    pageserver_default_tenant_config_compaction_algorithm: Optional[Dict[str, Any]],
-    pageserver_io_buffer_alignment: Optional[int],
-) -> NeonEnv:
-    disc = NeonEnvDiscriminants(
-        neon_binpath=neon_binpath,
-        # FIXME: There's no difference in e.g. having pageserver_virtual_file_io_engine=None, and
-        # explicitly specifying whatever the default is. We could share those envs.
-        pageserver_virtual_file_io_engine=pageserver_virtual_file_io_engine,
-        pageserver_aux_file_policy=pageserver_aux_file_policy,
-        pageserver_default_tenant_config_compaction_algorithm=pageserver_default_tenant_config_compaction_algorithm,
-        pageserver_io_buffer_alignment=pageserver_io_buffer_alignment,
-    )
-    return shared_environments.get_or_create(disc)
-
-
-class NeonTestTenant:
-    """
-    An object representing a single Neon tenant, in a shared environment
-
-    Notable functions and fields:
-
-    endpoints - A factory object for creating postgres compute nodes.
-
-    tenant_id - tenant ID of the initial tenant created in the repository
-
-    initial_timeline - timeline ID of the "main" branch
-
-    create_branch() - branch a new timeline from an existing one, returns
-        the new timeline id
-
-    create_timeline() - initializes a new timeline by running initdb, returns
-        the new timeline id
-    """
-
-    def __init__(self, env: NeonEnv):
-        self.tenant_id = TenantId.generate()
-        self.initial_timeline = TimelineId.generate()
-        self.created = False
-
-        self.endpoints = TenantEndpointFactory(self)
-        self.env = env
-
-    def create(
-        self,
-        conf: Optional[Dict[str, Any]] = None,
-        shard_count: Optional[int] = None,
-        shard_stripe_size: Optional[int] = None,
-        placement_policy: Optional[str] = None,
-        aux_file_policy: Optional[AuxFileStore] = None,
-    ):
-        assert not self.created
-        self.env.create_tenant(
-            tenant_id=self.tenant_id,
-            timeline_id=self.initial_timeline,
-            conf=conf,
-            shard_count=shard_count,
-            shard_stripe_size=shard_stripe_size,
-            placement_policy=placement_policy,
-            set_default=False,
-            aux_file_policy=aux_file_policy,
-        )
-        self.created = True
-
-    # Todo: this could be imeplemented
-    # def config_tenant(self, conf: Dict[str, str]):
-
-    def create_branch(
-        self,
-        new_branch_name: str = DEFAULT_BRANCH_NAME,
-        ancestor_branch_name: Optional[str] = None,
-        ancestor_start_lsn: Optional[Lsn] = None,
-        new_timeline_id: Optional[TimelineId] = None,
-    ) -> TimelineId:
-        return self.env.create_branch(
-            new_branch_name=new_branch_name,
-            tenant_id=self.tenant_id,
-            ancestor_branch_name=ancestor_branch_name,
-            ancestor_start_lsn=ancestor_start_lsn,
-            new_timeline_id=new_timeline_id,
-        )
-
-    def create_timeline(
-        self,
-        new_branch_name: str,
-        timeline_id: Optional[TimelineId] = None,
-    ) -> TimelineId:
-        return self.env.create_timeline(new_branch_name=new_branch_name, timeline_id=timeline_id)
-
-
-class TenantEndpointFactory:
-    """An object representing multiple compute endpoints of a single tenant."""
-
-    def __init__(self, tenant: NeonTestTenant):
-        self.tenant = tenant
-        self.num_instances: int = 0
-        self.endpoints: List[Endpoint] = []
-
-    def create(
-        self,
-        branch_name: str,
-        endpoint_id: Optional[str] = None,
-        lsn: Optional[Lsn] = None,
-        hot_standby: bool = False,
-        config_lines: Optional[List[str]] = None,
-    ) -> Endpoint:
-        ep = Endpoint(
-            self.tenant.env,
-            tenant_id=self.tenant.tenant_id,
-            pg_port=self.tenant.env.port_distributor.get_port(),
-            http_port=self.tenant.env.port_distributor.get_port(),
-        )
-
-        endpoint_id = endpoint_id or self.tenant.env.generate_endpoint_id()
-
-        self.num_instances += 1
-        self.endpoints.append(ep)
-
-        return ep.create(
-            branch_name=branch_name,
-            endpoint_id=endpoint_id,
-            lsn=lsn,
-            hot_standby=hot_standby,
-            config_lines=config_lines,
-        )
-
-    # FIXME: extra args for start
-    # remote_ext_config: Optional[str] = None,
-    # basebackup_request_tries: Optional[int] = None,
-    def create_start(self, *args, **kwargs):
-        ep = self.create(*args, **kwargs)
-        ep.start()
-        return ep
-
-    def stop_all(self, fail_on_error=True) -> "TenantEndpointFactory":
-        exception = None
-        for ep in self.endpoints:
-            try:
-                ep.stop()
-            except Exception as e:
-                log.error(f"Failed to stop endpoint {ep.endpoint_id}: {e}")
-                exception = e
-
-        if fail_on_error and exception is not None:
-            raise exception
-
-        return self
-
-    def new_replica(
-        self, origin: Endpoint, endpoint_id: str, config_lines: Optional[List[str]] = None
-    ):
-        branch_name = origin.branch_name
-        assert origin in self.endpoints
-        assert branch_name is not None
-
-        return self.create(
-            branch_name=branch_name,
-            endpoint_id=endpoint_id,
-            lsn=None,
-            hot_standby=True,
-            config_lines=config_lines,
-        )
-
-    def new_replica_start(
-        self, origin: Endpoint, endpoint_id: str, config_lines: Optional[List[str]] = None
-    ):
-        branch_name = origin.branch_name
-        assert origin in self.endpoints
-        assert branch_name is not None
-
-        return self.create_start(
-            branch_name=branch_name,
-            endpoint_id=endpoint_id,
-            lsn=None,
-            hot_standby=True,
-            config_lines=config_lines,
-        )

test_runner/fixtures/pageserver/remote_storage.py

@@ -7,7 +7,7 @@ from pathlib import Path
 from typing import Any, List, Tuple
 
 from fixtures.common_types import TenantId, TimelineId
-from fixtures.neon_fixtures import NeonEnv
+from fixtures.neon_fixtures import NeonEnv, Pagectl
 from fixtures.pageserver.common_types import (
     InvalidFileName,
     parse_layer_file_name,
@@ -35,7 +35,7 @@ def duplicate_one_tenant(env: NeonEnv, template_tenant: TenantId, new_tenant: Te
         for file in tl.iterdir():
             shutil.copy2(file, dst_tl_dir)
             if "__" in file.name:
-                env.pagectl.raw_cli(
+                Pagectl(env).raw_cli(
                     [
                         "layer",
                         "rewrite-summary",

test_runner/performance/pageserver/interactive/test_many_small_tenants.py

@@ -53,7 +53,7 @@ def setup_env(
             "checkpoint_distance": 268435456,
             "image_creation_threshold": 3,
         }
-        template_tenant, template_timeline = env.create_tenant(set_default=True)
+        template_tenant, template_timeline = env.neon_cli.create_tenant(set_default=True)
         env.pageserver.tenant_detach(template_tenant)
         env.pageserver.tenant_attach(template_tenant, config)
         ep = env.endpoints.create_start("main", tenant_id=template_tenant)

test_runner/performance/pageserver/pagebench/test_large_slru_basebackup.py

@@ -81,7 +81,7 @@ def setup_tenant_template(env: NeonEnv, n_txns: int):
         "image_creation_threshold": 3,
     }
 
-    template_tenant, template_timeline = env.create_tenant(set_default=True)
+    template_tenant, template_timeline = env.neon_cli.create_tenant(set_default=True)
     env.pageserver.tenant_detach(template_tenant)
     env.pageserver.tenant_attach(template_tenant, config)
 

test_runner/performance/pageserver/pagebench/test_pageserver_max_throughput_getpage_at_latest_lsn.py

@@ -162,7 +162,7 @@ def setup_tenant_template(env: NeonEnv, pg_bin: PgBin, scale: int):
         "checkpoint_distance": 268435456,
         "image_creation_threshold": 3,
     }
-    template_tenant, template_timeline = env.create_tenant(set_default=True)
+    template_tenant, template_timeline = env.neon_cli.create_tenant(set_default=True)
     env.pageserver.tenant_detach(template_tenant)
     env.pageserver.tenant_attach(template_tenant, config)
     ps_http = env.pageserver.http_client()

test_runner/performance/test_branch_creation.py

@@ -41,7 +41,7 @@ def test_branch_creation_heavy_write(neon_compare: NeonCompare, n_branches: int)
     pg_bin = neon_compare.pg_bin
 
     # Use aggressive GC and checkpoint settings, so GC and compaction happen more often during the test
-    tenant, _ = env.create_tenant(
+    tenant, _ = env.neon_cli.create_tenant(
         conf={
             "gc_period": "5 s",
             "gc_horizon": f"{4 * 1024 ** 2}",
@@ -64,7 +64,7 @@ def test_branch_creation_heavy_write(neon_compare: NeonCompare, n_branches: int)
 
         endpoint.stop()
 
-    env.create_branch("b0", tenant_id=tenant)
+    env.neon_cli.create_branch("b0", tenant_id=tenant)
 
     threads: List[threading.Thread] = []
     threads.append(threading.Thread(target=run_pgbench, args=("b0",), daemon=True))
@@ -78,7 +78,7 @@ def test_branch_creation_heavy_write(neon_compare: NeonCompare, n_branches: int)
         p = random.randint(0, i)
 
         timer = timeit.default_timer()
-        env.create_branch(f"b{i + 1}", ancestor_branch_name=f"b{p}", tenant_id=tenant)
+        env.neon_cli.create_branch(f"b{i + 1}", f"b{p}", tenant_id=tenant)
         dur = timeit.default_timer() - timer
 
         log.info(f"Creating branch b{i+1} took {dur}s")
@@ -104,7 +104,7 @@ def test_branch_creation_many(neon_compare: NeonCompare, n_branches: int, shape:
     # seed the prng so we will measure the same structure every time
     rng = random.Random("2024-02-29")
 
-    env.create_branch("b0")
+    env.neon_cli.create_branch("b0")
 
     endpoint = env.endpoints.create_start("b0")
     neon_compare.pg_bin.run_capture(["pgbench", "-i", "-I", "dtGvp", "-s10", endpoint.connstr()])
@@ -121,7 +121,7 @@ def test_branch_creation_many(neon_compare: NeonCompare, n_branches: int, shape:
 
         timer = timeit.default_timer()
         # each of these uploads to remote storage before completion
-        env.create_branch(f"b{i + 1}", ancestor_branch_name=parent)
+        env.neon_cli.create_branch(f"b{i + 1}", parent)
         dur = timeit.default_timer() - timer
         branch_creation_durations.append(dur)
 
@@ -222,7 +222,7 @@ def wait_and_record_startup_metrics(
 def test_branch_creation_many_relations(neon_compare: NeonCompare):
     env = neon_compare.env
 
-    timeline_id = env.create_branch("root")
+    timeline_id = env.neon_cli.create_branch("root")
 
     endpoint = env.endpoints.create_start("root")
     with closing(endpoint.connect()) as conn:
@@ -238,7 +238,7 @@ def test_branch_creation_many_relations(neon_compare: NeonCompare):
     )
 
     with neon_compare.record_duration("create_branch_time_not_busy_root"):
-        env.create_branch("child_not_busy", ancestor_branch_name="root")
+        env.neon_cli.create_branch("child_not_busy", "root")
 
     # run a concurrent insertion to make the ancestor "busy" during the branch creation
     thread = threading.Thread(
@@ -247,6 +247,6 @@ def test_branch_creation_many_relations(neon_compare: NeonCompare):
     thread.start()
 
     with neon_compare.record_duration("create_branch_time_busy_root"):
-        env.create_branch("child_busy", ancestor_branch_name="root")
+        env.neon_cli.create_branch("child_busy", "root")
 
     thread.join()

test_runner/performance/test_branching.py

@@ -41,7 +41,7 @@ def test_compare_child_and_root_pgbench_perf(neon_compare: NeonCompare):
         )
         neon_compare.zenbenchmark.record_pg_bench_result(branch, res)
 
-    env.create_branch("root")
+    env.neon_cli.create_branch("root")
     endpoint_root = env.endpoints.create_start("root")
     pg_bin.run_capture(["pgbench", "-i", "-I", "dtGvp", endpoint_root.connstr(), "-s10"])
 
@@ -55,14 +55,14 @@ def test_compare_child_and_root_pgbench_perf(neon_compare: NeonCompare):
 
 def test_compare_child_and_root_write_perf(neon_compare: NeonCompare):
     env = neon_compare.env
-    env.create_branch("root")
+    env.neon_cli.create_branch("root")
     endpoint_root = env.endpoints.create_start("root")
 
     endpoint_root.safe_psql(
         "CREATE TABLE foo(key serial primary key, t text default 'foooooooooooooooooooooooooooooooooooooooooooooooooooo')",
     )
 
-    env.create_branch("child", ancestor_branch_name="root")
+    env.neon_cli.create_branch("child", "root")
     endpoint_child = env.endpoints.create_start("child")
 
     with neon_compare.record_duration("root_run_duration"):
@@ -73,7 +73,7 @@ def test_compare_child_and_root_write_perf(neon_compare: NeonCompare):
 
 def test_compare_child_and_root_read_perf(neon_compare: NeonCompare):
     env = neon_compare.env
-    env.create_branch("root")
+    env.neon_cli.create_branch("root")
     endpoint_root = env.endpoints.create_start("root")
 
     endpoint_root.safe_psql_many(
@@ -83,7 +83,7 @@ def test_compare_child_and_root_read_perf(neon_compare: NeonCompare):
         ]
     )
 
-    env.create_branch("child", ancestor_branch_name="root")
+    env.neon_cli.create_branch("child", "root")
     endpoint_child = env.endpoints.create_start("child")
 
     with neon_compare.record_duration("root_run_duration"):

test_runner/performance/test_bulk_tenant_create.py

@@ -26,8 +26,10 @@ def test_bulk_tenant_create(
     for i in range(tenants_count):
         start = timeit.default_timer()
 
-        tenant, _ = env.create_tenant()
-        env.create_timeline(f"test_bulk_tenant_create_{tenants_count}_{i}", tenant_id=tenant)
+        tenant, _ = env.neon_cli.create_tenant()
+        env.neon_cli.create_timeline(
+            f"test_bulk_tenant_create_{tenants_count}_{i}", tenant_id=tenant
+        )
 
         # FIXME: We used to start new safekeepers here. Did that make sense? Should we do it now?
         # if use_safekeepers == 'with_sa':

test_runner/performance/test_bulk_update.py

@@ -16,7 +16,7 @@ def test_bulk_update(neon_env_builder: NeonEnvBuilder, zenbenchmark, fillfactor)
     env = neon_env_builder.init_start()
     n_records = 1000000
 
-    timeline_id = env.create_branch("test_bulk_update")
+    timeline_id = env.neon_cli.create_branch("test_bulk_update")
     tenant_id = env.initial_tenant
     endpoint = env.endpoints.create_start("test_bulk_update")
     cur = endpoint.connect().cursor()

test_runner/performance/test_compaction.py

@@ -17,7 +17,7 @@ def test_compaction(neon_compare: NeonCompare):
     env = neon_compare.env
     pageserver_http = env.pageserver.http_client()
 
-    tenant_id, timeline_id = env.create_tenant(
+    tenant_id, timeline_id = env.neon_cli.create_tenant(
         conf={
             # Disable background GC and compaction, we'll run compaction manually.
             "gc_period": "0s",
@@ -68,7 +68,7 @@ def test_compaction_l0_memory(neon_compare: NeonCompare):
     env = neon_compare.env
     pageserver_http = env.pageserver.http_client()
 
-    tenant_id, timeline_id = env.create_tenant(
+    tenant_id, timeline_id = env.neon_cli.create_tenant(
         conf={
             # Initially disable compaction so that we will build up a stack of L0s
             "compaction_period": "0s",

test_runner/performance/test_gc_feedback.py

@@ -11,7 +11,7 @@ def gc_feedback_impl(neon_env_builder: NeonEnvBuilder, zenbenchmark: NeonBenchma
     env = neon_env_builder.init_start()
     client = env.pageserver.http_client()
 
-    tenant_id, _ = env.create_tenant(
+    tenant_id, _ = env.neon_cli.create_tenant(
         conf={
             # disable default GC and compaction
             "gc_period": "1000 m",
@@ -63,7 +63,7 @@ def gc_feedback_impl(neon_env_builder: NeonEnvBuilder, zenbenchmark: NeonBenchma
             log.info(f"Physical storage size {physical_size}")
         if mode == "with_snapshots":
             if step == n_steps / 2:
-                env.create_branch("child")
+                env.neon_cli.create_branch("child")
 
     max_num_of_deltas_above_image = 0
     max_total_num_of_deltas = 0

test_runner/performance/test_layer_map.py

@@ -15,7 +15,7 @@ def test_layer_map(neon_env_builder: NeonEnvBuilder, zenbenchmark):
     # We want to have a lot of lot of layer files to exercise the layer map. Disable
     # GC, and make checkpoint_distance very small, so that we get a lot of small layer
     # files.
-    tenant, timeline = env.create_tenant(
+    tenant, timeline = env.neon_cli.create_tenant(
         conf={
             "gc_period": "0s",
             "checkpoint_distance": "16384",

test_runner/performance/test_lazy_startup.py

@@ -33,7 +33,7 @@ def test_lazy_startup(slru: str, neon_env_builder: NeonEnvBuilder, zenbenchmark:
     env = neon_env_builder.init_start()
 
     lazy_slru_download = "true" if slru == "lazy" else "false"
-    tenant, _ = env.create_tenant(
+    tenant, _ = env.neon_cli.create_tenant(
         conf={
             "lazy_slru_download": lazy_slru_download,
         }

test_runner/performance/test_sharding_autosplit.py

@@ -85,7 +85,7 @@ def test_sharding_autosplit(neon_env_builder: NeonEnvBuilder, pg_bin: PgBin):
     tenants = {}
     for tenant_id in set(TenantId.generate() for _i in range(0, tenant_count)):
         timeline_id = TimelineId.generate()
-        env.create_tenant(tenant_id, timeline_id, conf=tenant_conf)
+        env.neon_cli.create_tenant(tenant_id, timeline_id, conf=tenant_conf)
         endpoint = env.endpoints.create("main", tenant_id=tenant_id)
         tenants[tenant_id] = TenantState(timeline_id, endpoint)
         endpoint.start()

test_runner/performance/test_startup.py

@@ -27,7 +27,7 @@ def test_startup_simple(neon_env_builder: NeonEnvBuilder, zenbenchmark: NeonBenc
     neon_env_builder.num_safekeepers = 3
     env = neon_env_builder.init_start()
 
-    env.create_branch("test_startup")
+    env.neon_cli.create_branch("test_startup")
 
     endpoint = None
 

test_runner/regress/test_ancestor_branch.py

@@ -12,7 +12,7 @@ def test_ancestor_branch(neon_env_builder: NeonEnvBuilder):
     pageserver_http = env.pageserver.http_client()
 
     # Override defaults: 4M checkpoint_distance, disable background compaction and gc.
-    tenant, _ = env.create_tenant(
+    tenant, _ = env.neon_cli.create_tenant(
         conf={
             "checkpoint_distance": "4194304",
             "gc_period": "0s",
@@ -45,9 +45,7 @@ def test_ancestor_branch(neon_env_builder: NeonEnvBuilder):
     log.info(f"LSN after 100k rows: {lsn_100}")
 
     # Create branch1.
-    env.create_branch(
-        "branch1", ancestor_branch_name="main", ancestor_start_lsn=lsn_100, tenant_id=tenant
-    )
+    env.neon_cli.create_branch("branch1", "main", tenant_id=tenant, ancestor_start_lsn=lsn_100)
     endpoint_branch1 = env.endpoints.create_start("branch1", tenant_id=tenant)
 
     branch1_cur = endpoint_branch1.connect().cursor()
@@ -69,9 +67,7 @@ def test_ancestor_branch(neon_env_builder: NeonEnvBuilder):
     log.info(f"LSN after 200k rows: {lsn_200}")
 
     # Create branch2.
-    env.create_branch(
-        "branch2", ancestor_branch_name="branch1", ancestor_start_lsn=lsn_200, tenant_id=tenant
-    )
+    env.neon_cli.create_branch("branch2", "branch1", tenant_id=tenant, ancestor_start_lsn=lsn_200)
     endpoint_branch2 = env.endpoints.create_start("branch2", tenant_id=tenant)
     branch2_cur = endpoint_branch2.connect().cursor()
 

test_runner/regress/test_attach_tenant_config.py

@@ -41,7 +41,7 @@ def negative_env(neon_env_builder: NeonEnvBuilder) -> Generator[NegativeTests, N
     assert isinstance(env.pageserver_remote_storage, LocalFsStorage)
 
     ps_http = env.pageserver.http_client()
-    (tenant_id, _) = env.create_tenant()
+    (tenant_id, _) = env.neon_cli.create_tenant()
     assert ps_http.tenant_config(tenant_id).tenant_specific_overrides == {}
     config_pre_detach = ps_http.tenant_config(tenant_id)
     assert tenant_id in [TenantId(t["id"]) for t in ps_http.tenant_list()]
@@ -109,7 +109,7 @@ def test_empty_config(positive_env: NeonEnv, content_type: Optional[str]):
     """
     env = positive_env
     ps_http = env.pageserver.http_client()
-    (tenant_id, _) = env.create_tenant()
+    (tenant_id, _) = env.neon_cli.create_tenant()
     assert ps_http.tenant_config(tenant_id).tenant_specific_overrides == {}
     config_pre_detach = ps_http.tenant_config(tenant_id)
     assert tenant_id in [TenantId(t["id"]) for t in ps_http.tenant_list()]
@@ -182,7 +182,7 @@ def test_fully_custom_config(positive_env: NeonEnv):
         fully_custom_config.keys()
     ), "ensure we cover all config options"
 
-    (tenant_id, _) = env.create_tenant()
+    (tenant_id, _) = env.neon_cli.create_tenant()
     ps_http.set_tenant_config(tenant_id, fully_custom_config)
     our_tenant_config = ps_http.tenant_config(tenant_id)
     assert our_tenant_config.tenant_specific_overrides == fully_custom_config

test_runner/regress/test_auth.py

@@ -76,7 +76,7 @@ def test_compute_auth_to_pageserver(neon_env_builder: NeonEnvBuilder):
     env = neon_env_builder.init_start()
 
     branch = "test_compute_auth_to_pageserver"
-    env.create_branch(branch)
+    env.neon_cli.create_branch(branch)
     endpoint = env.endpoints.create_start(branch)
 
     with closing(endpoint.connect()) as conn:
@@ -186,7 +186,7 @@ def test_auth_failures(neon_env_builder: NeonEnvBuilder, auth_enabled: bool):
     env = neon_env_builder.init_start()
 
     branch = f"test_auth_failures_auth_enabled_{auth_enabled}"
-    timeline_id = env.create_branch(branch)
+    timeline_id = env.neon_cli.create_branch(branch)
     env.endpoints.create_start(branch)
 
     tenant_token = env.auth_keys.generate_tenant_token(env.initial_tenant)

test_runner/regress/test_backpressure.py

@@ -98,7 +98,7 @@ def check_backpressure(endpoint: Endpoint, stop_event: threading.Event, polling_
 def test_backpressure_received_lsn_lag(neon_env_builder: NeonEnvBuilder):
     env = neon_env_builder.init_start()
     # Create a branch for us
-    env.create_branch("test_backpressure")
+    env.neon_cli.create_branch("test_backpressure")
 
     endpoint = env.endpoints.create(
         "test_backpressure", config_lines=["max_replication_write_lag=30MB"]

test_runner/regress/test_bad_connection.py

@@ -22,7 +22,7 @@ def test_compute_pageserver_connection_stress(neon_env_builder: NeonEnvBuilder):
     pageserver_http = env.pageserver.http_client()
     pageserver_http.configure_failpoints(("simulated-bad-compute-connection", "50%return(15)"))
 
-    env.create_branch("test_compute_pageserver_connection_stress")
+    env.neon_cli.create_branch("test_compute_pageserver_connection_stress")
     endpoint = env.endpoints.create_start("test_compute_pageserver_connection_stress")
 
     pg_conn = endpoint.connect()

test_runner/regress/test_branch_and_gc.py

@@ -53,7 +53,7 @@ def test_branch_and_gc(neon_simple_env: NeonEnv, build_type: str):
     env = neon_simple_env
     pageserver_http_client = env.pageserver.http_client()
 
-    tenant, timeline_main = env.create_tenant(
+    tenant, timeline_main = env.neon_cli.create_tenant(
         conf={
             # disable background GC
             "gc_period": "0s",
@@ -90,7 +90,7 @@ def test_branch_and_gc(neon_simple_env: NeonEnv, build_type: str):
     pageserver_http_client.timeline_checkpoint(tenant, timeline_main)
     pageserver_http_client.timeline_gc(tenant, timeline_main, lsn2 - lsn1 + 1024)
 
-    env.create_branch(
+    env.neon_cli.create_branch(
         "test_branch", ancestor_branch_name="main", ancestor_start_lsn=lsn1, tenant_id=tenant
     )
     endpoint_branch = env.endpoints.create_start("test_branch", tenant_id=tenant)
@@ -127,7 +127,7 @@ def test_branch_creation_before_gc(neon_simple_env: NeonEnv):
     env.storage_controller.allowed_errors.extend(error_regexes)
 
     # Disable background GC but set the `pitr_interval` to be small, so GC can delete something
-    tenant, _ = env.create_tenant(
+    tenant, _ = env.neon_cli.create_tenant(
         conf={
             # disable background GC
             "gc_period": "0s",
@@ -145,7 +145,7 @@ def test_branch_creation_before_gc(neon_simple_env: NeonEnv):
         }
     )
 
-    b0 = env.create_branch("b0", tenant_id=tenant)
+    b0 = env.neon_cli.create_branch("b0", tenant_id=tenant)
     endpoint0 = env.endpoints.create_start("b0", tenant_id=tenant)
     res = endpoint0.safe_psql_many(
         queries=[
@@ -176,7 +176,7 @@ def test_branch_creation_before_gc(neon_simple_env: NeonEnv):
 
     # The starting LSN is invalid as the corresponding record is scheduled to be removed by in-queue GC.
     with pytest.raises(Exception, match="invalid branch start lsn: .*"):
-        env.create_branch("b1", ancestor_branch_name="b0", ancestor_start_lsn=lsn, tenant_id=tenant)
+        env.neon_cli.create_branch("b1", "b0", tenant_id=tenant, ancestor_start_lsn=lsn)
     # retry the same with the HTTP API, so that we can inspect the status code
     with pytest.raises(TimelineCreate406):
         new_timeline_id = TimelineId.generate()

test_runner/regress/test_branch_behind.py

@@ -23,7 +23,7 @@ def test_branch_behind(neon_env_builder: NeonEnvBuilder):
     env.storage_controller.allowed_errors.extend(error_regexes)
 
     # Branch at the point where only 100 rows were inserted
-    branch_behind_timeline_id = env.create_branch("test_branch_behind")
+    branch_behind_timeline_id = env.neon_cli.create_branch("test_branch_behind")
     endpoint_main = env.endpoints.create_start("test_branch_behind")
 
     main_cur = endpoint_main.connect().cursor()
@@ -58,10 +58,8 @@ def test_branch_behind(neon_env_builder: NeonEnvBuilder):
     log.info(f"LSN after 200100 rows: {lsn_b}")
 
     # Branch at the point where only 100 rows were inserted
-    env.create_branch(
-        "test_branch_behind_hundred",
-        ancestor_branch_name="test_branch_behind",
-        ancestor_start_lsn=lsn_a,
+    env.neon_cli.create_branch(
+        "test_branch_behind_hundred", "test_branch_behind", ancestor_start_lsn=lsn_a
     )
 
     # Insert many more rows. This generates enough WAL to fill a few segments.
@@ -77,10 +75,8 @@ def test_branch_behind(neon_env_builder: NeonEnvBuilder):
     log.info(f"LSN after 400100 rows: {lsn_c}")
 
     # Branch at the point where only 200100 rows were inserted
-    env.create_branch(
-        "test_branch_behind_more",
-        ancestor_branch_name="test_branch_behind",
-        ancestor_start_lsn=lsn_b,
+    env.neon_cli.create_branch(
+        "test_branch_behind_more", "test_branch_behind", ancestor_start_lsn=lsn_b
     )
 
     endpoint_hundred = env.endpoints.create_start("test_branch_behind_hundred")
@@ -101,17 +97,15 @@ def test_branch_behind(neon_env_builder: NeonEnvBuilder):
     pageserver_http = env.pageserver.http_client()
 
     # branch at segment boundary
-    env.create_branch(
-        "test_branch_segment_boundary",
-        ancestor_branch_name="test_branch_behind",
-        ancestor_start_lsn=Lsn("0/3000000"),
+    env.neon_cli.create_branch(
+        "test_branch_segment_boundary", "test_branch_behind", ancestor_start_lsn=Lsn("0/3000000")
     )
     endpoint = env.endpoints.create_start("test_branch_segment_boundary")
     assert endpoint.safe_psql("SELECT 1")[0][0] == 1
 
     # branch at pre-initdb lsn (from main branch)
     with pytest.raises(Exception, match="invalid branch start lsn: .*"):
-        env.create_branch("test_branch_preinitdb", ancestor_start_lsn=Lsn("0/42"))
+        env.neon_cli.create_branch("test_branch_preinitdb", ancestor_start_lsn=Lsn("0/42"))
     # retry the same with the HTTP API, so that we can inspect the status code
     with pytest.raises(TimelineCreate406):
         new_timeline_id = TimelineId.generate()
@@ -122,10 +116,8 @@ def test_branch_behind(neon_env_builder: NeonEnvBuilder):
 
     # branch at pre-ancestor lsn
     with pytest.raises(Exception, match="less than timeline ancestor lsn"):
-        env.create_branch(
-            "test_branch_preinitdb",
-            ancestor_branch_name="test_branch_behind",
-            ancestor_start_lsn=Lsn("0/42"),
+        env.neon_cli.create_branch(
+            "test_branch_preinitdb", "test_branch_behind", ancestor_start_lsn=Lsn("0/42")
         )
     # retry the same with the HTTP API, so that we can inspect the status code
     with pytest.raises(TimelineCreate406):
@@ -147,10 +139,8 @@ def test_branch_behind(neon_env_builder: NeonEnvBuilder):
     print_gc_result(gc_result)
     with pytest.raises(Exception, match="invalid branch start lsn: .*"):
         # this gced_lsn is pretty random, so if gc is disabled this woudln't fail
-        env.create_branch(
-            "test_branch_create_fail",
-            ancestor_branch_name="test_branch_behind",
-            ancestor_start_lsn=gced_lsn,
+        env.neon_cli.create_branch(
+            "test_branch_create_fail", "test_branch_behind", ancestor_start_lsn=gced_lsn
         )
     # retry the same with the HTTP API, so that we can inspect the status code
     with pytest.raises(TimelineCreate406):

test_runner/regress/test_branching.py

@@ -38,7 +38,7 @@ def test_branching_with_pgbench(
     env = neon_simple_env
 
     # Use aggressive GC and checkpoint settings, so that we also exercise GC during the test
-    tenant, _ = env.create_tenant(
+    tenant, _ = env.neon_cli.create_tenant(
         conf={
             "gc_period": "5 s",
             "gc_horizon": f"{1024 ** 2}",
@@ -55,7 +55,7 @@ def test_branching_with_pgbench(
         pg_bin.run_capture(["pgbench", "-i", "-I", "dtGvp", f"-s{scale}", connstr])
         pg_bin.run_capture(["pgbench", "-T15", connstr])
 
-    env.create_branch("b0", tenant_id=tenant)
+    env.neon_cli.create_branch("b0", tenant_id=tenant)
     endpoints: List[Endpoint] = []
     endpoints.append(env.endpoints.create_start("b0", tenant_id=tenant))
 
@@ -84,9 +84,9 @@ def test_branching_with_pgbench(
             threads = []
 
         if ty == "cascade":
-            env.create_branch(f"b{i + 1}", ancestor_branch_name=f"b{i}", tenant_id=tenant)
+            env.neon_cli.create_branch(f"b{i + 1}", f"b{i}", tenant_id=tenant)
         else:
-            env.create_branch(f"b{i + 1}", ancestor_branch_name="b0", tenant_id=tenant)
+            env.neon_cli.create_branch(f"b{i + 1}", "b0", tenant_id=tenant)
 
         endpoints.append(env.endpoints.create_start(f"b{i + 1}", tenant_id=tenant))
 
@@ -120,7 +120,7 @@ def test_branching_unnormalized_start_lsn(neon_simple_env: NeonEnv, pg_bin: PgBi
 
     env = neon_simple_env
 
-    env.create_branch("b0")
+    env.neon_cli.create_branch("b0")
     endpoint0 = env.endpoints.create_start("b0")
 
     pg_bin.run_capture(["pgbench", "-i", endpoint0.connstr()])
@@ -133,7 +133,7 @@ def test_branching_unnormalized_start_lsn(neon_simple_env: NeonEnv, pg_bin: PgBi
     start_lsn = Lsn((int(curr_lsn) - XLOG_BLCKSZ) // XLOG_BLCKSZ * XLOG_BLCKSZ)
 
     log.info(f"Branching b1 from b0 starting at lsn {start_lsn}...")
-    env.create_branch("b1", ancestor_branch_name="b0", ancestor_start_lsn=start_lsn)
+    env.neon_cli.create_branch("b1", "b0", ancestor_start_lsn=start_lsn)
     endpoint1 = env.endpoints.create_start("b1")
 
     pg_bin.run_capture(["pgbench", "-i", endpoint1.connstr()])
@@ -173,7 +173,7 @@ def test_cannot_create_endpoint_on_non_uploaded_timeline(neon_env_builder: NeonE
 
         wait_until_paused(env, "before-upload-index-pausable")
 
-        env.neon_cli.mappings_map_branch(initial_branch, env.initial_tenant, env.initial_timeline)
+        env.neon_cli.map_branch(initial_branch, env.initial_tenant, env.initial_timeline)
 
         with pytest.raises(RuntimeError, match="ERROR: Not found: Timeline"):
             env.endpoints.create_start(
@@ -432,7 +432,9 @@ def test_branching_while_stuck_find_gc_cutoffs(neon_env_builder: NeonEnvBuilder)
 
         wait_until_paused(env, failpoint)
 
-        env.create_branch("branch", ancestor_branch_name="main")
+        env.neon_cli.create_branch(
+            tenant_id=env.initial_tenant, ancestor_branch_name="main", new_branch_name="branch"
+        )
 
         client.configure_failpoints((failpoint, "off"))
 

test_runner/regress/test_broken_timeline.py

@@ -34,7 +34,7 @@ def test_local_corruption(neon_env_builder: NeonEnvBuilder):
     tenant_timelines: List[Tuple[TenantId, TimelineId, Endpoint]] = []
 
     for _ in range(3):
-        tenant_id, timeline_id = env.create_tenant()
+        tenant_id, timeline_id = env.neon_cli.create_tenant()
 
         endpoint = env.endpoints.create_start("main", tenant_id=tenant_id)
         with endpoint.cursor() as cur:
@@ -84,11 +84,13 @@ def test_local_corruption(neon_env_builder: NeonEnvBuilder):
 def test_create_multiple_timelines_parallel(neon_simple_env: NeonEnv):
     env = neon_simple_env
 
-    tenant_id, _ = env.create_tenant()
+    tenant_id, _ = env.neon_cli.create_tenant()
 
     with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
         futures = [
-            executor.submit(env.create_timeline, f"test-create-multiple-timelines-{i}", tenant_id)
+            executor.submit(
+                env.neon_cli.create_timeline, f"test-create-multiple-timelines-{i}", tenant_id
+            )
             for i in range(4)
         ]
         for future in futures:
@@ -109,7 +111,7 @@ def test_timeline_init_break_before_checkpoint(neon_env_builder: NeonEnvBuilder)
     tenant_id = env.initial_tenant
 
     timelines_dir = env.pageserver.timeline_dir(tenant_id)
-    old_tenant_timelines = env.neon_cli.timeline_list(tenant_id)
+    old_tenant_timelines = env.neon_cli.list_timelines(tenant_id)
     initial_timeline_dirs = [d for d in timelines_dir.iterdir()]
 
     # Introduce failpoint during timeline init (some intermediate files are on disk), before it's checkpointed.
@@ -121,7 +123,7 @@ def test_timeline_init_break_before_checkpoint(neon_env_builder: NeonEnvBuilder)
     env.pageserver.restart(immediate=True)
 
     # Creating the timeline didn't finish. The other timelines on tenant should still be present and work normally.
-    new_tenant_timelines = env.neon_cli.timeline_list(tenant_id)
+    new_tenant_timelines = env.neon_cli.list_timelines(tenant_id)
     assert (
         new_tenant_timelines == old_tenant_timelines
     ), f"Pageserver after restart should ignore non-initialized timelines for tenant {tenant_id}"
@@ -149,11 +151,11 @@ def test_timeline_init_break_before_checkpoint_recreate(
         ]
     )
 
-    env.create_tenant(env.initial_tenant)
+    env.neon_cli.create_tenant(env.initial_tenant)
     tenant_id = env.initial_tenant
 
     timelines_dir = env.pageserver.timeline_dir(tenant_id)
-    old_tenant_timelines = env.neon_cli.timeline_list(tenant_id)
+    old_tenant_timelines = env.neon_cli.list_timelines(tenant_id)
     initial_timeline_dirs = [d for d in timelines_dir.iterdir()]
 
     # Some fixed timeline ID (like control plane does)
@@ -174,7 +176,7 @@ def test_timeline_init_break_before_checkpoint_recreate(
     env.pageserver.restart(immediate=True)
 
     # Creating the timeline didn't finish. The other timelines on tenant should still be present and work normally.
-    new_tenant_timelines = env.neon_cli.timeline_list(tenant_id)
+    new_tenant_timelines = env.neon_cli.list_timelines(tenant_id)
     assert (
         new_tenant_timelines == old_tenant_timelines
     ), f"Pageserver after restart should ignore non-initialized timelines for tenant {tenant_id}"
@@ -199,7 +201,7 @@ def test_timeline_create_break_after_dir_creation(neon_env_builder: NeonEnvBuild
     tenant_id = env.initial_tenant
 
     timelines_dir = env.pageserver.timeline_dir(tenant_id)
-    old_tenant_timelines = env.neon_cli.timeline_list(tenant_id)
+    old_tenant_timelines = env.neon_cli.list_timelines(tenant_id)
     initial_timeline_dirs = [d for d in timelines_dir.iterdir()]
 
     # Introduce failpoint when creating a new timeline, right after creating its directory
@@ -209,7 +211,7 @@ def test_timeline_create_break_after_dir_creation(neon_env_builder: NeonEnvBuild
 
     # Creating the timeline didn't finish. The other timelines on tenant should still be present and work normally.
     # "New" timeline is not present in the list, allowing pageserver to retry the same request
-    new_tenant_timelines = env.neon_cli.timeline_list(tenant_id)
+    new_tenant_timelines = env.neon_cli.list_timelines(tenant_id)
     assert (
         new_tenant_timelines == old_tenant_timelines
     ), f"Pageserver after restart should ignore non-initialized timelines for tenant {tenant_id}"

test_runner/regress/test_change_pageserver.py

@@ -34,7 +34,7 @@ def test_change_pageserver(neon_env_builder: NeonEnvBuilder, make_httpserver):
         ignore_notify
     )
 
-    env.create_branch("test_change_pageserver")
+    env.neon_cli.create_branch("test_change_pageserver")
     endpoint = env.endpoints.create_start("test_change_pageserver")
 
     # Put this tenant into a dual-attached state

test_runner/regress/test_clog_truncate.py

@@ -2,14 +2,16 @@ import os
 import time
 
 from fixtures.log_helper import log
-from fixtures.neon_tenant import NeonTestTenant
+from fixtures.neon_fixtures import NeonEnv
 from fixtures.utils import query_scalar
 
 
 #
 # Test compute node start after clog truncation
 #
-def test_clog_truncate(neon_tenant: NeonTestTenant):
+def test_clog_truncate(neon_simple_env: NeonEnv):
+    env = neon_simple_env
+
     # set aggressive autovacuum to make sure that truncation will happen
     config = [
         "autovacuum_max_workers=10",
@@ -21,7 +23,7 @@ def test_clog_truncate(neon_tenant: NeonTestTenant):
         "autovacuum_freeze_max_age=100000",
     ]
 
-    endpoint = neon_tenant.endpoints.create_start("main", config_lines=config)
+    endpoint = env.endpoints.create_start("main", config_lines=config)
 
     # Install extension containing function needed for test
     endpoint.safe_psql("CREATE EXTENSION neon_test_utils")
@@ -54,12 +56,10 @@ def test_clog_truncate(neon_tenant: NeonTestTenant):
 
     # create new branch after clog truncation and start a compute node on it
     log.info(f"create branch at lsn_after_truncation {lsn_after_truncation}")
-    neon_tenant.create_branch(
-        "test_clog_truncate_new",
-        ancestor_branch_name="main",
-        ancestor_start_lsn=lsn_after_truncation,
+    env.neon_cli.create_branch(
+        "test_clog_truncate_new", "main", ancestor_start_lsn=lsn_after_truncation
     )
-    endpoint2 = neon_tenant.endpoints.create_start("test_clog_truncate_new")
+    endpoint2 = env.endpoints.create_start("test_clog_truncate_new")
 
     # check that new node doesn't contain truncated segment
     pg_xact_0000_path_new = os.path.join(endpoint2.pg_xact_dir_path(), "0000")

test_runner/regress/test_close_fds.py

@@ -23,7 +23,7 @@ def test_lsof_pageserver_pid(neon_simple_env: NeonEnv):
     env = neon_simple_env
 
     def start_workload():
-        env.create_branch("test_lsof_pageserver_pid")
+        env.neon_cli.create_branch("test_lsof_pageserver_pid")
         endpoint = env.endpoints.create_start("test_lsof_pageserver_pid")
         with closing(endpoint.connect()) as conn:
             with conn.cursor() as cur:

test_runner/regress/test_compatibility.py

@@ -517,7 +517,7 @@ def test_historic_storage_formats(
     assert metadata_summary["tenant_count"] >= 1
     assert metadata_summary["timeline_count"] >= 1
 
-    env.neon_cli.tenant_import(dataset.tenant_id)
+    env.neon_cli.import_tenant(dataset.tenant_id)
 
     # Discover timelines
     timelines = env.pageserver.http_client().timeline_list(dataset.tenant_id)

test_runner/regress/test_compute_catalog.py

@@ -1,9 +1,11 @@
 import requests
-from fixtures.neon_tenant import NeonTestTenant
+from fixtures.neon_fixtures import NeonEnv
 
 
-def test_compute_catalog(neon_tenant: NeonTestTenant):
-    endpoint = neon_tenant.endpoints.create_start("main", config_lines=["log_min_messages=debug1"])
+def test_compute_catalog(neon_simple_env: NeonEnv):
+    env = neon_simple_env
+
+    endpoint = env.endpoints.create_start("main", config_lines=["log_min_messages=debug1"])
     client = endpoint.http_client()
 
     objects = client.dbs_and_roles()

test_runner/regress/test_config.py

@@ -38,7 +38,7 @@ def test_safekeepers_reconfigure_reorder(
 ):
     neon_env_builder.num_safekeepers = 3
     env = neon_env_builder.init_start()
-    env.create_branch("test_safekeepers_reconfigure_reorder")
+    env.neon_cli.create_branch("test_safekeepers_reconfigure_reorder")
 
     endpoint = env.endpoints.create_start("test_safekeepers_reconfigure_reorder")
 

test_runner/regress/test_crafted_wal_end.py

@@ -1,7 +1,6 @@
 import pytest
 from fixtures.log_helper import log
-from fixtures.neon_cli import WalCraft
-from fixtures.neon_fixtures import NeonEnvBuilder
+from fixtures.neon_fixtures import NeonEnvBuilder, WalCraft
 
 # Restart nodes with WAL end having specially crafted shape, like last record
 # crossing segment boundary, to test decoding issues.
@@ -19,7 +18,7 @@ from fixtures.neon_fixtures import NeonEnvBuilder
 )
 def test_crafted_wal_end(neon_env_builder: NeonEnvBuilder, wal_type: str):
     env = neon_env_builder.init_start()
-    env.create_branch("test_crafted_wal_end")
+    env.neon_cli.create_branch("test_crafted_wal_end")
     env.pageserver.allowed_errors.extend(
         [
             # seems like pageserver stop triggers these
@@ -28,7 +27,7 @@ def test_crafted_wal_end(neon_env_builder: NeonEnvBuilder, wal_type: str):
     )
 
     endpoint = env.endpoints.create("test_crafted_wal_end")
-    wal_craft = WalCraft(extra_env=None, binpath=env.neon_binpath)
+    wal_craft = WalCraft(env)
     endpoint.config(wal_craft.postgres_config())
     endpoint.start()
     res = endpoint.safe_psql_many(

test_runner/regress/test_createdropdb.py

@@ -4,7 +4,6 @@ import pathlib
 import pytest
 from fixtures.log_helper import log
 from fixtures.neon_fixtures import NeonEnv, check_restored_datadir_content
-from fixtures.neon_tenant import NeonTestTenant
 from fixtures.pg_version import PgVersion
 from fixtures.utils import query_scalar
 
@@ -13,17 +12,18 @@ from fixtures.utils import query_scalar
 # Test CREATE DATABASE when there have been relmapper changes
 #
 @pytest.mark.parametrize("strategy", ["file_copy", "wal_log"])
-def test_createdb(neon_tenant: NeonTestTenant, strategy: str, pg_version: PgVersion):
-    if pg_version == PgVersion.V14 and strategy == "wal_log":
+def test_createdb(neon_simple_env: NeonEnv, strategy: str):
+    env = neon_simple_env
+    if env.pg_version == PgVersion.V14 and strategy == "wal_log":
         pytest.skip("wal_log strategy not supported on PostgreSQL 14")
 
-    endpoint = neon_tenant.endpoints.create_start("main")
+    endpoint = env.endpoints.create_start("main")
 
     with endpoint.cursor() as cur:
         # Cause a 'relmapper' change in the original branch
         cur.execute("VACUUM FULL pg_class")
 
-        if pg_version == PgVersion.V14:
+        if env.pg_version == PgVersion.V14:
             cur.execute("CREATE DATABASE foodb")
         else:
             cur.execute(f"CREATE DATABASE foodb STRATEGY={strategy}")
@@ -31,8 +31,8 @@ def test_createdb(neon_tenant: NeonTestTenant, strategy: str, pg_version: PgVers
         lsn = query_scalar(cur, "SELECT pg_current_wal_insert_lsn()")
 
     # Create a branch
-    neon_tenant.create_branch("test_createdb2", ancestor_branch_name="main", ancestor_start_lsn=lsn)
-    endpoint2 = neon_tenant.endpoints.create_start("test_createdb2")
+    env.neon_cli.create_branch("test_createdb2", "main", ancestor_start_lsn=lsn)
+    endpoint2 = env.endpoints.create_start("test_createdb2")
 
     # Test that you can connect to the new database on both branches
     for db in (endpoint, endpoint2):
@@ -77,14 +77,10 @@ def test_dropdb(neon_simple_env: NeonEnv, test_output_dir):
         lsn_after_drop = query_scalar(cur, "SELECT pg_current_wal_insert_lsn()")
 
     # Create two branches before and after database drop.
-    env.create_branch(
-        "test_before_dropdb", ancestor_branch_name="main", ancestor_start_lsn=lsn_before_drop
-    )
+    env.neon_cli.create_branch("test_before_dropdb", "main", ancestor_start_lsn=lsn_before_drop)
     endpoint_before = env.endpoints.create_start("test_before_dropdb")
 
-    env.create_branch(
-        "test_after_dropdb", ancestor_branch_name="main", ancestor_start_lsn=lsn_after_drop
-    )
+    env.neon_cli.create_branch("test_after_dropdb", "main", ancestor_start_lsn=lsn_after_drop)
     endpoint_after = env.endpoints.create_start("test_after_dropdb")
 
     # Test that database exists on the branch before drop

test_runner/regress/test_createuser.py

@@ -1,12 +1,13 @@
-from fixtures.neon_tenant import NeonTestTenant
+from fixtures.neon_fixtures import NeonEnv
 from fixtures.utils import query_scalar
 
 
 #
 # Test CREATE USER to check shared catalog restore
 #
-def test_createuser(neon_tenant: NeonTestTenant):
-    endpoint = neon_tenant.endpoints.create_start("main")
+def test_createuser(neon_simple_env: NeonEnv):
+    env = neon_simple_env
+    endpoint = env.endpoints.create_start("main")
 
     with endpoint.cursor() as cur:
         # Cause a 'relmapper' change in the original branch
@@ -17,10 +18,8 @@ def test_createuser(neon_tenant: NeonTestTenant):
         lsn = query_scalar(cur, "SELECT pg_current_wal_insert_lsn()")
 
     # Create a branch
-    neon_tenant.create_branch(
-        "test_createuser2", ancestor_branch_name="main", ancestor_start_lsn=lsn
-    )
-    endpoint2 = neon_tenant.endpoints.create_start("test_createuser2")
+    env.neon_cli.create_branch("test_createuser2", "main", ancestor_start_lsn=lsn)
+    endpoint2 = env.endpoints.create_start("test_createuser2")
 
     # Test that you can connect to new branch as a new user
     assert endpoint2.safe_psql("select current_user", user="testuser") == [("testuser",)]

test_runner/regress/test_ddl_forwarding.py

@@ -4,8 +4,7 @@ from typing import Any, Dict, List, Optional, Tuple, Type
 import psycopg2
 import pytest
 from fixtures.log_helper import log
-from fixtures.neon_fixtures import VanillaPostgres
-from fixtures.neon_tenant import NeonTestTenant
+from fixtures.neon_fixtures import NeonEnv, VanillaPostgres
 from pytest_httpserver import HTTPServer
 from werkzeug.wrappers.request import Request
 from werkzeug.wrappers.response import Response
@@ -289,8 +288,9 @@ def assert_db_connlimit(endpoint: Any, db_name: str, connlimit: int, msg: str):
 # 2. User can ignore, then compute_ctl will drop invalid databases
 #    automatically during full configuration
 # Here we test the latter. The first one is tested in test_ddl_forwarding
-def test_ddl_forwarding_invalid_db(neon_tenant: NeonTestTenant):
-    endpoint = neon_tenant.endpoints.create_start(
+def test_ddl_forwarding_invalid_db(neon_simple_env: NeonEnv):
+    env = neon_simple_env
+    endpoint = env.endpoints.create_start(
         "main",
         # Some non-existent url
         config_lines=["neon.console_url=http://localhost:9999/unknown/api/v0/roles_and_databases"],

test_runner/regress/test_disk_usage_eviction.py

@@ -59,11 +59,11 @@ def test_min_resident_size_override_handling(
     env.pageserver.stop()
     env.pageserver.start()
 
-    tenant_id, _ = env.create_tenant()
+    tenant_id, _ = env.neon_cli.create_tenant()
     assert_overrides(tenant_id, config_level_override)
 
     # Also ensure that specifying the paramter to create_tenant works, in addition to http-level recconfig.
-    tenant_id, _ = env.create_tenant(conf={"min_resident_size_override": "100"})
+    tenant_id, _ = env.neon_cli.create_tenant(conf={"min_resident_size_override": "100"})
     assert_config(tenant_id, 100, 100)
     ps_http.set_tenant_config(tenant_id, {})
     assert_config(tenant_id, None, config_level_override)
@@ -280,7 +280,7 @@ def _eviction_env(
 def pgbench_init_tenant(
     layer_size: int, scale: int, env: NeonEnv, pg_bin: PgBin
 ) -> Tuple[TenantId, TimelineId]:
-    tenant_id, timeline_id = env.create_tenant(
+    tenant_id, timeline_id = env.neon_cli.create_tenant(
         conf={
             "gc_period": "0s",
             "compaction_period": "0s",

test_runner/regress/test_download_extensions.py

@@ -81,7 +81,7 @@ def test_remote_extensions(
     # Start a compute node with remote_extension spec
     # and check that it can download the extensions and use them to CREATE EXTENSION.
     env = neon_env_builder_local.init_start()
-    env.create_branch("test_remote_extensions")
+    env.neon_cli.create_branch("test_remote_extensions")
     endpoint = env.endpoints.create(
         "test_remote_extensions",
         config_lines=["log_min_messages=debug3"],