a little more type-safety, a little more verbose...

.github/workflows/_build-and-test-locally.yml

@@ -53,6 +53,20 @@ jobs:
       BUILD_TAG: ${{ inputs.build-tag }}
 
     steps:
+      - name: Fix git ownership
+        run: |
+          # Workaround for `fatal: detected dubious ownership in repository at ...`
+          #
+          # Use both ${{ github.workspace }} and ${GITHUB_WORKSPACE} because they're different on host and in containers
+          #   Ref https://github.com/actions/checkout/issues/785
+          #
+          git config --global --add safe.directory ${{ github.workspace }}
+          git config --global --add safe.directory ${GITHUB_WORKSPACE}
+          for r in 14 15 16 17; do
+            git config --global --add safe.directory "${{ github.workspace }}/vendor/postgres-v$r"
+            git config --global --add safe.directory "${GITHUB_WORKSPACE}/vendor/postgres-v$r"
+          done
+
       - uses: actions/checkout@v4
         with:
           submodules: true

.github/workflows/build_and_test.yml

@@ -1078,6 +1078,20 @@ jobs:
     runs-on: [ self-hosted, small ]
     container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/ansible:latest
     steps:
+      - name: Fix git ownership
+        run: |
+          # Workaround for `fatal: detected dubious ownership in repository at ...`
+          #
+          # Use both ${{ github.workspace }} and ${GITHUB_WORKSPACE} because they're different on host and in containers
+          #   Ref https://github.com/actions/checkout/issues/785
+          #
+          git config --global --add safe.directory ${{ github.workspace }}
+          git config --global --add safe.directory ${GITHUB_WORKSPACE}
+          for r in 14 15 16 17; do
+            git config --global --add safe.directory "${{ github.workspace }}/vendor/postgres-v$r"
+            git config --global --add safe.directory "${GITHUB_WORKSPACE}/vendor/postgres-v$r"
+          done
+
       - uses: actions/checkout@v4
 
       - name: Trigger deploy workflow

.gitignore

@@ -6,8 +6,6 @@ __pycache__/
 test_output/
 .vscode
 .idea
-*.swp
-tags
 neon.iml
 /.neon
 /integration_tests/.neon

Cargo.lock

@@ -6272,7 +6272,7 @@ dependencies = [
 [[package]]
 name = "tokio-epoll-uring"
 version = "0.1.0"
-source = "git+https://github.com/neondatabase/tokio-epoll-uring.git?branch=main#cb2dcea2058034bc209e7917b01c5097712a3168"
+source = "git+https://github.com/neondatabase/tokio-epoll-uring.git?branch=main#08ccfa94ff5507727bf4d8d006666b5b192e04c6"
 dependencies = [
  "futures",
  "nix 0.26.4",
@@ -6788,7 +6788,7 @@ dependencies = [
 [[package]]
 name = "uring-common"
 version = "0.1.0"
-source = "git+https://github.com/neondatabase/tokio-epoll-uring.git?branch=main#cb2dcea2058034bc209e7917b01c5097712a3168"
+source = "git+https://github.com/neondatabase/tokio-epoll-uring.git?branch=main#08ccfa94ff5507727bf4d8d006666b5b192e04c6"
 dependencies = [
  "bytes",
  "io-uring",

control_plane/src/bin/neon_local.rs

@@ -1073,10 +1073,10 @@ async fn handle_tenant(subcmd: &TenantCmd, env: &mut local_env::LocalEnv) -> any
                     tenant_id,
                     TimelineCreateRequest {
                         new_timeline_id,
-                        mode: pageserver_api::models::TimelineCreateRequestMode::Bootstrap {
-                            existing_initdb_timeline_id: None,
-                            pg_version: Some(args.pg_version),
-                        },
+                        ancestor_timeline_id: None,
+                        ancestor_start_lsn: None,
+                        existing_initdb_timeline_id: None,
+                        pg_version: Some(args.pg_version),
                     },
                 )
                 .await?;
@@ -1133,10 +1133,10 @@ async fn handle_timeline(cmd: &TimelineCmd, env: &mut local_env::LocalEnv) -> Re
             let storage_controller = StorageController::from_env(env);
             let create_req = TimelineCreateRequest {
                 new_timeline_id,
-                mode: pageserver_api::models::TimelineCreateRequestMode::Bootstrap {
-                    existing_initdb_timeline_id: None,
-                    pg_version: Some(args.pg_version),
-                },
+                ancestor_timeline_id: None,
+                existing_initdb_timeline_id: None,
+                ancestor_start_lsn: None,
+                pg_version: Some(args.pg_version),
             };
             let timeline_info = storage_controller
                 .tenant_timeline_create(tenant_id, create_req)
@@ -1189,11 +1189,10 @@ async fn handle_timeline(cmd: &TimelineCmd, env: &mut local_env::LocalEnv) -> Re
             let storage_controller = StorageController::from_env(env);
             let create_req = TimelineCreateRequest {
                 new_timeline_id,
-                mode: pageserver_api::models::TimelineCreateRequestMode::Branch {
-                    ancestor_timeline_id,
-                    ancestor_start_lsn: start_lsn,
-                    pg_version: None,
-                },
+                ancestor_timeline_id: Some(ancestor_timeline_id),
+                existing_initdb_timeline_id: None,
+                ancestor_start_lsn: start_lsn,
+                pg_version: None,
             };
             let timeline_info = storage_controller
                 .tenant_timeline_create(tenant_id, create_req)

control_plane/src/pageserver.rs

@@ -529,6 +529,28 @@ impl PageServerNode {
         Ok(self.http_client.list_timelines(*tenant_shard_id).await?)
     }
 
+    pub async fn timeline_create(
+        &self,
+        tenant_shard_id: TenantShardId,
+        new_timeline_id: TimelineId,
+        ancestor_start_lsn: Option<Lsn>,
+        ancestor_timeline_id: Option<TimelineId>,
+        pg_version: Option<u32>,
+        existing_initdb_timeline_id: Option<TimelineId>,
+    ) -> anyhow::Result<TimelineInfo> {
+        let req = models::TimelineCreateRequest {
+            new_timeline_id,
+            ancestor_start_lsn,
+            ancestor_timeline_id,
+            pg_version,
+            existing_initdb_timeline_id,
+        };
+        Ok(self
+            .http_client
+            .timeline_create(tenant_shard_id, &req)
+            .await?)
+    }
+
     /// Import a basebackup prepared using either:
     /// a) `pg_basebackup -F tar`, or
     /// b) The `fullbackup` pageserver endpoint

libs/metrics/src/lib.rs

@@ -19,7 +19,6 @@ use once_cell::sync::Lazy;
 use prometheus::core::{
     Atomic, AtomicU64, Collector, GenericCounter, GenericCounterVec, GenericGauge, GenericGaugeVec,
 };
-pub use prometheus::local::LocalHistogram;
 pub use prometheus::opts;
 pub use prometheus::register;
 pub use prometheus::Error;

libs/pageserver_api/src/models.rs

@@ -211,30 +211,13 @@ pub enum TimelineState {
 #[derive(Serialize, Deserialize, Clone)]
 pub struct TimelineCreateRequest {
     pub new_timeline_id: TimelineId,
-    #[serde(flatten)]
-    pub mode: TimelineCreateRequestMode,
-}
-
-#[derive(Serialize, Deserialize, Clone)]
-#[serde(untagged)]
-pub enum TimelineCreateRequestMode {
-    Branch {
-        ancestor_timeline_id: TimelineId,
-        #[serde(default)]
-        ancestor_start_lsn: Option<Lsn>,
-        // TODO: cplane sets this, but, the branching code always
-        // inherits the ancestor's pg_version. Earlier code wasn't
-        // using a flattened enum, so, it was an accepted field, and
-        // we continue to accept it by having it here.
-        pg_version: Option<u32>,
-    },
-    // NB: Bootstrap is all-optional, and thus the serde(untagged) will cause serde to stop at Bootstrap.
-    // (serde picks the first matching enum variant, in declaration order).
-    Bootstrap {
-        #[serde(default)]
-        existing_initdb_timeline_id: Option<TimelineId>,
-        pg_version: Option<u32>,
-    },
+    #[serde(default)]
+    pub ancestor_timeline_id: Option<TimelineId>,
+    #[serde(default)]
+    pub existing_initdb_timeline_id: Option<TimelineId>,
+    #[serde(default)]
+    pub ancestor_start_lsn: Option<Lsn>,
+    pub pg_version: Option<u32>,
 }
 
 #[derive(Serialize, Deserialize, Clone)]
@@ -1068,12 +1051,6 @@ pub mod virtual_file {
     }
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct ScanDisposableKeysResponse {
-    pub disposable_count: usize,
-    pub not_disposable_count: usize,
-}
-
 // Wrapped in libpq CopyData
 #[derive(PartialEq, Eq, Debug)]
 pub enum PagestreamFeMessage {

libs/remote_storage/src/local_fs.rs

@@ -357,20 +357,22 @@ impl RemoteStorage for LocalFs {
                 .list_recursive(prefix)
                 .await
                 .map_err(DownloadError::Other)?;
-            let mut objects = Vec::with_capacity(keys.len());
-            for key in keys {
-                let path = key.with_base(&self.storage_root);
-                let metadata = file_metadata(&path).await?;
-                if metadata.is_dir() {
-                    continue;
-                }
-                objects.push(ListingObject {
-                    key: key.clone(),
-                    last_modified: metadata.modified()?,
-                    size: metadata.len(),
-                });
-            }
-            let objects = objects;
+            let objects = keys
+                .into_iter()
+                .filter_map(|k| {
+                    let path = k.with_base(&self.storage_root);
+                    if path.is_dir() {
+                        None
+                    } else {
+                        Some(ListingObject {
+                            key: k.clone(),
+                            // LocalFs is just for testing, so just specify a dummy time
+                            last_modified: SystemTime::now(),
+                            size: 0,
+                        })
+                    }
+                })
+                .collect();
 
             if let ListingMode::NoDelimiter = mode {
                 result.keys = objects;
@@ -408,8 +410,9 @@ impl RemoteStorage for LocalFs {
                     } else {
                         result.keys.push(ListingObject {
                             key: RemotePath::from_string(&relative_key).unwrap(),
-                            last_modified: object.last_modified,
-                            size: object.size,
+                            // LocalFs is just for testing
+                            last_modified: SystemTime::now(),
+                            size: 0,
                         });
                     }
                 }

pageserver/src/http/openapi_spec.yml

@@ -597,10 +597,6 @@ paths:
         Create a timeline. Returns new timeline id on success.
         Recreating the same timeline will succeed if the parameters match the existing timeline.
         If no pg_version is specified, assume DEFAULT_PG_VERSION hardcoded in the pageserver.
-
-        To ensure durability, the caller must retry the creation until success.
-        Just because the timeline is visible via other endpoints does not mean it is durable.
-        Future versions may stop showing timelines that are not yet durable.
       requestBody:
         content:
           application/json:

pageserver/src/http/routes.rs

@@ -38,7 +38,6 @@ use pageserver_api::models::TenantShardSplitRequest;
 use pageserver_api::models::TenantShardSplitResponse;
 use pageserver_api::models::TenantSorting;
 use pageserver_api::models::TimelineArchivalConfigRequest;
-use pageserver_api::models::TimelineCreateRequestMode;
 use pageserver_api::models::TimelinesInfoAndOffloaded;
 use pageserver_api::models::TopTenantShardItem;
 use pageserver_api::models::TopTenantShardsRequest;
@@ -86,7 +85,6 @@ use crate::tenant::timeline::Timeline;
 use crate::tenant::GetTimelineError;
 use crate::tenant::OffloadedTimeline;
 use crate::tenant::{LogicalSizeCalculationCause, PageReconstructError};
-use crate::DEFAULT_PG_VERSION;
 use crate::{disk_usage_eviction_task, tenant};
 use pageserver_api::models::{
     StatusResponse, TenantConfigRequest, TenantInfo, TimelineCreateRequest, TimelineGcRequest,
@@ -549,26 +547,6 @@ async fn timeline_create_handler(
     check_permission(&request, Some(tenant_shard_id.tenant_id))?;
 
     let new_timeline_id = request_data.new_timeline_id;
-    // fill in the default pg_version if not provided & convert request into domain model
-    let params: tenant::CreateTimelineParams = match request_data.mode {
-        TimelineCreateRequestMode::Bootstrap {
-            existing_initdb_timeline_id,
-            pg_version,
-        } => tenant::CreateTimelineParams::Bootstrap(tenant::CreateTimelineParamsBootstrap {
-            new_timeline_id,
-            existing_initdb_timeline_id,
-            pg_version: pg_version.unwrap_or(DEFAULT_PG_VERSION),
-        }),
-        TimelineCreateRequestMode::Branch {
-            ancestor_timeline_id,
-            ancestor_start_lsn,
-            pg_version: _,
-        } => tenant::CreateTimelineParams::Branch(tenant::CreateTimelineParamsBranch {
-            new_timeline_id,
-            ancestor_timeline_id,
-            ancestor_start_lsn,
-        }),
-    };
 
     let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Error);
 
@@ -581,12 +559,22 @@ async fn timeline_create_handler(
 
         tenant.wait_to_become_active(ACTIVE_TENANT_TIMEOUT).await?;
 
-        // earlier versions of the code had pg_version and ancestor_lsn in the span
-        // => continue to provide that information, but, through a log message that doesn't require us to destructure
-        tracing::info!(?params, "creating timeline");
+        if let Some(ancestor_id) = request_data.ancestor_timeline_id.as_ref() {
+            tracing::info!(%ancestor_id, "starting to branch");
+        } else {
+            tracing::info!("bootstrapping");
+        }
 
         match tenant
-            .create_timeline(params, state.broker_client.clone(), &ctx)
+            .create_timeline(
+                new_timeline_id,
+                request_data.ancestor_timeline_id,
+                request_data.ancestor_start_lsn,
+                request_data.pg_version.unwrap_or(crate::DEFAULT_PG_VERSION),
+                request_data.existing_initdb_timeline_id,
+                state.broker_client.clone(),
+                &ctx,
+            )
             .await
         {
             Ok(new_timeline) => {
@@ -637,6 +625,8 @@ async fn timeline_create_handler(
         tenant_id = %tenant_shard_id.tenant_id,
         shard_id = %tenant_shard_id.shard_slug(),
         timeline_id = %new_timeline_id,
+        lsn=?request_data.ancestor_start_lsn,
+        pg_version=?request_data.pg_version
     ))
     .await
 }
@@ -1293,99 +1283,6 @@ async fn layer_map_info_handler(
     json_response(StatusCode::OK, layer_map_info)
 }
 
-#[instrument(skip_all, fields(tenant_id, shard_id, timeline_id, layer_name))]
-async fn timeline_layer_scan_disposable_keys(
-    request: Request<Body>,
-    cancel: CancellationToken,
-) -> Result<Response<Body>, ApiError> {
-    let tenant_shard_id: TenantShardId = parse_request_param(&request, "tenant_shard_id")?;
-    let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
-    let layer_name: LayerName = parse_request_param(&request, "layer_name")?;
-
-    tracing::Span::current().record(
-        "tenant_id",
-        tracing::field::display(&tenant_shard_id.tenant_id),
-    );
-    tracing::Span::current().record(
-        "shard_id",
-        tracing::field::display(tenant_shard_id.shard_slug()),
-    );
-    tracing::Span::current().record("timeline_id", tracing::field::display(&timeline_id));
-    tracing::Span::current().record("layer_name", tracing::field::display(&layer_name));
-
-    let state = get_state(&request);
-
-    check_permission(&request, Some(tenant_shard_id.tenant_id))?;
-
-    // technically the timeline need not be active for this scan to complete
-    let timeline =
-        active_timeline_of_active_tenant(&state.tenant_manager, tenant_shard_id, timeline_id)
-            .await?;
-
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
-
-    let guard = timeline.layers.read().await;
-    let Some(layer) = guard.try_get_from_key(&layer_name.clone().into()) else {
-        return Err(ApiError::NotFound(
-            anyhow::anyhow!("Layer {tenant_shard_id}/{timeline_id}/{layer_name} not found").into(),
-        ));
-    };
-
-    let resident_layer = layer
-        .download_and_keep_resident()
-        .await
-        .map_err(|err| match err {
-            tenant::storage_layer::layer::DownloadError::TimelineShutdown
-            | tenant::storage_layer::layer::DownloadError::DownloadCancelled => {
-                ApiError::ShuttingDown
-            }
-            tenant::storage_layer::layer::DownloadError::ContextAndConfigReallyDeniesDownloads
-            | tenant::storage_layer::layer::DownloadError::DownloadRequired
-            | tenant::storage_layer::layer::DownloadError::NotFile(_)
-            | tenant::storage_layer::layer::DownloadError::DownloadFailed
-            | tenant::storage_layer::layer::DownloadError::PreStatFailed(_) => {
-                ApiError::InternalServerError(err.into())
-            }
-            #[cfg(test)]
-            tenant::storage_layer::layer::DownloadError::Failpoint(_) => {
-                ApiError::InternalServerError(err.into())
-            }
-        })?;
-
-    let keys = resident_layer
-        .load_keys(&ctx)
-        .await
-        .map_err(ApiError::InternalServerError)?;
-
-    let shard_identity = timeline.get_shard_identity();
-
-    let mut disposable_count = 0;
-    let mut not_disposable_count = 0;
-    let cancel = cancel.clone();
-    for (i, key) in keys.into_iter().enumerate() {
-        if shard_identity.is_key_disposable(&key) {
-            disposable_count += 1;
-            tracing::debug!(key = %key, key.dbg=?key, "disposable key");
-        } else {
-            not_disposable_count += 1;
-        }
-        #[allow(clippy::collapsible_if)]
-        if i % 10000 == 0 {
-            if cancel.is_cancelled() || timeline.cancel.is_cancelled() || timeline.is_stopping() {
-                return Err(ApiError::ShuttingDown);
-            }
-        }
-    }
-
-    json_response(
-        StatusCode::OK,
-        pageserver_api::models::ScanDisposableKeysResponse {
-            disposable_count,
-            not_disposable_count,
-        },
-    )
-}
-
 async fn layer_download_handler(
     request: Request<Body>,
     _cancel: CancellationToken,
@@ -3248,10 +3145,6 @@ pub fn make_router(
             "/v1/tenant/:tenant_shard_id/timeline/:timeline_id/layer/:layer_file_name",
             |r| api_handler(r, evict_timeline_layer_handler),
         )
-        .post(
-            "/v1/tenant/:tenant_shard_id/timeline/:timeline_id/layer/:layer_name/scan_disposable_keys",
-            |r| testing_api_handler("timeline_layer_scan_disposable_keys", r, timeline_layer_scan_disposable_keys),
-        )
         .post(
             "/v1/tenant/:tenant_shard_id/timeline/:timeline_id/block_gc",
             |r| api_handler(r, timeline_gc_blocking_handler),

pageserver/src/metrics.rs

@@ -3040,111 +3040,13 @@ impl<F: Future<Output = Result<O, E>>, O, E> Future for MeasuredRemoteOp<F> {
 }
 
 pub mod tokio_epoll_uring {
-    use std::{
-        collections::HashMap,
-        sync::{Arc, Mutex},
-    };
-
-    use metrics::{register_histogram, register_int_counter, Histogram, LocalHistogram, UIntGauge};
+    use metrics::{register_int_counter, UIntGauge};
     use once_cell::sync::Lazy;
 
-    /// Shared storage for tokio-epoll-uring thread local metrics.
-    pub(crate) static THREAD_LOCAL_METRICS_STORAGE: Lazy<ThreadLocalMetricsStorage> =
-        Lazy::new(|| {
-            let slots_submission_queue_depth = register_histogram!(
-                "pageserver_tokio_epoll_uring_slots_submission_queue_depth",
-                "The slots waiters queue depth of each tokio_epoll_uring system",
-                vec![1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0, 1024.0],
-            )
-            .expect("failed to define a metric");
-            ThreadLocalMetricsStorage {
-                observers: Mutex::new(HashMap::new()),
-                slots_submission_queue_depth,
-            }
-        });
-
-    pub struct ThreadLocalMetricsStorage {
-        /// List of thread local metrics observers.
-        observers: Mutex<HashMap<u64, Arc<ThreadLocalMetrics>>>,
-        /// A histogram shared between all thread local systems
-        /// for collecting slots submission queue depth.
-        slots_submission_queue_depth: Histogram,
-    }
-
-    /// Each thread-local [`tokio_epoll_uring::System`] gets one of these as its
-    /// [`tokio_epoll_uring::metrics::PerSystemMetrics`] generic.
-    ///
-    /// The System makes observations into [`Self`] and periodically, the collector
-    /// comes along and flushes [`Self`] into the shared storage [`THREAD_LOCAL_METRICS_STORAGE`].
-    ///
-    /// [`LocalHistogram`] is `!Send`, so, we need to put it behind a [`Mutex`].
-    /// But except for the periodic flush, the lock is uncontended so there's no waiting
-    /// for cache coherence protocol to get an exclusive cache line.
-    pub struct ThreadLocalMetrics {
-        /// Local observer of thread local tokio-epoll-uring system's slots waiters queue depth.
-        slots_submission_queue_depth: Mutex<LocalHistogram>,
-    }
-
-    impl ThreadLocalMetricsStorage {
-        /// Registers a new thread local system. Returns a thread local metrics observer.
-        pub fn register_system(&self, id: u64) -> Arc<ThreadLocalMetrics> {
-            let per_system_metrics = Arc::new(ThreadLocalMetrics::new(
-                self.slots_submission_queue_depth.local(),
-            ));
-            let mut g = self.observers.lock().unwrap();
-            g.insert(id, Arc::clone(&per_system_metrics));
-            per_system_metrics
-        }
-
-        /// Removes metrics observer for a thread local system.
-        /// This should be called before dropping a thread local system.
-        pub fn remove_system(&self, id: u64) {
-            let mut g = self.observers.lock().unwrap();
-            g.remove(&id);
-        }
-
-        /// Flush all thread local metrics to the shared storage.
-        pub fn flush_thread_local_metrics(&self) {
-            let g = self.observers.lock().unwrap();
-            g.values().for_each(|local| {
-                local.flush();
-            });
-        }
-    }
-
-    impl ThreadLocalMetrics {
-        pub fn new(slots_submission_queue_depth: LocalHistogram) -> Self {
-            ThreadLocalMetrics {
-                slots_submission_queue_depth: Mutex::new(slots_submission_queue_depth),
-            }
-        }
-
-        /// Flushes the thread local metrics to shared aggregator.
-        pub fn flush(&self) {
-            let Self {
-                slots_submission_queue_depth,
-            } = self;
-            slots_submission_queue_depth.lock().unwrap().flush();
-        }
-    }
-
-    impl tokio_epoll_uring::metrics::PerSystemMetrics for ThreadLocalMetrics {
-        fn observe_slots_submission_queue_depth(&self, queue_depth: u64) {
-            let Self {
-                slots_submission_queue_depth,
-            } = self;
-            slots_submission_queue_depth
-                .lock()
-                .unwrap()
-                .observe(queue_depth as f64);
-        }
-    }
-
     pub struct Collector {
         descs: Vec<metrics::core::Desc>,
         systems_created: UIntGauge,
         systems_destroyed: UIntGauge,
-        thread_local_metrics_storage: &'static ThreadLocalMetricsStorage,
     }
 
     impl metrics::core::Collector for Collector {
@@ -3154,7 +3056,7 @@ pub mod tokio_epoll_uring {
 
         fn collect(&self) -> Vec<metrics::proto::MetricFamily> {
             let mut mfs = Vec::with_capacity(Self::NMETRICS);
-            let tokio_epoll_uring::metrics::GlobalMetrics {
+            let tokio_epoll_uring::metrics::Metrics {
                 systems_created,
                 systems_destroyed,
             } = tokio_epoll_uring::metrics::global();
@@ -3162,21 +3064,12 @@ pub mod tokio_epoll_uring {
             mfs.extend(self.systems_created.collect());
             self.systems_destroyed.set(systems_destroyed);
             mfs.extend(self.systems_destroyed.collect());
-
-            self.thread_local_metrics_storage
-                .flush_thread_local_metrics();
-
-            mfs.extend(
-                self.thread_local_metrics_storage
-                    .slots_submission_queue_depth
-                    .collect(),
-            );
             mfs
         }
     }
 
     impl Collector {
-        const NMETRICS: usize = 3;
+        const NMETRICS: usize = 2;
 
         #[allow(clippy::new_without_default)]
         pub fn new() -> Self {
@@ -3208,7 +3101,6 @@ pub mod tokio_epoll_uring {
                 descs,
                 systems_created,
                 systems_destroyed,
-                thread_local_metrics_storage: &THREAD_LOCAL_METRICS_STORAGE,
             }
         }
     }
@@ -3568,7 +3460,6 @@ pub fn preinitialize_metrics() {
     Lazy::force(&RECONSTRUCT_TIME);
     Lazy::force(&BASEBACKUP_QUERY_TIME);
     Lazy::force(&COMPUTE_COMMANDS_COUNTERS);
-    Lazy::force(&tokio_epoll_uring::THREAD_LOCAL_METRICS_STORAGE);
 
     tenant_throttling::preinitialize_global_metrics();
 }

pageserver/src/pgdatadir_mapping.rs

@@ -1506,42 +1506,35 @@ impl<'a> DatadirModification<'a> {
         Ok(())
     }
 
-    /// Drop some relations
-    pub(crate) async fn put_rel_drops(
-        &mut self,
-        drop_relations: HashMap<(u32, u32), Vec<RelTag>>,
-        ctx: &RequestContext,
-    ) -> anyhow::Result<()> {
-        for ((spc_node, db_node), rel_tags) in drop_relations {
-            let dir_key = rel_dir_to_key(spc_node, db_node);
-            let buf = self.get(dir_key, ctx).await?;
-            let mut dir = RelDirectory::des(&buf)?;
+    /// Drop a relation.
+    pub async fn put_rel_drop(&mut self, rel: RelTag, ctx: &RequestContext) -> anyhow::Result<()> {
+        anyhow::ensure!(rel.relnode != 0, RelationError::InvalidRelnode);
 
-            let mut dirty = false;
-            for rel_tag in rel_tags {
-                if dir.rels.remove(&(rel_tag.relnode, rel_tag.forknum)) {
-                    dirty = true;
+        // Remove it from the directory entry
+        let dir_key = rel_dir_to_key(rel.spcnode, rel.dbnode);
+        let buf = self.get(dir_key, ctx).await?;
+        let mut dir = RelDirectory::des(&buf)?;
 
-                    // update logical size
-                    let size_key = rel_size_to_key(rel_tag);
-                    let old_size = self.get(size_key, ctx).await?.get_u32_le();
-                    self.pending_nblocks -= old_size as i64;
+        self.pending_directory_entries
+            .push((DirectoryKind::Rel, dir.rels.len()));
 
-                    // Remove entry from relation size cache
-                    self.tline.remove_cached_rel_size(&rel_tag);
-
-                    // Delete size entry, as well as all blocks
-                    self.delete(rel_key_range(rel_tag));
-                }
-            }
-
-            if dirty {
-                self.put(dir_key, Value::Image(Bytes::from(RelDirectory::ser(&dir)?)));
-                self.pending_directory_entries
-                    .push((DirectoryKind::Rel, dir.rels.len()));
-            }
+        if dir.rels.remove(&(rel.relnode, rel.forknum)) {
+            self.put(dir_key, Value::Image(Bytes::from(RelDirectory::ser(&dir)?)));
+        } else {
+            warn!("dropped rel {} did not exist in rel directory", rel);
         }
 
+        // update logical size
+        let size_key = rel_size_to_key(rel);
+        let old_size = self.get(size_key, ctx).await?.get_u32_le();
+        self.pending_nblocks -= old_size as i64;
+
+        // Remove enty from relation size cache
+        self.tline.remove_cached_rel_size(&rel);
+
+        // Delete size entry, as well as all blocks
+        self.delete(rel_key_range(rel));
+
         Ok(())
     }
 

pageserver/src/tenant.rs

@@ -294,11 +294,11 @@ pub struct Tenant {
 
     /// During timeline creation, we first insert the TimelineId to the
     /// creating map, then `timelines`, then remove it from the creating map.
-    /// **Lock order**: if acquiring all (or a subset), acquire them in order `timelines`, `timelines_offloaded`, `timelines_creating`
+    /// **Lock order**: if acquiring both, acquire`timelines` before `timelines_creating`
     timelines_creating: std::sync::Mutex<HashSet<TimelineId>>,
 
     /// Possibly offloaded and archived timelines
-    /// **Lock order**: if acquiring all (or a subset), acquire them in order `timelines`, `timelines_offloaded`, `timelines_creating`
+    /// **Lock order**: if acquiring both, acquire`timelines` before `timelines_offloaded`
     timelines_offloaded: Mutex<HashMap<TimelineId, Arc<OffloadedTimeline>>>,
 
     // This mutex prevents creation of new timelines during GC.
@@ -584,40 +584,30 @@ impl OffloadedTimeline {
     }
 }
 
-impl fmt::Debug for OffloadedTimeline {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "OffloadedTimeline<{}>", self.timeline_id)
-    }
-}
-
 #[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
 pub enum MaybeOffloaded {
     Yes,
     No,
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone)]
 pub enum TimelineOrOffloaded {
     Timeline(Arc<Timeline>),
     Offloaded(Arc<OffloadedTimeline>),
 }
 
 impl TimelineOrOffloaded {
-    pub fn arc_ref(&self) -> TimelineOrOffloadedArcRef<'_> {
+    pub fn tenant_shard_id(&self) -> TenantShardId {
         match self {
-            TimelineOrOffloaded::Timeline(timeline) => {
-                TimelineOrOffloadedArcRef::Timeline(timeline)
-            }
-            TimelineOrOffloaded::Offloaded(offloaded) => {
-                TimelineOrOffloadedArcRef::Offloaded(offloaded)
-            }
+            TimelineOrOffloaded::Timeline(timeline) => timeline.tenant_shard_id,
+            TimelineOrOffloaded::Offloaded(offloaded) => offloaded.tenant_shard_id,
         }
     }
-    pub fn tenant_shard_id(&self) -> TenantShardId {
-        self.arc_ref().tenant_shard_id()
-    }
     pub fn timeline_id(&self) -> TimelineId {
-        self.arc_ref().timeline_id()
+        match self {
+            TimelineOrOffloaded::Timeline(timeline) => timeline.timeline_id,
+            TimelineOrOffloaded::Offloaded(offloaded) => offloaded.timeline_id,
+        }
     }
     pub fn delete_progress(&self) -> &Arc<tokio::sync::Mutex<DeleteTimelineFlow>> {
         match self {
@@ -625,7 +615,7 @@ impl TimelineOrOffloaded {
             TimelineOrOffloaded::Offloaded(offloaded) => &offloaded.delete_progress,
         }
     }
-    fn remote_client_maybe_construct(&self, tenant: &Tenant) -> Arc<RemoteTimelineClient> {
+    pub fn remote_client_maybe_construct(&self, tenant: &Tenant) -> Arc<RemoteTimelineClient> {
         match self {
             TimelineOrOffloaded::Timeline(timeline) => timeline.remote_client.clone(),
             TimelineOrOffloaded::Offloaded(offloaded) => match offloaded.remote_client.clone() {
@@ -642,38 +632,6 @@ impl TimelineOrOffloaded {
     }
 }
 
-pub enum TimelineOrOffloadedArcRef<'a> {
-    Timeline(&'a Arc<Timeline>),
-    Offloaded(&'a Arc<OffloadedTimeline>),
-}
-
-impl TimelineOrOffloadedArcRef<'_> {
-    pub fn tenant_shard_id(&self) -> TenantShardId {
-        match self {
-            TimelineOrOffloadedArcRef::Timeline(timeline) => timeline.tenant_shard_id,
-            TimelineOrOffloadedArcRef::Offloaded(offloaded) => offloaded.tenant_shard_id,
-        }
-    }
-    pub fn timeline_id(&self) -> TimelineId {
-        match self {
-            TimelineOrOffloadedArcRef::Timeline(timeline) => timeline.timeline_id,
-            TimelineOrOffloadedArcRef::Offloaded(offloaded) => offloaded.timeline_id,
-        }
-    }
-}
-
-impl<'a> From<&'a Arc<Timeline>> for TimelineOrOffloadedArcRef<'a> {
-    fn from(timeline: &'a Arc<Timeline>) -> Self {
-        Self::Timeline(timeline)
-    }
-}
-
-impl<'a> From<&'a Arc<OffloadedTimeline>> for TimelineOrOffloadedArcRef<'a> {
-    fn from(timeline: &'a Arc<OffloadedTimeline>) -> Self {
-        Self::Offloaded(timeline)
-    }
-}
-
 #[derive(Debug, thiserror::Error, PartialEq, Eq)]
 pub enum GetTimelineError {
     #[error("Timeline is shutting down")]
@@ -779,99 +737,6 @@ impl Debug for SetStoppingError {
     }
 }
 
-/// Arguments to [`Tenant::create_timeline`].
-///
-/// Not usable as an idempotency key for timeline creation because if [`CreateTimelineParamsBranch::ancestor_start_lsn`]
-/// is `None`, the result of the timeline create call is not deterministic.
-///
-/// See [`CreateTimelineIdempotency`] for an idempotency key.
-#[derive(Debug)]
-pub(crate) enum CreateTimelineParams {
-    Bootstrap(CreateTimelineParamsBootstrap),
-    Branch(CreateTimelineParamsBranch),
-}
-
-#[derive(Debug)]
-pub(crate) struct CreateTimelineParamsBootstrap {
-    pub(crate) new_timeline_id: TimelineId,
-    pub(crate) existing_initdb_timeline_id: Option<TimelineId>,
-    pub(crate) pg_version: u32,
-}
-
-/// NB: See comment on [`CreateTimelineIdempotency::Branch`] for why there's no `pg_version` here.
-#[derive(Debug)]
-pub(crate) struct CreateTimelineParamsBranch {
-    pub(crate) new_timeline_id: TimelineId,
-    pub(crate) ancestor_timeline_id: TimelineId,
-    pub(crate) ancestor_start_lsn: Option<Lsn>,
-}
-
-/// What is used to determine idempotency of a [`Tenant::create_timeline`] call in  [`Tenant::start_creating_timeline`].
-///
-/// Each [`Timeline`] object holds [`Self`] as an immutable property in [`Timeline::create_idempotency`].
-///
-/// We lower timeline creation requests to [`Self`], and then use [`PartialEq::eq`] to compare [`Timeline::create_idempotency`] with the request.
-/// If they are equal, we return a reference to the existing timeline, otherwise it's an idempotency conflict.
-///
-/// There is special treatment for [`Self::FailWithConflict`] to always return an idempotency conflict.
-/// It would be nice to have more advanced derive macros to make that special treatment declarative.
-///
-/// Notes:
-/// - Unlike [`CreateTimelineParams`], ancestor LSN is fixed, so, branching will be at a deterministic LSN.
-/// - We make some trade-offs though, e.g., [`CreateTimelineParamsBootstrap::existing_initdb_timeline_id`]
-///   is not considered for idempotency. We can improve on this over time if we deem it necessary.
-///
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub(crate) enum CreateTimelineIdempotency {
-    /// NB: special treatment, see comment in [`Self`].
-    FailWithConflict,
-    Bootstrap {
-        pg_version: u32,
-    },
-    /// NB: branches always have the same `pg_version` as their ancestor.
-    /// While [`pageserver_api::models::TimelineCreateRequestMode::Branch::pg_version`]
-    /// exists as a field, and is set by cplane, it has always been ignored by pageserver when
-    /// determining the child branch pg_version.
-    Branch {
-        ancestor_timeline_id: TimelineId,
-        ancestor_start_lsn: Lsn,
-    },
-}
-
-/// What is returned by [`Tenant::start_creating_timeline`].
-#[must_use]
-enum StartCreatingTimelineResult<'t> {
-    CreateGuard(TimelineCreateGuard<'t>),
-    Idempotent(Arc<Timeline>),
-}
-
-/// What is returned by [`Tenant::create_timeline`].
-enum CreateTimelineResult {
-    Created(Arc<Timeline>),
-    Idempotent(Arc<Timeline>),
-}
-
-impl CreateTimelineResult {
-    fn discriminant(&self) -> &'static str {
-        match self {
-            Self::Created(_) => "Created",
-            Self::Idempotent(_) => "Idempotent",
-        }
-    }
-    fn timeline(&self) -> &Arc<Timeline> {
-        match self {
-            Self::Created(t) | Self::Idempotent(t) => t,
-        }
-    }
-    /// Unit test timelines aren't activated, test has to do it if it needs to.
-    #[cfg(test)]
-    fn into_timeline_for_test(self) -> Arc<Timeline> {
-        match self {
-            Self::Created(t) | Self::Idempotent(t) => t,
-        }
-    }
-}
-
 #[derive(thiserror::Error, Debug)]
 pub enum CreateTimelineError {
     #[error("creation of timeline with the given ID is in progress")]
@@ -1011,24 +876,12 @@ impl Tenant {
     ) -> anyhow::Result<()> {
         let tenant_id = self.tenant_shard_id;
 
-        let idempotency = if metadata.ancestor_timeline().is_none() {
-            CreateTimelineIdempotency::Bootstrap {
-                pg_version: metadata.pg_version(),
-            }
-        } else {
-            CreateTimelineIdempotency::Branch {
-                ancestor_timeline_id: metadata.ancestor_timeline().unwrap(),
-                ancestor_start_lsn: metadata.ancestor_lsn(),
-            }
-        };
-
         let timeline = self.create_timeline_struct(
             timeline_id,
             &metadata,
             ancestor.clone(),
             resources,
             CreateTimelineCause::Load,
-            idempotency.clone(),
         )?;
         let disk_consistent_lsn = timeline.get_disk_consistent_lsn();
         anyhow::ensure!(
@@ -1821,8 +1674,6 @@ impl Tenant {
     }
 
     /// Loads the specified (offloaded) timeline from S3 and attaches it as a loaded timeline
-    ///
-    /// Counterpart to [`offload_timeline`].
     async fn unoffload_timeline(
         self: &Arc<Self>,
         timeline_id: TimelineId,
@@ -1831,24 +1682,6 @@ impl Tenant {
     ) -> Result<Arc<Timeline>, TimelineArchivalError> {
         info!("unoffloading timeline");
         let cancel = self.cancel.clone();
-
-        // Protect against concurrent attempts to use this TimelineId
-        // We don't care much about idempotency, as it's ensured a layer above.
-        let allow_offloaded = true;
-        let _create_guard = self
-            .create_timeline_create_guard(
-                timeline_id,
-                CreateTimelineIdempotency::FailWithConflict,
-                allow_offloaded,
-            )
-            .map_err(|err| match err {
-                TimelineExclusionError::AlreadyCreating => TimelineArchivalError::AlreadyInProgress,
-                TimelineExclusionError::AlreadyExists { .. } => {
-                    TimelineArchivalError::Other(anyhow::anyhow!("Timeline already exists"))
-                }
-                TimelineExclusionError::Other(e) => TimelineArchivalError::Other(e),
-            })?;
-
         let timeline_preload = self
             .load_timeline_metadata(timeline_id, self.remote_storage.clone(), cancel.clone())
             .await;
@@ -2115,17 +1948,16 @@ impl Tenant {
         self.timelines.lock().unwrap().keys().cloned().collect()
     }
 
-    /// This is used by tests & import-from-basebackup.
+    /// This is used to create the initial 'main' timeline during bootstrapping,
+    /// or when importing a new base backup. The caller is expected to load an
+    /// initial image of the datadir to the new timeline after this.
     ///
-    /// The returned [`UninitializedTimeline`] contains no data nor metadata and it is in
-    /// a state that will fail [`Tenant::load_remote_timeline`] because `disk_consistent_lsn=Lsn(0)`.
+    /// Until that happens, the on-disk state is invalid (disk_consistent_lsn=Lsn(0))
+    /// and the timeline will fail to load at a restart.
     ///
-    /// The caller is responsible for getting the timeline into a state that will be accepted
-    /// by [`Tenant::load_remote_timeline`] / [`Tenant::attach`].
-    /// Then they may call [`UninitializedTimeline::finish_creation`] to add the timeline
-    /// to the [`Tenant::timelines`].
-    ///
-    /// Tests should use `Tenant::create_test_timeline` to set up the minimum required metadata keys.
+    /// For tests, use `DatadirModification::init_empty_test_timeline` + `commit` to setup the
+    /// minimum amount of keys required to get a writable timeline.
+    /// (Without it, `put` might fail due to `repartition` failing.)
     pub(crate) async fn create_empty_timeline(
         &self,
         new_timeline_id: TimelineId,
@@ -2139,15 +1971,7 @@ impl Tenant {
         );
 
         // Protect against concurrent attempts to use this TimelineId
-        let create_guard = match self
-            .start_creating_timeline(new_timeline_id, CreateTimelineIdempotency::FailWithConflict)
-            .await?
-        {
-            StartCreatingTimelineResult::CreateGuard(guard) => guard,
-            StartCreatingTimelineResult::Idempotent(_) => {
-                unreachable!("FailWithConflict implies we get an error instead")
-            }
-        };
+        let create_guard = self.create_timeline_create_guard(new_timeline_id)?;
 
         let new_metadata = TimelineMetadata::new(
             // Initialize disk_consistent LSN to 0, The caller must import some data to
@@ -2266,7 +2090,11 @@ impl Tenant {
     #[allow(clippy::too_many_arguments)]
     pub(crate) async fn create_timeline(
         self: &Arc<Tenant>,
-        params: CreateTimelineParams,
+        new_timeline_id: TimelineId,
+        ancestor_timeline_id: Option<TimelineId>,
+        mut ancestor_start_lsn: Option<Lsn>,
+        pg_version: u32,
+        load_existing_initdb: Option<TimelineId>,
         broker_client: storage_broker::BrokerClientChannel,
         ctx: &RequestContext,
     ) -> Result<Arc<Timeline>, CreateTimelineError> {
@@ -2285,25 +2113,54 @@ impl Tenant {
             .enter()
             .map_err(|_| CreateTimelineError::ShuttingDown)?;
 
-        let result: CreateTimelineResult = match params {
-            CreateTimelineParams::Bootstrap(CreateTimelineParamsBootstrap {
-                new_timeline_id,
-                existing_initdb_timeline_id,
-                pg_version,
-            }) => {
-                self.bootstrap_timeline(
-                    new_timeline_id,
-                    pg_version,
-                    existing_initdb_timeline_id,
-                    ctx,
-                )
-                .await?
+        // Get exclusive access to the timeline ID: this ensures that it does not already exist,
+        // and that no other creation attempts will be allowed in while we are working.
+        let create_guard = match self.create_timeline_create_guard(new_timeline_id) {
+            Ok(m) => m,
+            Err(TimelineExclusionError::AlreadyCreating) => {
+                // Creation is in progress, we cannot create it again, and we cannot
+                // check if this request matches the existing one, so caller must try
+                // again later.
+                return Err(CreateTimelineError::AlreadyCreating);
             }
-            CreateTimelineParams::Branch(CreateTimelineParamsBranch {
-                new_timeline_id,
-                ancestor_timeline_id,
-                mut ancestor_start_lsn,
-            }) => {
+            Err(TimelineExclusionError::Other(e)) => {
+                return Err(CreateTimelineError::Other(e));
+            }
+            Err(TimelineExclusionError::AlreadyExists(existing)) => {
+                debug!("timeline {new_timeline_id} already exists");
+
+                // Idempotency: creating the same timeline twice is not an error, unless
+                // the second creation has different parameters.
+                if existing.get_ancestor_timeline_id() != ancestor_timeline_id
+                    || existing.pg_version != pg_version
+                    || (ancestor_start_lsn.is_some()
+                        && ancestor_start_lsn != Some(existing.get_ancestor_lsn()))
+                {
+                    return Err(CreateTimelineError::Conflict);
+                }
+
+                // Wait for uploads to complete, so that when we return Ok, the timeline
+                // is known to be durable on remote storage. Just like we do at the end of
+                // this function, after we have created the timeline ourselves.
+                //
+                // We only really care that the initial version of `index_part.json` has
+                // been uploaded. That's enough to remember that the timeline
+                // exists. However, there is no function to wait specifically for that so
+                // we just wait for all in-progress uploads to finish.
+                existing
+                    .remote_client
+                    .wait_completion()
+                    .await
+                    .context("wait for timeline uploads to complete")?;
+
+                return Ok(existing);
+            }
+        };
+
+        pausable_failpoint!("timeline-creation-after-uninit");
+
+        let loaded_timeline = match ancestor_timeline_id {
+            Some(ancestor_timeline_id) => {
                 let ancestor_timeline = self
                     .get_timeline(ancestor_timeline_id, false)
                     .context("Cannot branch off the timeline that's not present in pageserver")?;
@@ -2350,48 +2207,43 @@ impl Tenant {
                         })?;
                 }
 
-                self.branch_timeline(&ancestor_timeline, new_timeline_id, ancestor_start_lsn, ctx)
-                    .await?
+                self.branch_timeline(
+                    &ancestor_timeline,
+                    new_timeline_id,
+                    ancestor_start_lsn,
+                    create_guard,
+                    ctx,
+                )
+                .await?
+            }
+            None => {
+                self.bootstrap_timeline(
+                    new_timeline_id,
+                    pg_version,
+                    load_existing_initdb,
+                    create_guard,
+                    ctx,
+                )
+                .await?
             }
         };
 
         // At this point we have dropped our guard on [`Self::timelines_creating`], and
         // the timeline is visible in [`Self::timelines`], but it is _not_ durable yet.  We must
-        // not send a success to the caller until it is.  The same applies to idempotent retries.
-        //
-        // TODO: the timeline is already visible in [`Self::timelines`]; a caller could incorrectly
-        // assume that, because they can see the timeline via API, that the creation is done and
-        // that it is durable. Ideally, we would keep the timeline hidden (in [`Self::timelines_creating`])
-        // until it is durable, e.g., by extending the time we hold the creation guard. This also
-        // interacts with UninitializedTimeline and is generally a bit tricky.
-        //
-        // To re-emphasize: the only correct way to create a timeline is to repeat calling the
-        // creation API until it returns success. Only then is durability guaranteed.
-        info!(creation_result=%result.discriminant(), "waiting for timeline to be durable");
-        result
-            .timeline()
+        // not send a success to the caller until it is.  The same applies to handling retries,
+        // see the handling of [`TimelineExclusionError::AlreadyExists`] above.
+        let kind = ancestor_timeline_id
+            .map(|_| "branched")
+            .unwrap_or("bootstrapped");
+        loaded_timeline
             .remote_client
             .wait_completion()
             .await
-            .context("wait for timeline initial uploads to complete")?;
+            .with_context(|| format!("wait for {} timeline initial uploads to complete", kind))?;
 
-        // The creating task is responsible for activating the timeline.
-        // We do this after `wait_completion()` so that we don't spin up tasks that start
-        // doing stuff before the IndexPart is durable in S3, which is done by the previous section.
-        let activated_timeline = match result {
-            CreateTimelineResult::Created(timeline) => {
-                timeline.activate(self.clone(), broker_client, None, ctx);
-                timeline
-            }
-            CreateTimelineResult::Idempotent(timeline) => {
-                info!(
-                    "request was deemed idempotent, activation will be done by the creating task"
-                );
-                timeline
-            }
-        };
+        loaded_timeline.activate(self.clone(), broker_client, None, ctx);
 
-        Ok(activated_timeline)
+        Ok(loaded_timeline)
     }
 
     pub(crate) async fn delete_timeline(
@@ -3048,58 +2900,33 @@ impl Tenant {
         &self,
         child_shards: &Vec<TenantShardId>,
     ) -> anyhow::Result<()> {
-        let (timelines, offloaded) = {
-            let timelines = self.timelines.lock().unwrap();
-            let offloaded = self.timelines_offloaded.lock().unwrap();
-            (timelines.clone(), offloaded.clone())
-        };
-        let timelines_iter = timelines
-            .values()
-            .map(TimelineOrOffloadedArcRef::<'_>::from)
-            .chain(
-                offloaded
-                    .values()
-                    .map(TimelineOrOffloadedArcRef::<'_>::from),
-            );
-        for timeline in timelines_iter {
+        let timelines = self.timelines.lock().unwrap().clone();
+        for timeline in timelines.values() {
             // We do not block timeline creation/deletion during splits inside the pageserver: it is up to higher levels
             // to ensure that they do not start a split if currently in the process of doing these.
 
-            let timeline_id = timeline.timeline_id();
-
-            if let TimelineOrOffloadedArcRef::Timeline(timeline) = timeline {
-                // Upload an index from the parent: this is partly to provide freshness for the
-                // child tenants that will copy it, and partly for general ease-of-debugging: there will
-                // always be a parent shard index in the same generation as we wrote the child shard index.
-                tracing::info!(%timeline_id, "Uploading index");
-                timeline
-                    .remote_client
-                    .schedule_index_upload_for_file_changes()?;
-                timeline.remote_client.wait_completion().await?;
-            }
-
-            let remote_client = match timeline {
-                TimelineOrOffloadedArcRef::Timeline(timeline) => timeline.remote_client.clone(),
-                TimelineOrOffloadedArcRef::Offloaded(offloaded) => {
-                    let remote_client = self
-                        .build_timeline_client(offloaded.timeline_id, self.remote_storage.clone());
-                    Arc::new(remote_client)
-                }
-            };
+            // Upload an index from the parent: this is partly to provide freshness for the
+            // child tenants that will copy it, and partly for general ease-of-debugging: there will
+            // always be a parent shard index in the same generation as we wrote the child shard index.
+            tracing::info!(timeline_id=%timeline.timeline_id, "Uploading index");
+            timeline
+                .remote_client
+                .schedule_index_upload_for_file_changes()?;
+            timeline.remote_client.wait_completion().await?;
 
             // Shut down the timeline's remote client: this means that the indices we write
             // for child shards will not be invalidated by the parent shard deleting layers.
-            tracing::info!(%timeline_id, "Shutting down remote storage client");
-            remote_client.shutdown().await;
+            tracing::info!(timeline_id=%timeline.timeline_id, "Shutting down remote storage client");
+            timeline.remote_client.shutdown().await;
 
             // Download methods can still be used after shutdown, as they don't flow through the remote client's
             // queue.  In principal the RemoteTimelineClient could provide this without downloading it, but this
             // operation is rare, so it's simpler to just download it (and robustly guarantees that the index
             // we use here really is the remotely persistent one).
-            tracing::info!(%timeline_id, "Downloading index_part from parent");
-            let result = remote_client
+            tracing::info!(timeline_id=%timeline.timeline_id, "Downloading index_part from parent");
+            let result = timeline.remote_client
                 .download_index_file(&self.cancel)
-                .instrument(info_span!("download_index_file", tenant_id=%self.tenant_shard_id.tenant_id, shard_id=%self.tenant_shard_id.shard_slug(), %timeline_id))
+                .instrument(info_span!("download_index_file", tenant_id=%self.tenant_shard_id.tenant_id, shard_id=%self.tenant_shard_id.shard_slug(), timeline_id=%timeline.timeline_id))
                 .await?;
             let index_part = match result {
                 MaybeDeletedIndexPart::Deleted(_) => {
@@ -3109,11 +2936,11 @@ impl Tenant {
             };
 
             for child_shard in child_shards {
-                tracing::info!(%timeline_id, "Uploading index_part for child {}", child_shard.to_index());
+                tracing::info!(timeline_id=%timeline.timeline_id, "Uploading index_part for child {}", child_shard.to_index());
                 upload_index_part(
                     &self.remote_storage,
                     child_shard,
-                    &timeline_id,
+                    &timeline.timeline_id,
                     self.generation,
                     &index_part,
                     &self.cancel,
@@ -3122,6 +2949,8 @@ impl Tenant {
             }
         }
 
+        // TODO: also copy index files of offloaded timelines
+
         let tenant_manifest = self.tenant_manifest();
         // TODO: generation support
         let generation = remote_timeline_client::TENANT_MANIFEST_GENERATION;
@@ -3404,7 +3233,6 @@ impl Tenant {
         ancestor: Option<Arc<Timeline>>,
         resources: TimelineResources,
         cause: CreateTimelineCause,
-        create_idempotency: CreateTimelineIdempotency,
     ) -> anyhow::Result<Arc<Timeline>> {
         let state = match cause {
             CreateTimelineCause::Load => {
@@ -3434,7 +3262,6 @@ impl Tenant {
             pg_version,
             state,
             self.attach_wal_lag_cooldown.clone(),
-            create_idempotency,
             self.cancel.child_token(),
         );
 
@@ -3920,16 +3747,16 @@ impl Tenant {
     /// timeline background tasks are launched, except the flush loop.
     #[cfg(test)]
     async fn branch_timeline_test(
-        self: &Arc<Self>,
+        &self,
         src_timeline: &Arc<Timeline>,
         dst_id: TimelineId,
         ancestor_lsn: Option<Lsn>,
         ctx: &RequestContext,
     ) -> Result<Arc<Timeline>, CreateTimelineError> {
+        let create_guard = self.create_timeline_create_guard(dst_id).unwrap();
         let tl = self
-            .branch_timeline_impl(src_timeline, dst_id, ancestor_lsn, ctx)
-            .await?
-            .into_timeline_for_test();
+            .branch_timeline_impl(src_timeline, dst_id, ancestor_lsn, create_guard, ctx)
+            .await?;
         tl.set_state(TimelineState::Active);
         Ok(tl)
     }
@@ -3938,7 +3765,7 @@ impl Tenant {
     #[cfg(test)]
     #[allow(clippy::too_many_arguments)]
     pub async fn branch_timeline_test_with_layers(
-        self: &Arc<Self>,
+        &self,
         src_timeline: &Arc<Timeline>,
         dst_id: TimelineId,
         ancestor_lsn: Option<Lsn>,
@@ -3986,24 +3813,28 @@ impl Tenant {
     }
 
     /// Branch an existing timeline.
+    ///
+    /// The caller is responsible for activating the returned timeline.
     async fn branch_timeline(
-        self: &Arc<Self>,
+        &self,
         src_timeline: &Arc<Timeline>,
         dst_id: TimelineId,
         start_lsn: Option<Lsn>,
+        timeline_create_guard: TimelineCreateGuard<'_>,
         ctx: &RequestContext,
-    ) -> Result<CreateTimelineResult, CreateTimelineError> {
-        self.branch_timeline_impl(src_timeline, dst_id, start_lsn, ctx)
+    ) -> Result<Arc<Timeline>, CreateTimelineError> {
+        self.branch_timeline_impl(src_timeline, dst_id, start_lsn, timeline_create_guard, ctx)
             .await
     }
 
     async fn branch_timeline_impl(
-        self: &Arc<Self>,
+        &self,
         src_timeline: &Arc<Timeline>,
         dst_id: TimelineId,
         start_lsn: Option<Lsn>,
+        timeline_create_guard: TimelineCreateGuard<'_>,
         _ctx: &RequestContext,
-    ) -> Result<CreateTimelineResult, CreateTimelineError> {
+    ) -> Result<Arc<Timeline>, CreateTimelineError> {
         let src_id = src_timeline.timeline_id;
 
         // We will validate our ancestor LSN in this function.  Acquire the GC lock so that
@@ -4018,23 +3849,6 @@ impl Tenant {
             lsn
         });
 
-        // we finally have determined the ancestor_start_lsn, so we can get claim exclusivity now
-        let timeline_create_guard = match self
-            .start_creating_timeline(
-                dst_id,
-                CreateTimelineIdempotency::Branch {
-                    ancestor_timeline_id: src_timeline.timeline_id,
-                    ancestor_start_lsn: start_lsn,
-                },
-            )
-            .await?
-        {
-            StartCreatingTimelineResult::CreateGuard(guard) => guard,
-            StartCreatingTimelineResult::Idempotent(timeline) => {
-                return Ok(CreateTimelineResult::Idempotent(timeline));
-            }
-        };
-
         // Ensure that `start_lsn` is valid, i.e. the LSN is within the PITR
         // horizon on the source timeline
         //
@@ -4120,92 +3934,28 @@ impl Tenant {
             .schedule_index_upload_for_full_metadata_update(&metadata)
             .context("branch initial metadata upload")?;
 
-        // Callers are responsible to wait for uploads to complete and for activating the timeline.
-
-        Ok(CreateTimelineResult::Created(new_timeline))
+        Ok(new_timeline)
     }
 
     /// For unit tests, make this visible so that other modules can directly create timelines
     #[cfg(test)]
     #[tracing::instrument(skip_all, fields(tenant_id=%self.tenant_shard_id.tenant_id, shard_id=%self.tenant_shard_id.shard_slug(), %timeline_id))]
     pub(crate) async fn bootstrap_timeline_test(
-        self: &Arc<Self>,
+        &self,
         timeline_id: TimelineId,
         pg_version: u32,
         load_existing_initdb: Option<TimelineId>,
         ctx: &RequestContext,
     ) -> anyhow::Result<Arc<Timeline>> {
-        self.bootstrap_timeline(timeline_id, pg_version, load_existing_initdb, ctx)
-            .await
-            .map_err(anyhow::Error::new)
-            .map(|r| r.into_timeline_for_test())
-    }
-
-    /// Get exclusive access to the timeline ID for creation.
-    ///
-    /// Timeline-creating code paths must use this function before making changes
-    /// to in-memory or persistent state.
-    ///
-    /// The `state` parameter is a description of the timeline creation operation
-    /// we intend to perform.
-    /// If the timeline was already created in the meantime, we check whether this
-    /// request conflicts or is idempotent , based on `state`.
-    async fn start_creating_timeline(
-        &self,
-        new_timeline_id: TimelineId,
-        idempotency: CreateTimelineIdempotency,
-    ) -> Result<StartCreatingTimelineResult<'_>, CreateTimelineError> {
-        let allow_offloaded = false;
-        match self.create_timeline_create_guard(new_timeline_id, idempotency, allow_offloaded) {
-            Ok(create_guard) => {
-                pausable_failpoint!("timeline-creation-after-uninit");
-                Ok(StartCreatingTimelineResult::CreateGuard(create_guard))
-            }
-            Err(TimelineExclusionError::AlreadyCreating) => {
-                // Creation is in progress, we cannot create it again, and we cannot
-                // check if this request matches the existing one, so caller must try
-                // again later.
-                Err(CreateTimelineError::AlreadyCreating)
-            }
-            Err(TimelineExclusionError::Other(e)) => Err(CreateTimelineError::Other(e)),
-            Err(TimelineExclusionError::AlreadyExists {
-                existing: TimelineOrOffloaded::Offloaded(_existing),
-                ..
-            }) => {
-                info!("timeline already exists but is offloaded");
-                Err(CreateTimelineError::Conflict)
-            }
-            Err(TimelineExclusionError::AlreadyExists {
-                existing: TimelineOrOffloaded::Timeline(existing),
-                arg,
-            }) => {
-                {
-                    let existing = &existing.create_idempotency;
-                    let _span = info_span!("idempotency_check", ?existing, ?arg).entered();
-                    debug!("timeline already exists");
-
-                    match (existing, &arg) {
-                        // FailWithConflict => no idempotency check
-                        (CreateTimelineIdempotency::FailWithConflict, _)
-                        | (_, CreateTimelineIdempotency::FailWithConflict) => {
-                            warn!("timeline already exists, failing request");
-                            return Err(CreateTimelineError::Conflict);
-                        }
-                        // Idempotent <=> CreateTimelineIdempotency is identical
-                        (x, y) if x == y => {
-                            info!("timeline already exists and idempotency matches, succeeding request");
-                            // fallthrough
-                        }
-                        (_, _) => {
-                            warn!("idempotency conflict, failing request");
-                            return Err(CreateTimelineError::Conflict);
-                        }
-                    }
-                }
-
-                Ok(StartCreatingTimelineResult::Idempotent(existing))
-            }
-        }
+        let create_guard = self.create_timeline_create_guard(timeline_id).unwrap();
+        self.bootstrap_timeline(
+            timeline_id,
+            pg_version,
+            load_existing_initdb,
+            create_guard,
+            ctx,
+        )
+        .await
     }
 
     async fn upload_initdb(
@@ -4259,26 +4009,16 @@ impl Tenant {
 
     /// - run initdb to init temporary instance and get bootstrap data
     /// - after initialization completes, tar up the temp dir and upload it to S3.
+    ///
+    /// The caller is responsible for activating the returned timeline.
     async fn bootstrap_timeline(
-        self: &Arc<Self>,
+        &self,
         timeline_id: TimelineId,
         pg_version: u32,
         load_existing_initdb: Option<TimelineId>,
+        timeline_create_guard: TimelineCreateGuard<'_>,
         ctx: &RequestContext,
-    ) -> Result<CreateTimelineResult, CreateTimelineError> {
-        let timeline_create_guard = match self
-            .start_creating_timeline(
-                timeline_id,
-                CreateTimelineIdempotency::Bootstrap { pg_version },
-            )
-            .await?
-        {
-            StartCreatingTimelineResult::CreateGuard(guard) => guard,
-            StartCreatingTimelineResult::Idempotent(timeline) => {
-                return Ok(CreateTimelineResult::Idempotent(timeline))
-            }
-        };
-
+    ) -> anyhow::Result<Arc<Timeline>> {
         // create a `tenant/{tenant_id}/timelines/basebackup-{timeline_id}.{TEMP_FILE_SUFFIX}/`
         // temporary directory for basebackup files for the given timeline.
 
@@ -4342,9 +4082,7 @@ impl Tenant {
                 .context("extract initdb tar")?;
         } else {
             // Init temporarily repo to get bootstrap data, this creates a directory in the `pgdata_path` path
-            run_initdb(self.conf, &pgdata_path, pg_version, &self.cancel)
-                .await
-                .context("run initdb")?;
+            run_initdb(self.conf, &pgdata_path, pg_version, &self.cancel).await?;
 
             // Upload the created data dir to S3
             if self.tenant_shard_id().is_shard_zero() {
@@ -4398,9 +4136,7 @@ impl Tenant {
         })?;
 
         fail::fail_point!("before-checkpoint-new-timeline", |_| {
-            Err(CreateTimelineError::Other(anyhow::anyhow!(
-                "failpoint before-checkpoint-new-timeline"
-            )))
+            anyhow::bail!("failpoint before-checkpoint-new-timeline");
         });
 
         unfinished_timeline
@@ -4415,9 +4151,7 @@ impl Tenant {
         // All done!
         let timeline = raw_timeline.finish_creation()?;
 
-        // Callers are responsible to wait for uploads to complete and for activating the timeline.
-
-        Ok(CreateTimelineResult::Created(timeline))
+        Ok(timeline)
     }
 
     fn build_timeline_remote_client(&self, timeline_id: TimelineId) -> RemoteTimelineClient {
@@ -4467,7 +4201,6 @@ impl Tenant {
                 ancestor,
                 resources,
                 CreateTimelineCause::Load,
-                create_guard.idempotency.clone(),
             )
             .context("Failed to create timeline data structure")?;
 
@@ -4505,26 +4238,15 @@ impl Tenant {
 
     /// Get a guard that provides exclusive access to the timeline directory, preventing
     /// concurrent attempts to create the same timeline.
-    ///
-    /// The `allow_offloaded` parameter controls whether to tolerate the existence of
-    /// offloaded timelines or not.
     fn create_timeline_create_guard(
         &self,
         timeline_id: TimelineId,
-        idempotency: CreateTimelineIdempotency,
-        allow_offloaded: bool,
     ) -> Result<TimelineCreateGuard, TimelineExclusionError> {
         let tenant_shard_id = self.tenant_shard_id;
 
         let timeline_path = self.conf.timeline_path(&tenant_shard_id, &timeline_id);
 
-        let create_guard = TimelineCreateGuard::new(
-            self,
-            timeline_id,
-            timeline_path.clone(),
-            idempotency,
-            allow_offloaded,
-        )?;
+        let create_guard = TimelineCreateGuard::new(self, timeline_id, timeline_path.clone())?;
 
         // At this stage, we have got exclusive access to in-memory state for this timeline ID
         // for creation.
@@ -5160,10 +4882,7 @@ mod tests {
             .await
         {
             Ok(_) => panic!("duplicate timeline creation should fail"),
-            Err(e) => assert_eq!(
-                e.to_string(),
-                "timeline already exists with different parameters".to_string()
-            ),
+            Err(e) => assert_eq!(e.to_string(), "Already exists".to_string()),
         }
 
         Ok(())

pageserver/src/tenant/remote_timeline_client.rs

@@ -1278,14 +1278,10 @@ impl RemoteTimelineClient {
         let fut = {
             let mut guard = self.upload_queue.lock().unwrap();
             let upload_queue = match &mut *guard {
-                UploadQueue::Stopped(_) => {
-                    scopeguard::ScopeGuard::into_inner(sg);
-                    return;
-                }
+                UploadQueue::Stopped(_) => return,
                 UploadQueue::Uninitialized => {
                     // transition into Stopped state
                     self.stop_impl(&mut guard);
-                    scopeguard::ScopeGuard::into_inner(sg);
                     return;
                 }
                 UploadQueue::Initialized(ref mut init) => init,

pageserver/src/tenant/storage_layer.rs

@@ -1,6 +1,5 @@
 //! Common traits and structs for layers
 
-pub mod batch_split_writer;
 pub mod delta_layer;
 pub mod filter_iterator;
 pub mod image_layer;
@@ -9,6 +8,7 @@ pub(crate) mod layer;
 mod layer_desc;
 mod layer_name;
 pub mod merge_iterator;
+pub mod split_writer;
 
 use crate::context::{AccessStatsBehavior, RequestContext};
 use crate::repository::Value;

pageserver/src/tenant/storage_layer/delta_layer.rs

@@ -1084,7 +1084,7 @@ impl DeltaLayerInner {
         }
     }
 
-    pub(crate) async fn index_entries<'a>(
+    pub(super) async fn load_keys<'a>(
         &'a self,
         ctx: &RequestContext,
     ) -> Result<Vec<DeltaEntry<'a>>> {
@@ -1346,7 +1346,7 @@ impl DeltaLayerInner {
 
         tree_reader.dump().await?;
 
-        let keys = self.index_entries(ctx).await?;
+        let keys = self.load_keys(ctx).await?;
 
         async fn dump_blob(val: &ValueRef<'_>, ctx: &RequestContext) -> anyhow::Result<String> {
             let buf = val.load_raw(ctx).await?;
@@ -1453,16 +1453,6 @@ impl DeltaLayerInner {
             ),
         }
     }
-
-    /// NB: not super efficient, but not terrible either. Should prob be an iterator.
-    //
-    // We're reusing the index traversal logical in plan_reads; would be nice to
-    // factor that out.
-    pub(crate) async fn load_keys(&self, ctx: &RequestContext) -> anyhow::Result<Vec<Key>> {
-        self.index_entries(ctx)
-            .await
-            .map(|entries| entries.into_iter().map(|entry| entry.key).collect())
-    }
 }
 
 /// A set of data associated with a delta layer key and its value

pageserver/src/tenant/storage_layer/image_layer.rs

@@ -673,21 +673,6 @@ impl ImageLayerInner {
             ),
         }
     }
-
-    /// NB: not super efficient, but not terrible either. Should prob be an iterator.
-    //
-    // We're reusing the index traversal logical in plan_reads; would be nice to
-    // factor that out.
-    pub(crate) async fn load_keys(&self, ctx: &RequestContext) -> anyhow::Result<Vec<Key>> {
-        let plan = self
-            .plan_reads(KeySpace::single(self.key_range.clone()), None, ctx)
-            .await?;
-        Ok(plan
-            .into_iter()
-            .flat_map(|read| read.blobs_at)
-            .map(|(_, blob_meta)| blob_meta.key)
-            .collect())
-    }
 }
 
 /// A builder object for constructing a new image layer.
@@ -1024,7 +1009,7 @@ impl ImageLayerWriter {
         self.inner.take().unwrap().finish(ctx, None).await
     }
 
-    /// Finish writing the image layer with an end key, used in [`super::batch_split_writer::SplitImageLayerWriter`]. The end key determines the end of the image layer's covered range and is exclusive.
+    /// Finish writing the image layer with an end key, used in [`super::split_writer::SplitImageLayerWriter`]. The end key determines the end of the image layer's covered range and is exclusive.
     pub(super) async fn finish_with_end_key(
         mut self,
         end_key: Key,

pageserver/src/tenant/storage_layer/layer.rs

@@ -19,7 +19,7 @@ use crate::task_mgr::TaskKind;
 use crate::tenant::timeline::{CompactionError, GetVectoredError};
 use crate::tenant::{remote_timeline_client::LayerFileMetadata, Timeline};
 
-use super::delta_layer::{self};
+use super::delta_layer::{self, DeltaEntry};
 use super::image_layer::{self};
 use super::{
     AsLayerDesc, ImageLayerWriter, LayerAccessStats, LayerAccessStatsReset, LayerName,
@@ -1841,22 +1841,23 @@ impl ResidentLayer {
     pub(crate) async fn load_keys<'a>(
         &'a self,
         ctx: &RequestContext,
-    ) -> anyhow::Result<Vec<pageserver_api::key::Key>> {
+    ) -> anyhow::Result<Vec<DeltaEntry<'a>>> {
         use LayerKind::*;
 
         let owner = &self.owner.0;
-        let inner = self.downloaded.get(owner, ctx).await?;
+        match self.downloaded.get(owner, ctx).await? {
+            Delta(ref d) => {
+                // this is valid because the DownloadedLayer::kind is a OnceCell, not a
+                // Mutex<OnceCell>, so we cannot go and deinitialize the value with OnceCell::take
+                // while it's being held.
+                self.owner.record_access(ctx);
 
-        // this is valid because the DownloadedLayer::kind is a OnceCell, not a
-        // Mutex<OnceCell>, so we cannot go and deinitialize the value with OnceCell::take
-        // while it's being held.
-        self.owner.record_access(ctx);
-
-        let res = match inner {
-            Delta(ref d) => delta_layer::DeltaLayerInner::load_keys(d, ctx).await,
-            Image(ref i) => image_layer::ImageLayerInner::load_keys(i, ctx).await,
-        };
-        res.with_context(|| format!("Layer index is corrupted for {self}"))
+                delta_layer::DeltaLayerInner::load_keys(d, ctx)
+                    .await
+                    .with_context(|| format!("Layer index is corrupted for {self}"))
+            }
+            Image(_) => anyhow::bail!(format!("cannot load_keys on a image layer {self}")),
+        }
     }
 
     /// Read all they keys in this layer which match the ShardIdentity, and write them all to

pageserver/src/tenant/storage_layer/layer_desc.rs

@@ -57,34 +57,6 @@ impl std::fmt::Display for PersistentLayerKey {
     }
 }
 
-impl From<ImageLayerName> for PersistentLayerKey {
-    fn from(image_layer_name: ImageLayerName) -> Self {
-        Self {
-            key_range: image_layer_name.key_range,
-            lsn_range: PersistentLayerDesc::image_layer_lsn_range(image_layer_name.lsn),
-            is_delta: false,
-        }
-    }
-}
-
-impl From<DeltaLayerName> for PersistentLayerKey {
-    fn from(delta_layer_name: DeltaLayerName) -> Self {
-        Self {
-            key_range: delta_layer_name.key_range,
-            lsn_range: delta_layer_name.lsn_range,
-            is_delta: true,
-        }
-    }
-}
-
-impl From<LayerName> for PersistentLayerKey {
-    fn from(layer_name: LayerName) -> Self {
-        match layer_name {
-            LayerName::Image(i) => i.into(),
-            LayerName::Delta(d) => d.into(),
-        }
-    }
-}
 impl PersistentLayerDesc {
     pub fn key(&self) -> PersistentLayerKey {
         PersistentLayerKey {

pageserver/src/tenant/storage_layer/split_writer.rs

@@ -12,154 +12,41 @@ use super::{
     DeltaLayerWriter, ImageLayerWriter, PersistentLayerDesc, PersistentLayerKey, ResidentLayer,
 };
 
-pub(crate) enum BatchWriterResult {
+pub(crate) enum SplitWriterResult {
     Produced(ResidentLayer),
     Discarded(PersistentLayerKey),
 }
 
 #[cfg(test)]
-impl BatchWriterResult {
+impl SplitWriterResult {
     fn into_resident_layer(self) -> ResidentLayer {
         match self {
-            BatchWriterResult::Produced(layer) => layer,
-            BatchWriterResult::Discarded(_) => panic!("unexpected discarded layer"),
+            SplitWriterResult::Produced(layer) => layer,
+            SplitWriterResult::Discarded(_) => panic!("unexpected discarded layer"),
         }
     }
 
     fn into_discarded_layer(self) -> PersistentLayerKey {
         match self {
-            BatchWriterResult::Produced(_) => panic!("unexpected produced layer"),
-            BatchWriterResult::Discarded(layer) => layer,
+            SplitWriterResult::Produced(_) => panic!("unexpected produced layer"),
+            SplitWriterResult::Discarded(layer) => layer,
         }
     }
 }
 
-enum LayerWriterWrapper {
-    Image(ImageLayerWriter),
-    Delta(DeltaLayerWriter),
-}
-
-/// An layer writer that takes unfinished layers and finish them atomically.
-#[must_use]
-pub struct BatchLayerWriter {
-    generated_layer_writers: Vec<(LayerWriterWrapper, PersistentLayerKey)>,
-    conf: &'static PageServerConf,
-}
-
-impl BatchLayerWriter {
-    pub async fn new(conf: &'static PageServerConf) -> anyhow::Result<Self> {
-        Ok(Self {
-            generated_layer_writers: Vec::new(),
-            conf,
-        })
-    }
-
-    pub fn add_unfinished_image_writer(
-        &mut self,
-        writer: ImageLayerWriter,
-        key_range: Range<Key>,
-        lsn: Lsn,
-    ) {
-        self.generated_layer_writers.push((
-            LayerWriterWrapper::Image(writer),
-            PersistentLayerKey {
-                key_range,
-                lsn_range: PersistentLayerDesc::image_layer_lsn_range(lsn),
-                is_delta: false,
-            },
-        ));
-    }
-
-    pub fn add_unfinished_delta_writer(
-        &mut self,
-        writer: DeltaLayerWriter,
-        key_range: Range<Key>,
-        lsn_range: Range<Lsn>,
-    ) {
-        self.generated_layer_writers.push((
-            LayerWriterWrapper::Delta(writer),
-            PersistentLayerKey {
-                key_range,
-                lsn_range,
-                is_delta: true,
-            },
-        ));
-    }
-
-    pub(crate) async fn finish_with_discard_fn<D, F>(
-        self,
-        tline: &Arc<Timeline>,
-        ctx: &RequestContext,
-        discard_fn: D,
-    ) -> anyhow::Result<Vec<BatchWriterResult>>
-    where
-        D: Fn(&PersistentLayerKey) -> F,
-        F: Future<Output = bool>,
-    {
-        let Self {
-            generated_layer_writers,
-            ..
-        } = self;
-        let clean_up_layers = |generated_layers: Vec<BatchWriterResult>| {
-            for produced_layer in generated_layers {
-                if let BatchWriterResult::Produced(resident_layer) = produced_layer {
-                    let layer: Layer = resident_layer.into();
-                    layer.delete_on_drop();
-                }
-            }
-        };
-        // BEGIN: catch every error and do the recovery in the below section
-        let mut generated_layers: Vec<BatchWriterResult> = Vec::new();
-        for (inner, layer_key) in generated_layer_writers {
-            if discard_fn(&layer_key).await {
-                generated_layers.push(BatchWriterResult::Discarded(layer_key));
-            } else {
-                let res = match inner {
-                    LayerWriterWrapper::Delta(writer) => {
-                        writer.finish(layer_key.key_range.end, ctx).await
-                    }
-                    LayerWriterWrapper::Image(writer) => {
-                        writer
-                            .finish_with_end_key(layer_key.key_range.end, ctx)
-                            .await
-                    }
-                };
-                let layer = match res {
-                    Ok((desc, path)) => {
-                        match Layer::finish_creating(self.conf, tline, desc, &path) {
-                            Ok(layer) => layer,
-                            Err(e) => {
-                                tokio::fs::remove_file(&path).await.ok();
-                                clean_up_layers(generated_layers);
-                                return Err(e);
-                            }
-                        }
-                    }
-                    Err(e) => {
-                        // Image/DeltaLayerWriter::finish will clean up the temporary layer if anything goes wrong,
-                        // so we don't need to remove the layer we just failed to create by ourselves.
-                        clean_up_layers(generated_layers);
-                        return Err(e);
-                    }
-                };
-                generated_layers.push(BatchWriterResult::Produced(layer));
-            }
-        }
-        // END: catch every error and do the recovery in the above section
-        Ok(generated_layers)
-    }
-}
-
 /// An image writer that takes images and produces multiple image layers.
+///
+/// The interface does not guarantee atomicity (i.e., if the image layer generation
+/// fails, there might be leftover files to be cleaned up)
 #[must_use]
 pub struct SplitImageLayerWriter {
     inner: ImageLayerWriter,
     target_layer_size: u64,
-    lsn: Lsn,
+    generated_layer_writers: Vec<(ImageLayerWriter, PersistentLayerKey)>,
     conf: &'static PageServerConf,
     timeline_id: TimelineId,
     tenant_shard_id: TenantShardId,
-    batches: BatchLayerWriter,
+    lsn: Lsn,
     start_key: Key,
 }
 
@@ -184,10 +71,10 @@ impl SplitImageLayerWriter {
                 ctx,
             )
             .await?,
+            generated_layer_writers: Vec::new(),
             conf,
             timeline_id,
             tenant_shard_id,
-            batches: BatchLayerWriter::new(conf).await?,
             lsn,
             start_key,
         })
@@ -215,13 +102,16 @@ impl SplitImageLayerWriter {
                 ctx,
             )
             .await?;
+            let layer_key = PersistentLayerKey {
+                key_range: self.start_key..key,
+                lsn_range: PersistentLayerDesc::image_layer_lsn_range(self.lsn),
+                is_delta: false,
+            };
             let prev_image_writer = std::mem::replace(&mut self.inner, next_image_writer);
-            self.batches.add_unfinished_image_writer(
-                prev_image_writer,
-                self.start_key..key,
-                self.lsn,
-            );
             self.start_key = key;
+
+            self.generated_layer_writers
+                .push((prev_image_writer, layer_key));
         }
         self.inner.put_image(key, img, ctx).await
     }
@@ -232,18 +122,64 @@ impl SplitImageLayerWriter {
         ctx: &RequestContext,
         end_key: Key,
         discard_fn: D,
-    ) -> anyhow::Result<Vec<BatchWriterResult>>
+    ) -> anyhow::Result<Vec<SplitWriterResult>>
     where
         D: Fn(&PersistentLayerKey) -> F,
         F: Future<Output = bool>,
     {
         let Self {
-            mut batches, inner, ..
+            mut generated_layer_writers,
+            inner,
+            ..
         } = self;
         if inner.num_keys() != 0 {
-            batches.add_unfinished_image_writer(inner, self.start_key..end_key, self.lsn);
+            let layer_key = PersistentLayerKey {
+                key_range: self.start_key..end_key,
+                lsn_range: PersistentLayerDesc::image_layer_lsn_range(self.lsn),
+                is_delta: false,
+            };
+            generated_layer_writers.push((inner, layer_key));
         }
-        batches.finish_with_discard_fn(tline, ctx, discard_fn).await
+        let clean_up_layers = |generated_layers: Vec<SplitWriterResult>| {
+            for produced_layer in generated_layers {
+                if let SplitWriterResult::Produced(image_layer) = produced_layer {
+                    let layer: Layer = image_layer.into();
+                    layer.delete_on_drop();
+                }
+            }
+        };
+        // BEGIN: catch every error and do the recovery in the below section
+        let mut generated_layers = Vec::new();
+        for (inner, layer_key) in generated_layer_writers {
+            if discard_fn(&layer_key).await {
+                generated_layers.push(SplitWriterResult::Discarded(layer_key));
+            } else {
+                let layer = match inner
+                    .finish_with_end_key(layer_key.key_range.end, ctx)
+                    .await
+                {
+                    Ok((desc, path)) => {
+                        match Layer::finish_creating(self.conf, tline, desc, &path) {
+                            Ok(layer) => layer,
+                            Err(e) => {
+                                tokio::fs::remove_file(&path).await.ok();
+                                clean_up_layers(generated_layers);
+                                return Err(e);
+                            }
+                        }
+                    }
+                    Err(e) => {
+                        // ImageLayerWriter::finish will clean up the temporary layer if anything goes wrong,
+                        // so we don't need to remove the layer we just failed to create by ourselves.
+                        clean_up_layers(generated_layers);
+                        return Err(e);
+                    }
+                };
+                generated_layers.push(SplitWriterResult::Produced(layer));
+            }
+        }
+        // END: catch every error and do the recovery in the above section
+        Ok(generated_layers)
     }
 
     #[cfg(test)]
@@ -252,7 +188,7 @@ impl SplitImageLayerWriter {
         tline: &Arc<Timeline>,
         ctx: &RequestContext,
         end_key: Key,
-    ) -> anyhow::Result<Vec<BatchWriterResult>> {
+    ) -> anyhow::Result<Vec<SplitWriterResult>> {
         self.finish_with_discard_fn(tline, ctx, end_key, |_| async { false })
             .await
     }
@@ -260,6 +196,9 @@ impl SplitImageLayerWriter {
 
 /// A delta writer that takes key-lsn-values and produces multiple delta layers.
 ///
+/// The interface does not guarantee atomicity (i.e., if the delta layer generation fails,
+/// there might be leftover files to be cleaned up).
+///
 /// Note that if updates of a single key exceed the target size limit, all of the updates will be batched
 /// into a single file. This behavior might change in the future. For reference, the legacy compaction algorithm
 /// will split them into multiple files based on size.
@@ -267,12 +206,12 @@ impl SplitImageLayerWriter {
 pub struct SplitDeltaLayerWriter {
     inner: Option<(Key, DeltaLayerWriter)>,
     target_layer_size: u64,
+    generated_layer_writers: Vec<(DeltaLayerWriter, PersistentLayerKey)>,
     conf: &'static PageServerConf,
     timeline_id: TimelineId,
     tenant_shard_id: TenantShardId,
     lsn_range: Range<Lsn>,
     last_key_written: Key,
-    batches: BatchLayerWriter,
 }
 
 impl SplitDeltaLayerWriter {
@@ -286,12 +225,12 @@ impl SplitDeltaLayerWriter {
         Ok(Self {
             target_layer_size,
             inner: None,
+            generated_layer_writers: Vec::new(),
             conf,
             timeline_id,
             tenant_shard_id,
             lsn_range,
             last_key_written: Key::MIN,
-            batches: BatchLayerWriter::new(conf).await?,
         })
     }
 
@@ -340,11 +279,13 @@ impl SplitDeltaLayerWriter {
                 .await?;
                 let (start_key, prev_delta_writer) =
                     std::mem::replace(&mut self.inner, Some((key, next_delta_writer))).unwrap();
-                self.batches.add_unfinished_delta_writer(
-                    prev_delta_writer,
-                    start_key..key,
-                    self.lsn_range.clone(),
-                );
+                let layer_key = PersistentLayerKey {
+                    key_range: start_key..key,
+                    lsn_range: self.lsn_range.clone(),
+                    is_delta: true,
+                };
+                self.generated_layer_writers
+                    .push((prev_delta_writer, layer_key));
             } else if inner.estimated_size() >= S3_UPLOAD_LIMIT {
                 // We have to produce a very large file b/c a key is updated too often.
                 anyhow::bail!(
@@ -364,25 +305,64 @@ impl SplitDeltaLayerWriter {
         tline: &Arc<Timeline>,
         ctx: &RequestContext,
         discard_fn: D,
-    ) -> anyhow::Result<Vec<BatchWriterResult>>
+    ) -> anyhow::Result<Vec<SplitWriterResult>>
     where
         D: Fn(&PersistentLayerKey) -> F,
         F: Future<Output = bool>,
     {
         let Self {
-            mut batches, inner, ..
+            mut generated_layer_writers,
+            inner,
+            ..
         } = self;
         if let Some((start_key, writer)) = inner {
             if writer.num_keys() != 0 {
                 let end_key = self.last_key_written.next();
-                batches.add_unfinished_delta_writer(
-                    writer,
-                    start_key..end_key,
-                    self.lsn_range.clone(),
-                );
+                let layer_key = PersistentLayerKey {
+                    key_range: start_key..end_key,
+                    lsn_range: self.lsn_range.clone(),
+                    is_delta: true,
+                };
+                generated_layer_writers.push((writer, layer_key));
             }
         }
-        batches.finish_with_discard_fn(tline, ctx, discard_fn).await
+        let clean_up_layers = |generated_layers: Vec<SplitWriterResult>| {
+            for produced_layer in generated_layers {
+                if let SplitWriterResult::Produced(delta_layer) = produced_layer {
+                    let layer: Layer = delta_layer.into();
+                    layer.delete_on_drop();
+                }
+            }
+        };
+        // BEGIN: catch every error and do the recovery in the below section
+        let mut generated_layers = Vec::new();
+        for (inner, layer_key) in generated_layer_writers {
+            if discard_fn(&layer_key).await {
+                generated_layers.push(SplitWriterResult::Discarded(layer_key));
+            } else {
+                let layer = match inner.finish(layer_key.key_range.end, ctx).await {
+                    Ok((desc, path)) => {
+                        match Layer::finish_creating(self.conf, tline, desc, &path) {
+                            Ok(layer) => layer,
+                            Err(e) => {
+                                tokio::fs::remove_file(&path).await.ok();
+                                clean_up_layers(generated_layers);
+                                return Err(e);
+                            }
+                        }
+                    }
+                    Err(e) => {
+                        // DeltaLayerWriter::finish will clean up the temporary layer if anything goes wrong,
+                        // so we don't need to remove the layer we just failed to create by ourselves.
+                        clean_up_layers(generated_layers);
+                        return Err(e);
+                    }
+                };
+                generated_layers.push(SplitWriterResult::Produced(layer));
+            }
+        }
+        // END: catch every error and do the recovery in the above section
+        Ok(generated_layers)
     }
 
     #[cfg(test)]
@@ -390,7 +370,7 @@ impl SplitDeltaLayerWriter {
         self,
         tline: &Arc<Timeline>,
         ctx: &RequestContext,
-    ) -> anyhow::Result<Vec<BatchWriterResult>> {
+    ) -> anyhow::Result<Vec<SplitWriterResult>> {
         self.finish_with_discard_fn(tline, ctx, |_| async { false })
             .await
     }

pageserver/src/tenant/timeline.rs

@@ -424,9 +424,6 @@ pub struct Timeline {
     pub(crate) handles: handle::PerTimelineState<crate::page_service::TenantManagerTypes>,
 
     pub(crate) attach_wal_lag_cooldown: Arc<OnceLock<WalLagCooldown>>,
-
-    /// Cf. [`crate::tenant::CreateTimelineIdempotency`].
-    pub(crate) create_idempotency: crate::tenant::CreateTimelineIdempotency,
 }
 
 pub type TimelineDeleteProgress = Arc<tokio::sync::Mutex<DeleteTimelineFlow>>;
@@ -2139,7 +2136,6 @@ impl Timeline {
         pg_version: u32,
         state: TimelineState,
         attach_wal_lag_cooldown: Arc<OnceLock<WalLagCooldown>>,
-        create_idempotency: crate::tenant::CreateTimelineIdempotency,
         cancel: CancellationToken,
     ) -> Arc<Self> {
         let disk_consistent_lsn = metadata.disk_consistent_lsn();
@@ -2278,8 +2274,6 @@ impl Timeline {
                 handles: Default::default(),
 
                 attach_wal_lag_cooldown,
-
-                create_idempotency,
             };
 
             result.repartition_threshold =

pageserver/src/tenant/timeline/compaction.rs

@@ -32,11 +32,11 @@ use crate::page_cache;
 use crate::statvfs::Statvfs;
 use crate::tenant::checks::check_valid_layermap;
 use crate::tenant::remote_timeline_client::WaitCompletionError;
-use crate::tenant::storage_layer::batch_split_writer::{
-    BatchWriterResult, SplitDeltaLayerWriter, SplitImageLayerWriter,
-};
 use crate::tenant::storage_layer::filter_iterator::FilterIterator;
 use crate::tenant::storage_layer::merge_iterator::MergeIterator;
+use crate::tenant::storage_layer::split_writer::{
+    SplitDeltaLayerWriter, SplitImageLayerWriter, SplitWriterResult,
+};
 use crate::tenant::storage_layer::{
     AsLayerDesc, PersistentLayerDesc, PersistentLayerKey, ValueReconstructState,
 };
@@ -834,12 +834,7 @@ impl Timeline {
                 if self.cancel.is_cancelled() {
                     return Err(CompactionError::ShuttingDown);
                 }
-                let delta = l.get_as_delta(ctx).await.map_err(CompactionError::Other)?;
-                let keys = delta
-                    .index_entries(ctx)
-                    .await
-                    .map_err(CompactionError::Other)?;
-                all_keys.extend(keys);
+                all_keys.extend(l.load_keys(ctx).await.map_err(CompactionError::Other)?);
             }
             // The current stdlib sorting implementation is designed in a way where it is
             // particularly fast where the slice is made up of sorted sub-ranges.
@@ -2043,11 +2038,11 @@ impl Timeline {
         let produced_image_layers_len = produced_image_layers.len();
         for action in produced_delta_layers {
             match action {
-                BatchWriterResult::Produced(layer) => {
+                SplitWriterResult::Produced(layer) => {
                     stat.produce_delta_layer(layer.layer_desc().file_size());
                     compact_to.push(layer);
                 }
-                BatchWriterResult::Discarded(l) => {
+                SplitWriterResult::Discarded(l) => {
                     keep_layers.insert(l);
                     stat.discard_delta_layer();
                 }
@@ -2055,11 +2050,11 @@ impl Timeline {
         }
         for action in produced_image_layers {
             match action {
-                BatchWriterResult::Produced(layer) => {
+                SplitWriterResult::Produced(layer) => {
                     stat.produce_image_layer(layer.layer_desc().file_size());
                     compact_to.push(layer);
                 }
-                BatchWriterResult::Discarded(l) => {
+                SplitWriterResult::Discarded(l) => {
                     keep_layers.insert(l);
                     stat.discard_image_layer();
                 }
@@ -2443,7 +2438,7 @@ impl CompactionDeltaLayer<TimelineAdaptor> for ResidentDeltaLayer {
     type DeltaEntry<'a> = DeltaEntry<'a>;
 
     async fn load_keys<'a>(&self, ctx: &RequestContext) -> anyhow::Result<Vec<DeltaEntry<'_>>> {
-        self.0.get_as_delta(ctx).await?.index_entries(ctx).await
+        self.0.load_keys(ctx).await
     }
 }
 

pageserver/src/tenant/timeline/delete.rs

@@ -313,7 +313,6 @@ impl DeleteTimelineFlow {
                 // Important. We dont pass ancestor above because it can be missing.
                 // Thus we need to skip the validation here.
                 CreateTimelineCause::Delete,
-                crate::tenant::CreateTimelineIdempotency::FailWithConflict, // doesn't matter what we put here
             )
             .context("create_timeline_struct")?;
 

pageserver/src/tenant/timeline/layer_manager.rs

@@ -45,16 +45,13 @@ impl LayerManager {
     pub(crate) fn get_from_key(&self, key: &PersistentLayerKey) -> Layer {
         // The assumption for the `expect()` is that all code maintains the following invariant:
         // A layer's descriptor is present in the LayerMap => the LayerFileManager contains a layer for the descriptor.
-        self.try_get_from_key(key)
+        self.layers()
+            .get(key)
             .with_context(|| format!("get layer from key: {key}"))
             .expect("not found")
             .clone()
     }
 
-    pub(crate) fn try_get_from_key(&self, key: &PersistentLayerKey) -> Option<&Layer> {
-        self.layers().get(key)
-    }
-
     pub(crate) fn get_from_desc(&self, desc: &PersistentLayerDesc) -> Layer {
         self.get_from_key(&desc.key())
     }

pageserver/src/tenant/timeline/uninit.rs

@@ -5,11 +5,7 @@ use camino::Utf8PathBuf;
 use tracing::{error, info, info_span};
 use utils::{fs_ext, id::TimelineId, lsn::Lsn};
 
-use crate::{
-    context::RequestContext,
-    import_datadir,
-    tenant::{CreateTimelineIdempotency, Tenant, TimelineOrOffloaded},
-};
+use crate::{context::RequestContext, import_datadir, tenant::Tenant};
 
 use super::Timeline;
 
@@ -169,17 +165,13 @@ pub(crate) struct TimelineCreateGuard<'t> {
     owning_tenant: &'t Tenant,
     timeline_id: TimelineId,
     pub(crate) timeline_path: Utf8PathBuf,
-    pub(crate) idempotency: CreateTimelineIdempotency,
 }
 
 /// Errors when acquiring exclusive access to a timeline ID for creation
 #[derive(thiserror::Error, Debug)]
 pub(crate) enum TimelineExclusionError {
     #[error("Already exists")]
-    AlreadyExists {
-        existing: TimelineOrOffloaded,
-        arg: CreateTimelineIdempotency,
-    },
+    AlreadyExists(Arc<Timeline>),
     #[error("Already creating")]
     AlreadyCreating,
 
@@ -193,42 +185,27 @@ impl<'t> TimelineCreateGuard<'t> {
         owning_tenant: &'t Tenant,
         timeline_id: TimelineId,
         timeline_path: Utf8PathBuf,
-        idempotency: CreateTimelineIdempotency,
-        allow_offloaded: bool,
     ) -> Result<Self, TimelineExclusionError> {
         // Lock order: this is the only place we take both locks.  During drop() we only
         // lock creating_timelines
         let timelines = owning_tenant.timelines.lock().unwrap();
-        let timelines_offloaded = owning_tenant.timelines_offloaded.lock().unwrap();
         let mut creating_timelines: std::sync::MutexGuard<
             '_,
             std::collections::HashSet<TimelineId>,
         > = owning_tenant.timelines_creating.lock().unwrap();
 
         if let Some(existing) = timelines.get(&timeline_id) {
-            return Err(TimelineExclusionError::AlreadyExists {
-                existing: TimelineOrOffloaded::Timeline(existing.clone()),
-                arg: idempotency,
-            });
+            Err(TimelineExclusionError::AlreadyExists(existing.clone()))
+        } else if creating_timelines.contains(&timeline_id) {
+            Err(TimelineExclusionError::AlreadyCreating)
+        } else {
+            creating_timelines.insert(timeline_id);
+            Ok(Self {
+                owning_tenant,
+                timeline_id,
+                timeline_path,
+            })
         }
-        if !allow_offloaded {
-            if let Some(existing) = timelines_offloaded.get(&timeline_id) {
-                return Err(TimelineExclusionError::AlreadyExists {
-                    existing: TimelineOrOffloaded::Offloaded(existing.clone()),
-                    arg: idempotency,
-                });
-            }
-        }
-        if creating_timelines.contains(&timeline_id) {
-            return Err(TimelineExclusionError::AlreadyCreating);
-        }
-        creating_timelines.insert(timeline_id);
-        Ok(Self {
-            owning_tenant,
-            timeline_id,
-            timeline_path,
-            idempotency,
-        })
     }
 }
 

pageserver/src/virtual_file/io_engine/tokio_epoll_uring_ext.rs

@@ -16,24 +16,18 @@ use tokio_epoll_uring::{System, SystemHandle};
 
 use crate::virtual_file::on_fatal_io_error;
 
-use crate::metrics::tokio_epoll_uring::{self as metrics, THREAD_LOCAL_METRICS_STORAGE};
+use crate::metrics::tokio_epoll_uring as metrics;
 
 #[derive(Clone)]
 struct ThreadLocalState(Arc<ThreadLocalStateInner>);
 
 struct ThreadLocalStateInner {
-    cell: tokio::sync::OnceCell<SystemHandle<metrics::ThreadLocalMetrics>>,
+    cell: tokio::sync::OnceCell<SystemHandle>,
     launch_attempts: AtomicU32,
     /// populated through fetch_add from [`THREAD_LOCAL_STATE_ID`]
     thread_local_state_id: u64,
 }
 
-impl Drop for ThreadLocalStateInner {
-    fn drop(&mut self) {
-        THREAD_LOCAL_METRICS_STORAGE.remove_system(self.thread_local_state_id);
-    }
-}
-
 impl ThreadLocalState {
     pub fn new() -> Self {
         Self(Arc::new(ThreadLocalStateInner {
@@ -77,8 +71,7 @@ pub async fn thread_local_system() -> Handle {
                         &fake_cancel,
                     )
                     .await;
-                    let per_system_metrics = metrics::THREAD_LOCAL_METRICS_STORAGE.register_system(inner.thread_local_state_id);
-                    let res = System::launch_with_metrics(per_system_metrics)
+                    let res = System::launch()
                     // this might move us to another executor thread => loop outside the get_or_try_init, not inside it
                     .await;
                     match res {
@@ -93,7 +86,6 @@ pub async fn thread_local_system() -> Handle {
                                 emit_launch_failure_process_stats();
                             });
                             metrics::THREAD_LOCAL_LAUNCH_FAILURES.inc();
-                            metrics::THREAD_LOCAL_METRICS_STORAGE.remove_system(inner.thread_local_state_id);
                             Err(())
                         }
                         // abort the process instead of panicking because pageserver usually becomes half-broken if we panic somewhere.
@@ -123,7 +115,7 @@ fn emit_launch_failure_process_stats() {
     // number of threads
     // rss / system memory usage generally
 
-    let tokio_epoll_uring::metrics::GlobalMetrics {
+    let tokio_epoll_uring::metrics::Metrics {
         systems_created,
         systems_destroyed,
     } = tokio_epoll_uring::metrics::global();
@@ -190,7 +182,7 @@ fn emit_launch_failure_process_stats() {
 pub struct Handle(ThreadLocalState);
 
 impl std::ops::Deref for Handle {
-    type Target = SystemHandle<metrics::ThreadLocalMetrics>;
+    type Target = SystemHandle;
 
     fn deref(&self) -> &Self::Target {
         self.0

pgxn/neon/Makefile

@@ -8,7 +8,6 @@ OBJS = \
 	file_cache.o \
 	hll.o \
 	libpagestore.o \
-	logical_replication_monitor.o \
 	neon.o \
 	neon_pgversioncompat.o \
 	neon_perf_counters.o \

pgxn/neon/logical_replication_monitor.c

@@ -1,253 +0,0 @@
-#include <limits.h>
-#include <string.h>
-#include <dirent.h>
-#include <signal.h>
-
-#include "postgres.h"
-
-#include "miscadmin.h"
-#include "postmaster/bgworker.h"
-#include "postmaster/interrupt.h"
-#include "replication/slot.h"
-#include "storage/fd.h"
-#include "storage/procsignal.h"
-#include "tcop/tcopprot.h"
-#include "utils/guc.h"
-#include "utils/wait_event.h"
-
-#include "logical_replication_monitor.h"
-
-#define LS_MONITOR_CHECK_INTERVAL 10000 /* ms */
-
-static int	logical_replication_max_snap_files = 300;
-
-PGDLLEXPORT void LogicalSlotsMonitorMain(Datum main_arg);
-
-static int
-LsnDescComparator(const void *a, const void *b)
-{
-	XLogRecPtr	lsn1 = *((const XLogRecPtr *) a);
-	XLogRecPtr	lsn2 = *((const XLogRecPtr *) b);
-
-	if (lsn1 < lsn2)
-		return 1;
-	else if (lsn1 == lsn2)
-		return 0;
-	else
-		return -1;
-}
-
-/*
- * Look at .snap files and calculate minimum allowed restart_lsn of slot so that
- * next gc would leave not more than logical_replication_max_snap_files; all
- * slots having lower restart_lsn should be dropped.
- */
-static XLogRecPtr
-get_num_snap_files_lsn_threshold(void)
-{
-	DIR		   *dirdesc;
-	struct dirent *de;
-	char	   *snap_path = "pg_logical/snapshots/";
-	int			lsns_allocated = 1024;
-	int			lsns_num = 0;
-	XLogRecPtr *lsns;
-	XLogRecPtr	cutoff;
-
-	if (logical_replication_max_snap_files < 0)
-		return 0;
-
-	lsns = palloc(sizeof(XLogRecPtr) * lsns_allocated);
-
-	/* find all .snap files and get their lsns */
-	dirdesc = AllocateDir(snap_path);
-	while ((de = ReadDir(dirdesc, snap_path)) != NULL)
-	{
-		XLogRecPtr	lsn;
-		uint32		hi;
-		uint32		lo;
-
-		if (strcmp(de->d_name, ".") == 0 ||
-			strcmp(de->d_name, "..") == 0)
-			continue;
-
-		if (sscanf(de->d_name, "%X-%X.snap", &hi, &lo) != 2)
-		{
-			ereport(LOG,
-					(errmsg("could not parse file name as .snap file \"%s\"", de->d_name)));
-			continue;
-		}
-
-		lsn = ((uint64) hi) << 32 | lo;
-		elog(DEBUG5, "found snap file %X/%X", LSN_FORMAT_ARGS(lsn));
-		if (lsns_allocated == lsns_num)
-		{
-			lsns_allocated *= 2;
-			lsns = repalloc(lsns, sizeof(XLogRecPtr) * lsns_allocated);
-		}
-		lsns[lsns_num++] = lsn;
-	}
-	/* sort by lsn desc */
-	qsort(lsns, lsns_num, sizeof(XLogRecPtr), LsnDescComparator);
-	/* and take cutoff at logical_replication_max_snap_files */
-	if (logical_replication_max_snap_files > lsns_num)
-		cutoff = 0;
-	/* have less files than cutoff */
-	else
-	{
-		cutoff = lsns[logical_replication_max_snap_files - 1];
-		elog(LOG, "ls_monitor: dropping logical slots with restart_lsn lower %X/%X, found %d .snap files, limit is %d",
-			 LSN_FORMAT_ARGS(cutoff), lsns_num, logical_replication_max_snap_files);
-	}
-	pfree(lsns);
-	FreeDir(dirdesc);
-	return cutoff;
-}
-
-void
-InitLogicalReplicationMonitor(void)
-{
-	BackgroundWorker bgw;
-
-	DefineCustomIntVariable(
-							"neon.logical_replication_max_snap_files",
-							"Maximum allowed logical replication .snap files. When exceeded, slots are dropped until the limit is met. -1 disables the limit.",
-							NULL,
-							&logical_replication_max_snap_files,
-							300, -1, INT_MAX,
-							PGC_SIGHUP,
-							0,
-							NULL, NULL, NULL);
-
-	memset(&bgw, 0, sizeof(bgw));
-	bgw.bgw_flags = BGWORKER_SHMEM_ACCESS;
-	bgw.bgw_start_time = BgWorkerStart_RecoveryFinished;
-	snprintf(bgw.bgw_library_name, BGW_MAXLEN, "neon");
-	snprintf(bgw.bgw_function_name, BGW_MAXLEN, "LogicalSlotsMonitorMain");
-	snprintf(bgw.bgw_name, BGW_MAXLEN, "Logical replication monitor");
-	snprintf(bgw.bgw_type, BGW_MAXLEN, "Logical replication monitor");
-	bgw.bgw_restart_time = 5;
-	bgw.bgw_notify_pid = 0;
-	bgw.bgw_main_arg = (Datum) 0;
-
-	RegisterBackgroundWorker(&bgw);
-}
-
-/*
- * Unused logical replication slots pins WAL and prevents deletion of snapshots.
- * WAL bloat is guarded by max_slot_wal_keep_size; this bgw removes slots which
- * need too many .snap files.
- */
-void
-LogicalSlotsMonitorMain(Datum main_arg)
-{
-	/* Establish signal handlers. */
-	pqsignal(SIGUSR1, procsignal_sigusr1_handler);
-	pqsignal(SIGHUP, SignalHandlerForConfigReload);
-	pqsignal(SIGTERM, die);
-
-	BackgroundWorkerUnblockSignals();
-
-	for (;;)
-	{
-		XLogRecPtr	cutoff_lsn;
-
-		/* In case of a SIGHUP, just reload the configuration. */
-		if (ConfigReloadPending)
-		{
-			ConfigReloadPending = false;
-			ProcessConfigFile(PGC_SIGHUP);
-		}
-
-		/*
-		 * If there are too many .snap files, just drop all logical slots to
-		 * prevent aux files bloat.
-		 */
-		cutoff_lsn = get_num_snap_files_lsn_threshold();
-		if (cutoff_lsn > 0)
-		{
-			for (int i = 0; i < max_replication_slots; i++)
-			{
-				char		slot_name[NAMEDATALEN];
-				ReplicationSlot *s = &ReplicationSlotCtl->replication_slots[i];
-				XLogRecPtr	restart_lsn;
-
-				/* find the name */
-				LWLockAcquire(ReplicationSlotControlLock, LW_SHARED);
-				/* Consider only logical repliction slots */
-				if (!s->in_use || !SlotIsLogical(s))
-				{
-					LWLockRelease(ReplicationSlotControlLock);
-					continue;
-				}
-
-				/* do we need to drop it? */
-				SpinLockAcquire(&s->mutex);
-				restart_lsn = s->data.restart_lsn;
-				SpinLockRelease(&s->mutex);
-				if (restart_lsn >= cutoff_lsn)
-				{
-					LWLockRelease(ReplicationSlotControlLock);
-					continue;
-				}
-
-				strlcpy(slot_name, s->data.name.data, NAMEDATALEN);
-				elog(LOG, "ls_monitor: dropping slot %s with restart_lsn %X/%X below horizon %X/%X",
-					 slot_name, LSN_FORMAT_ARGS(restart_lsn), LSN_FORMAT_ARGS(cutoff_lsn));
-				LWLockRelease(ReplicationSlotControlLock);
-
-				/* now try to drop it, killing owner before if any */
-				for (;;)
-				{
-					pid_t		active_pid;
-
-					SpinLockAcquire(&s->mutex);
-					active_pid = s->active_pid;
-					SpinLockRelease(&s->mutex);
-
-					if (active_pid == 0)
-					{
-						/*
-						 * Slot is releasted, try to drop it. Though of course
-						 * it could have been reacquired, so drop can ERROR
-						 * out. Similarly it could have been dropped in the
-						 * meanwhile.
-						 *
-						 * In principle we could remove pg_try/pg_catch, that
-						 * would restart the whole bgworker.
-						 */
-						ConditionVariableCancelSleep();
-						PG_TRY();
-						{
-							ReplicationSlotDrop(slot_name, true);
-							elog(LOG, "ls_monitor: slot %s dropped", slot_name);
-						}
-						PG_CATCH();
-						{
-							/* log ERROR and reset elog stack */
-							EmitErrorReport();
-							FlushErrorState();
-							elog(LOG, "ls_monitor: failed to drop slot %s", slot_name);
-						}
-						PG_END_TRY();
-						break;
-					}
-					else
-					{
-						/* kill the owner and wait for release */
-						elog(LOG, "ls_monitor: killing slot %s owner %d", slot_name, active_pid);
-						(void) kill(active_pid, SIGTERM);
-						/* We shouldn't get stuck, but to be safe add timeout. */
-						ConditionVariableTimedSleep(&s->active_cv, 1000, WAIT_EVENT_REPLICATION_SLOT_DROP);
-					}
-				}
-			}
-		}
-
-		(void) WaitLatch(MyLatch,
-						 WL_LATCH_SET | WL_EXIT_ON_PM_DEATH | WL_TIMEOUT,
-						 LS_MONITOR_CHECK_INTERVAL,
-						 PG_WAIT_EXTENSION);
-		ResetLatch(MyLatch);
-		CHECK_FOR_INTERRUPTS();
-	}
-}

pgxn/neon/logical_replication_monitor.h

@@ -1,6 +0,0 @@
-#ifndef __NEON_LOGICAL_REPLICATION_MONITOR_H__
-#define __NEON_LOGICAL_REPLICATION_MONITOR_H__
-
-void InitLogicalReplicationMonitor(void);
-
-#endif

pgxn/neon/neon.c

@@ -14,22 +14,32 @@
 #include "miscadmin.h"
 #include "access/subtrans.h"
 #include "access/twophase.h"
+#include "access/xact.h"
 #include "access/xlog.h"
+#include "storage/buf_internals.h"
+#include "storage/bufmgr.h"
+#include "catalog/pg_type.h"
+#include "postmaster/bgworker.h"
+#include "postmaster/interrupt.h"
 #include "replication/logical.h"
 #include "replication/slot.h"
 #include "replication/walsender.h"
 #include "storage/proc.h"
+#include "storage/procsignal.h"
+#include "tcop/tcopprot.h"
 #include "funcapi.h"
 #include "access/htup_details.h"
 #include "utils/builtins.h"
 #include "utils/pg_lsn.h"
 #include "utils/guc.h"
 #include "utils/guc_tables.h"
+#include "utils/wait_event.h"
 
 #include "extension_server.h"
 #include "neon.h"
+#include "walproposer.h"
+#include "pagestore_client.h"
 #include "control_plane_connector.h"
-#include "logical_replication_monitor.h"
 #include "walsender_hooks.h"
 #if PG_MAJORVERSION_NUM >= 16
 #include "storage/ipc.h"
@@ -38,6 +48,7 @@
 PG_MODULE_MAGIC;
 void		_PG_init(void);
 
+static int	logical_replication_max_snap_files = 300;
 
 static int  running_xacts_overflow_policy;
 
@@ -71,6 +82,237 @@ static const struct config_enum_entry running_xacts_overflow_policies[] = {
 	{NULL, 0, false}
 };
 
+static void
+InitLogicalReplicationMonitor(void)
+{
+	BackgroundWorker bgw;
+
+	DefineCustomIntVariable(
+							"neon.logical_replication_max_snap_files",
+							"Maximum allowed logical replication .snap files. When exceeded, slots are dropped until the limit is met. -1 disables the limit.",
+							NULL,
+							&logical_replication_max_snap_files,
+							300, -1, INT_MAX,
+							PGC_SIGHUP,
+							0,
+							NULL, NULL, NULL);
+
+	memset(&bgw, 0, sizeof(bgw));
+	bgw.bgw_flags = BGWORKER_SHMEM_ACCESS;
+	bgw.bgw_start_time = BgWorkerStart_RecoveryFinished;
+	snprintf(bgw.bgw_library_name, BGW_MAXLEN, "neon");
+	snprintf(bgw.bgw_function_name, BGW_MAXLEN, "LogicalSlotsMonitorMain");
+	snprintf(bgw.bgw_name, BGW_MAXLEN, "Logical replication monitor");
+	snprintf(bgw.bgw_type, BGW_MAXLEN, "Logical replication monitor");
+	bgw.bgw_restart_time = 5;
+	bgw.bgw_notify_pid = 0;
+	bgw.bgw_main_arg = (Datum) 0;
+
+	RegisterBackgroundWorker(&bgw);
+}
+
+static int
+LsnDescComparator(const void *a, const void *b)
+{
+	XLogRecPtr	lsn1 = *((const XLogRecPtr *) a);
+	XLogRecPtr	lsn2 = *((const XLogRecPtr *) b);
+
+	if (lsn1 < lsn2)
+		return 1;
+	else if (lsn1 == lsn2)
+		return 0;
+	else
+		return -1;
+}
+
+/*
+ * Look at .snap files and calculate minimum allowed restart_lsn of slot so that
+ * next gc would leave not more than logical_replication_max_snap_files; all
+ * slots having lower restart_lsn should be dropped.
+ */
+static XLogRecPtr
+get_num_snap_files_lsn_threshold(void)
+{
+	DIR		   *dirdesc;
+	struct dirent *de;
+	char	   *snap_path = "pg_logical/snapshots/";
+	int			lsns_allocated = 1024;
+	int			lsns_num = 0;
+	XLogRecPtr *lsns;
+	XLogRecPtr	cutoff;
+
+	if (logical_replication_max_snap_files < 0)
+		return 0;
+
+	lsns = palloc(sizeof(XLogRecPtr) * lsns_allocated);
+
+	/* find all .snap files and get their lsns */
+	dirdesc = AllocateDir(snap_path);
+	while ((de = ReadDir(dirdesc, snap_path)) != NULL)
+	{
+		XLogRecPtr	lsn;
+		uint32		hi;
+		uint32		lo;
+
+		if (strcmp(de->d_name, ".") == 0 ||
+			strcmp(de->d_name, "..") == 0)
+			continue;
+
+		if (sscanf(de->d_name, "%X-%X.snap", &hi, &lo) != 2)
+		{
+			ereport(LOG,
+					(errmsg("could not parse file name as .snap file \"%s\"", de->d_name)));
+			continue;
+		}
+
+		lsn = ((uint64) hi) << 32 | lo;
+		elog(DEBUG5, "found snap file %X/%X", LSN_FORMAT_ARGS(lsn));
+		if (lsns_allocated == lsns_num)
+		{
+			lsns_allocated *= 2;
+			lsns = repalloc(lsns, sizeof(XLogRecPtr) * lsns_allocated);
+		}
+		lsns[lsns_num++] = lsn;
+	}
+	/* sort by lsn desc */
+	qsort(lsns, lsns_num, sizeof(XLogRecPtr), LsnDescComparator);
+	/* and take cutoff at logical_replication_max_snap_files */
+	if (logical_replication_max_snap_files > lsns_num)
+		cutoff = 0;
+	/* have less files than cutoff */
+	else
+	{
+		cutoff = lsns[logical_replication_max_snap_files - 1];
+		elog(LOG, "ls_monitor: dropping logical slots with restart_lsn lower %X/%X, found %d .snap files, limit is %d",
+			 LSN_FORMAT_ARGS(cutoff), lsns_num, logical_replication_max_snap_files);
+	}
+	pfree(lsns);
+	FreeDir(dirdesc);
+	return cutoff;
+}
+
+#define LS_MONITOR_CHECK_INTERVAL 10000 /* ms */
+
+/*
+ * Unused logical replication slots pins WAL and prevents deletion of snapshots.
+ * WAL bloat is guarded by max_slot_wal_keep_size; this bgw removes slots which
+ * need too many .snap files.
+ */
+PGDLLEXPORT void
+LogicalSlotsMonitorMain(Datum main_arg)
+{
+	/* Establish signal handlers. */
+	pqsignal(SIGUSR1, procsignal_sigusr1_handler);
+	pqsignal(SIGHUP, SignalHandlerForConfigReload);
+	pqsignal(SIGTERM, die);
+
+	BackgroundWorkerUnblockSignals();
+
+	for (;;)
+	{
+		XLogRecPtr	cutoff_lsn;
+
+		/* In case of a SIGHUP, just reload the configuration. */
+		if (ConfigReloadPending)
+		{
+			ConfigReloadPending = false;
+			ProcessConfigFile(PGC_SIGHUP);
+		}
+
+		/*
+		 * If there are too many .snap files, just drop all logical slots to
+		 * prevent aux files bloat.
+		 */
+		cutoff_lsn = get_num_snap_files_lsn_threshold();
+		if (cutoff_lsn > 0)
+		{
+			for (int i = 0; i < max_replication_slots; i++)
+			{
+				char		slot_name[NAMEDATALEN];
+				ReplicationSlot *s = &ReplicationSlotCtl->replication_slots[i];
+				XLogRecPtr	restart_lsn;
+
+				/* find the name */
+				LWLockAcquire(ReplicationSlotControlLock, LW_SHARED);
+				/* Consider only logical repliction slots */
+				if (!s->in_use || !SlotIsLogical(s))
+				{
+					LWLockRelease(ReplicationSlotControlLock);
+					continue;
+				}
+
+				/* do we need to drop it? */
+				SpinLockAcquire(&s->mutex);
+				restart_lsn = s->data.restart_lsn;
+				SpinLockRelease(&s->mutex);
+				if (restart_lsn >= cutoff_lsn)
+				{
+					LWLockRelease(ReplicationSlotControlLock);
+					continue;
+				}
+
+				strlcpy(slot_name, s->data.name.data, NAMEDATALEN);
+				elog(LOG, "ls_monitor: dropping slot %s with restart_lsn %X/%X below horizon %X/%X",
+					 slot_name, LSN_FORMAT_ARGS(restart_lsn), LSN_FORMAT_ARGS(cutoff_lsn));
+				LWLockRelease(ReplicationSlotControlLock);
+
+				/* now try to drop it, killing owner before if any */
+				for (;;)
+				{
+					pid_t		active_pid;
+
+					SpinLockAcquire(&s->mutex);
+					active_pid = s->active_pid;
+					SpinLockRelease(&s->mutex);
+
+					if (active_pid == 0)
+					{
+						/*
+						 * Slot is releasted, try to drop it. Though of course
+						 * it could have been reacquired, so drop can ERROR
+						 * out. Similarly it could have been dropped in the
+						 * meanwhile.
+						 *
+						 * In principle we could remove pg_try/pg_catch, that
+						 * would restart the whole bgworker.
+						 */
+						ConditionVariableCancelSleep();
+						PG_TRY();
+						{
+							ReplicationSlotDrop(slot_name, true);
+							elog(LOG, "ls_monitor: slot %s dropped", slot_name);
+						}
+						PG_CATCH();
+						{
+							/* log ERROR and reset elog stack */
+							EmitErrorReport();
+							FlushErrorState();
+							elog(LOG, "ls_monitor: failed to drop slot %s", slot_name);
+						}
+						PG_END_TRY();
+						break;
+					}
+					else
+					{
+						/* kill the owner and wait for release */
+						elog(LOG, "ls_monitor: killing slot %s owner %d", slot_name, active_pid);
+						(void) kill(active_pid, SIGTERM);
+						/* We shouldn't get stuck, but to be safe add timeout. */
+						ConditionVariableTimedSleep(&s->active_cv, 1000, WAIT_EVENT_REPLICATION_SLOT_DROP);
+					}
+				}
+			}
+		}
+
+		(void) WaitLatch(MyLatch,
+						 WL_LATCH_SET | WL_EXIT_ON_PM_DEATH | WL_TIMEOUT,
+						 LS_MONITOR_CHECK_INTERVAL,
+						 PG_WAIT_EXTENSION);
+		ResetLatch(MyLatch);
+		CHECK_FOR_INTERRUPTS();
+	}
+}
+
 /*
  * XXX: These private to procarray.c, but we need them here.
  */
@@ -425,6 +667,7 @@ _PG_init(void)
 	SlotFuncs_Custom_XLogReaderRoutines = NeonOnDemandXLogReaderRoutines;
 
 	InitLogicalReplicationMonitor();
+
 	InitControlPlaneConnector();
 
 	pg_init_extension_server();

proxy/src/auth/backend/jwt.rs

@@ -5,7 +5,6 @@ use std::time::{Duration, SystemTime};
 use arc_swap::ArcSwapOption;
 use dashmap::DashMap;
 use jose_jwk::crypto::KeyInfo;
-use reqwest::{redirect, Client};
 use serde::de::Visitor;
 use serde::{Deserialize, Deserializer};
 use signature::Verifier;
@@ -25,7 +24,6 @@ const MIN_RENEW: Duration = Duration::from_secs(30);
 const AUTO_RENEW: Duration = Duration::from_secs(300);
 const MAX_RENEW: Duration = Duration::from_secs(3600);
 const MAX_JWK_BODY_SIZE: usize = 64 * 1024;
-const JWKS_USER_AGENT: &str = "neon-proxy";
 
 /// How to get the JWT auth rules
 pub(crate) trait FetchAuthRules: Clone + Send + Sync + 'static {
@@ -52,6 +50,7 @@ pub(crate) struct AuthRule {
     pub(crate) role_names: Vec<RoleNameInt>,
 }
 
+#[derive(Default)]
 pub struct JwkCache {
     client: reqwest::Client,
 
@@ -358,20 +357,6 @@ impl JwkCache {
     }
 }
 
-impl Default for JwkCache {
-    fn default() -> Self {
-        let client = Client::builder()
-            .user_agent(JWKS_USER_AGENT)
-            .redirect(redirect::Policy::none())
-            .build()
-            .expect("using &str and standard redirect::Policy");
-        JwkCache {
-            client,
-            map: DashMap::default(),
-        }
-    }
-}
-
 fn verify_ec_signature(data: &[u8], sig: &[u8], key: &jose_jwk::Ec) -> Result<(), JwtError> {
     use ecdsa::Signature;
     use signature::Verifier;

proxy/src/auth/backend/mod.rs

@@ -21,10 +21,7 @@ use crate::auth::{self, validate_password_and_exchange, AuthError, ComputeUserIn
 use crate::cache::Cached;
 use crate::config::AuthenticationConfig;
 use crate::context::RequestMonitoring;
-use crate::control_plane::errors::GetAuthInfoError;
-use crate::control_plane::provider::{
-    CachedAllowedIps, CachedNodeInfo, CachedRoleSecret, ControlPlaneBackend,
-};
+use crate::control_plane::provider::{CachedNodeInfo, ControlPlaneBackend};
 use crate::control_plane::{self, Api, AuthSecret};
 use crate::intern::EndpointIdInt;
 use crate::metrics::Metrics;
@@ -35,38 +32,19 @@ use crate::stream::Stream;
 use crate::types::{EndpointCacheKey, EndpointId, RoleName};
 use crate::{scram, stream};
 
-/// Alternative to [`std::borrow::Cow`] but doesn't need `T: ToOwned` as we don't need that functionality
-pub enum MaybeOwned<'a, T> {
-    Owned(T),
-    Borrowed(&'a T),
-}
-
-impl<T> std::ops::Deref for MaybeOwned<'_, T> {
-    type Target = T;
-
-    fn deref(&self) -> &Self::Target {
-        match self {
-            MaybeOwned::Owned(t) => t,
-            MaybeOwned::Borrowed(t) => t,
-        }
-    }
-}
-
-/// This type serves two purposes:
-///
-/// * When `T` is `()`, it's just a regular auth backend selector
-///   which we use in [`crate::config::ProxyConfig`].
-///
-/// * However, when we substitute `T` with [`ComputeUserInfoMaybeEndpoint`],
-///   this helps us provide the credentials only to those auth
-///   backends which require them for the authentication process.
-pub enum Backend<'a, T> {
+/// The [crate::serverless] module can authenticate either using control-plane
+/// to get authentication state, or by using JWKs stored in the filesystem.
+#[derive(Clone, Copy)]
+pub enum ServerlessBackend<'a> {
     /// Cloud API (V2).
-    ControlPlane(MaybeOwned<'a, ControlPlaneBackend>, T),
+    ControlPlane(&'a ControlPlaneBackend),
     /// Local proxy uses configured auth credentials and does not wake compute
-    Local(MaybeOwned<'a, LocalBackend>),
+    Local(&'a LocalBackend),
 }
 
+#[cfg(test)]
+use crate::control_plane::provider::{CachedAllowedIps, CachedRoleSecret};
+
 #[cfg(test)]
 pub(crate) trait TestBackend: Send + Sync + 'static {
     fn wake_compute(&self) -> Result<CachedNodeInfo, control_plane::errors::WakeComputeError>;
@@ -83,56 +61,20 @@ impl Clone for Box<dyn TestBackend> {
     }
 }
 
-impl std::fmt::Display for Backend<'_, ()> {
+impl std::fmt::Display for ControlPlaneBackend {
     fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            Self::ControlPlane(api, ()) => match &**api {
-                ControlPlaneBackend::Management(endpoint) => fmt
-                    .debug_tuple("ControlPlane::Management")
-                    .field(&endpoint.url())
-                    .finish(),
-                #[cfg(any(test, feature = "testing"))]
-                ControlPlaneBackend::PostgresMock(endpoint) => fmt
-                    .debug_tuple("ControlPlane::PostgresMock")
-                    .field(&endpoint.url())
-                    .finish(),
-                #[cfg(test)]
-                ControlPlaneBackend::Test(_) => fmt.debug_tuple("ControlPlane::Test").finish(),
-            },
-            Self::Local(_) => fmt.debug_tuple("Local").finish(),
-        }
-    }
-}
-
-impl<T> Backend<'_, T> {
-    /// Very similar to [`std::option::Option::as_ref`].
-    /// This helps us pass structured config to async tasks.
-    pub(crate) fn as_ref(&self) -> Backend<'_, &T> {
-        match self {
-            Self::ControlPlane(c, x) => Backend::ControlPlane(MaybeOwned::Borrowed(c), x),
-            Self::Local(l) => Backend::Local(MaybeOwned::Borrowed(l)),
-        }
-    }
-}
-
-impl<'a, T> Backend<'a, T> {
-    /// Very similar to [`std::option::Option::map`].
-    /// Maps [`Backend<T>`] to [`Backend<R>`] by applying
-    /// a function to a contained value.
-    pub(crate) fn map<R>(self, f: impl FnOnce(T) -> R) -> Backend<'a, R> {
-        match self {
-            Self::ControlPlane(c, x) => Backend::ControlPlane(c, f(x)),
-            Self::Local(l) => Backend::Local(l),
-        }
-    }
-}
-impl<'a, T, E> Backend<'a, Result<T, E>> {
-    /// Very similar to [`std::option::Option::transpose`].
-    /// This is most useful for error handling.
-    pub(crate) fn transpose(self) -> Result<Backend<'a, T>, E> {
-        match self {
-            Self::ControlPlane(c, x) => x.map(|x| Backend::ControlPlane(c, x)),
-            Self::Local(l) => Ok(Backend::Local(l)),
+            ControlPlaneBackend::Management(endpoint) => fmt
+                .debug_tuple("ControlPlane::Management")
+                .field(&endpoint.url())
+                .finish(),
+            #[cfg(any(test, feature = "testing"))]
+            ControlPlaneBackend::PostgresMock(endpoint) => fmt
+                .debug_tuple("ControlPlane::PostgresMock")
+                .field(&endpoint.url())
+                .finish(),
+            #[cfg(test)]
+            ControlPlaneBackend::Test(_) => fmt.debug_tuple("ControlPlane::Test").finish(),
         }
     }
 }
@@ -399,96 +341,79 @@ async fn authenticate_with_secret(
     classic::authenticate(ctx, info, client, config, secret).await
 }
 
-impl<'a> Backend<'a, ComputeUserInfoMaybeEndpoint> {
-    /// Get username from the credentials.
-    pub(crate) fn get_user(&self) -> &str {
-        match self {
-            Self::ControlPlane(_, user_info) => &user_info.user,
-            Self::Local(_) => "local",
-        }
-    }
-
-    /// Authenticate the client via the requested backend, possibly using credentials.
+impl ControlPlaneBackend {
     #[tracing::instrument(fields(allow_cleartext = allow_cleartext), skip_all)]
     pub(crate) async fn authenticate(
-        self,
+        &self,
         ctx: &RequestMonitoring,
+        user_info: ComputeUserInfoMaybeEndpoint,
         client: &mut stream::PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
         allow_cleartext: bool,
         config: &'static AuthenticationConfig,
         endpoint_rate_limiter: Arc<EndpointRateLimiter>,
-    ) -> auth::Result<Backend<'a, ComputeCredentials>> {
-        let res = match self {
-            Self::ControlPlane(api, user_info) => {
-                info!(
-                    user = &*user_info.user,
-                    project = user_info.endpoint(),
-                    "performing authentication using the console"
-                );
+    ) -> auth::Result<ControlPlaneComputeBackend> {
+        info!(
+            user = &*user_info.user,
+            project = user_info.endpoint(),
+            "performing authentication using the console"
+        );
 
-                let credentials = auth_quirks(
-                    ctx,
-                    &*api,
-                    user_info,
-                    client,
-                    allow_cleartext,
-                    config,
-                    endpoint_rate_limiter,
-                )
-                .await?;
-                Backend::ControlPlane(api, credentials)
-            }
-            Self::Local(_) => {
-                return Err(auth::AuthError::bad_auth_method("invalid for local proxy"))
-            }
-        };
+        let credentials = auth_quirks(
+            ctx,
+            self,
+            user_info,
+            client,
+            allow_cleartext,
+            config,
+            endpoint_rate_limiter,
+        )
+        .await?;
 
         info!("user successfully authenticated");
-        Ok(res)
+        Ok(ControlPlaneComputeBackend {
+            api: self,
+            creds: credentials,
+        })
+    }
+
+    pub(crate) fn attach_to_credentials(
+        &self,
+        creds: ComputeCredentials,
+    ) -> ControlPlaneComputeBackend {
+        ControlPlaneComputeBackend { api: self, creds }
     }
 }
 
-impl Backend<'_, ComputeUserInfo> {
-    pub(crate) async fn get_role_secret(
-        &self,
-        ctx: &RequestMonitoring,
-    ) -> Result<CachedRoleSecret, GetAuthInfoError> {
-        match self {
-            Self::ControlPlane(api, user_info) => api.get_role_secret(ctx, user_info).await,
-            Self::Local(_) => Ok(Cached::new_uncached(None)),
-        }
-    }
-
-    pub(crate) async fn get_allowed_ips_and_secret(
-        &self,
-        ctx: &RequestMonitoring,
-    ) -> Result<(CachedAllowedIps, Option<CachedRoleSecret>), GetAuthInfoError> {
-        match self {
-            Self::ControlPlane(api, user_info) => {
-                api.get_allowed_ips_and_secret(ctx, user_info).await
-            }
-            Self::Local(_) => Ok((Cached::new_uncached(Arc::new(vec![])), None)),
-        }
-    }
+pub struct ControlPlaneComputeBackend<'a> {
+    api: &'a ControlPlaneBackend,
+    creds: ComputeCredentials,
 }
 
 #[async_trait::async_trait]
-impl ComputeConnectBackend for Backend<'_, ComputeCredentials> {
+impl ComputeConnectBackend for ControlPlaneComputeBackend<'static> {
     async fn wake_compute(
         &self,
         ctx: &RequestMonitoring,
     ) -> Result<CachedNodeInfo, control_plane::errors::WakeComputeError> {
-        match self {
-            Self::ControlPlane(api, creds) => api.wake_compute(ctx, &creds.info).await,
-            Self::Local(local) => Ok(Cached::new_uncached(local.node_info.clone())),
-        }
+        self.api.wake_compute(ctx, &self.creds.info).await
     }
 
     fn get_keys(&self) -> &ComputeCredentialKeys {
-        match self {
-            Self::ControlPlane(_, creds) => &creds.keys,
-            Self::Local(_) => &ComputeCredentialKeys::None,
-        }
+        &self.creds.keys
+    }
+}
+
+#[async_trait::async_trait]
+impl ComputeConnectBackend for LocalBackend {
+    async fn wake_compute(
+        &self,
+        _ctx: &RequestMonitoring,
+    ) -> Result<CachedNodeInfo, control_plane::errors::WakeComputeError> {
+        Ok(Cached::new_uncached(self.node_info.clone()))
+    }
+
+    fn get_keys(&self) -> &ComputeCredentialKeys {
+        &ComputeCredentialKeys::None
     }
 }
 

proxy/src/auth/mod.rs

@@ -1,7 +1,7 @@
 //! Client authentication mechanisms.
 
 pub mod backend;
-pub use backend::Backend;
+pub use backend::ServerlessBackend;
 
 mod credentials;
 pub(crate) use credentials::{

proxy/src/bin/local_proxy.rs

@@ -203,7 +203,7 @@ async fn main() -> anyhow::Result<()> {
 
     let task = serverless::task_main(
         config,
-        auth_backend,
+        auth::ServerlessBackend::Local(auth_backend),
         http_listener,
         shutdown.clone(),
         Arc::new(CancellationHandlerMain::new(
@@ -295,12 +295,8 @@ fn build_config(args: &LocalProxyCliArgs) -> anyhow::Result<&'static ProxyConfig
 }
 
 /// auth::Backend is created at proxy startup, and lives forever.
-fn build_auth_backend(
-    args: &LocalProxyCliArgs,
-) -> anyhow::Result<&'static auth::Backend<'static, ()>> {
-    let auth_backend = proxy::auth::Backend::Local(proxy::auth::backend::MaybeOwned::Owned(
-        LocalBackend::new(args.postgres, args.compute_ctl.clone()),
-    ));
+fn build_auth_backend(args: &LocalProxyCliArgs) -> anyhow::Result<&'static LocalBackend> {
+    let auth_backend = LocalBackend::new(args.postgres, args.compute_ctl.clone());
 
     Ok(Box::leak(Box::new(auth_backend)))
 }

proxy/src/bin/proxy.rs

@@ -13,13 +13,14 @@ use aws_config::web_identity_token::WebIdentityTokenCredentialsProvider;
 use aws_config::Region;
 use futures::future::Either;
 use proxy::auth::backend::jwt::JwkCache;
-use proxy::auth::backend::{AuthRateLimiter, ConsoleRedirectBackend, MaybeOwned};
+use proxy::auth::backend::{AuthRateLimiter, ConsoleRedirectBackend};
 use proxy::cancellation::{CancelMap, CancellationHandler};
 use proxy::config::{
     self, remote_storage_from_toml, AuthenticationConfig, CacheOptions, HttpConfig,
     ProjectInfoCacheOptions, ProxyConfig, ProxyProtocolV2,
 };
 use proxy::context::parquet::ParquetUploadArgs;
+use proxy::control_plane::provider::ControlPlaneBackend;
 use proxy::http::health_server::AppMetrics;
 use proxy::metrics::Metrics;
 use proxy::rate_limiter::{
@@ -467,7 +468,7 @@ async fn main() -> anyhow::Result<()> {
             if let Some(serverless_listener) = serverless_listener {
                 client_tasks.spawn(serverless::task_main(
                     config,
-                    auth_backend,
+                    auth::ServerlessBackend::ControlPlane(auth_backend),
                     serverless_listener,
                     cancellation_token.clone(),
                     cancellation_handler.clone(),
@@ -515,40 +516,38 @@ async fn main() -> anyhow::Result<()> {
         ));
     }
 
-    if let Either::Left(auth::Backend::ControlPlane(api, _)) = &auth_backend {
-        if let proxy::control_plane::provider::ControlPlaneBackend::Management(api) = &**api {
-            match (redis_notifications_client, regional_redis_client.clone()) {
-                (None, None) => {}
-                (client1, client2) => {
-                    let cache = api.caches.project_info.clone();
-                    if let Some(client) = client1 {
-                        maintenance_tasks.spawn(notifications::task_main(
-                            client,
-                            cache.clone(),
-                            cancel_map.clone(),
-                            args.region.clone(),
-                        ));
-                    }
-                    if let Some(client) = client2 {
-                        maintenance_tasks.spawn(notifications::task_main(
-                            client,
-                            cache.clone(),
-                            cancel_map.clone(),
-                            args.region.clone(),
-                        ));
-                    }
-                    maintenance_tasks.spawn(async move { cache.clone().gc_worker().await });
+    if let Either::Left(ControlPlaneBackend::Management(api)) = &auth_backend {
+        match (redis_notifications_client, regional_redis_client.clone()) {
+            (None, None) => {}
+            (client1, client2) => {
+                let cache = api.caches.project_info.clone();
+                if let Some(client) = client1 {
+                    maintenance_tasks.spawn(notifications::task_main(
+                        client,
+                        cache.clone(),
+                        cancel_map.clone(),
+                        args.region.clone(),
+                    ));
                 }
+                if let Some(client) = client2 {
+                    maintenance_tasks.spawn(notifications::task_main(
+                        client,
+                        cache.clone(),
+                        cancel_map.clone(),
+                        args.region.clone(),
+                    ));
+                }
+                maintenance_tasks.spawn(async move { cache.clone().gc_worker().await });
             }
-            if let Some(regional_redis_client) = regional_redis_client {
-                let cache = api.caches.endpoints_cache.clone();
-                let con = regional_redis_client;
-                let span = tracing::info_span!("endpoints_cache");
-                maintenance_tasks.spawn(
-                    async move { cache.do_read(con, cancellation_token.clone()).await }
-                        .instrument(span),
-                );
-            }
+        }
+        if let Some(regional_redis_client) = regional_redis_client {
+            let cache = api.caches.endpoints_cache.clone();
+            let con = regional_redis_client;
+            let span = tracing::info_span!("endpoints_cache");
+            maintenance_tasks.spawn(
+                async move { cache.do_read(con, cancellation_token.clone()).await }
+                    .instrument(span),
+            );
         }
     }
 
@@ -694,7 +693,7 @@ fn build_config(args: &ProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> {
 /// auth::Backend is created at proxy startup, and lives forever.
 fn build_auth_backend(
     args: &ProxyCliArgs,
-) -> anyhow::Result<Either<&'static auth::Backend<'static, ()>, &'static ConsoleRedirectBackend>> {
+) -> anyhow::Result<Either<&'static ControlPlaneBackend, &'static ConsoleRedirectBackend>> {
     match &args.auth_backend {
         AuthBackendType::Console => {
             let wake_compute_cache_config: CacheOptions = args.wake_compute_cache.parse()?;
@@ -744,8 +743,7 @@ fn build_auth_backend(
                 locks,
                 wake_compute_endpoint_rate_limiter,
             );
-            let api = control_plane::provider::ControlPlaneBackend::Management(api);
-            let auth_backend = auth::Backend::ControlPlane(MaybeOwned::Owned(api), ());
+            let auth_backend = control_plane::provider::ControlPlaneBackend::Management(api);
 
             let config = Box::leak(Box::new(auth_backend));
 
@@ -756,9 +754,7 @@ fn build_auth_backend(
         AuthBackendType::Postgres => {
             let url = args.auth_endpoint.parse()?;
             let api = control_plane::provider::mock::Api::new(url, !args.is_private_access_proxy);
-            let api = control_plane::provider::ControlPlaneBackend::PostgresMock(api);
-
-            let auth_backend = auth::Backend::ControlPlane(MaybeOwned::Owned(api), ());
+            let auth_backend = control_plane::provider::ControlPlaneBackend::PostgresMock(api);
 
             let config = Box::leak(Box::new(auth_backend));
 

proxy/src/proxy/connect_compute.rs

@@ -56,7 +56,7 @@ pub(crate) trait ConnectMechanism {
 }
 
 #[async_trait]
-pub(crate) trait ComputeConnectBackend {
+pub(crate) trait ComputeConnectBackend: Send + Sync + 'static {
     async fn wake_compute(
         &self,
         ctx: &RequestMonitoring,
@@ -98,10 +98,10 @@ impl ConnectMechanism for TcpMechanism<'_> {
 
 /// Try to connect to the compute node, retrying if necessary.
 #[tracing::instrument(skip_all)]
-pub(crate) async fn connect_to_compute<M: ConnectMechanism, B: ComputeConnectBackend>(
+pub(crate) async fn connect_to_compute<M: ConnectMechanism>(
     ctx: &RequestMonitoring,
     mechanism: &M,
-    user_info: &B,
+    user_info: &dyn ComputeConnectBackend,
     allow_self_signed_compute: bool,
     wake_compute_retry_config: RetryConfig,
     connect_to_compute_retry_config: RetryConfig,

proxy/src/proxy/mod.rs

@@ -26,6 +26,7 @@ use self::passthrough::ProxyPassthrough;
 use crate::cancellation::{self, CancellationHandlerMain, CancellationHandlerMainInternal};
 use crate::config::{ProxyConfig, ProxyProtocolV2, TlsConfig};
 use crate::context::RequestMonitoring;
+use crate::control_plane::provider::ControlPlaneBackend;
 use crate::error::ReportableError;
 use crate::metrics::{Metrics, NumClientConnectionsGuard};
 use crate::protocol2::read_proxy_protocol;
@@ -54,7 +55,7 @@ pub async fn run_until_cancelled<F: std::future::Future>(
 
 pub async fn task_main(
     config: &'static ProxyConfig,
-    auth_backend: &'static auth::Backend<'static, ()>,
+    auth_backend: &'static ControlPlaneBackend,
     listener: tokio::net::TcpListener,
     cancellation_token: CancellationToken,
     cancellation_handler: Arc<CancellationHandlerMain>,
@@ -241,7 +242,7 @@ impl ReportableError for ClientRequestError {
 #[allow(clippy::too_many_arguments)]
 pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
     config: &'static ProxyConfig,
-    auth_backend: &'static auth::Backend<'static, ()>,
+    auth_backend: &'static ControlPlaneBackend,
     ctx: &RequestMonitoring,
     cancellation_handler: Arc<CancellationHandlerMain>,
     stream: S,
@@ -282,20 +283,17 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
     let common_names = tls.map(|tls| &tls.common_names);
 
     // Extract credentials which we're going to use for auth.
-    let result = auth_backend
-        .as_ref()
-        .map(|()| auth::ComputeUserInfoMaybeEndpoint::parse(ctx, &params, hostname, common_names))
-        .transpose();
-
+    let result = auth::ComputeUserInfoMaybeEndpoint::parse(ctx, &params, hostname, common_names);
     let user_info = match result {
         Ok(user_info) => user_info,
         Err(e) => stream.throw_error(e).await?,
     };
 
-    let user = user_info.get_user().to_owned();
-    let user_info = match user_info
+    let user = user_info.user.clone();
+    let user_info = match auth_backend
         .authenticate(
             ctx,
+            user_info,
             &mut stream,
             mode.allow_cleartext(),
             &config.authentication_config,

proxy/src/proxy/tests/mod.rs

@@ -6,6 +6,7 @@ use std::time::Duration;
 
 use anyhow::{bail, Context};
 use async_trait::async_trait;
+use auth::backend::ControlPlaneComputeBackend;
 use http::StatusCode;
 use retry::{retry_after, ShouldRetryWakeCompute};
 use rstest::rstest;
@@ -19,7 +20,7 @@ use super::connect_compute::ConnectMechanism;
 use super::retry::CouldRetry;
 use super::*;
 use crate::auth::backend::{
-    ComputeCredentialKeys, ComputeCredentials, ComputeUserInfo, MaybeOwned, TestBackend,
+    ComputeCredentialKeys, ComputeCredentials, ComputeUserInfo, TestBackend,
 };
 use crate::config::{CertResolver, RetryConfig};
 use crate::control_plane::messages::{ControlPlaneError, Details, MetricsAuxInfo, Status};
@@ -566,19 +567,21 @@ fn helper_create_cached_node_info(cache: &'static NodeInfoCache) -> CachedNodeIn
 
 fn helper_create_connect_info(
     mechanism: &TestConnectMechanism,
-) -> auth::Backend<'static, ComputeCredentials> {
-    let user_info = auth::Backend::ControlPlane(
-        MaybeOwned::Owned(ControlPlaneBackend::Test(Box::new(mechanism.clone()))),
-        ComputeCredentials {
-            info: ComputeUserInfo {
-                endpoint: "endpoint".into(),
-                user: "user".into(),
-                options: NeonOptions::parse_options_raw(""),
-            },
-            keys: ComputeCredentialKeys::Password("password".into()),
+) -> ControlPlaneComputeBackend<'static> {
+    let api = Box::leak(Box::new(ControlPlaneBackend::Test(Box::new(
+        mechanism.clone(),
+    ))));
+
+    let creds = ComputeCredentials {
+        info: ComputeUserInfo {
+            endpoint: "endpoint".into(),
+            user: "user".into(),
+            options: NeonOptions::parse_options_raw(""),
         },
-    );
-    user_info
+        keys: ComputeCredentialKeys::Password("password".into()),
+    };
+
+    api.attach_to_credentials(creds)
 }
 
 #[tokio::test]

proxy/src/proxy/wake_compute.rs

@@ -11,10 +11,10 @@ use crate::metrics::{
 };
 use crate::proxy::retry::{retry_after, should_retry};
 
-pub(crate) async fn wake_compute<B: ComputeConnectBackend>(
+pub(crate) async fn wake_compute(
     num_retries: &mut u32,
     ctx: &RequestMonitoring,
-    api: &B,
+    api: &dyn ComputeConnectBackend,
     config: RetryConfig,
 ) -> Result<CachedNodeInfo, WakeComputeError> {
     let retry_type = RetryType::WakeCompute;

proxy/src/serverless/backend.rs

@@ -15,9 +15,9 @@ use super::conn_pool::poll_client;
 use super::conn_pool_lib::{Client, ConnInfo, GlobalConnPool};
 use super::http_conn_pool::{self, poll_http2_client, Send};
 use super::local_conn_pool::{self, LocalClient, LocalConnPool, EXT_NAME, EXT_SCHEMA, EXT_VERSION};
-use crate::auth::backend::local::StaticAuthRules;
+use crate::auth::backend::local::{LocalBackend, StaticAuthRules};
 use crate::auth::backend::{ComputeCredentials, ComputeUserInfo};
-use crate::auth::{self, check_peer_addr_is_in_list, AuthError};
+use crate::auth::{check_peer_addr_is_in_list, AuthError, ServerlessBackend};
 use crate::compute;
 use crate::compute_ctl::{
     ComputeCtlError, ExtensionInstallRequest, Privilege, SetRoleGrantsRequest,
@@ -26,11 +26,11 @@ use crate::config::ProxyConfig;
 use crate::context::RequestMonitoring;
 use crate::control_plane::errors::{GetAuthInfoError, WakeComputeError};
 use crate::control_plane::locks::ApiLocks;
-use crate::control_plane::provider::ApiLockError;
-use crate::control_plane::CachedNodeInfo;
+use crate::control_plane::provider::{ApiLockError, ControlPlaneBackend};
+use crate::control_plane::{Api, CachedNodeInfo};
 use crate::error::{ErrorKind, ReportableError, UserFacingError};
 use crate::intern::EndpointIdInt;
-use crate::proxy::connect_compute::ConnectMechanism;
+use crate::proxy::connect_compute::{ComputeConnectBackend, ConnectMechanism};
 use crate::proxy::retry::{CouldRetry, ShouldRetryWakeCompute};
 use crate::rate_limiter::EndpointRateLimiter;
 use crate::types::{EndpointId, Host};
@@ -41,7 +41,6 @@ pub(crate) struct PoolingBackend {
     pub(crate) pool: Arc<GlobalConnPool<tokio_postgres::Client>>,
 
     pub(crate) config: &'static ProxyConfig,
-    pub(crate) auth_backend: &'static crate::auth::Backend<'static, ()>,
     pub(crate) endpoint_rate_limiter: Arc<EndpointRateLimiter>,
 }
 
@@ -49,12 +48,13 @@ impl PoolingBackend {
     pub(crate) async fn authenticate_with_password(
         &self,
         ctx: &RequestMonitoring,
+        auth_backend: &ControlPlaneBackend,
         user_info: &ComputeUserInfo,
         password: &[u8],
     ) -> Result<ComputeCredentials, AuthError> {
-        let user_info = user_info.clone();
-        let backend = self.auth_backend.as_ref().map(|()| user_info.clone());
-        let (allowed_ips, maybe_secret) = backend.get_allowed_ips_and_secret(ctx).await?;
+        let (allowed_ips, maybe_secret) = auth_backend
+            .get_allowed_ips_and_secret(ctx, user_info)
+            .await?;
         if self.config.authentication_config.ip_allowlist_check_enabled
             && !check_peer_addr_is_in_list(&ctx.peer_addr(), &allowed_ips)
         {
@@ -68,7 +68,7 @@ impl PoolingBackend {
         }
         let cached_secret = match maybe_secret {
             Some(secret) => secret,
-            None => backend.get_role_secret(ctx).await?,
+            None => auth_backend.get_role_secret(ctx, user_info).await?,
         };
 
         let secret = match cached_secret.value.clone() {
@@ -103,7 +103,7 @@ impl PoolingBackend {
             }
         };
         res.map(|key| ComputeCredentials {
-            info: user_info,
+            info: user_info.clone(),
             keys: key,
         })
     }
@@ -111,11 +111,12 @@ impl PoolingBackend {
     pub(crate) async fn authenticate_with_jwt(
         &self,
         ctx: &RequestMonitoring,
+        auth_backend: ServerlessBackend<'static>,
         user_info: &ComputeUserInfo,
         jwt: String,
     ) -> Result<ComputeCredentials, AuthError> {
-        match &self.auth_backend {
-            crate::auth::Backend::ControlPlane(console, ()) => {
+        match auth_backend {
+            ServerlessBackend::ControlPlane(console) => {
                 self.config
                     .authentication_config
                     .jwks_cache
@@ -123,7 +124,7 @@ impl PoolingBackend {
                         ctx,
                         user_info.endpoint.clone(),
                         &user_info.user,
-                        &**console,
+                        console,
                         &jwt,
                     )
                     .await
@@ -134,7 +135,7 @@ impl PoolingBackend {
                     keys: crate::auth::backend::ComputeCredentialKeys::None,
                 })
             }
-            crate::auth::Backend::Local(_) => {
+            ServerlessBackend::Local(_) => {
                 let keys = self
                     .config
                     .authentication_config
@@ -164,6 +165,7 @@ impl PoolingBackend {
     pub(crate) async fn connect_to_compute(
         &self,
         ctx: &RequestMonitoring,
+        auth_backend: ServerlessBackend<'static>,
         conn_info: ConnInfo,
         keys: ComputeCredentials,
         force_new: bool,
@@ -182,7 +184,14 @@ impl PoolingBackend {
         let conn_id = uuid::Uuid::new_v4();
         tracing::Span::current().record("conn_id", display(conn_id));
         info!(%conn_id, "pool: opening a new connection '{conn_info}'");
-        let backend = self.auth_backend.as_ref().map(|()| keys);
+
+        let api = match auth_backend {
+            ServerlessBackend::ControlPlane(cplane) => {
+                &cplane.attach_to_credentials(keys) as &dyn ComputeConnectBackend
+            }
+            ServerlessBackend::Local(local_proxy) => local_proxy as &dyn ComputeConnectBackend,
+        };
+
         crate::proxy::connect_compute::connect_to_compute(
             ctx,
             &TokioMechanism {
@@ -191,7 +200,7 @@ impl PoolingBackend {
                 pool: self.pool.clone(),
                 locks: &self.config.connect_compute_locks,
             },
-            &backend,
+            api,
             false, // do not allow self signed compute for http flow
             self.config.wake_compute_retry_config,
             self.config.connect_to_compute_retry_config,
@@ -204,6 +213,7 @@ impl PoolingBackend {
     pub(crate) async fn connect_to_local_proxy(
         &self,
         ctx: &RequestMonitoring,
+        auth_backend: &'static ControlPlaneBackend,
         conn_info: ConnInfo,
     ) -> Result<http_conn_pool::Client<Send>, HttpConnError> {
         info!("pool: looking for an existing connection");
@@ -214,7 +224,8 @@ impl PoolingBackend {
         let conn_id = uuid::Uuid::new_v4();
         tracing::Span::current().record("conn_id", display(conn_id));
         info!(%conn_id, "pool: opening a new connection '{conn_info}'");
-        let backend = self.auth_backend.as_ref().map(|()| ComputeCredentials {
+
+        let backend = auth_backend.attach_to_credentials(ComputeCredentials {
             info: ComputeUserInfo {
                 user: conn_info.user_info.user.clone(),
                 endpoint: EndpointId::from(format!("{}-local-proxy", conn_info.user_info.endpoint)),
@@ -249,26 +260,20 @@ impl PoolingBackend {
     pub(crate) async fn connect_to_local_postgres(
         &self,
         ctx: &RequestMonitoring,
+        auth_backend: &LocalBackend,
         conn_info: ConnInfo,
     ) -> Result<LocalClient<tokio_postgres::Client>, HttpConnError> {
         if let Some(client) = self.local_pool.get(ctx, &conn_info)? {
             return Ok(client);
         }
 
-        let local_backend = match &self.auth_backend {
-            auth::Backend::ControlPlane(_, ()) => {
-                unreachable!("only local_proxy can connect to local postgres")
-            }
-            auth::Backend::Local(local) => local,
-        };
-
         if !self.local_pool.initialized(&conn_info) {
             // only install and grant usage one at a time.
-            let _permit = local_backend.initialize.acquire().await.unwrap();
+            let _permit = auth_backend.initialize.acquire().await.unwrap();
 
             // check again for race
             if !self.local_pool.initialized(&conn_info) {
-                local_backend
+                auth_backend
                     .compute_ctl
                     .install_extension(&ExtensionInstallRequest {
                         extension: EXT_NAME,
@@ -277,7 +282,7 @@ impl PoolingBackend {
                     })
                     .await?;
 
-                local_backend
+                auth_backend
                     .compute_ctl
                     .grant_role(&SetRoleGrantsRequest {
                         schema: EXT_SCHEMA,
@@ -295,7 +300,7 @@ impl PoolingBackend {
         tracing::Span::current().record("conn_id", display(conn_id));
         info!(%conn_id, "local_pool: opening a new connection '{conn_info}'");
 
-        let mut node_info = local_backend.node_info.clone();
+        let mut node_info = auth_backend.node_info.clone();
 
         let (key, jwk) = create_random_jwk();
 

proxy/src/serverless/mod.rs

@@ -32,7 +32,6 @@ use hyper_util::rt::TokioExecutor;
 use hyper_util::server::conn::auto::Builder;
 use rand::rngs::StdRng;
 use rand::SeedableRng;
-use sql_over_http::{uuid_to_header_value, NEON_REQUEST_ID};
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::net::{TcpListener, TcpStream};
 use tokio::time::timeout;
@@ -42,6 +41,7 @@ use tokio_util::task::TaskTracker;
 use tracing::{info, warn, Instrument};
 use utils::http::error::ApiError;
 
+use crate::auth::ServerlessBackend;
 use crate::cancellation::CancellationHandlerMain;
 use crate::config::ProxyConfig;
 use crate::context::RequestMonitoring;
@@ -56,7 +56,7 @@ pub(crate) const SERVERLESS_DRIVER_SNI: &str = "api";
 
 pub async fn task_main(
     config: &'static ProxyConfig,
-    auth_backend: &'static crate::auth::Backend<'static, ()>,
+    auth_backend: ServerlessBackend<'static>,
     ws_listener: TcpListener,
     cancellation_token: CancellationToken,
     cancellation_handler: Arc<CancellationHandlerMain>,
@@ -112,7 +112,6 @@ pub async fn task_main(
         local_pool,
         pool: Arc::clone(&conn_pool),
         config,
-        auth_backend,
         endpoint_rate_limiter: Arc::clone(&endpoint_rate_limiter),
     });
     let tls_acceptor: Arc<dyn MaybeTlsAcceptor> = match config.tls_config.as_ref() {
@@ -185,6 +184,7 @@ pub async fn task_main(
 
                 Box::pin(connection_handler(
                     config,
+                    auth_backend,
                     backend,
                     connections2,
                     cancellation_handler,
@@ -290,6 +290,7 @@ async fn connection_startup(
 #[allow(clippy::too_many_arguments)]
 async fn connection_handler(
     config: &'static ProxyConfig,
+    auth_backend: ServerlessBackend<'static>,
     backend: Arc<PoolingBackend>,
     connections: TaskTracker,
     cancellation_handler: Arc<CancellationHandlerMain>,
@@ -310,18 +311,7 @@ async fn connection_handler(
         hyper_util::rt::TokioIo::new(conn),
         hyper::service::service_fn(move |req: hyper::Request<Incoming>| {
             // First HTTP request shares the same session ID
-            let mut session_id = session_id.take().unwrap_or_else(uuid::Uuid::new_v4);
-
-            if matches!(backend.auth_backend, crate::auth::Backend::Local(_)) {
-                // take session_id from request, if given.
-                if let Some(id) = req
-                    .headers()
-                    .get(&NEON_REQUEST_ID)
-                    .and_then(|id| uuid::Uuid::try_parse_ascii(id.as_bytes()).ok())
-                {
-                    session_id = id;
-                }
-            }
+            let session_id = session_id.take().unwrap_or_else(uuid::Uuid::new_v4);
 
             // Cancel the current inflight HTTP request if the requets stream is closed.
             // This is slightly different to `_cancel_connection` in that
@@ -335,6 +325,7 @@ async fn connection_handler(
                 request_handler(
                     req,
                     config,
+                    auth_backend,
                     backend.clone(),
                     connections.clone(),
                     cancellation_handler.clone(),
@@ -347,15 +338,8 @@ async fn connection_handler(
                 .map_ok_or_else(api_error_into_response, |r| r),
             );
             async move {
-                let mut res = handler.await;
+                let res = handler.await;
                 cancel_request.disarm();
-
-                // add the session ID to the response
-                if let Ok(resp) = &mut res {
-                    resp.headers_mut()
-                        .append(&NEON_REQUEST_ID, uuid_to_header_value(session_id));
-                }
-
                 res
             }
         }),
@@ -381,6 +365,7 @@ async fn connection_handler(
 async fn request_handler(
     mut request: hyper::Request<Incoming>,
     config: &'static ProxyConfig,
+    auth_backend: ServerlessBackend<'static>,
     backend: Arc<PoolingBackend>,
     ws_connections: TaskTracker,
     cancellation_handler: Arc<CancellationHandlerMain>,
@@ -401,6 +386,10 @@ async fn request_handler(
     if config.http_config.accept_websockets
         && framed_websockets::upgrade::is_upgrade_request(&request)
     {
+        let ServerlessBackend::ControlPlane(auth_backend) = auth_backend else {
+            return json_response(StatusCode::BAD_REQUEST, "query is not supported");
+        };
+
         let ctx = RequestMonitoring::new(
             session_id,
             peer_addr,
@@ -418,7 +407,7 @@ async fn request_handler(
             async move {
                 if let Err(e) = websocket::serve_websocket(
                     config,
-                    backend.auth_backend,
+                    auth_backend,
                     ctx,
                     websocket,
                     cancellation_handler,
@@ -444,9 +433,16 @@ async fn request_handler(
         );
         let span = ctx.span();
 
-        sql_over_http::handle(config, ctx, request, backend, http_cancellation_token)
-            .instrument(span)
-            .await
+        sql_over_http::handle(
+            config,
+            ctx,
+            request,
+            auth_backend,
+            backend,
+            http_cancellation_token,
+        )
+        .instrument(span)
+        .await
     } else if request.uri().path() == "/sql" && *request.method() == Method::OPTIONS {
         Response::builder()
             .header("Allow", "OPTIONS, POST")

proxy/src/serverless/sql_over_http.rs

@@ -23,7 +23,6 @@ use typed_json::json;
 use url::Url;
 use urlencoding;
 use utils::http::error::ApiError;
-use uuid::Uuid;
 
 use super::backend::{LocalProxyConnError, PoolingBackend};
 use super::conn_pool::{AuthData, ConnInfoWithAuth};
@@ -31,10 +30,11 @@ use super::conn_pool_lib::{self, ConnInfo};
 use super::http_util::json_response;
 use super::json::{json_to_pg_text, pg_text_row_to_json, JsonConversionError};
 use super::local_conn_pool;
-use crate::auth::backend::{ComputeCredentialKeys, ComputeUserInfo};
-use crate::auth::{endpoint_sni, ComputeUserInfoParseError};
+use crate::auth::backend::{ComputeCredentialKeys, ComputeCredentials, ComputeUserInfo};
+use crate::auth::{endpoint_sni, ComputeUserInfoParseError, ServerlessBackend};
 use crate::config::{AuthenticationConfig, HttpConfig, ProxyConfig, TlsConfig};
 use crate::context::RequestMonitoring;
+use crate::control_plane::provider::ControlPlaneBackend;
 use crate::error::{ErrorKind, ReportableError, UserFacingError};
 use crate::metrics::{HttpDirection, Metrics};
 use crate::proxy::{run_until_cancelled, NeonOptions};
@@ -64,8 +64,6 @@ enum Payload {
     Batch(BatchQueryData),
 }
 
-pub(super) static NEON_REQUEST_ID: HeaderName = HeaderName::from_static("neon-request-id");
-
 static CONN_STRING: HeaderName = HeaderName::from_static("neon-connection-string");
 static RAW_TEXT_OUTPUT: HeaderName = HeaderName::from_static("neon-raw-text-output");
 static ARRAY_MODE: HeaderName = HeaderName::from_static("neon-array-mode");
@@ -243,10 +241,11 @@ pub(crate) async fn handle(
     config: &'static ProxyConfig,
     ctx: RequestMonitoring,
     request: Request<Incoming>,
+    auth_backend: ServerlessBackend<'static>,
     backend: Arc<PoolingBackend>,
     cancel: CancellationToken,
 ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, ApiError> {
-    let result = handle_inner(cancel, config, &ctx, request, backend).await;
+    let result = handle_inner(cancel, config, &ctx, request, auth_backend, backend).await;
 
     let mut response = match result {
         Ok(r) => {
@@ -501,6 +500,7 @@ async fn handle_inner(
     config: &'static ProxyConfig,
     ctx: &RequestMonitoring,
     request: Request<Incoming>,
+    auth_backend: ServerlessBackend<'static>,
     backend: Arc<PoolingBackend>,
 ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, SqlOverHttpError> {
     let _requeset_gauge = Metrics::get()
@@ -525,7 +525,11 @@ async fn handle_inner(
 
     match conn_info.auth {
         AuthData::Jwt(jwt) if config.authentication_config.is_auth_broker => {
-            handle_auth_broker_inner(ctx, request, conn_info.conn_info, jwt, backend).await
+            let ServerlessBackend::ControlPlane(cplane) = auth_backend else {
+                panic!("auth_broker must be configured with a control-plane auth backend.")
+            };
+
+            handle_auth_broker_inner(ctx, request, conn_info.conn_info, jwt, cplane, backend).await
         }
         auth => {
             handle_db_inner(
@@ -535,6 +539,7 @@ async fn handle_inner(
                 request,
                 conn_info.conn_info,
                 auth,
+                auth_backend,
                 backend,
             )
             .await
@@ -542,6 +547,7 @@ async fn handle_inner(
     }
 }
 
+#[allow(clippy::too_many_arguments)]
 async fn handle_db_inner(
     cancel: CancellationToken,
     config: &'static ProxyConfig,
@@ -549,6 +555,7 @@ async fn handle_db_inner(
     request: Request<Incoming>,
     conn_info: ConnInfo,
     auth: AuthData,
+    auth_backend: ServerlessBackend<'static>,
     backend: Arc<PoolingBackend>,
 ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, SqlOverHttpError> {
     //
@@ -591,45 +598,58 @@ async fn handle_db_inner(
         .map_err(SqlOverHttpError::from),
     );
 
-    let authenticate_and_connect = Box::pin(
-        async {
-            let is_local_proxy = matches!(backend.auth_backend, crate::auth::Backend::Local(_));
+    let authenticate_and_connect = Box::pin(async {
+        let creds = match auth {
+            AuthData::Password(pw) => {
+                let ServerlessBackend::ControlPlane(cplane) = auth_backend else {
+                    return Err(SqlOverHttpError::ConnInfo(
+                        ConnInfoError::MissingCredentials(Credentials::BearerJwt),
+                    ));
+                };
 
-            let keys = match auth {
-                AuthData::Password(pw) => {
-                    backend
-                        .authenticate_with_password(ctx, &conn_info.user_info, &pw)
-                        .await?
-                }
-                AuthData::Jwt(jwt) => {
-                    backend
-                        .authenticate_with_jwt(ctx, &conn_info.user_info, jwt)
-                        .await?
-                }
-            };
+                backend
+                    .authenticate_with_password(ctx, cplane, &conn_info.user_info, &pw)
+                    .await
+                    .map_err(HttpConnError::from)?
+            }
+            AuthData::Jwt(jwt) => backend
+                .authenticate_with_jwt(ctx, auth_backend, &conn_info.user_info, jwt)
+                .await
+                .map_err(HttpConnError::from)?,
+        };
 
-            let client = match keys.keys {
-                ComputeCredentialKeys::JwtPayload(payload) if is_local_proxy => {
-                    let mut client = backend.connect_to_local_postgres(ctx, conn_info).await?;
-                    let (cli_inner, _dsc) = client.client_inner();
-                    cli_inner.set_jwt_session(&payload).await?;
-                    Client::Local(client)
-                }
-                _ => {
-                    let client = backend
-                        .connect_to_compute(ctx, conn_info, keys, !allow_pool)
-                        .await?;
-                    Client::Remote(client)
-                }
-            };
+        let client = match (creds.keys, auth_backend) {
+            (ComputeCredentialKeys::JwtPayload(payload), ServerlessBackend::Local(local)) => {
+                let mut client = backend
+                    .connect_to_local_postgres(ctx, local, conn_info)
+                    .await?;
+                let (cli_inner, _dsc) = client.client_inner();
+                cli_inner.set_jwt_session(&payload).await?;
+                Client::Local(client)
+            }
+            (keys, auth_backend) => {
+                let client = backend
+                    .connect_to_compute(
+                        ctx,
+                        auth_backend,
+                        conn_info,
+                        ComputeCredentials {
+                            keys,
+                            info: creds.info,
+                        },
+                        !allow_pool,
+                    )
+                    .await
+                    .map_err(HttpConnError::from)?;
+                Client::Remote(client)
+            }
+        };
 
-            // not strictly necessary to mark success here,
-            // but it's just insurance for if we forget it somewhere else
-            ctx.success();
-            Ok::<_, HttpConnError>(client)
-        }
-        .map_err(SqlOverHttpError::from),
-    );
+        // not strictly necessary to mark success here,
+        // but it's just insurance for if we forget it somewhere else
+        ctx.success();
+        Ok::<_, SqlOverHttpError>(client)
+    });
 
     let (payload, mut client) = match run_until_cancelled(
         // Run both operations in parallel
@@ -709,25 +729,27 @@ static HEADERS_TO_FORWARD: &[&HeaderName] = &[
     &TXN_DEFERRABLE,
 ];
 
-pub(crate) fn uuid_to_header_value(id: Uuid) -> HeaderValue {
-    let mut uuid = [0; uuid::fmt::Hyphenated::LENGTH];
-    HeaderValue::from_str(id.as_hyphenated().encode_lower(&mut uuid[..]))
-        .expect("uuid hyphenated format should be all valid header characters")
-}
-
 async fn handle_auth_broker_inner(
     ctx: &RequestMonitoring,
     request: Request<Incoming>,
     conn_info: ConnInfo,
     jwt: String,
+    auth_backend: &'static ControlPlaneBackend,
     backend: Arc<PoolingBackend>,
 ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, SqlOverHttpError> {
     backend
-        .authenticate_with_jwt(ctx, &conn_info.user_info, jwt)
+        .authenticate_with_jwt(
+            ctx,
+            ServerlessBackend::ControlPlane(auth_backend),
+            &conn_info.user_info,
+            jwt,
+        )
         .await
         .map_err(HttpConnError::from)?;
 
-    let mut client = backend.connect_to_local_proxy(ctx, conn_info).await?;
+    let mut client = backend
+        .connect_to_local_proxy(ctx, auth_backend, conn_info)
+        .await?;
 
     let local_proxy_uri = ::http::Uri::from_static("http://proxy.local/sql");
 
@@ -741,7 +763,6 @@ async fn handle_auth_broker_inner(
             req = req.header(h, hv);
         }
     }
-    req = req.header(&NEON_REQUEST_ID, uuid_to_header_value(ctx.session_id()));
 
     let req = req
         .body(body)

proxy/src/serverless/websocket.rs

@@ -15,6 +15,7 @@ use tracing::warn;
 use crate::cancellation::CancellationHandlerMain;
 use crate::config::ProxyConfig;
 use crate::context::RequestMonitoring;
+use crate::control_plane::provider::ControlPlaneBackend;
 use crate::error::{io_error, ReportableError};
 use crate::metrics::Metrics;
 use crate::proxy::{handle_client, ClientMode, ErrorSource};
@@ -125,7 +126,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin> AsyncBufRead for WebSocketRw<S> {
 
 pub(crate) async fn serve_websocket(
     config: &'static ProxyConfig,
-    auth_backend: &'static crate::auth::Backend<'static, ()>,
+    auth_backend: &'static ControlPlaneBackend,
     ctx: RequestMonitoring,
     websocket: OnUpgrade,
     cancellation_handler: Arc<CancellationHandlerMain>,

safekeeper/src/bin/safekeeper.rs

@@ -193,8 +193,6 @@ struct Args {
     /// Usually, timeline eviction has to wait for `partial_backup_timeout` before being eligible for eviction,
     /// but if a timeline is un-evicted and then _not_ written to, it would immediately flap to evicting again,
     /// if it weren't for `eviction_min_resident` preventing that.
-    ///
-    /// Also defines interval for eviction retries.
     #[arg(long, value_parser = humantime::parse_duration, default_value = DEFAULT_EVICTION_MIN_RESIDENT)]
     eviction_min_resident: Duration,
 }

safekeeper/src/control_file.rs

@@ -14,10 +14,12 @@ use std::path::Path;
 use std::time::Instant;
 
 use crate::control_file_upgrade::downgrade_v9_to_v8;
-use crate::control_file_upgrade::upgrade_control_file;
 use crate::metrics::PERSIST_CONTROL_FILE_SECONDS;
 use crate::state::{EvictionState, TimelinePersistentState};
-use utils::bin_ser::LeSer;
+use crate::{control_file_upgrade::upgrade_control_file, timeline::get_timeline_dir};
+use utils::{bin_ser::LeSer, id::TenantTimelineId};
+
+use crate::SafeKeeperConf;
 
 pub const SK_MAGIC: u32 = 0xcafeceefu32;
 pub const SK_FORMAT_VERSION: u32 = 9;
@@ -52,12 +54,13 @@ pub struct FileStorage {
 
 impl FileStorage {
     /// Initialize storage by loading state from disk.
-    pub fn restore_new(timeline_dir: &Utf8Path, no_sync: bool) -> Result<FileStorage> {
-        let state = Self::load_control_file_from_dir(timeline_dir)?;
+    pub fn restore_new(ttid: &TenantTimelineId, conf: &SafeKeeperConf) -> Result<FileStorage> {
+        let timeline_dir = get_timeline_dir(conf, ttid);
+        let state = Self::load_control_file_from_dir(&timeline_dir)?;
 
         Ok(FileStorage {
-            timeline_dir: timeline_dir.to_path_buf(),
-            no_sync,
+            timeline_dir,
+            no_sync: conf.no_sync,
             state,
             last_persist_at: Instant::now(),
         })
@@ -68,16 +71,16 @@ impl FileStorage {
     /// Note: we normally call this in temp directory for atomic init, so
     /// interested in FileStorage as a result only in tests.
     pub async fn create_new(
-        timeline_dir: &Utf8Path,
+        dir: Utf8PathBuf,
+        conf: &SafeKeeperConf,
         state: TimelinePersistentState,
-        no_sync: bool,
     ) -> Result<FileStorage> {
         // we don't support creating new timelines in offloaded state
         assert!(matches!(state.eviction_state, EvictionState::Present));
 
         let mut store = FileStorage {
-            timeline_dir: timeline_dir.to_path_buf(),
-            no_sync,
+            timeline_dir: dir,
+            no_sync: conf.no_sync,
             state: state.clone(),
             last_persist_at: Instant::now(),
         };
@@ -236,46 +239,89 @@ mod test {
     use tokio::fs;
     use utils::lsn::Lsn;
 
-    const NO_SYNC: bool = true;
+    fn stub_conf() -> SafeKeeperConf {
+        let workdir = camino_tempfile::tempdir().unwrap().into_path();
+        SafeKeeperConf {
+            workdir,
+            ..SafeKeeperConf::dummy()
+        }
+    }
 
-    #[tokio::test]
-    async fn test_read_write_safekeeper_state() -> anyhow::Result<()> {
-        let tempdir = camino_tempfile::tempdir()?;
-        let mut state = TimelinePersistentState::empty();
-        let mut storage = FileStorage::create_new(tempdir.path(), state.clone(), NO_SYNC).await?;
+    async fn load_from_control_file(
+        conf: &SafeKeeperConf,
+        ttid: &TenantTimelineId,
+    ) -> Result<(FileStorage, TimelinePersistentState)> {
+        let timeline_dir = get_timeline_dir(conf, ttid);
+        fs::create_dir_all(&timeline_dir)
+            .await
+            .expect("failed to create timeline dir");
+        Ok((
+            FileStorage::restore_new(ttid, conf)?,
+            FileStorage::load_control_file_from_dir(&timeline_dir)?,
+        ))
+    }
 
-        // Make a change.
-        state.commit_lsn = Lsn(42);
-        storage.persist(&state).await?;
-
-        // Reload the state. It should match the previously persisted state.
-        let loaded_state = FileStorage::load_control_file_from_dir(tempdir.path())?;
-        assert_eq!(loaded_state, state);
-        Ok(())
+    async fn create(
+        conf: &SafeKeeperConf,
+        ttid: &TenantTimelineId,
+    ) -> Result<(FileStorage, TimelinePersistentState)> {
+        let timeline_dir = get_timeline_dir(conf, ttid);
+        fs::create_dir_all(&timeline_dir)
+            .await
+            .expect("failed to create timeline dir");
+        let state = TimelinePersistentState::empty();
+        let storage = FileStorage::create_new(timeline_dir, conf, state.clone()).await?;
+        Ok((storage, state))
     }
 
     #[tokio::test]
-    async fn test_safekeeper_state_checksum_mismatch() -> anyhow::Result<()> {
-        let tempdir = camino_tempfile::tempdir()?;
-        let mut state = TimelinePersistentState::empty();
-        let mut storage = FileStorage::create_new(tempdir.path(), state.clone(), NO_SYNC).await?;
-
-        // Make a change.
-        state.commit_lsn = Lsn(42);
-        storage.persist(&state).await?;
-
-        // Change the first byte to fail checksum validation.
-        let ctrl_path = tempdir.path().join(CONTROL_FILE_NAME);
-        let mut data = fs::read(&ctrl_path).await?;
-        data[0] += 1;
-        fs::write(&ctrl_path, &data).await?;
-
-        // Loading the file should fail checksum validation.
-        if let Err(err) = FileStorage::load_control_file_from_dir(tempdir.path()) {
-            assert!(err.to_string().contains("control file checksum mismatch"))
-        } else {
-            panic!("expected checksum error")
+    async fn test_read_write_safekeeper_state() {
+        let conf = stub_conf();
+        let ttid = TenantTimelineId::generate();
+        {
+            let (mut storage, mut state) =
+                create(&conf, &ttid).await.expect("failed to create state");
+            // change something
+            state.commit_lsn = Lsn(42);
+            storage
+                .persist(&state)
+                .await
+                .expect("failed to persist state");
+        }
+
+        let (_, state) = load_from_control_file(&conf, &ttid)
+            .await
+            .expect("failed to read state");
+        assert_eq!(state.commit_lsn, Lsn(42));
+    }
+
+    #[tokio::test]
+    async fn test_safekeeper_state_checksum_mismatch() {
+        let conf = stub_conf();
+        let ttid = TenantTimelineId::generate();
+        {
+            let (mut storage, mut state) =
+                create(&conf, &ttid).await.expect("failed to read state");
+
+            // change something
+            state.commit_lsn = Lsn(42);
+            storage
+                .persist(&state)
+                .await
+                .expect("failed to persist state");
+        }
+        let control_path = get_timeline_dir(&conf, &ttid).join(CONTROL_FILE_NAME);
+        let mut data = fs::read(&control_path).await.unwrap();
+        data[0] += 1; // change the first byte of the file to fail checksum validation
+        fs::write(&control_path, &data)
+            .await
+            .expect("failed to write control file");
+
+        match load_from_control_file(&conf, &ttid).await {
+            Err(err) => assert!(err
+                .to_string()
+                .contains("safekeeper control file checksum mismatch")),
+            Ok(_) => panic!("expected error"),
         }
-        Ok(())
     }
 }

safekeeper/src/copy_timeline.rs

@@ -154,7 +154,7 @@ pub async fn handle_request(request: Request) -> Result<()> {
     new_state.peer_horizon_lsn = request.until_lsn;
     new_state.backup_lsn = new_backup_lsn;
 
-    FileStorage::create_new(&tli_dir_path, new_state.clone(), conf.no_sync).await?;
+    FileStorage::create_new(tli_dir_path.clone(), conf, new_state.clone()).await?;
 
     // now we have a ready timeline in a temp directory
     validate_temp_timeline(conf, request.destination_ttid, &tli_dir_path).await?;

safekeeper/src/lib.rs

@@ -113,7 +113,6 @@ impl SafeKeeperConf {
 
 impl SafeKeeperConf {
     #[cfg(test)]
-    #[allow(unused)]
     fn dummy() -> Self {
         SafeKeeperConf {
             workdir: Utf8PathBuf::from("./"),

safekeeper/src/state.rs

@@ -143,8 +143,8 @@ impl TimelinePersistentState {
         TimelinePersistentState::new(
             &TenantTimelineId::empty(),
             ServerInfo {
-                pg_version: 170000, /* Postgres server version (major * 10000) */
-                system_id: 0,       /* Postgres system identifier */
+                pg_version: 17, /* Postgres server version */
+                system_id: 0,   /* Postgres system identifier */
                 wal_seg_size: 16 * 1024 * 1024,
             },
             vec![],

safekeeper/src/timeline.rs

@@ -328,19 +328,15 @@ impl SharedState {
     /// Restore SharedState from control file. If file doesn't exist, bails out.
     fn restore(conf: &SafeKeeperConf, ttid: &TenantTimelineId) -> Result<Self> {
         let timeline_dir = get_timeline_dir(conf, ttid);
-        let control_store = control_file::FileStorage::restore_new(&timeline_dir, conf.no_sync)?;
+        let control_store = control_file::FileStorage::restore_new(ttid, conf)?;
         if control_store.server.wal_seg_size == 0 {
             bail!(TimelineError::UninitializedWalSegSize(*ttid));
         }
 
         let sk = match control_store.eviction_state {
             EvictionState::Present => {
-                let wal_store = wal_storage::PhysicalStorage::new(
-                    ttid,
-                    &timeline_dir,
-                    &control_store,
-                    conf.no_sync,
-                )?;
+                let wal_store =
+                    wal_storage::PhysicalStorage::new(ttid, timeline_dir, conf, &control_store)?;
                 StateSK::Loaded(SafeKeeper::new(
                     TimelineState::new(control_store),
                     wal_store,
@@ -1050,9 +1046,9 @@ impl ManagerTimeline {
         // trying to restore WAL storage
         let wal_store = wal_storage::PhysicalStorage::new(
             &self.ttid,
-            &self.timeline_dir,
+            self.timeline_dir.clone(),
+            &conf,
             shared.sk.state(),
-            conf.no_sync,
         )?;
 
         // updating control file

safekeeper/src/timeline_eviction.rs

@@ -66,15 +66,15 @@ impl Manager {
         ready
     }
 
-    /// Evict the timeline to remote storage. Returns whether the eviction was successful.
+    /// Evict the timeline to remote storage.
     #[instrument(name = "evict_timeline", skip_all)]
-    pub(crate) async fn evict_timeline(&mut self) -> bool {
+    pub(crate) async fn evict_timeline(&mut self) {
         assert!(!self.is_offloaded);
         let partial_backup_uploaded = match &self.partial_backup_uploaded {
             Some(p) => p.clone(),
             None => {
                 warn!("no partial backup uploaded, skipping eviction");
-                return false;
+                return;
             }
         };
 
@@ -91,12 +91,11 @@ impl Manager {
 
         if let Err(e) = do_eviction(self, &partial_backup_uploaded).await {
             warn!("failed to evict timeline: {:?}", e);
-            return false;
+            return;
         }
 
         info!("successfully evicted timeline");
         NUM_EVICTED_TIMELINES.inc();
-        true
     }
 
     /// Attempt to restore evicted timeline from remote storage; it must be

safekeeper/src/timeline_manager.rs

@@ -297,12 +297,7 @@ pub async fn main_task(
                 match mgr.global_rate_limiter.try_acquire_eviction() {
                     Some(_permit) => {
                         mgr.set_status(Status::EvictTimeline);
-                        if !mgr.evict_timeline().await {
-                            // eviction failed, try again later
-                            mgr.evict_not_before =
-                                Instant::now() + rand_duration(&mgr.conf.eviction_min_resident);
-                            update_next_event(&mut next_event, mgr.evict_not_before);
-                        }
+                        mgr.evict_timeline().await;
                     }
                     None => {
                         // we can't evict timeline now, will try again later

safekeeper/src/timelines_global_map.rs

@@ -244,7 +244,7 @@ impl GlobalTimelines {
         // immediately initialize first WAL segment as well.
         let state =
             TimelinePersistentState::new(&ttid, server_info, vec![], commit_lsn, local_start_lsn)?;
-        control_file::FileStorage::create_new(&tmp_dir_path, state, conf.no_sync).await?;
+        control_file::FileStorage::create_new(tmp_dir_path.clone(), &conf, state).await?;
         let timeline = GlobalTimelines::load_temp_timeline(ttid, &tmp_dir_path, true).await?;
         Ok(timeline)
     }
@@ -596,7 +596,7 @@ pub async fn validate_temp_timeline(
         bail!("wal_seg_size is not set");
     }
 
-    let wal_store = wal_storage::PhysicalStorage::new(&ttid, path, &control_store, conf.no_sync)?;
+    let wal_store = wal_storage::PhysicalStorage::new(&ttid, path.clone(), conf, &control_store)?;
 
     let commit_lsn = control_store.commit_lsn;
     let flush_lsn = wal_store.flush_lsn();

safekeeper/src/wal_storage.rs

@@ -29,6 +29,7 @@ use crate::metrics::{
 };
 use crate::state::TimelinePersistentState;
 use crate::wal_backup::{read_object, remote_timeline_path};
+use crate::SafeKeeperConf;
 use postgres_ffi::waldecoder::WalStreamDecoder;
 use postgres_ffi::XLogFileName;
 use postgres_ffi::XLOG_BLCKSZ;
@@ -86,9 +87,7 @@ pub trait Storage {
 pub struct PhysicalStorage {
     metrics: WalStorageMetrics,
     timeline_dir: Utf8PathBuf,
-
-    /// Disables fsync if true.
-    no_sync: bool,
+    conf: SafeKeeperConf,
 
     /// Size of WAL segment in bytes.
     wal_seg_size: usize,
@@ -152,9 +151,9 @@ impl PhysicalStorage {
     /// the disk. Otherwise, all LSNs are set to zero.
     pub fn new(
         ttid: &TenantTimelineId,
-        timeline_dir: &Utf8Path,
+        timeline_dir: Utf8PathBuf,
+        conf: &SafeKeeperConf,
         state: &TimelinePersistentState,
-        no_sync: bool,
     ) -> Result<PhysicalStorage> {
         let wal_seg_size = state.server.wal_seg_size as usize;
 
@@ -199,8 +198,8 @@ impl PhysicalStorage {
 
         Ok(PhysicalStorage {
             metrics: WalStorageMetrics::default(),
-            timeline_dir: timeline_dir.to_path_buf(),
-            no_sync,
+            timeline_dir,
+            conf: conf.clone(),
             wal_seg_size,
             pg_version: state.server.pg_version,
             system_id: state.server.system_id,
@@ -225,7 +224,7 @@ impl PhysicalStorage {
 
     /// Call fdatasync if config requires so.
     async fn fdatasync_file(&mut self, file: &File) -> Result<()> {
-        if !self.no_sync {
+        if !self.conf.no_sync {
             self.metrics
                 .observe_flush_seconds(time_io_closure(file.sync_data()).await?);
         }
@@ -264,7 +263,9 @@ impl PhysicalStorage {
 
             // Note: this doesn't get into observe_flush_seconds metric. But
             // segment init should be separate metric, if any.
-            if let Err(e) = durable_rename(&tmp_path, &wal_file_partial_path, !self.no_sync).await {
+            if let Err(e) =
+                durable_rename(&tmp_path, &wal_file_partial_path, !self.conf.no_sync).await
+            {
                 // Probably rename succeeded, but fsync of it failed. Remove
                 // the file then to avoid using it.
                 remove_file(wal_file_partial_path)

storage_controller/src/reconciler.rs

@@ -450,9 +450,6 @@ impl Reconciler {
         }
     }
 
-    /// This function does _not_ mutate any state, so it is cancellation safe.
-    ///
-    /// This function does not respect [`Self::cancel`], callers should handle that.
     async fn await_lsn(
         &self,
         tenant_shard_id: TenantShardId,
@@ -573,10 +570,8 @@ impl Reconciler {
 
         if let Some(baseline) = baseline_lsns {
             tracing::info!("🕑 Waiting for LSN to catch up...");
-            tokio::select! {
-                r = self.await_lsn(self.tenant_shard_id, &dest_ps, baseline) => {r?;}
-                _ = self.cancel.cancelled() => {return Err(ReconcileError::Cancel)}
-            };
+            self.await_lsn(self.tenant_shard_id, &dest_ps, baseline)
+                .await?;
         }
 
         tracing::info!("🔁 Notifying compute to use pageserver {dest_ps}");

storage_controller/src/service.rs

@@ -3130,11 +3130,9 @@ impl Service {
             .await?;
 
             // Propagate the LSN that shard zero picked, if caller didn't provide one
-            match &mut create_req.mode {
-                models::TimelineCreateRequestMode::Branch { ancestor_start_lsn, .. } if ancestor_start_lsn.is_none() => {
-                    *ancestor_start_lsn = timeline_info.ancestor_lsn;
-                },
-                _ => {}
+            if create_req.ancestor_timeline_id.is_some() && create_req.ancestor_start_lsn.is_none()
+            {
+                create_req.ancestor_start_lsn = timeline_info.ancestor_lsn;
             }
 
             // Create timeline on remaining shards with number >0

test_runner/fixtures/metrics.py

@@ -150,7 +150,6 @@ PAGESERVER_GLOBAL_METRICS: tuple[str, ...] = (
     counter("pageserver_tenant_throttling_count_accounted_finish_global"),
     counter("pageserver_tenant_throttling_wait_usecs_sum_global"),
     counter("pageserver_tenant_throttling_count_global"),
-    *histogram("pageserver_tokio_epoll_uring_slots_submission_queue_depth"),
 )
 
 PAGESERVER_PER_TENANT_METRICS: tuple[str, ...] = (

test_runner/fixtures/neon_fixtures.py

@@ -44,14 +44,7 @@ from urllib3.util.retry import Retry
 
 from fixtures import overlayfs
 from fixtures.auth_tokens import AuthKeys, TokenScope
-from fixtures.common_types import (
-    Lsn,
-    NodeId,
-    TenantId,
-    TenantShardId,
-    TimelineArchivalState,
-    TimelineId,
-)
+from fixtures.common_types import Lsn, NodeId, TenantId, TenantShardId, TimelineId
 from fixtures.endpoint.http import EndpointHttpClient
 from fixtures.log_helper import log
 from fixtures.metrics import Metrics, MetricsGetter, parse_metrics
@@ -61,11 +54,7 @@ from fixtures.pageserver.allowed_errors import (
     DEFAULT_STORAGE_CONTROLLER_ALLOWED_ERRORS,
 )
 from fixtures.pageserver.common_types import LayerName, parse_layer_file_name
-from fixtures.pageserver.http import (
-    HistoricLayerInfo,
-    PageserverHttpClient,
-    ScanDisposableKeysResponse,
-)
+from fixtures.pageserver.http import PageserverHttpClient
 from fixtures.pageserver.utils import (
     wait_for_last_record_lsn,
 )
@@ -2143,24 +2132,6 @@ class NeonStorageController(MetricsGetter, LogUtils):
         response.raise_for_status()
         return response.json()
 
-    def timeline_archival_config(
-        self,
-        tenant_id: TenantId,
-        timeline_id: TimelineId,
-        state: TimelineArchivalState,
-    ):
-        config = {"state": state.value}
-        log.info(
-            f"requesting timeline archival config {config} for tenant {tenant_id} and timeline {timeline_id}"
-        )
-        res = self.request(
-            "PUT",
-            f"{self.api}/v1/tenant/{tenant_id}/timeline/{timeline_id}/archival_config",
-            json=config,
-            headers=self.headers(TokenScope.ADMIN),
-        )
-        return res.json()
-
     def configure_failpoints(self, config_strings: tuple[str, str] | list[tuple[str, str]]):
         if isinstance(config_strings, tuple):
             pairs = [config_strings]
@@ -2674,51 +2645,6 @@ class NeonPageserver(PgProtocol, LogUtils):
         layers = self.list_layers(tenant_id, timeline_id)
         return layer_name in [parse_layer_file_name(p.name) for p in layers]
 
-    def timeline_scan_no_disposable_keys(
-        self, tenant_shard_id: TenantShardId, timeline_id: TimelineId
-    ) -> TimelineAssertNoDisposableKeysResult:
-        """
-        Scan all keys in all layers of the tenant/timeline for disposable keys.
-        Disposable keys are keys that are present in a layer referenced by the shard
-        but are not going to be accessed by the shard.
-        For example, after shard split, the child shards will reference the parent's layer
-        files until new data is ingested and/or compaction rewrites the layers.
-        """
-
-        ps_http = self.http_client()
-        tally = ScanDisposableKeysResponse(0, 0)
-        per_layer = []
-        with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
-            futs = []
-            shard_layer_map = ps_http.layer_map_info(tenant_shard_id, timeline_id)
-            for layer in shard_layer_map.historic_layers:
-
-                def do_layer(
-                    shard_ps_http: PageserverHttpClient,
-                    tenant_shard_id: TenantShardId,
-                    timeline_id: TimelineId,
-                    layer: HistoricLayerInfo,
-                ) -> tuple[HistoricLayerInfo, ScanDisposableKeysResponse]:
-                    return (
-                        layer,
-                        shard_ps_http.timeline_layer_scan_disposable_keys(
-                            tenant_shard_id, timeline_id, layer.layer_file_name
-                        ),
-                    )
-
-                futs.append(executor.submit(do_layer, ps_http, tenant_shard_id, timeline_id, layer))
-            for fut in futs:
-                layer, result = fut.result()
-                tally += result
-                per_layer.append((layer, result))
-        return TimelineAssertNoDisposableKeysResult(tally, per_layer)
-
-
-@dataclass
-class TimelineAssertNoDisposableKeysResult:
-    tally: ScanDisposableKeysResponse
-    per_layer: list[tuple[HistoricLayerInfo, ScanDisposableKeysResponse]]
-
 
 class PgBin:
     """A helper class for executing postgres binaries"""

test_runner/fixtures/pageserver/http.py

@@ -129,26 +129,6 @@ class LayerMapInfo:
         return set(x.layer_file_name for x in self.historic_layers)
 
 
-@dataclass
-class ScanDisposableKeysResponse:
-    disposable_count: int
-    not_disposable_count: int
-
-    def __add__(self, b):
-        a = self
-        assert isinstance(a, ScanDisposableKeysResponse)
-        assert isinstance(b, ScanDisposableKeysResponse)
-        return ScanDisposableKeysResponse(
-            a.disposable_count + b.disposable_count, a.not_disposable_count + b.not_disposable_count
-        )
-
-    @classmethod
-    def from_json(cls, d: dict[str, Any]) -> ScanDisposableKeysResponse:
-        disposable_count = d["disposable_count"]
-        not_disposable_count = d["not_disposable_count"]
-        return ScanDisposableKeysResponse(disposable_count, not_disposable_count)
-
-
 @dataclass
 class TenantConfig:
     tenant_specific_overrides: dict[str, Any]
@@ -162,19 +142,6 @@ class TenantConfig:
         )
 
 
-@dataclass
-class TimelinesInfoAndOffloaded:
-    timelines: list[dict[str, Any]]
-    offloaded: list[dict[str, Any]]
-
-    @classmethod
-    def from_json(cls, d: dict[str, Any]) -> TimelinesInfoAndOffloaded:
-        return TimelinesInfoAndOffloaded(
-            timelines=d["timelines"],
-            offloaded=d["offloaded"],
-        )
-
-
 class PageserverHttpClient(requests.Session, MetricsGetter):
     def __init__(
         self,
@@ -497,18 +464,6 @@ class PageserverHttpClient(requests.Session, MetricsGetter):
         assert isinstance(res_json, list)
         return res_json
 
-    def timeline_and_offloaded_list(
-        self,
-        tenant_id: Union[TenantId, TenantShardId],
-    ) -> TimelinesInfoAndOffloaded:
-        res = self.get(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline_and_offloaded",
-        )
-        self.verbose_error(res)
-        res_json = res.json()
-        assert isinstance(res_json, dict)
-        return TimelinesInfoAndOffloaded.from_json(res_json)
-
     def timeline_create(
         self,
         pg_version: PgVersion,
@@ -521,13 +476,12 @@ class PageserverHttpClient(requests.Session, MetricsGetter):
     ) -> dict[Any, Any]:
         body: dict[str, Any] = {
             "new_timeline_id": str(new_timeline_id),
+            "ancestor_start_lsn": str(ancestor_start_lsn) if ancestor_start_lsn else None,
+            "ancestor_timeline_id": str(ancestor_timeline_id) if ancestor_timeline_id else None,
+            "existing_initdb_timeline_id": str(existing_initdb_timeline_id)
+            if existing_initdb_timeline_id
+            else None,
         }
-        if ancestor_timeline_id:
-            body["ancestor_timeline_id"] = str(ancestor_timeline_id)
-        if ancestor_start_lsn:
-            body["ancestor_start_lsn"] = str(ancestor_start_lsn)
-        if existing_initdb_timeline_id:
-            body["existing_initdb_timeline_id"] = str(existing_initdb_timeline_id)
         if pg_version != PgVersion.NOT_SET:
             body["pg_version"] = int(pg_version)
 
@@ -925,16 +879,6 @@ class PageserverHttpClient(requests.Session, MetricsGetter):
         self.verbose_error(res)
         return LayerMapInfo.from_json(res.json())
 
-    def timeline_layer_scan_disposable_keys(
-        self, tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId, layer_name: str
-    ) -> ScanDisposableKeysResponse:
-        res = self.post(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/layer/{layer_name}/scan_disposable_keys",
-        )
-        self.verbose_error(res)
-        assert res.status_code == 200
-        return ScanDisposableKeysResponse.from_json(res.json())
-
     def download_layer(
         self, tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId, layer_name: str
     ):

test_runner/regress/test_pg_regress.py

@@ -3,13 +3,10 @@
 #
 from __future__ import annotations
 
-import os
-from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from typing import TYPE_CHECKING, cast
 
 import pytest
-from fixtures.log_helper import log
 from fixtures.neon_fixtures import (
     Endpoint,
     NeonEnv,
@@ -327,97 +324,3 @@ def test_sql_regress(
         pg_bin.run(pg_regress_command, env=env_vars, cwd=runpath)
 
     post_checks(env, test_output_dir, DBNAME, endpoint)
-
-
-@pytest.mark.skipif(os.environ.get("BUILD_TYPE") == "debug", reason="only run with release build")
-def test_tx_abort_with_many_relations(
-    neon_env_builder: NeonEnvBuilder,
-):
-    """
-    This is not a pg_regress test as such, but perhaps it should be -- this test exercises postgres
-    behavior when aborting a transaction with lots of relations.
-
-    Reproducer for https://github.com/neondatabase/neon/issues/9505
-    """
-
-    env = neon_env_builder.init_start()
-    ep = env.endpoints.create_start(
-        "main",
-        tenant_id=env.initial_tenant,
-        config_lines=[
-            "shared_buffers=1000MB",
-            "max_locks_per_transaction=16384",
-        ],
-    )
-
-    # How many relations: this number is tuned to be long enough to take tens of seconds
-    # if the rollback code path is buggy, tripping the test's timeout.
-    n = 4000
-
-    def create():
-        # Create many relations
-        log.info(f"Creating {n} relations...")
-        ep.safe_psql_many(
-            [
-                "BEGIN",
-                f"""DO $$
-            DECLARE
-                i INT;
-                table_name TEXT;
-            BEGIN
-                FOR i IN 1..{n} LOOP
-                    table_name := 'table_' || i;
-                    EXECUTE 'CREATE TABLE IF NOT EXISTS ' || table_name || ' (id SERIAL PRIMARY KEY, data TEXT)';
-                END LOOP;
-            END $$;
-            """,
-                "COMMIT",
-            ]
-        )
-
-    def truncate():
-        # Truncate relations, then roll back the transaction containing the truncations
-        log.info(f"Truncating {n} relations...")
-        ep.safe_psql_many(
-            [
-                "BEGIN",
-                f"""DO $$
-            DECLARE
-                i INT;
-                table_name TEXT;
-            BEGIN
-                FOR i IN 1..{n} LOOP
-                    table_name := 'table_' || i;
-                    EXECUTE 'TRUNCATE ' || table_name ;
-                END LOOP;
-            END $$;
-            """,
-            ]
-        )
-
-    def rollback_and_wait():
-        log.info(f"Rolling back after truncating {n} relations...")
-        ep.safe_psql("ROLLBACK")
-
-        # Restart the endpoint: this ensures that we can read back what we just wrote, i.e. pageserver
-        # ingest has caught up.
-        ep.stop()
-        log.info(f"Starting endpoint after truncating {n} relations...")
-        ep.start()
-        log.info(f"Started endpoint after truncating {n} relations...")
-
-    # Actual create & truncate phases may be slow, these involves lots of WAL records.  We do not
-    # apply a special timeout, they are expected to complete within general test timeout
-    create()
-    truncate()
-
-    # Run in a thread because the failure case is to take pathologically long time, and we don't want
-    # to block the test executor on that.
-    with ThreadPoolExecutor(max_workers=1) as exec:
-        try:
-            # Rollback phase should be fast: this is one WAL record that we should process efficiently
-            fut = exec.submit(rollback_and_wait)
-            fut.result(timeout=5)
-        except:
-            exec.shutdown(wait=False, cancel_futures=True)
-            raise

test_runner/regress/test_readonly_node.py

@@ -169,24 +169,23 @@ def test_readonly_node_gc(neon_env_builder: NeonEnvBuilder):
             )
         return last_flush_lsn
 
-    def trigger_gc_and_select(env: NeonEnv, ep_static: Endpoint, ctx: str):
+    def trigger_gc_and_select(env: NeonEnv, ep_static: Endpoint):
         """
         Trigger GC manually on all pageservers. Then run an `SELECT` query.
         """
         for shard, ps in tenant_get_shards(env, env.initial_tenant):
             client = ps.http_client()
             gc_result = client.timeline_gc(shard, env.initial_timeline, 0)
-            # Note: cannot assert on `layers_removed` here because it could be layers
-            # not guarded by the lease. Rely on successful execution of the query instead.
             log.info(f"{gc_result=}")
 
+            assert (
+                gc_result["layers_removed"] == 0
+            ), "No layers should be removed, old layers are guarded by leases."
+
         with ep_static.cursor() as cur:
-            # Following query should succeed if pages are properly guarded by leases.
             cur.execute("SELECT count(*) FROM t0")
             assert cur.fetchone() == (ROW_COUNT,)
 
-        log.info(f"`SELECT` query succeed after GC, {ctx=}")
-
     # Insert some records on main branch
     with env.endpoints.create_start("main") as ep_main:
         with ep_main.cursor() as cur:
@@ -211,9 +210,9 @@ def test_readonly_node_gc(neon_env_builder: NeonEnvBuilder):
             # Wait for static compute to renew lease at least once.
             time.sleep(LSN_LEASE_LENGTH / 2)
 
-            generate_updates_on_main(env, ep_main, 3, end=100)
+            generate_updates_on_main(env, ep_main, i, end=100)
 
-            trigger_gc_and_select(env, ep_static, ctx="Before pageservers restart")
+            trigger_gc_and_select(env, ep_static)
 
             # Trigger Pageserver restarts
             for ps in env.pageservers:
@@ -222,7 +221,7 @@ def test_readonly_node_gc(neon_env_builder: NeonEnvBuilder):
                 time.sleep(LSN_LEASE_LENGTH / 2)
                 ps.start()
 
-            trigger_gc_and_select(env, ep_static, ctx="After pageservers restart")
+            trigger_gc_and_select(env, ep_static)
 
             # Reconfigure pageservers
             env.pageservers[0].stop()
@@ -231,7 +230,7 @@ def test_readonly_node_gc(neon_env_builder: NeonEnvBuilder):
             )
             env.storage_controller.reconcile_until_idle()
 
-            trigger_gc_and_select(env, ep_static, ctx="After putting pageserver 0 offline")
+            trigger_gc_and_select(env, ep_static)
 
         # Do some update so we can increment latest_gc_cutoff
         generate_updates_on_main(env, ep_main, i, end=100)

test_runner/regress/test_sharding.py

@@ -3,11 +3,11 @@ from __future__ import annotations
 import os
 import time
 from collections import defaultdict
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 import pytest
 import requests
-from fixtures.common_types import Lsn, TenantId, TenantShardId, TimelineArchivalState, TimelineId
+from fixtures.common_types import Lsn, TenantId, TenantShardId, TimelineId
 from fixtures.compute_reconfigure import ComputeReconfigure
 from fixtures.log_helper import log
 from fixtures.neon_fixtures import (
@@ -188,9 +188,7 @@ def test_sharding_split_unsharded(
         "compact-shard-ancestors-persistent",
     ],
 )
-def test_sharding_split_compaction(
-    neon_env_builder: NeonEnvBuilder, failpoint: Optional[str], build_type: str
-):
+def test_sharding_split_compaction(neon_env_builder: NeonEnvBuilder, failpoint: Optional[str]):
     """
     Test that after a split, we clean up parent layer data in the child shards via compaction.
     """
@@ -324,19 +322,9 @@ def test_sharding_split_compaction(
             # Physical size should shrink because layers are smaller
             assert detail_after["current_physical_size"] < detail_before["current_physical_size"]
 
-    # Validate filtering compaction actually happened
+    # Validate size statistics
     for shard in shards:
         ps = env.get_tenant_pageserver(shard)
-
-        log.info("scan all layer files for disposable keys, there shouldn't be any")
-        result = ps.timeline_scan_no_disposable_keys(shard, timeline_id)
-        tally = result.tally
-        raw_page_count = tally.not_disposable_count + tally.disposable_count
-        assert tally.not_disposable_count > (
-            raw_page_count // 2
-        ), "compaction doesn't rewrite layers that are >=50pct local"
-
-        log.info("check sizes")
         timeline_info = ps.http_client().timeline_detail(shard, timeline_id)
         reported_size = timeline_info["current_physical_size"]
         layer_paths = ps.list_layers(shard, timeline_id)
@@ -365,145 +353,6 @@ def test_sharding_split_compaction(
     workload.validate()
 
 
-def test_sharding_split_offloading(neon_env_builder: NeonEnvBuilder):
-    """
-    Test that during a split, we don't miss archived and offloaded timelines.
-    """
-
-    TENANT_CONF = {
-        # small checkpointing and compaction targets to ensure we generate many upload operations
-        "checkpoint_distance": 128 * 1024,
-        "compaction_threshold": 1,
-        "compaction_target_size": 128 * 1024,
-        # no PITR horizon, we specify the horizon when we request on-demand GC
-        "pitr_interval": "3600s",
-        # disable background compaction, GC and offloading. We invoke it manually when we want it to happen.
-        "gc_period": "0s",
-        "compaction_period": "0s",
-        # Disable automatic creation of image layers, as we will create them explicitly when we want them
-        "image_creation_threshold": 9999,
-        "image_layer_creation_check_threshold": 0,
-        "lsn_lease_length": "0s",
-    }
-
-    neon_env_builder.storage_controller_config = {
-        # Default neon_local uses a small timeout: use a longer one to tolerate longer pageserver restarts.
-        "max_offline": "30s",
-        "max_warming_up": "300s",
-    }
-
-    env = neon_env_builder.init_start(initial_tenant_conf=TENANT_CONF)
-    tenant_id = env.initial_tenant
-    timeline_id_main = env.initial_timeline
-
-    # Check that we created with an unsharded TenantShardId: this is the default,
-    # but check it in case we change the default in future
-    assert env.storage_controller.inspect(TenantShardId(tenant_id, 0, 0)) is not None
-
-    workload_main = Workload(env, tenant_id, timeline_id_main, branch_name="main")
-    workload_main.init()
-    workload_main.write_rows(256)
-    workload_main.validate()
-    workload_main.stop()
-
-    # Create two timelines, archive one, offload the other
-    timeline_id_archived = env.create_branch("archived_not_offloaded")
-    timeline_id_offloaded = env.create_branch("archived_offloaded")
-
-    def timeline_id_set_for(list: list[dict[str, Any]]) -> set[TimelineId]:
-        return set(
-            map(
-                lambda t: TimelineId(t["timeline_id"]),
-                list,
-            )
-        )
-
-    expected_offloaded_set = {timeline_id_offloaded}
-    expected_timeline_set = {timeline_id_main, timeline_id_archived}
-
-    with env.get_tenant_pageserver(tenant_id).http_client() as http_client:
-        http_client.timeline_archival_config(
-            tenant_id, timeline_id_archived, TimelineArchivalState.ARCHIVED
-        )
-        http_client.timeline_archival_config(
-            tenant_id, timeline_id_offloaded, TimelineArchivalState.ARCHIVED
-        )
-        http_client.timeline_offload(tenant_id, timeline_id_offloaded)
-        list = http_client.timeline_and_offloaded_list(tenant_id)
-        assert timeline_id_set_for(list.offloaded) == expected_offloaded_set
-        assert timeline_id_set_for(list.timelines) == expected_timeline_set
-
-        # Do a full image layer generation before splitting
-        http_client.timeline_checkpoint(
-            tenant_id, timeline_id_main, force_image_layer_creation=True, wait_until_uploaded=True
-        )
-
-    # Split one shard into two
-    shards = env.storage_controller.tenant_shard_split(tenant_id, shard_count=2)
-
-    # Let all shards move into their stable locations, so that during subsequent steps we
-    # don't have reconciles in progress (simpler to reason about what messages we expect in logs)
-    env.storage_controller.reconcile_until_idle()
-
-    # Check we got the shard IDs we expected
-    assert env.storage_controller.inspect(TenantShardId(tenant_id, 0, 2)) is not None
-    assert env.storage_controller.inspect(TenantShardId(tenant_id, 1, 2)) is not None
-
-    workload_main.validate()
-    workload_main.stop()
-
-    env.storage_controller.consistency_check()
-
-    # Ensure each shard has the same list of timelines and offloaded timelines
-    for shard in shards:
-        ps = env.get_tenant_pageserver(shard)
-
-        list = ps.http_client().timeline_and_offloaded_list(shard)
-        assert timeline_id_set_for(list.offloaded) == expected_offloaded_set
-        assert timeline_id_set_for(list.timelines) == expected_timeline_set
-
-        ps.http_client().timeline_compact(shard, timeline_id_main)
-
-    # Check that we can still read all the data
-    workload_main.validate()
-
-    # Force a restart, which requires the state to be persisted.
-    env.pageserver.stop()
-    env.pageserver.start()
-
-    # Ensure each shard has the same list of timelines and offloaded timelines
-    for shard in shards:
-        ps = env.get_tenant_pageserver(shard)
-
-        list = ps.http_client().timeline_and_offloaded_list(shard)
-        assert timeline_id_set_for(list.offloaded) == expected_offloaded_set
-        assert timeline_id_set_for(list.timelines) == expected_timeline_set
-
-        ps.http_client().timeline_compact(shard, timeline_id_main)
-
-    # Compaction shouldn't make anything unreadable
-    workload_main.validate()
-
-    # Do sharded unarchival
-    env.storage_controller.timeline_archival_config(
-        tenant_id, timeline_id_offloaded, TimelineArchivalState.UNARCHIVED
-    )
-    env.storage_controller.timeline_archival_config(
-        tenant_id, timeline_id_archived, TimelineArchivalState.UNARCHIVED
-    )
-
-    for shard in shards:
-        ps = env.get_tenant_pageserver(shard)
-
-        list = ps.http_client().timeline_and_offloaded_list(shard)
-        assert timeline_id_set_for(list.offloaded) == set()
-        assert timeline_id_set_for(list.timelines) == {
-            timeline_id_main,
-            timeline_id_archived,
-            timeline_id_offloaded,
-        }
-
-
 def test_sharding_split_smoke(
     neon_env_builder: NeonEnvBuilder,
 ):

test_runner/regress/test_tenant_relocation.py

@@ -435,9 +435,7 @@ def test_emergency_relocate_with_branches_slow_replay(
 
     # This fail point will pause the WAL ingestion on the main branch, after the
     # the first insert
-    pageserver_http.configure_failpoints(
-        [("pageserver-wal-ingest-logical-message-sleep", "return(5000)")]
-    )
+    pageserver_http.configure_failpoints([("wal-ingest-logical-message-sleep", "return(5000)")])
 
     # Attach and wait a few seconds to give it time to load the tenants, attach to the
     # safekeepers, and to stream and ingest the WAL up to the pause-point.
@@ -455,13 +453,11 @@ def test_emergency_relocate_with_branches_slow_replay(
         assert cur.fetchall() == [("before pause",), ("after pause",)]
 
     # Sanity check that the failpoint was reached
-    env.pageserver.assert_log_contains(
-        'failpoint "pageserver-wal-ingest-logical-message-sleep": sleep done'
-    )
+    env.pageserver.assert_log_contains('failpoint "wal-ingest-logical-message-sleep": sleep done')
     assert time.time() - before_attach_time > 5
 
     # Clean up
-    pageserver_http.configure_failpoints(("pageserver-wal-ingest-logical-message-sleep", "off"))
+    pageserver_http.configure_failpoints(("wal-ingest-logical-message-sleep", "off"))
 
 
 # Simulate hard crash of pageserver and re-attach a tenant with a branch
@@ -585,9 +581,7 @@ def test_emergency_relocate_with_branches_createdb(
     # bug reproduced easily even without this, as there is always some delay between
     # loading the timeline and establishing the connection to the safekeeper to stream and
     # ingest the WAL, but let's make this less dependent on accidental timing.
-    pageserver_http.configure_failpoints(
-        [("pageserver-wal-ingest-logical-message-sleep", "return(5000)")]
-    )
+    pageserver_http.configure_failpoints([("wal-ingest-logical-message-sleep", "return(5000)")])
     before_attach_time = time.time()
     env.pageserver.tenant_attach(tenant_id)
 
@@ -596,10 +590,8 @@ def test_emergency_relocate_with_branches_createdb(
         assert query_scalar(cur, "SELECT count(*) FROM test_migrate_one") == 200
 
     # Sanity check that the failpoint was reached
-    env.pageserver.assert_log_contains(
-        'failpoint "pageserver-wal-ingest-logical-message-sleep": sleep done'
-    )
+    env.pageserver.assert_log_contains('failpoint "wal-ingest-logical-message-sleep": sleep done')
     assert time.time() - before_attach_time > 5
 
     # Clean up
-    pageserver_http.configure_failpoints(("pageserver-wal-ingest-logical-message-sleep", "off"))
+    pageserver_http.configure_failpoints(("wal-ingest-logical-message-sleep", "off"))

test_runner/regress/test_wal_receiver.py

@@ -1,12 +1,11 @@
 from __future__ import annotations
 
-import os
+import time
 from typing import TYPE_CHECKING
 
 from fixtures.common_types import Lsn, TenantId
 from fixtures.log_helper import log
 from fixtures.neon_fixtures import NeonEnv, NeonEnvBuilder
-from fixtures.utils import wait_until
 
 if TYPE_CHECKING:
     from typing import Any
@@ -20,10 +19,6 @@ def test_pageserver_lsn_wait_error_start(neon_env_builder: NeonEnvBuilder):
     env = neon_env_builder.init_start()
     env.pageserver.http_client()
 
-    # In this test we force 'Timed out while waiting for WAL record error' while
-    # fetching basebackup and don't want any retries.
-    os.environ["NEON_COMPUTE_TESTING_BASEBACKUP_RETRIES"] = "1"
-
     tenant_id, timeline_id = env.create_tenant()
     expected_timeout_error = f"Timed out while waiting for WAL record at LSN {future_lsn} to arrive"
     env.pageserver.allowed_errors.append(f".*{expected_timeout_error}.*")
@@ -54,14 +49,11 @@ def test_pageserver_lsn_wait_error_start(neon_env_builder: NeonEnvBuilder):
 def test_pageserver_lsn_wait_error_safekeeper_stop(neon_env_builder: NeonEnvBuilder):
     # Trigger WAL wait timeout faster
     def customize_pageserver_toml(ps_cfg: dict[str, Any]):
-        ps_cfg["wait_lsn_timeout"] = "2s"
+        ps_cfg["wait_lsn_timeout"] = "1s"
         tenant_config = ps_cfg.setdefault("tenant_config", {})
         tenant_config["walreceiver_connect_timeout"] = "2s"
         tenant_config["lagging_wal_timeout"] = "2s"
 
-    # In this test we force 'Timed out while waiting for WAL record error' while
-    # fetching basebackup and don't want any retries.
-    os.environ["NEON_COMPUTE_TESTING_BASEBACKUP_RETRIES"] = "1"
     neon_env_builder.pageserver_config_override = customize_pageserver_toml
 
     # Have notable SK ids to ensure we check logs for their presence, not some other random numbers
@@ -72,6 +64,7 @@ def test_pageserver_lsn_wait_error_safekeeper_stop(neon_env_builder: NeonEnvBuil
 
     tenant_id, timeline_id = env.create_tenant()
 
+    elements_to_insert = 1_000_000
     expected_timeout_error = f"Timed out while waiting for WAL record at LSN {future_lsn} to arrive"
     env.pageserver.allowed_errors.append(f".*{expected_timeout_error}.*")
     # we configure wait_lsn_timeout to a shorter value than the lagging_wal_timeout / walreceiver_connect_timeout
@@ -81,50 +74,45 @@ def test_pageserver_lsn_wait_error_safekeeper_stop(neon_env_builder: NeonEnvBuil
         ".*ingesting record with timestamp lagging more than wait_lsn_timeout.*"
     )
 
-    insert_test_elements(env, tenant_id, start=0, count=1)
+    insert_test_elements(env, tenant_id, start=0, count=elements_to_insert)
 
-    def all_sks_in_wareceiver_state():
-        try:
-            trigger_wait_lsn_timeout(env, tenant_id)
-        except Exception as e:
-            exception_string = str(e)
+    try:
+        trigger_wait_lsn_timeout(env, tenant_id)
+    except Exception as e:
+        exception_string = str(e)
+        assert expected_timeout_error in exception_string, "Should time out during waiting for WAL"
+
+        for safekeeper in env.safekeepers:
             assert (
-                expected_timeout_error in exception_string
-            ), "Should time out during waiting for WAL"
-
-            for safekeeper in env.safekeepers:
-                assert (
-                    str(safekeeper.id) in exception_string
-                ), f"Should have safekeeper {safekeeper.id} printed in walreceiver state after WAL wait timeout"
-
-    wait_until(60, 0.5, all_sks_in_wareceiver_state)
+                str(safekeeper.id) in exception_string
+            ), f"Should have safekeeper {safekeeper.id} printed in walreceiver state after WAL wait timeout"
 
     stopped_safekeeper = env.safekeepers[-1]
     stopped_safekeeper_id = stopped_safekeeper.id
     log.info(f"Stopping safekeeper {stopped_safekeeper.id}")
     stopped_safekeeper.stop()
+    # sleep until stopped safekeeper is removed from candidates
+    time.sleep(2)
 
-    def all_but_stopped_sks_in_wareceiver_state():
-        try:
-            trigger_wait_lsn_timeout(env, tenant_id)
-        except Exception as e:
-            # Strip out the part before stdout, as it contains full command with the list of all safekeepers
-            exception_string = str(e).split("stdout", 1)[-1]
-            assert (
-                expected_timeout_error in exception_string
-            ), "Should time out during waiting for WAL"
+    # Spend some more time inserting, to ensure SKs report updated statuses and walreceiver in PS have time to update its connection stats.
+    insert_test_elements(env, tenant_id, start=elements_to_insert + 1, count=elements_to_insert)
 
-            for safekeeper in env.safekeepers:
-                if safekeeper.id == stopped_safekeeper_id:
-                    assert (
-                        str(safekeeper.id) not in exception_string
-                    ), f"Should not have stopped safekeeper {safekeeper.id} printed in walreceiver state after 2nd WAL wait timeout"
-                else:
-                    assert (
-                        str(safekeeper.id) in exception_string
-                    ), f"Should have safekeeper {safekeeper.id} printed in walreceiver state after 2nd WAL wait timeout"
+    try:
+        trigger_wait_lsn_timeout(env, tenant_id)
+    except Exception as e:
+        # Strip out the part before stdout, as it contains full command with the list of all safekeepers
+        exception_string = str(e).split("stdout", 1)[-1]
+        assert expected_timeout_error in exception_string, "Should time out during waiting for WAL"
 
-    wait_until(60, 0.5, all_but_stopped_sks_in_wareceiver_state)
+        for safekeeper in env.safekeepers:
+            if safekeeper.id == stopped_safekeeper_id:
+                assert (
+                    str(safekeeper.id) not in exception_string
+                ), f"Should not have stopped safekeeper {safekeeper.id} printed in walreceiver state after 2nd WAL wait timeout"
+            else:
+                assert (
+                    str(safekeeper.id) in exception_string
+                ), f"Should have safekeeper {safekeeper.id} printed in walreceiver state after 2nd WAL wait timeout"
 
 
 def insert_test_elements(env: NeonEnv, tenant_id: TenantId, start: int, count: int):