keep holding layer map lock inside compact_level0_phase1

This is preliminary work for/from #4220 (async `Layer::get_value_reconstruct_data`).

That plan involves turning Timeline::layers into a `tokio::sync::RwLock`.

Before this patch, `compact_level0_phase1` would drop and re-acquire
`Timeline::layers` (`read().unwrap()`).

We can't have that if we switch to an async RwLock because
tokio::sync::RwLock's read guard is `!Send`.
Which makes the compaction task future `!Send`.
Which doesn't compile because the compaction task runs inside a `task_mgr`
task, and `task_mgr` requires the task futures to be `Send`.

I think the performance impact of holding the lock for longer will be minimal
as the number of L0s should be bounded, and so, any of the prep work that
we're doing in the section where we previously didn't hold the lock should
be fast anyways.

The far worse part of this function in terms of lock contention / performance
is the fact that we hold the read lock while we write out the L1s.

But that's a pre-existing condition => not changing it here.

Cargo.lock

@@ -110,12 +110,6 @@ dependencies = [
  "backtrace",
 ]
 
-[[package]]
-name = "arc-swap"
-version = "1.6.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bddcadddf5e9015d310179a59bb28c4d4b9920ad0f11e8e14dbadf654890c9a6"
-
 [[package]]
 name = "archery"
 version = "0.5.0"
@@ -2548,7 +2542,6 @@ name = "pageserver"
 version = "0.1.0"
 dependencies = [
  "anyhow",
- "arc-swap",
  "async-stream",
  "async-trait",
  "byteorder",

Cargo.toml

@@ -32,7 +32,6 @@ license = "Apache-2.0"
 ## All dependency versions, used in the project
 [workspace.dependencies]
 anyhow = { version = "1.0", features = ["backtrace"] }
-arc-swap = "1.6"
 async-stream = "0.3"
 async-trait = "0.1"
 atty = "0.2.14"

compute_tools/src/compute.rs

@@ -370,6 +370,11 @@ impl ComputeNode {
         // 'Close' connection
         drop(client);
 
+        info!(
+            "finished configuration of compute for project {}",
+            spec.cluster.cluster_id.as_deref().unwrap_or("None")
+        );
+
         Ok(())
     }
 
@@ -422,22 +427,22 @@ impl ComputeNode {
     #[instrument(skip(self))]
     pub fn start_compute(&self) -> Result<std::process::Child> {
         let compute_state = self.state.lock().unwrap().clone();
-        let pspec = compute_state.pspec.as_ref().expect("spec must be set");
+        let spec = compute_state.pspec.as_ref().expect("spec must be set");
         info!(
             "starting compute for project {}, operation {}, tenant {}, timeline {}",
-            pspec.spec.cluster.cluster_id.as_deref().unwrap_or("None"),
-            pspec.spec.operation_uuid.as_deref().unwrap_or("None"),
-            pspec.tenant_id,
-            pspec.timeline_id,
+            spec.spec.cluster.cluster_id.as_deref().unwrap_or("None"),
+            spec.spec.operation_uuid.as_deref().unwrap_or("None"),
+            spec.tenant_id,
+            spec.timeline_id,
         );
 
         self.prepare_pgdata(&compute_state)?;
 
         let start_time = Utc::now();
 
-        let pg = self.start_postgres(pspec.storage_auth_token.clone())?;
+        let pg = self.start_postgres(spec.storage_auth_token.clone())?;
 
-        if pspec.spec.mode == ComputeMode::Primary && !pspec.spec.skip_pg_catalog_updates {
+        if spec.spec.mode == ComputeMode::Primary {
             self.apply_config(&compute_state)?;
         }
 
@@ -457,11 +462,6 @@ impl ComputeNode {
         }
         self.set_status(ComputeStatus::Running);
 
-        info!(
-            "finished configuration of compute for project {}",
-            pspec.spec.cluster.cluster_id.as_deref().unwrap_or("None")
-        );
-
         Ok(pg)
     }
 

control_plane/src/endpoint.rs

@@ -450,7 +450,6 @@ impl Endpoint {
 
         // Create spec file
         let spec = ComputeSpec {
-            skip_pg_catalog_updates: false,
             format_version: 1.0,
             operation_uuid: None,
             cluster: Cluster {

libs/compute_api/src/spec.rs

@@ -27,12 +27,6 @@ pub struct ComputeSpec {
     pub cluster: Cluster,
     pub delta_operations: Option<Vec<DeltaOp>>,
 
-    /// An optinal hint that can be passed to speed up startup time if we know
-    /// that no pg catalog mutations (like role creation, database creation,
-    /// extension creation) need to be done on the actual database to start.
-    #[serde(default)] // Default false
-    pub skip_pg_catalog_updates: bool,
-
     // Information needed to connect to the storage layer.
     //
     // `tenant_id`, `timeline_id` and `pageserver_connstring` are always needed.

pageserver/Cargo.toml

@@ -12,7 +12,6 @@ testing = ["fail/failpoints"]
 
 [dependencies]
 anyhow.workspace = true
-arc-swap.workspace = true
 async-stream.workspace = true
 async-trait.workspace = true
 byteorder.workspace = true

pageserver/benches/bench_layer_map.rs

@@ -1,23 +1,22 @@
 use pageserver::keyspace::{KeyPartitioning, KeySpace};
 use pageserver::repository::Key;
 use pageserver::tenant::layer_map::LayerMap;
-use pageserver::tenant::storage_layer::{tests::LayerDescriptor, Layer, LayerFileName};
-use pageserver::tenant::storage_layer::{PersistentLayer, PersistentLayerDesc};
+use pageserver::tenant::storage_layer::{Layer, LayerDescriptor, LayerFileName};
 use rand::prelude::{SeedableRng, SliceRandom, StdRng};
 use std::cmp::{max, min};
 use std::fs::File;
 use std::io::{BufRead, BufReader};
 use std::path::PathBuf;
 use std::str::FromStr;
+use std::sync::Arc;
 use std::time::Instant;
-use utils::id::{TenantId, TimelineId};
 
 use utils::lsn::Lsn;
 
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 
-fn build_layer_map(filename_dump: PathBuf) -> LayerMap {
-    let mut layer_map = LayerMap::default();
+fn build_layer_map(filename_dump: PathBuf) -> LayerMap<LayerDescriptor> {
+    let mut layer_map = LayerMap::<LayerDescriptor>::default();
 
     let mut min_lsn = Lsn(u64::MAX);
     let mut max_lsn = Lsn(0);
@@ -34,7 +33,7 @@ fn build_layer_map(filename_dump: PathBuf) -> LayerMap {
         min_lsn = min(min_lsn, lsn_range.start);
         max_lsn = max(max_lsn, Lsn(lsn_range.end.0 - 1));
 
-        updates.insert_historic(layer.layer_desc().clone());
+        updates.insert_historic(layer.get_persistent_layer_desc(), Arc::new(layer));
     }
 
     println!("min: {min_lsn}, max: {max_lsn}");
@@ -44,7 +43,7 @@ fn build_layer_map(filename_dump: PathBuf) -> LayerMap {
 }
 
 /// Construct a layer map query pattern for benchmarks
-fn uniform_query_pattern(layer_map: &LayerMap) -> Vec<(Key, Lsn)> {
+fn uniform_query_pattern(layer_map: &LayerMap<LayerDescriptor>) -> Vec<(Key, Lsn)> {
     // For each image layer we query one of the pages contained, at LSN right
     // before the image layer was created. This gives us a somewhat uniform
     // coverage of both the lsn and key space because image layers have
@@ -70,7 +69,7 @@ fn uniform_query_pattern(layer_map: &LayerMap) -> Vec<(Key, Lsn)> {
 
 // Construct a partitioning for testing get_difficulty map when we
 // don't have an exact result of `collect_keyspace` to work with.
-fn uniform_key_partitioning(layer_map: &LayerMap, _lsn: Lsn) -> KeyPartitioning {
+fn uniform_key_partitioning(layer_map: &LayerMap<LayerDescriptor>, _lsn: Lsn) -> KeyPartitioning {
     let mut parts = Vec::new();
 
     // We add a partition boundary at the start of each image layer,
@@ -210,15 +209,13 @@ fn bench_sequential(c: &mut Criterion) {
     for i in 0..100_000 {
         let i32 = (i as u32) % 100;
         let zero = Key::from_hex("000000000000000000000000000000000000").unwrap();
-        let layer = LayerDescriptor::from(PersistentLayerDesc::new_img(
-            TenantId::generate(),
-            TimelineId::generate(),
-            zero.add(10 * i32)..zero.add(10 * i32 + 1),
-            Lsn(i),
-            false,
-            0,
-        ));
-        updates.insert_historic(layer.layer_desc().clone());
+        let layer = LayerDescriptor {
+            key: zero.add(10 * i32)..zero.add(10 * i32 + 1),
+            lsn: Lsn(i)..Lsn(i + 1),
+            is_incremental: false,
+            short_id: format!("Layer {}", i),
+        };
+        updates.insert_historic(layer.get_persistent_layer_desc(), Arc::new(layer));
     }
     updates.flush();
     println!("Finished layer map init in {:?}", now.elapsed());

pageserver/src/disk_usage_eviction_task.rs

@@ -516,7 +516,7 @@ async fn collect_eviction_candidates(
             if !tl.is_active() {
                 continue;
             }
-            let info = tl.get_local_layers_for_disk_usage_eviction().await;
+            let info = tl.get_local_layers_for_disk_usage_eviction();
             debug!(tenant_id=%tl.tenant_id, timeline_id=%tl.timeline_id, "timeline resident layers count: {}", info.resident_layers.len());
             tenant_candidates.extend(
                 info.resident_layers

pageserver/src/http/routes.rs

@@ -215,7 +215,7 @@ async fn build_timeline_info(
 ) -> anyhow::Result<TimelineInfo> {
     crate::tenant::debug_assert_current_span_has_tenant_and_timeline_id();
 
-    let mut info = build_timeline_info_common(timeline, ctx).await?;
+    let mut info = build_timeline_info_common(timeline, ctx)?;
     if include_non_incremental_logical_size {
         // XXX we should be using spawn_ondemand_logical_size_calculation here.
         // Otherwise, if someone deletes the timeline / detaches the tenant while
@@ -233,7 +233,7 @@ async fn build_timeline_info(
     Ok(info)
 }
 
-async fn build_timeline_info_common(
+fn build_timeline_info_common(
     timeline: &Arc<Timeline>,
     ctx: &RequestContext,
 ) -> anyhow::Result<TimelineInfo> {
@@ -264,7 +264,7 @@ async fn build_timeline_info_common(
             None
         }
     };
-    let current_physical_size = Some(timeline.layer_size_sum().await);
+    let current_physical_size = Some(timeline.layer_size_sum());
     let state = timeline.current_state();
     let remote_consistent_lsn = timeline.get_remote_consistent_lsn().unwrap_or(Lsn(0));
 
@@ -330,7 +330,6 @@ async fn timeline_create_handler(
             Ok(Some(new_timeline)) => {
                 // Created. Construct a TimelineInfo for it.
                 let timeline_info = build_timeline_info_common(&new_timeline, &ctx)
-                    .await
                     .map_err(ApiError::InternalServerError)?;
                 json_response(StatusCode::CREATED, timeline_info)
             }
@@ -592,7 +591,7 @@ async fn tenant_status(
         // Calculate total physical size of all timelines
         let mut current_physical_size = 0;
         for timeline in tenant.list_timelines().iter() {
-            current_physical_size += timeline.layer_size_sum().await;
+            current_physical_size += timeline.layer_size_sum();
         }
 
         let state = tenant.current_state();
@@ -702,7 +701,7 @@ async fn layer_map_info_handler(
     check_permission(&request, Some(tenant_id))?;
 
     let timeline = active_timeline_of_active_tenant(tenant_id, timeline_id).await?;
-    let layer_map_info = timeline.layer_map_info(reset).await;
+    let layer_map_info = timeline.layer_map_info(reset);
 
     json_response(StatusCode::OK, layer_map_info)
 }

pageserver/src/import_datadir.rs

@@ -75,12 +75,12 @@ pub async fn import_timeline_from_postgres_datadir(
             {
                 pg_control = Some(control_file);
             }
-            modification.flush().await?;
+            modification.flush()?;
         }
     }
 
     // We're done importing all the data files.
-    modification.commit().await?;
+    modification.commit()?;
 
     // We expect the Postgres server to be shut down cleanly.
     let pg_control = pg_control.context("pg_control file not found")?;
@@ -359,7 +359,7 @@ pub async fn import_basebackup_from_tar(
                     // We found the pg_control file.
                     pg_control = Some(res);
                 }
-                modification.flush().await?;
+                modification.flush()?;
             }
             tokio_tar::EntryType::Directory => {
                 debug!("directory {:?}", file_path);
@@ -377,7 +377,7 @@ pub async fn import_basebackup_from_tar(
     // sanity check: ensure that pg_control is loaded
     let _pg_control = pg_control.context("pg_control file not found")?;
 
-    modification.commit().await?;
+    modification.commit()?;
     Ok(())
 }
 
@@ -594,7 +594,7 @@ async fn import_file(
         // zenith.signal is not necessarily the last file, that we handle
         // but it is ok to call `finish_write()`, because final `modification.commit()`
         // will update lsn once more to the final one.
-        let writer = modification.tline.writer().await;
+        let writer = modification.tline.writer();
         writer.finish_write(prev_lsn);
 
         debug!("imported zenith signal {}", prev_lsn);

pageserver/src/page_service.rs

@@ -489,9 +489,7 @@ impl PageServerHandler {
         // Create empty timeline
         info!("creating new timeline");
         let tenant = get_active_tenant_with_timeout(tenant_id, &ctx).await?;
-        let timeline = tenant
-            .create_empty_timeline(timeline_id, base_lsn, pg_version, &ctx)
-            .await?;
+        let timeline = tenant.create_empty_timeline(timeline_id, base_lsn, pg_version, &ctx)?;
 
         // TODO mark timeline as not ready until it reaches end_lsn.
         // We might have some wal to import as well, and we should prevent compute

pageserver/src/pgdatadir_mapping.rs

@@ -1122,7 +1122,7 @@ impl<'a> DatadirModification<'a> {
     /// retains all the metadata, but data pages are flushed. That's again OK
     /// for bulk import, where you are just loading data pages and won't try to
     /// modify the same pages twice.
-    pub async fn flush(&mut self) -> anyhow::Result<()> {
+    pub fn flush(&mut self) -> anyhow::Result<()> {
         // Unless we have accumulated a decent amount of changes, it's not worth it
         // to scan through the pending_updates list.
         let pending_nblocks = self.pending_nblocks;
@@ -1130,20 +1130,19 @@ impl<'a> DatadirModification<'a> {
             return Ok(());
         }
 
-        let writer = self.tline.writer().await;
+        let writer = self.tline.writer();
 
         // Flush relation and  SLRU data blocks, keep metadata.
-        let mut retained_pending_updates = HashMap::new();
-        for (key, value) in self.pending_updates.drain() {
-            if is_rel_block_key(key) || is_slru_block_key(key) {
-                // This bails out on first error without modifying pending_updates.
-                // That's Ok, cf this function's doc comment.
-                writer.put(key, self.lsn, &value).await?;
+        let mut result: anyhow::Result<()> = Ok(());
+        self.pending_updates.retain(|&key, value| {
+            if result.is_ok() && (is_rel_block_key(key) || is_slru_block_key(key)) {
+                result = writer.put(key, self.lsn, value);
+                false
             } else {
-                retained_pending_updates.insert(key, value);
+                true
             }
-        }
-        self.pending_updates.extend(retained_pending_updates);
+        });
+        result?;
 
         if pending_nblocks != 0 {
             writer.update_current_logical_size(pending_nblocks * i64::from(BLCKSZ));
@@ -1158,17 +1157,17 @@ impl<'a> DatadirModification<'a> {
     /// underlying timeline.
     /// All the modifications in this atomic update are stamped by the specified LSN.
     ///
-    pub async fn commit(&mut self) -> anyhow::Result<()> {
-        let writer = self.tline.writer().await;
+    pub fn commit(&mut self) -> anyhow::Result<()> {
+        let writer = self.tline.writer();
         let lsn = self.lsn;
         let pending_nblocks = self.pending_nblocks;
         self.pending_nblocks = 0;
 
         for (key, value) in self.pending_updates.drain() {
-            writer.put(key, lsn, &value).await?;
+            writer.put(key, lsn, &value)?;
         }
         for key_range in self.pending_deletions.drain(..) {
-            writer.delete(key_range, lsn).await?;
+            writer.delete(key_range, lsn)?;
         }
 
         writer.finish_write(lsn);

pageserver/src/tenant.rs

@@ -11,7 +11,7 @@
 //! parent timeline, and the last LSN that has been written to disk.
 //!
 
-use anyhow::{bail, ensure, Context};
+use anyhow::{bail, Context};
 use futures::FutureExt;
 use pageserver_api::models::TimelineState;
 use remote_storage::DownloadError;
@@ -85,9 +85,7 @@ pub mod blob_io;
 pub mod block_io;
 pub mod disk_btree;
 pub(crate) mod ephemeral_file;
-pub mod layer_cache;
 pub mod layer_map;
-pub mod layer_map_mgr;
 pub mod manifest;
 
 pub mod metadata;
@@ -188,13 +186,18 @@ struct TimelineUninitMark {
 }
 
 impl UninitializedTimeline<'_> {
-    /// Finish timeline creation: insert it into the Tenant's timelines map and remove the
-    /// uninit mark file.
+    /// Ensures timeline data is valid, loads it into pageserver's memory and removes
+    /// uninit mark file on success.
     ///
     /// This function launches the flush loop if not already done.
     ///
     /// The caller is responsible for activating the timeline (function `.activate()`).
-    fn finish_creation(mut self) -> anyhow::Result<Arc<Timeline>> {
+    fn initialize_with_lock(
+        mut self,
+        _ctx: &RequestContext,
+        timelines: &mut HashMap<TimelineId, Arc<Timeline>>,
+        load_layer_map: bool,
+    ) -> anyhow::Result<Arc<Timeline>> {
         let timeline_id = self.timeline_id;
         let tenant_id = self.owning_tenant.tenant_id;
 
@@ -202,19 +205,25 @@ impl UninitializedTimeline<'_> {
             format!("No timeline for initalization found for {tenant_id}/{timeline_id}")
         })?;
 
-        // Check that the caller initialized disk_consistent_lsn
         let new_disk_consistent_lsn = new_timeline.get_disk_consistent_lsn();
-        ensure!(
-            new_disk_consistent_lsn.is_valid(),
-            "new timeline {tenant_id}/{timeline_id} has invalid disk_consistent_lsn"
-        );
+        // TODO it would be good to ensure that, but apparently a lot of our testing is dependend on that at least
+        // ensure!(new_disk_consistent_lsn.is_valid(),
+        //     "Timeline {tenant_id}/{timeline_id} has invalid disk_consistent_lsn and cannot be initialized");
 
-        let mut timelines = self.owning_tenant.timelines.lock().unwrap();
         match timelines.entry(timeline_id) {
             Entry::Occupied(_) => anyhow::bail!(
                 "Found freshly initialized timeline {tenant_id}/{timeline_id} in the tenant map"
             ),
             Entry::Vacant(v) => {
+                if load_layer_map {
+                    new_timeline
+                        .load_layer_map(new_disk_consistent_lsn)
+                        .with_context(|| {
+                            format!(
+                                "Failed to load layermap for timeline {tenant_id}/{timeline_id}"
+                            )
+                        })?;
+                }
                 uninit_mark.remove_uninit_mark().with_context(|| {
                     format!(
                         "Failed to remove uninit mark file for timeline {tenant_id}/{timeline_id}"
@@ -243,10 +252,9 @@ impl UninitializedTimeline<'_> {
             .await
             .context("Failed to import basebackup")?;
 
-        // Flush the new layer files to disk, before we make the timeline as available to
-        // the outside world.
-        //
         // Flush loop needs to be spawned in order to be able to flush.
+        // We want to run proper checkpoint before we mark timeline as available to outside world
+        // Thus spawning flush loop manually and skipping flush_loop setup in initialize_with_lock
         raw_timeline.maybe_spawn_flush_loop();
 
         fail::fail_point!("before-checkpoint-new-timeline", |_| {
@@ -258,9 +266,10 @@ impl UninitializedTimeline<'_> {
             .await
             .context("Failed to flush after basebackup import")?;
 
-        // All the data has been imported. Insert the Timeline into the tenant's timelines
-        // map and remove the uninit mark file.
-        let tl = self.finish_creation()?;
+        // Initialize without loading the layer map. We started with an empty layer map, and already
+        // updated it for the layers that we created during the import.
+        let mut timelines = self.owning_tenant.timelines.lock().unwrap();
+        let tl = self.initialize_with_lock(ctx, &mut timelines, false)?;
         tl.activate(broker_client, None, ctx);
         Ok(tl)
     }
@@ -303,6 +312,15 @@ fn cleanup_timeline_directory(uninit_mark: TimelineUninitMark) {
 }
 
 impl TimelineUninitMark {
+    /// Useful for initializing timelines, existing on disk after the restart.
+    pub fn dummy() -> Self {
+        Self {
+            uninit_mark_deleted: true,
+            uninit_mark_path: PathBuf::new(),
+            timeline_path: PathBuf::new(),
+        }
+    }
+
     fn new(uninit_mark_path: PathBuf, timeline_path: PathBuf) -> Self {
         Self {
             uninit_mark_deleted: false,
@@ -496,7 +514,7 @@ impl Tenant {
         ancestor: Option<Arc<Timeline>>,
         first_save: bool,
         init_order: Option<&InitializationOrder>,
-        _ctx: &RequestContext,
+        ctx: &RequestContext,
     ) -> anyhow::Result<()> {
         let tenant_id = self.tenant_id;
 
@@ -507,39 +525,54 @@ impl Tenant {
         .context("merge_local_remote_metadata")?
         .to_owned();
 
-        let timeline = self.create_timeline_struct(
-            timeline_id,
-            up_to_date_metadata,
-            ancestor.clone(),
-            remote_client,
-            init_order,
-        )?;
-        let new_disk_consistent_lsn = timeline.get_disk_consistent_lsn();
-        anyhow::ensure!(
-            new_disk_consistent_lsn.is_valid(),
-            "Timeline {tenant_id}/{timeline_id} has invalid disk_consistent_lsn"
-        );
-        timeline
-            .load_layer_map(new_disk_consistent_lsn)
-            .await
-            .with_context(|| {
-                format!("Failed to load layermap for timeline {tenant_id}/{timeline_id}")
-            })?;
-
-        {
+        let timeline = {
             // avoiding holding it across awaits
             let mut timelines_accessor = self.timelines.lock().unwrap();
-            match timelines_accessor.entry(timeline_id) {
-                Entry::Occupied(_) => {
-                    // The uninit mark file acts as a lock that prevents another task from
-                    // initializing the timeline at the same time.
-                    unreachable!(
-                        "Timeline {tenant_id}/{timeline_id} already exists in the tenant map"
-                    );
-                }
-                Entry::Vacant(v) => {
-                    v.insert(Arc::clone(&timeline));
-                    timeline.maybe_spawn_flush_loop();
+            if timelines_accessor.contains_key(&timeline_id) {
+                anyhow::bail!(
+                    "Timeline {tenant_id}/{timeline_id} already exists in the tenant map"
+                );
+            }
+
+            let dummy_timeline = self.create_timeline_data(
+                timeline_id,
+                up_to_date_metadata,
+                ancestor.clone(),
+                remote_client,
+                init_order,
+            )?;
+
+            let timeline = UninitializedTimeline {
+                owning_tenant: self,
+                timeline_id,
+                raw_timeline: Some((dummy_timeline, TimelineUninitMark::dummy())),
+            };
+            // Do not start walreceiver here. We do need loaded layer map for reconcile_with_remote
+            // But we shouldnt start walreceiver before we have all the data locally, because working walreceiver
+            // will ingest data which may require looking at the layers which are not yet available locally
+            match timeline.initialize_with_lock(ctx, &mut timelines_accessor, true) {
+                Ok(new_timeline) => new_timeline,
+                Err(e) => {
+                    error!("Failed to initialize timeline {tenant_id}/{timeline_id}: {e:?}");
+                    // FIXME using None is a hack, it wont hurt, just ugly.
+                    //     Ideally initialize_with_lock error should return timeline in the error
+                    //     Or return ownership of itself completely so somethin like into_broken
+                    //     can be called directly on Uninitielized timeline
+                    //     also leades to redundant .clone
+                    let broken_timeline = self
+                        .create_timeline_data(
+                            timeline_id,
+                            up_to_date_metadata,
+                            ancestor.clone(),
+                            None,
+                            None,
+                        )
+                        .with_context(|| {
+                            format!("creating broken timeline data for {tenant_id}/{timeline_id}")
+                        })?;
+                    broken_timeline.set_broken(e.to_string());
+                    timelines_accessor.insert(timeline_id, broken_timeline);
+                    return Err(e);
                 }
             }
         };
@@ -557,10 +590,16 @@ impl Tenant {
                 .context("failed to reconcile with remote")?
         }
 
-        let layers = timeline.layer_mgr.read();
         // Sanity check: a timeline should have some content.
         anyhow::ensure!(
-            ancestor.is_some() || layers.iter_historic_layers().next().is_some(),
+            ancestor.is_some()
+                || timeline
+                    .layers
+                    .read()
+                    .unwrap()
+                    .iter_historic_layers()
+                    .next()
+                    .is_some(),
             "Timeline has no ancestor and no layer files"
         );
 
@@ -1125,14 +1164,14 @@ impl Tenant {
                                 .init_upload_queue_stopped_to_continue_deletion(&index_part)?;
 
                             let timeline = self
-                                .create_timeline_struct(
+                                .create_timeline_data(
                                     timeline_id,
                                     &local_metadata,
                                     ancestor,
                                     Some(remote_client),
                                     init_order,
                                 )
-                                .context("create_timeline_struct")?;
+                                .context("create_timeline_data")?;
 
                             let guard = Arc::clone(&timeline.delete_lock).lock_owned().await;
 
@@ -1241,7 +1280,7 @@ impl Tenant {
     /// 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 async fn create_empty_timeline(
+    pub fn create_empty_timeline(
         &self,
         new_timeline_id: TimelineId,
         initdb_lsn: Lsn,
@@ -1253,15 +1292,11 @@ impl Tenant {
             "Cannot create empty timelines on inactive tenant"
         );
 
-        let timeline_uninit_mark = {
-            let timelines: MutexGuard<'_, HashMap<TimelineId, Arc<Timeline>>> =
-                self.timelines.lock().unwrap();
-            self.create_timeline_uninit_mark(new_timeline_id, &timelines)?
-        };
+        let timelines = self.timelines.lock().unwrap();
+        let timeline_uninit_mark = self.create_timeline_uninit_mark(new_timeline_id, &timelines)?;
+        drop(timelines);
 
         let new_metadata = TimelineMetadata::new(
-            // Initialize disk_consistent LSN to 0, The caller must import some data to
-            // make it valid, before calling finish_creation()
             Lsn(0),
             None,
             None,
@@ -1270,14 +1305,13 @@ impl Tenant {
             initdb_lsn,
             pg_version,
         );
-        self.prepare_new_timeline(
+        self.prepare_timeline(
             new_timeline_id,
             &new_metadata,
             timeline_uninit_mark,
-            initdb_lsn,
+            true,
             None,
         )
-        .await
     }
 
     /// Helper for unit tests to create an emtpy timeline.
@@ -1286,16 +1320,14 @@ impl Tenant {
     // This makes the various functions which anyhow::ensure! for Active state work in tests.
     // Our current tests don't need the background loops.
     #[cfg(test)]
-    pub async fn create_test_timeline(
+    pub fn create_test_timeline(
         &self,
         new_timeline_id: TimelineId,
         initdb_lsn: Lsn,
         pg_version: u32,
         ctx: &RequestContext,
     ) -> anyhow::Result<Arc<Timeline>> {
-        let uninit_tl = self
-            .create_empty_timeline(new_timeline_id, initdb_lsn, pg_version, ctx)
-            .await?;
+        let uninit_tl = self.create_empty_timeline(new_timeline_id, initdb_lsn, pg_version, ctx)?;
         let tline = uninit_tl.raw_timeline().expect("we just created it");
         assert_eq!(tline.get_last_record_lsn(), Lsn(0));
 
@@ -1306,14 +1338,12 @@ impl Tenant {
             .context("init_empty_test_timeline")?;
         modification
             .commit()
-            .await
             .context("commit init_empty_test_timeline modification")?;
 
-        // Flush to disk so that uninit_tl's check for valid disk_consistent_lsn passes.
-        tline.maybe_spawn_flush_loop();
-        tline.freeze_and_flush().await.context("freeze_and_flush")?;
-
-        let tl = uninit_tl.finish_creation()?;
+        let mut timelines = self.timelines.lock().unwrap();
+        // load_layers=false because create_empty_timeline already did that what's necessary (set next_open_layer)
+        // and modification.init_empty() already created layers.
+        let tl = uninit_tl.initialize_with_lock(ctx, &mut timelines, false)?;
         // The non-test code would call tl.activate() here.
         tl.set_state(TimelineState::Active);
         Ok(tl)
@@ -1559,7 +1589,7 @@ impl Tenant {
             // No timeout here, GC & Compaction should be responsive to the
             // `TimelineState::Stopping` change.
             info!("waiting for layer_removal_cs.lock()");
-            let layer_removal_guard = timeline.lcache.delete_guard().await;
+            let layer_removal_guard = timeline.layer_removal_cs.lock().await;
             info!("got layer_removal_cs.lock(), deleting layer files");
 
             // NB: storage_sync upload tasks that reference these layers have been cancelled
@@ -2245,12 +2275,7 @@ impl Tenant {
         }
     }
 
-    /// Helper function to create a new Timeline struct.
-    ///
-    /// The returned Timeline is in Loading state. The caller is responsible for
-    /// initializing any on-disk state, and for inserting the Timeline to the 'timelines'
-    /// map.
-    fn create_timeline_struct(
+    fn create_timeline_data(
         &self,
         new_timeline_id: TimelineId,
         new_metadata: &TimelineMetadata,
@@ -2670,7 +2695,7 @@ impl Tenant {
         src_timeline: &Arc<Timeline>,
         dst_id: TimelineId,
         start_lsn: Option<Lsn>,
-        _ctx: &RequestContext,
+        ctx: &RequestContext,
     ) -> anyhow::Result<Arc<Timeline>> {
         let src_id = src_timeline.timeline_id;
 
@@ -2755,17 +2780,17 @@ impl Tenant {
             src_timeline.pg_version,
         );
 
-        let uninitialized_timeline = self
-            .prepare_new_timeline(
+        let new_timeline = {
+            let mut timelines = self.timelines.lock().unwrap();
+            self.prepare_timeline(
                 dst_id,
                 &metadata,
                 timeline_uninit_mark,
-                start_lsn + 1,
+                false,
                 Some(Arc::clone(src_timeline)),
-            )
-            .await?;
-
-        let new_timeline = uninitialized_timeline.finish_creation()?;
+            )?
+            .initialize_with_lock(ctx, &mut timelines, true)?
+        };
 
         // Root timeline gets its layers during creation and uploads them along with the metadata.
         // A branch timeline though, when created, can get no writes for some time, hence won't get any layers created.
@@ -2841,15 +2866,8 @@ impl Tenant {
             pgdata_lsn,
             pg_version,
         );
-        let raw_timeline = self
-            .prepare_new_timeline(
-                timeline_id,
-                &new_metadata,
-                timeline_uninit_mark,
-                pgdata_lsn,
-                None,
-            )
-            .await?;
+        let raw_timeline =
+            self.prepare_timeline(timeline_id, &new_metadata, timeline_uninit_mark, true, None)?;
 
         let tenant_id = raw_timeline.owning_tenant.tenant_id;
         let unfinished_timeline = raw_timeline.raw_timeline()?;
@@ -2865,10 +2883,10 @@ impl Tenant {
             format!("Failed to import pgdatadir for timeline {tenant_id}/{timeline_id}")
         })?;
 
-        // Flush the new layer files to disk, before we make the timeline as available to
+        // Flush the new layer files to disk, before we mark the timeline as available to
         // the outside world.
         //
-        // Flush loop needs to be spawned in order to be able to flush.
+        // Thus spawn flush loop manually and skip flush_loop setup in initialize_with_lock
         unfinished_timeline.maybe_spawn_flush_loop();
 
         fail::fail_point!("before-checkpoint-new-timeline", |_| {
@@ -2884,8 +2902,12 @@ impl Tenant {
                 )
             })?;
 
-        // All done!
-        let timeline = raw_timeline.finish_creation()?;
+        // Initialize the timeline without loading the layer map, because we already updated the layer
+        // map above, when we imported the datadir.
+        let timeline = {
+            let mut timelines = self.timelines.lock().unwrap();
+            raw_timeline.initialize_with_lock(ctx, &mut timelines, false)?
+        };
 
         info!(
             "created root timeline {} timeline.lsn {}",
@@ -2896,18 +2918,14 @@ impl Tenant {
         Ok(timeline)
     }
 
-    /// Creates intermediate timeline structure and its files.
-    ///
-    /// An empty layer map is initialized, and new data and WAL can be imported starting
-    /// at 'disk_consistent_lsn'. After any initial data has been imported, call
-    /// `finish_creation` to insert the Timeline into the timelines map and to remove the
-    /// uninit mark file.
-    async fn prepare_new_timeline(
+    /// Creates intermediate timeline structure and its files, without loading it into memory.
+    /// It's up to the caller to import the necesary data and import the timeline into memory.
+    fn prepare_timeline(
         &self,
         new_timeline_id: TimelineId,
         new_metadata: &TimelineMetadata,
         uninit_mark: TimelineUninitMark,
-        start_lsn: Lsn,
+        init_layers: bool,
         ancestor: Option<Arc<Timeline>>,
     ) -> anyhow::Result<UninitializedTimeline> {
         let tenant_id = self.tenant_id;
@@ -2925,27 +2943,33 @@ impl Tenant {
             None
         };
 
-        let timeline_struct = self
-            .create_timeline_struct(new_timeline_id, new_metadata, ancestor, remote_client, None)
-            .context("Failed to create timeline data structure")?;
-
-        timeline_struct.init_empty_layer_map(start_lsn).await?;
-
-        if let Err(e) =
-            self.create_timeline_files(&uninit_mark.timeline_path, new_timeline_id, new_metadata)
-        {
-            error!("Failed to create initial files for timeline {tenant_id}/{new_timeline_id}, cleaning up: {e:?}");
-            cleanup_timeline_directory(uninit_mark);
-            return Err(e);
+        match self.create_timeline_files(
+            &uninit_mark.timeline_path,
+            new_timeline_id,
+            new_metadata,
+            ancestor,
+            remote_client,
+        ) {
+            Ok(new_timeline) => {
+                if init_layers {
+                    new_timeline.layers.write().unwrap().next_open_layer_at =
+                        Some(new_timeline.initdb_lsn);
+                }
+                debug!(
+                    "Successfully created initial files for timeline {tenant_id}/{new_timeline_id}"
+                );
+                Ok(UninitializedTimeline {
+                    owning_tenant: self,
+                    timeline_id: new_timeline_id,
+                    raw_timeline: Some((new_timeline, uninit_mark)),
+                })
+            }
+            Err(e) => {
+                error!("Failed to create initial files for timeline {tenant_id}/{new_timeline_id}, cleaning up: {e:?}");
+                cleanup_timeline_directory(uninit_mark);
+                Err(e)
+            }
         }
-
-        debug!("Successfully created initial files for timeline {tenant_id}/{new_timeline_id}");
-
-        Ok(UninitializedTimeline {
-            owning_tenant: self,
-            timeline_id: new_timeline_id,
-            raw_timeline: Some((timeline_struct, uninit_mark)),
-        })
     }
 
     fn create_timeline_files(
@@ -2953,8 +2977,13 @@ impl Tenant {
         timeline_path: &Path,
         new_timeline_id: TimelineId,
         new_metadata: &TimelineMetadata,
-    ) -> anyhow::Result<()> {
-        crashsafe::create_dir(timeline_path).context("Failed to create timeline directory")?;
+        ancestor: Option<Arc<Timeline>>,
+        remote_client: Option<RemoteTimelineClient>,
+    ) -> anyhow::Result<Arc<Timeline>> {
+        let timeline_data = self
+            .create_timeline_data(new_timeline_id, new_metadata, ancestor, remote_client, None)
+            .context("Failed to create timeline data structure")?;
+        crashsafe::create_dir_all(timeline_path).context("Failed to create timeline directory")?;
 
         fail::fail_point!("after-timeline-uninit-mark-creation", |_| {
             anyhow::bail!("failpoint after-timeline-uninit-mark-creation");
@@ -2968,7 +2997,8 @@ impl Tenant {
             true,
         )
         .context("Failed to create timeline metadata")?;
-        Ok(())
+
+        Ok(timeline_data)
     }
 
     /// Attempts to create an uninit mark file for the timeline initialization.
@@ -3583,21 +3613,16 @@ mod tests {
     #[tokio::test]
     async fn test_basic() -> anyhow::Result<()> {
         let (tenant, ctx) = TenantHarness::create("test_basic")?.load().await;
-        let tline = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(0x08), DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let tline =
+            tenant.create_test_timeline(TIMELINE_ID, Lsn(0x08), DEFAULT_PG_VERSION, &ctx)?;
 
-        let writer = tline.writer().await;
-        writer
-            .put(*TEST_KEY, Lsn(0x10), &Value::Image(TEST_IMG("foo at 0x10")))
-            .await?;
+        let writer = tline.writer();
+        writer.put(*TEST_KEY, Lsn(0x10), &Value::Image(TEST_IMG("foo at 0x10")))?;
         writer.finish_write(Lsn(0x10));
         drop(writer);
 
-        let writer = tline.writer().await;
-        writer
-            .put(*TEST_KEY, Lsn(0x20), &Value::Image(TEST_IMG("foo at 0x20")))
-            .await?;
+        let writer = tline.writer();
+        writer.put(*TEST_KEY, Lsn(0x20), &Value::Image(TEST_IMG("foo at 0x20")))?;
         writer.finish_write(Lsn(0x20));
         drop(writer);
 
@@ -3622,14 +3647,9 @@ mod tests {
         let (tenant, ctx) = TenantHarness::create("no_duplicate_timelines")?
             .load()
             .await;
-        let _ = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let _ = tenant.create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)?;
 
-        match tenant
-            .create_empty_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)
-            .await
-        {
+        match tenant.create_empty_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx) {
             Ok(_) => panic!("duplicate timeline creation should fail"),
             Err(e) => assert_eq!(
                 e.to_string(),
@@ -3658,10 +3678,9 @@ mod tests {
         use std::str::from_utf8;
 
         let (tenant, ctx) = TenantHarness::create("test_branch")?.load().await;
-        let tline = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)
-            .await?;
-        let writer = tline.writer().await;
+        let tline =
+            tenant.create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)?;
+        let writer = tline.writer();
 
         #[allow(non_snake_case)]
         let TEST_KEY_A: Key = Key::from_hex("112222222233333333444444445500000001").unwrap();
@@ -3669,21 +3688,13 @@ mod tests {
         let TEST_KEY_B: Key = Key::from_hex("112222222233333333444444445500000002").unwrap();
 
         // Insert a value on the timeline
-        writer
-            .put(TEST_KEY_A, Lsn(0x20), &test_value("foo at 0x20"))
-            .await?;
-        writer
-            .put(TEST_KEY_B, Lsn(0x20), &test_value("foobar at 0x20"))
-            .await?;
+        writer.put(TEST_KEY_A, Lsn(0x20), &test_value("foo at 0x20"))?;
+        writer.put(TEST_KEY_B, Lsn(0x20), &test_value("foobar at 0x20"))?;
         writer.finish_write(Lsn(0x20));
 
-        writer
-            .put(TEST_KEY_A, Lsn(0x30), &test_value("foo at 0x30"))
-            .await?;
+        writer.put(TEST_KEY_A, Lsn(0x30), &test_value("foo at 0x30"))?;
         writer.finish_write(Lsn(0x30));
-        writer
-            .put(TEST_KEY_A, Lsn(0x40), &test_value("foo at 0x40"))
-            .await?;
+        writer.put(TEST_KEY_A, Lsn(0x40), &test_value("foo at 0x40"))?;
         writer.finish_write(Lsn(0x40));
 
         //assert_current_logical_size(&tline, Lsn(0x40));
@@ -3695,10 +3706,8 @@ mod tests {
         let newtline = tenant
             .get_timeline(NEW_TIMELINE_ID, true)
             .expect("Should have a local timeline");
-        let new_writer = newtline.writer().await;
-        new_writer
-            .put(TEST_KEY_A, Lsn(0x40), &test_value("bar at 0x40"))
-            .await?;
+        let new_writer = newtline.writer();
+        new_writer.put(TEST_KEY_A, Lsn(0x40), &test_value("bar at 0x40"))?;
         new_writer.finish_write(Lsn(0x40));
 
         // Check page contents on both branches
@@ -3724,46 +3733,38 @@ mod tests {
         let mut lsn = start_lsn;
         #[allow(non_snake_case)]
         {
-            let writer = tline.writer().await;
+            let writer = tline.writer();
             // Create a relation on the timeline
-            writer
-                .put(
-                    *TEST_KEY,
-                    lsn,
-                    &Value::Image(TEST_IMG(&format!("foo at {}", lsn))),
-                )
-                .await?;
+            writer.put(
+                *TEST_KEY,
+                lsn,
+                &Value::Image(TEST_IMG(&format!("foo at {}", lsn))),
+            )?;
             writer.finish_write(lsn);
             lsn += 0x10;
-            writer
-                .put(
-                    *TEST_KEY,
-                    lsn,
-                    &Value::Image(TEST_IMG(&format!("foo at {}", lsn))),
-                )
-                .await?;
+            writer.put(
+                *TEST_KEY,
+                lsn,
+                &Value::Image(TEST_IMG(&format!("foo at {}", lsn))),
+            )?;
             writer.finish_write(lsn);
             lsn += 0x10;
         }
         tline.freeze_and_flush().await?;
         {
-            let writer = tline.writer().await;
-            writer
-                .put(
-                    *TEST_KEY,
-                    lsn,
-                    &Value::Image(TEST_IMG(&format!("foo at {}", lsn))),
-                )
-                .await?;
+            let writer = tline.writer();
+            writer.put(
+                *TEST_KEY,
+                lsn,
+                &Value::Image(TEST_IMG(&format!("foo at {}", lsn))),
+            )?;
             writer.finish_write(lsn);
             lsn += 0x10;
-            writer
-                .put(
-                    *TEST_KEY,
-                    lsn,
-                    &Value::Image(TEST_IMG(&format!("foo at {}", lsn))),
-                )
-                .await?;
+            writer.put(
+                *TEST_KEY,
+                lsn,
+                &Value::Image(TEST_IMG(&format!("foo at {}", lsn))),
+            )?;
             writer.finish_write(lsn);
         }
         tline.freeze_and_flush().await
@@ -3775,9 +3776,8 @@ mod tests {
             TenantHarness::create("test_prohibit_branch_creation_on_garbage_collected_data")?
                 .load()
                 .await;
-        let tline = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let tline =
+            tenant.create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)?;
         make_some_layers(tline.as_ref(), Lsn(0x20)).await?;
 
         // this removes layers before lsn 40 (50 minus 10), so there are two remaining layers, image and delta for 31-50
@@ -3814,9 +3814,8 @@ mod tests {
                 .load()
                 .await;
 
-        let tline = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(0x50), DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let tline =
+            tenant.create_test_timeline(TIMELINE_ID, Lsn(0x50), DEFAULT_PG_VERSION, &ctx)?;
         // try to branch at lsn 0x25, should fail because initdb lsn is 0x50
         match tenant
             .branch_timeline_test(&tline, NEW_TIMELINE_ID, Some(Lsn(0x25)), &ctx)
@@ -3865,9 +3864,8 @@ mod tests {
             TenantHarness::create("test_get_branchpoints_from_an_inactive_timeline")?
                 .load()
                 .await;
-        let tline = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let tline =
+            tenant.create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)?;
         make_some_layers(tline.as_ref(), Lsn(0x20)).await?;
 
         tenant
@@ -3915,9 +3913,8 @@ mod tests {
             TenantHarness::create("test_retain_data_in_parent_which_is_needed_for_child")?
                 .load()
                 .await;
-        let tline = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let tline =
+            tenant.create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)?;
         make_some_layers(tline.as_ref(), Lsn(0x20)).await?;
 
         tenant
@@ -3940,9 +3937,8 @@ mod tests {
             TenantHarness::create("test_parent_keeps_data_forever_after_branching")?
                 .load()
                 .await;
-        let tline = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let tline =
+            tenant.create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)?;
         make_some_layers(tline.as_ref(), Lsn(0x20)).await?;
 
         tenant
@@ -3974,9 +3970,8 @@ mod tests {
         let harness = TenantHarness::create(TEST_NAME)?;
         {
             let (tenant, ctx) = harness.load().await;
-            let tline = tenant
-                .create_test_timeline(TIMELINE_ID, Lsn(0x7000), DEFAULT_PG_VERSION, &ctx)
-                .await?;
+            let tline =
+                tenant.create_test_timeline(TIMELINE_ID, Lsn(0x7000), DEFAULT_PG_VERSION, &ctx)?;
             make_some_layers(tline.as_ref(), Lsn(0x8000)).await?;
         }
 
@@ -3995,9 +3990,8 @@ mod tests {
         // create two timelines
         {
             let (tenant, ctx) = harness.load().await;
-            let tline = tenant
-                .create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)
-                .await?;
+            let tline =
+                tenant.create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)?;
 
             make_some_layers(tline.as_ref(), Lsn(0x20)).await?;
 
@@ -4034,9 +4028,8 @@ mod tests {
         let harness = TenantHarness::create(TEST_NAME)?;
         let (tenant, ctx) = harness.load().await;
 
-        let tline = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let tline =
+            tenant.create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)?;
         drop(tline);
         drop(tenant);
 
@@ -4074,44 +4067,35 @@ mod tests {
     #[tokio::test]
     async fn test_images() -> anyhow::Result<()> {
         let (tenant, ctx) = TenantHarness::create("test_images")?.load().await;
-        let tline = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(0x08), DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let tline =
+            tenant.create_test_timeline(TIMELINE_ID, Lsn(0x08), DEFAULT_PG_VERSION, &ctx)?;
 
-        let writer = tline.writer().await;
-        writer
-            .put(*TEST_KEY, Lsn(0x10), &Value::Image(TEST_IMG("foo at 0x10")))
-            .await?;
+        let writer = tline.writer();
+        writer.put(*TEST_KEY, Lsn(0x10), &Value::Image(TEST_IMG("foo at 0x10")))?;
         writer.finish_write(Lsn(0x10));
         drop(writer);
 
         tline.freeze_and_flush().await?;
         tline.compact(&ctx).await?;
 
-        let writer = tline.writer().await;
-        writer
-            .put(*TEST_KEY, Lsn(0x20), &Value::Image(TEST_IMG("foo at 0x20")))
-            .await?;
+        let writer = tline.writer();
+        writer.put(*TEST_KEY, Lsn(0x20), &Value::Image(TEST_IMG("foo at 0x20")))?;
         writer.finish_write(Lsn(0x20));
         drop(writer);
 
         tline.freeze_and_flush().await?;
         tline.compact(&ctx).await?;
 
-        let writer = tline.writer().await;
-        writer
-            .put(*TEST_KEY, Lsn(0x30), &Value::Image(TEST_IMG("foo at 0x30")))
-            .await?;
+        let writer = tline.writer();
+        writer.put(*TEST_KEY, Lsn(0x30), &Value::Image(TEST_IMG("foo at 0x30")))?;
         writer.finish_write(Lsn(0x30));
         drop(writer);
 
         tline.freeze_and_flush().await?;
         tline.compact(&ctx).await?;
 
-        let writer = tline.writer().await;
-        writer
-            .put(*TEST_KEY, Lsn(0x40), &Value::Image(TEST_IMG("foo at 0x40")))
-            .await?;
+        let writer = tline.writer();
+        writer.put(*TEST_KEY, Lsn(0x40), &Value::Image(TEST_IMG("foo at 0x40")))?;
         writer.finish_write(Lsn(0x40));
         drop(writer);
 
@@ -4149,9 +4133,8 @@ mod tests {
     #[tokio::test]
     async fn test_bulk_insert() -> anyhow::Result<()> {
         let (tenant, ctx) = TenantHarness::create("test_bulk_insert")?.load().await;
-        let tline = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(0x08), DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let tline =
+            tenant.create_test_timeline(TIMELINE_ID, Lsn(0x08), DEFAULT_PG_VERSION, &ctx)?;
 
         let mut lsn = Lsn(0x10);
 
@@ -4162,14 +4145,12 @@ mod tests {
         for _ in 0..50 {
             for _ in 0..10000 {
                 test_key.field6 = blknum;
-                let writer = tline.writer().await;
-                writer
-                    .put(
-                        test_key,
-                        lsn,
-                        &Value::Image(TEST_IMG(&format!("{} at {}", blknum, lsn))),
-                    )
-                    .await?;
+                let writer = tline.writer();
+                writer.put(
+                    test_key,
+                    lsn,
+                    &Value::Image(TEST_IMG(&format!("{} at {}", blknum, lsn))),
+                )?;
                 writer.finish_write(lsn);
                 drop(writer);
 
@@ -4195,9 +4176,8 @@ mod tests {
     #[tokio::test]
     async fn test_random_updates() -> anyhow::Result<()> {
         let (tenant, ctx) = TenantHarness::create("test_random_updates")?.load().await;
-        let tline = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let tline =
+            tenant.create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)?;
 
         const NUM_KEYS: usize = 1000;
 
@@ -4214,14 +4194,12 @@ mod tests {
         for blknum in 0..NUM_KEYS {
             lsn = Lsn(lsn.0 + 0x10);
             test_key.field6 = blknum as u32;
-            let writer = tline.writer().await;
-            writer
-                .put(
-                    test_key,
-                    lsn,
-                    &Value::Image(TEST_IMG(&format!("{} at {}", blknum, lsn))),
-                )
-                .await?;
+            let writer = tline.writer();
+            writer.put(
+                test_key,
+                lsn,
+                &Value::Image(TEST_IMG(&format!("{} at {}", blknum, lsn))),
+            )?;
             writer.finish_write(lsn);
             updated[blknum] = lsn;
             drop(writer);
@@ -4234,14 +4212,12 @@ mod tests {
                 lsn = Lsn(lsn.0 + 0x10);
                 let blknum = thread_rng().gen_range(0..NUM_KEYS);
                 test_key.field6 = blknum as u32;
-                let writer = tline.writer().await;
-                writer
-                    .put(
-                        test_key,
-                        lsn,
-                        &Value::Image(TEST_IMG(&format!("{} at {}", blknum, lsn))),
-                    )
-                    .await?;
+                let writer = tline.writer();
+                writer.put(
+                    test_key,
+                    lsn,
+                    &Value::Image(TEST_IMG(&format!("{} at {}", blknum, lsn))),
+                )?;
                 writer.finish_write(lsn);
                 drop(writer);
                 updated[blknum] = lsn;
@@ -4274,9 +4250,8 @@ mod tests {
         let (tenant, ctx) = TenantHarness::create("test_traverse_branches")?
             .load()
             .await;
-        let mut tline = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let mut tline =
+            tenant.create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)?;
 
         const NUM_KEYS: usize = 1000;
 
@@ -4293,14 +4268,12 @@ mod tests {
         for blknum in 0..NUM_KEYS {
             lsn = Lsn(lsn.0 + 0x10);
             test_key.field6 = blknum as u32;
-            let writer = tline.writer().await;
-            writer
-                .put(
-                    test_key,
-                    lsn,
-                    &Value::Image(TEST_IMG(&format!("{} at {}", blknum, lsn))),
-                )
-                .await?;
+            let writer = tline.writer();
+            writer.put(
+                test_key,
+                lsn,
+                &Value::Image(TEST_IMG(&format!("{} at {}", blknum, lsn))),
+            )?;
             writer.finish_write(lsn);
             updated[blknum] = lsn;
             drop(writer);
@@ -4321,14 +4294,12 @@ mod tests {
                 lsn = Lsn(lsn.0 + 0x10);
                 let blknum = thread_rng().gen_range(0..NUM_KEYS);
                 test_key.field6 = blknum as u32;
-                let writer = tline.writer().await;
-                writer
-                    .put(
-                        test_key,
-                        lsn,
-                        &Value::Image(TEST_IMG(&format!("{} at {}", blknum, lsn))),
-                    )
-                    .await?;
+                let writer = tline.writer();
+                writer.put(
+                    test_key,
+                    lsn,
+                    &Value::Image(TEST_IMG(&format!("{} at {}", blknum, lsn))),
+                )?;
                 println!("updating {} at {}", blknum, lsn);
                 writer.finish_write(lsn);
                 drop(writer);
@@ -4362,9 +4333,8 @@ mod tests {
         let (tenant, ctx) = TenantHarness::create("test_traverse_ancestors")?
             .load()
             .await;
-        let mut tline = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let mut tline =
+            tenant.create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)?;
 
         const NUM_KEYS: usize = 100;
         const NUM_TLINES: usize = 50;
@@ -4389,14 +4359,12 @@ mod tests {
                 lsn = Lsn(lsn.0 + 0x10);
                 let blknum = thread_rng().gen_range(0..NUM_KEYS);
                 test_key.field6 = blknum as u32;
-                let writer = tline.writer().await;
-                writer
-                    .put(
-                        test_key,
-                        lsn,
-                        &Value::Image(TEST_IMG(&format!("{} {} at {}", idx, blknum, lsn))),
-                    )
-                    .await?;
+                let writer = tline.writer();
+                writer.put(
+                    test_key,
+                    lsn,
+                    &Value::Image(TEST_IMG(&format!("{} {} at {}", idx, blknum, lsn))),
+                )?;
                 println!("updating [{}][{}] at {}", idx, blknum, lsn);
                 writer.finish_write(lsn);
                 drop(writer);
@@ -4429,9 +4397,8 @@ mod tests {
             .await;
 
         let initdb_lsn = Lsn(0x20);
-        let utline = tenant
-            .create_empty_timeline(TIMELINE_ID, initdb_lsn, DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let utline =
+            tenant.create_empty_timeline(TIMELINE_ID, initdb_lsn, DEFAULT_PG_VERSION, &ctx)?;
         let tline = utline.raw_timeline().unwrap();
 
         // Spawn flush loop now so that we can set the `expect_initdb_optimization`
@@ -4468,7 +4435,6 @@ mod tests {
             .context("init_empty_test_timeline")?;
         modification
             .commit()
-            .await
             .context("commit init_empty_test_timeline modification")?;
 
         // Do the flush. The flush code will check the expectations that we set above.

pageserver/src/tenant/layer_cache.rs

@@ -1,143 +0,0 @@
-use super::storage_layer::{PersistentLayer, PersistentLayerDesc, PersistentLayerKey, RemoteLayer};
-use super::Timeline;
-use crate::tenant::layer_map::{self, LayerMap};
-use anyhow::Result;
-use std::sync::{Mutex, Weak};
-use std::{collections::HashMap, sync::Arc};
-
-pub struct LayerCache {
-    /// Layer removal lock.
-    /// A lock to ensure that no layer of the timeline is removed concurrently by other tasks.
-    /// This lock is acquired in [`Timeline::gc`], [`Timeline::compact`],
-    /// and [`Tenant::delete_timeline`]. This is an `Arc<Mutex>` lock because we need an owned
-    /// lock guard in functions that will be spawned to tokio I/O pool (which requires `'static`).
-    pub layers_removal_lock: Arc<tokio::sync::Mutex<()>>,
-
-    /// We need this lock b/c we do not have any way to prevent GC/compaction from removing files in-use.
-    /// We need to do reference counting on Arc to prevent this from happening, and we can safely remove this lock.
-    pub layers_operation_lock: Arc<tokio::sync::RwLock<()>>,
-
-    /// Will be useful when we move evict / download to layer cache.
-    #[allow(unused)]
-    timeline: Weak<Timeline>,
-
-    mapping: Mutex<HashMap<PersistentLayerKey, Arc<dyn PersistentLayer>>>,
-}
-
-pub struct LayerInUseWrite(tokio::sync::OwnedRwLockWriteGuard<()>);
-
-pub struct LayerInUseRead(tokio::sync::OwnedRwLockReadGuard<()>);
-
-#[derive(Clone)]
-pub struct DeleteGuard(Arc<tokio::sync::OwnedMutexGuard<()>>);
-
-impl LayerCache {
-    pub fn new(timeline: Weak<Timeline>) -> Self {
-        Self {
-            layers_operation_lock: Arc::new(tokio::sync::RwLock::new(())),
-            layers_removal_lock: Arc::new(tokio::sync::Mutex::new(())),
-            mapping: Mutex::new(HashMap::new()),
-            timeline,
-        }
-    }
-
-    pub fn get_from_desc(&self, desc: &PersistentLayerDesc) -> Arc<dyn PersistentLayer> {
-        let guard = self.mapping.lock().unwrap();
-        guard.get(&desc.key()).expect("not found").clone()
-    }
-
-    /// This function is to mock the original behavior of `layers` lock in `Timeline`. Can be removed after we ensure
-    /// we won't delete files that are being read.
-    pub async fn layer_in_use_write(&self) -> LayerInUseWrite {
-        LayerInUseWrite(self.layers_operation_lock.clone().write_owned().await)
-    }
-
-    /// This function is to mock the original behavior of `layers` lock in `Timeline`. Can be removed after we ensure
-    /// we won't delete files that are being read.
-    pub async fn layer_in_use_read(&self) -> LayerInUseRead {
-        LayerInUseRead(self.layers_operation_lock.clone().read_owned().await)
-    }
-
-    /// Ensures only one of compaction / gc can happen at a time.
-    pub async fn delete_guard(&self) -> DeleteGuard {
-        DeleteGuard(Arc::new(
-            self.layers_removal_lock.clone().lock_owned().await,
-        ))
-    }
-
-    /// Should only be called when initializing the timeline. Bypass checks and layer operation lock.
-    pub fn remove_local_when_init(&self, layer: Arc<dyn PersistentLayer>) {
-        let mut guard = self.mapping.lock().unwrap();
-        guard.remove(&layer.layer_desc().key());
-    }
-
-    /// Should only be called when initializing the timeline. Bypass checks and layer operation lock.
-    pub fn populate_remote_when_init(&self, layer: Arc<RemoteLayer>) {
-        let mut guard = self.mapping.lock().unwrap();
-        guard.insert(layer.layer_desc().key(), layer);
-    }
-
-    /// Should only be called when initializing the timeline. Bypass checks and layer operation lock.
-    pub fn populate_local_when_init(&self, layer: Arc<dyn PersistentLayer>) {
-        let mut guard = self.mapping.lock().unwrap();
-        guard.insert(layer.layer_desc().key(), layer);
-    }
-
-    /// Called within read path.
-    pub fn replace_and_verify(
-        &self,
-        expected: Arc<dyn PersistentLayer>,
-        new: Arc<dyn PersistentLayer>,
-    ) -> Result<()> {
-        let mut guard = self.mapping.lock().unwrap();
-
-        use super::layer_map::LayerKey;
-        let key = LayerKey::from(&*expected);
-        let other = LayerKey::from(&*new);
-
-        let expected_l0 = LayerMap::is_l0(expected.layer_desc());
-        let new_l0 = LayerMap::is_l0(new.layer_desc());
-
-        fail::fail_point!("layermap-replace-notfound", |_| anyhow::bail!(
-            "replacing downloaded layer into layermap failed because layer was not found"
-        ));
-
-        anyhow::ensure!(
-            key == other,
-            "replacing downloaded layer into layermap failed because two layers have different keys: {key:?} != {other:?}"
-        );
-
-        anyhow::ensure!(
-             expected_l0 == new_l0,
-             "replacing downloaded layer into layermap failed because one layer is l0 while the other is not: {expected_l0} != {new_l0}"
-         );
-
-        if let Some(layer) = guard.get_mut(&expected.layer_desc().key()) {
-            anyhow::ensure!(
-                layer_map::compare_arced_layers(&expected, layer),
-                "replacing downloaded layer into layermap failed because another layer was found instead of expected, expected={expected:?}, new={new:?}",
-                expected = Arc::as_ptr(&expected),
-                new = Arc::as_ptr(layer),
-            );
-            *layer = new;
-            Ok(())
-        } else {
-            anyhow::bail!(
-                "replacing downloaded layer into layermap failed because layer was not found"
-            );
-        }
-    }
-
-    /// Called within write path. When compaction and image layer creation we will create new layers.
-    pub fn create_new_layer(&self, layer: Arc<dyn PersistentLayer>) {
-        let mut guard = self.mapping.lock().unwrap();
-        guard.insert(layer.layer_desc().key(), layer);
-    }
-
-    /// Called within write path. When GC and compaction we will remove layers and delete them on disk.
-    /// Will move logic to delete files here later.
-    pub fn delete_layer(&self, layer: Arc<dyn PersistentLayer>) {
-        let mut guard = self.mapping.lock().unwrap();
-        guard.remove(&layer.layer_desc().key());
-    }
-}

pageserver/src/tenant/layer_map.rs

@@ -51,23 +51,25 @@ use crate::keyspace::KeyPartitioning;
 use crate::repository::Key;
 use crate::tenant::storage_layer::InMemoryLayer;
 use crate::tenant::storage_layer::Layer;
+use anyhow::Context;
 use anyhow::Result;
+use std::collections::HashMap;
 use std::collections::VecDeque;
 use std::ops::Range;
 use std::sync::Arc;
 use utils::lsn::Lsn;
 
 use historic_layer_coverage::BufferedHistoricLayerCoverage;
-pub use historic_layer_coverage::{LayerKey, Replacement};
+pub use historic_layer_coverage::Replacement;
 
 use super::storage_layer::range_eq;
 use super::storage_layer::PersistentLayerDesc;
+use super::storage_layer::PersistentLayerKey;
 
 ///
 /// LayerMap tracks what layers exist on a timeline.
 ///
-#[derive(Default, Clone)]
-pub struct LayerMap {
+pub struct LayerMap<L: ?Sized> {
     //
     // 'open_layer' holds the current InMemoryLayer that is accepting new
     // records. If it is None, 'next_open_layer_at' will be set instead, indicating
@@ -93,6 +95,24 @@ pub struct LayerMap {
     /// L0 layers have key range Key::MIN..Key::MAX, and locating them using R-Tree search is very inefficient.
     /// So L0 layers are held in l0_delta_layers vector, in addition to the R-tree.
     l0_delta_layers: Vec<Arc<PersistentLayerDesc>>,
+
+    /// Mapping from persistent layer key to the actual layer object. Currently, it stores delta, image, and
+    /// remote layers. In future refactors, this will be eventually moved out of LayerMap into Timeline, and
+    /// RemoteLayer will be removed.
+    mapping: HashMap<PersistentLayerKey, Arc<L>>,
+}
+
+impl<L: ?Sized> Default for LayerMap<L> {
+    fn default() -> Self {
+        Self {
+            open_layer: None,
+            next_open_layer_at: None,
+            frozen_layers: VecDeque::default(),
+            l0_delta_layers: Vec::default(),
+            historic: BufferedHistoricLayerCoverage::default(),
+            mapping: HashMap::default(),
+        }
+    }
 }
 
 /// The primary update API for the layer map.
@@ -100,21 +120,24 @@ pub struct LayerMap {
 /// Batching historic layer insertions and removals is good for
 /// performance and this struct helps us do that correctly.
 #[must_use]
-pub struct BatchedUpdates<'a> {
+pub struct BatchedUpdates<'a, L: ?Sized + Layer> {
     // While we hold this exclusive reference to the layer map the type checker
     // will prevent us from accidentally reading any unflushed updates.
-    layer_map: &'a mut LayerMap,
+    layer_map: &'a mut LayerMap<L>,
 }
 
 /// Provide ability to batch more updates while hiding the read
 /// API so we don't accidentally read without flushing.
-impl BatchedUpdates<'_> {
+impl<L> BatchedUpdates<'_, L>
+where
+    L: ?Sized + Layer,
+{
     ///
     /// Insert an on-disk layer.
     ///
     // TODO remove the `layer` argument when `mapping` is refactored out of `LayerMap`
-    pub fn insert_historic(&mut self, layer_desc: PersistentLayerDesc) {
-        self.layer_map.insert_historic_noflush(layer_desc)
+    pub fn insert_historic(&mut self, layer_desc: PersistentLayerDesc, layer: Arc<L>) {
+        self.layer_map.insert_historic_noflush(layer_desc, layer)
     }
 
     ///
@@ -122,8 +145,31 @@ impl BatchedUpdates<'_> {
     ///
     /// This should be called when the corresponding file on disk has been deleted.
     ///
-    pub fn remove_historic(&mut self, layer_desc: PersistentLayerDesc) {
-        self.layer_map.remove_historic_noflush(layer_desc)
+    pub fn remove_historic(&mut self, layer_desc: PersistentLayerDesc, layer: Arc<L>) {
+        self.layer_map.remove_historic_noflush(layer_desc, layer)
+    }
+
+    /// Replaces existing layer iff it is the `expected`.
+    ///
+    /// If the expected layer has been removed it will not be inserted by this function.
+    ///
+    /// Returned `Replacement` describes succeeding in replacement or the reason why it could not
+    /// be done.
+    ///
+    /// TODO replacement can be done without buffering and rebuilding layer map updates.
+    ///      One way to do that is to add a layer of indirection for returned values, so
+    ///      that we can replace values only by updating a hashmap.
+    pub fn replace_historic(
+        &mut self,
+        expected_desc: PersistentLayerDesc,
+        expected: &Arc<L>,
+        new_desc: PersistentLayerDesc,
+        new: Arc<L>,
+    ) -> anyhow::Result<Replacement<Arc<L>>> {
+        fail::fail_point!("layermap-replace-notfound", |_| Ok(Replacement::NotFound));
+
+        self.layer_map
+            .replace_historic_noflush(expected_desc, expected, new_desc, new)
     }
 
     // We will flush on drop anyway, but this method makes it
@@ -139,19 +185,25 @@ impl BatchedUpdates<'_> {
 // than panic later or read without flushing.
 //
 // TODO maybe warn if flush hasn't explicitly been called
-impl Drop for BatchedUpdates<'_> {
+impl<L> Drop for BatchedUpdates<'_, L>
+where
+    L: ?Sized + Layer,
+{
     fn drop(&mut self) {
         self.layer_map.flush_updates();
     }
 }
 
 /// Return value of LayerMap::search
-pub struct SearchResult {
-    pub layer: Arc<PersistentLayerDesc>,
+pub struct SearchResult<L: ?Sized> {
+    pub layer: Arc<L>,
     pub lsn_floor: Lsn,
 }
 
-impl LayerMap {
+impl<L> LayerMap<L>
+where
+    L: ?Sized + Layer,
+{
     ///
     /// Find the latest layer (by lsn.end) that covers the given
     /// 'key', with lsn.start < 'end_lsn'.
@@ -183,7 +235,7 @@ impl LayerMap {
     /// NOTE: This only searches the 'historic' layers, *not* the
     /// 'open' and 'frozen' layers!
     ///
-    pub fn search(&self, key: Key, end_lsn: Lsn) -> Option<SearchResult> {
+    pub fn search(&self, key: Key, end_lsn: Lsn) -> Option<SearchResult<L>> {
         let version = self.historic.get().unwrap().get_version(end_lsn.0 - 1)?;
         let latest_delta = version.delta_coverage.query(key.to_i128());
         let latest_image = version.image_coverage.query(key.to_i128());
@@ -192,6 +244,7 @@ impl LayerMap {
             (None, None) => None,
             (None, Some(image)) => {
                 let lsn_floor = image.get_lsn_range().start;
+                let image = self.get_layer_from_mapping(&image.key()).clone();
                 Some(SearchResult {
                     layer: image,
                     lsn_floor,
@@ -199,6 +252,7 @@ impl LayerMap {
             }
             (Some(delta), None) => {
                 let lsn_floor = delta.get_lsn_range().start;
+                let delta = self.get_layer_from_mapping(&delta.key()).clone();
                 Some(SearchResult {
                     layer: delta,
                     lsn_floor,
@@ -209,6 +263,7 @@ impl LayerMap {
                 let image_is_newer = image.get_lsn_range().end >= delta.get_lsn_range().end;
                 let image_exact_match = img_lsn + 1 == end_lsn;
                 if image_is_newer || image_exact_match {
+                    let image = self.get_layer_from_mapping(&image.key()).clone();
                     Some(SearchResult {
                         layer: image,
                         lsn_floor: img_lsn,
@@ -216,6 +271,7 @@ impl LayerMap {
                 } else {
                     let lsn_floor =
                         std::cmp::max(delta.get_lsn_range().start, image.get_lsn_range().start + 1);
+                    let delta = self.get_layer_from_mapping(&delta.key()).clone();
                     Some(SearchResult {
                         layer: delta,
                         lsn_floor,
@@ -226,7 +282,7 @@ impl LayerMap {
     }
 
     /// Start a batch of updates, applied on drop
-    pub fn batch_update(&mut self) -> BatchedUpdates<'_> {
+    pub fn batch_update(&mut self) -> BatchedUpdates<'_, L> {
         BatchedUpdates { layer_map: self }
     }
 
@@ -236,32 +292,48 @@ impl LayerMap {
     /// Helper function for BatchedUpdates::insert_historic
     ///
     /// TODO(chi): remove L generic so that we do not need to pass layer object.
-    pub(self) fn insert_historic_noflush(&mut self, layer_desc: PersistentLayerDesc) {
+    pub(self) fn insert_historic_noflush(
+        &mut self,
+        layer_desc: PersistentLayerDesc,
+        layer: Arc<L>,
+    ) {
+        self.mapping.insert(layer_desc.key(), layer.clone());
+
         // TODO: See #3869, resulting #4088, attempted fix and repro #4094
 
-        if Self::is_l0(&layer_desc) {
+        if Self::is_l0(&layer) {
             self.l0_delta_layers.push(layer_desc.clone().into());
         }
 
         self.historic.insert(
-            historic_layer_coverage::LayerKey::from(&layer_desc),
+            historic_layer_coverage::LayerKey::from(&*layer),
             layer_desc.into(),
         );
     }
 
+    fn get_layer_from_mapping(&self, key: &PersistentLayerKey) -> &Arc<L> {
+        let layer = self
+            .mapping
+            .get(key)
+            .with_context(|| format!("{key:?}"))
+            .expect("inconsistent layer mapping");
+        layer
+    }
+
     ///
     /// Remove an on-disk layer from the map.
     ///
     /// Helper function for BatchedUpdates::remove_historic
     ///
-    pub fn remove_historic_noflush(&mut self, layer_desc: PersistentLayerDesc) {
+    pub fn remove_historic_noflush(&mut self, layer_desc: PersistentLayerDesc, layer: Arc<L>) {
         self.historic
-            .remove(historic_layer_coverage::LayerKey::from(&layer_desc));
-        let layer_key = layer_desc.key();
-        if Self::is_l0(&layer_desc) {
+            .remove(historic_layer_coverage::LayerKey::from(&*layer));
+        if Self::is_l0(&layer) {
             let len_before = self.l0_delta_layers.len();
             let mut l0_delta_layers = std::mem::take(&mut self.l0_delta_layers);
-            l0_delta_layers.retain(|other| other.key() != layer_key);
+            l0_delta_layers.retain(|other| {
+                !Self::compare_arced_layers(self.get_layer_from_mapping(&other.key()), &layer)
+            });
             self.l0_delta_layers = l0_delta_layers;
             // this assertion is related to use of Arc::ptr_eq in Self::compare_arced_layers,
             // there's a chance that the comparison fails at runtime due to it comparing (pointer,
@@ -272,6 +344,69 @@ impl LayerMap {
                 "failed to locate removed historic layer from l0_delta_layers"
             );
         }
+        self.mapping.remove(&layer_desc.key());
+    }
+
+    pub(self) fn replace_historic_noflush(
+        &mut self,
+        expected_desc: PersistentLayerDesc,
+        expected: &Arc<L>,
+        new_desc: PersistentLayerDesc,
+        new: Arc<L>,
+    ) -> anyhow::Result<Replacement<Arc<L>>> {
+        let key = historic_layer_coverage::LayerKey::from(&**expected);
+        let other = historic_layer_coverage::LayerKey::from(&*new);
+
+        let expected_l0 = Self::is_l0(expected);
+        let new_l0 = Self::is_l0(&new);
+
+        anyhow::ensure!(
+            key == other,
+            "expected and new must have equal LayerKeys: {key:?} != {other:?}"
+        );
+
+        anyhow::ensure!(
+            expected_l0 == new_l0,
+            "expected and new must both be l0 deltas or neither should be: {expected_l0} != {new_l0}"
+        );
+
+        let l0_index = if expected_l0 {
+            // find the index in case replace worked, we need to replace that as well
+            let pos = self.l0_delta_layers.iter().position(|slot| {
+                Self::compare_arced_layers(self.get_layer_from_mapping(&slot.key()), expected)
+            });
+
+            if pos.is_none() {
+                return Ok(Replacement::NotFound);
+            }
+            pos
+        } else {
+            None
+        };
+
+        let new_desc = Arc::new(new_desc);
+        let replaced = self.historic.replace(&key, new_desc.clone(), |existing| {
+            **existing == expected_desc
+        });
+
+        if let Replacement::Replaced { .. } = &replaced {
+            self.mapping.remove(&expected_desc.key());
+            self.mapping.insert(new_desc.key(), new);
+            if let Some(index) = l0_index {
+                self.l0_delta_layers[index] = new_desc;
+            }
+        }
+
+        let replaced = match replaced {
+            Replacement::Replaced { in_buffered } => Replacement::Replaced { in_buffered },
+            Replacement::NotFound => Replacement::NotFound,
+            Replacement::RemovalBuffered => Replacement::RemovalBuffered,
+            Replacement::Unexpected(x) => {
+                Replacement::Unexpected(self.get_layer_from_mapping(&x.key()).clone())
+            }
+        };
+
+        Ok(replaced)
     }
 
     /// Helper function for BatchedUpdates::drop.
@@ -319,8 +454,10 @@ impl LayerMap {
         Ok(true)
     }
 
-    pub fn iter_historic_layers(&self) -> impl '_ + Iterator<Item = Arc<PersistentLayerDesc>> {
-        self.historic.iter()
+    pub fn iter_historic_layers(&self) -> impl '_ + Iterator<Item = Arc<L>> {
+        self.historic
+            .iter()
+            .map(|x| self.get_layer_from_mapping(&x.key()).clone())
     }
 
     ///
@@ -335,7 +472,7 @@ impl LayerMap {
         &self,
         key_range: &Range<Key>,
         lsn: Lsn,
-    ) -> Result<Vec<(Range<Key>, Option<Arc<PersistentLayerDesc>>)>> {
+    ) -> Result<Vec<(Range<Key>, Option<Arc<L>>)>> {
         let version = match self.historic.get().unwrap().get_version(lsn.0) {
             Some(v) => v,
             None => return Ok(vec![]),
@@ -345,26 +482,36 @@ impl LayerMap {
         let end = key_range.end.to_i128();
 
         // Initialize loop variables
-        let mut coverage: Vec<(Range<Key>, Option<Arc<PersistentLayerDesc>>)> = vec![];
+        let mut coverage: Vec<(Range<Key>, Option<Arc<L>>)> = vec![];
         let mut current_key = start;
         let mut current_val = version.image_coverage.query(start);
 
         // Loop through the change events and push intervals
         for (change_key, change_val) in version.image_coverage.range(start..end) {
             let kr = Key::from_i128(current_key)..Key::from_i128(change_key);
-            coverage.push((kr, current_val.take()));
+            coverage.push((
+                kr,
+                current_val
+                    .take()
+                    .map(|l| self.get_layer_from_mapping(&l.key()).clone()),
+            ));
             current_key = change_key;
             current_val = change_val.clone();
         }
 
         // Add the final interval
         let kr = Key::from_i128(current_key)..Key::from_i128(end);
-        coverage.push((kr, current_val.take()));
+        coverage.push((
+            kr,
+            current_val
+                .take()
+                .map(|l| self.get_layer_from_mapping(&l.key()).clone()),
+        ));
 
         Ok(coverage)
     }
 
-    pub fn is_l0(layer: &PersistentLayerDesc) -> bool {
+    pub fn is_l0(layer: &L) -> bool {
         range_eq(&layer.get_key_range(), &(Key::MIN..Key::MAX))
     }
 
@@ -390,7 +537,7 @@ impl LayerMap {
     /// TODO The optimal number should probably be slightly higher than 1, but to
     ///      implement that we need to plumb a lot more context into this function
     ///      than just the current partition_range.
-    pub fn is_reimage_worthy(layer: &PersistentLayerDesc, partition_range: &Range<Key>) -> bool {
+    pub fn is_reimage_worthy(layer: &L, partition_range: &Range<Key>) -> bool {
         // Case 1
         if !Self::is_l0(layer) {
             return true;
@@ -448,7 +595,9 @@ impl LayerMap {
                     let kr = Key::from_i128(current_key)..Key::from_i128(change_key);
                     let lr = lsn.start..val.get_lsn_range().start;
                     if !kr.is_empty() {
-                        let base_count = Self::is_reimage_worthy(&val, key) as usize;
+                        let base_count =
+                            Self::is_reimage_worthy(self.get_layer_from_mapping(&val.key()), key)
+                                as usize;
                         let new_limit = limit.map(|l| l - base_count);
                         let max_stacked_deltas_underneath =
                             self.count_deltas(&kr, &lr, new_limit)?;
@@ -471,7 +620,9 @@ impl LayerMap {
                 let lr = lsn.start..val.get_lsn_range().start;
 
                 if !kr.is_empty() {
-                    let base_count = Self::is_reimage_worthy(&val, key) as usize;
+                    let base_count =
+                        Self::is_reimage_worthy(self.get_layer_from_mapping(&val.key()), key)
+                            as usize;
                     let new_limit = limit.map(|l| l - base_count);
                     let max_stacked_deltas_underneath = self.count_deltas(&kr, &lr, new_limit)?;
                     max_stacked_deltas = std::cmp::max(
@@ -621,8 +772,12 @@ impl LayerMap {
     }
 
     /// Return all L0 delta layers
-    pub fn get_level0_deltas(&self) -> Result<Vec<Arc<PersistentLayerDesc>>> {
-        Ok(self.l0_delta_layers.to_vec())
+    pub fn get_level0_deltas(&self) -> Result<Vec<Arc<L>>> {
+        Ok(self
+            .l0_delta_layers
+            .iter()
+            .map(|x| self.get_layer_from_mapping(&x.key()).clone())
+            .collect())
     }
 
     /// debugging function to print out the contents of the layer map
@@ -647,76 +802,97 @@ impl LayerMap {
         println!("End dump LayerMap");
         Ok(())
     }
-}
 
-/// Similar to `Arc::ptr_eq`, but only compares the object pointers, not vtables.
-///
-/// Returns `true` if the two `Arc` point to the same layer, false otherwise.
-///
-/// If comparing persistent layers, ALWAYS compare the layer descriptor key.
-#[inline(always)]
-pub fn compare_arced_layers<L: ?Sized>(left: &Arc<L>, right: &Arc<L>) -> bool {
-    // "dyn Trait" objects are "fat pointers" in that they have two components:
-    // - pointer to the object
-    // - pointer to the vtable
-    //
-    // rust does not provide a guarantee that these vtables are unique, but however
-    // `Arc::ptr_eq` as of writing (at least up to 1.67) uses a comparison where both the
-    // pointer and the vtable need to be equal.
-    //
-    // See: https://github.com/rust-lang/rust/issues/103763
-    //
-    // A future version of rust will most likely use this form below, where we cast each
-    // pointer into a pointer to unit, which drops the inaccessible vtable pointer, making it
-    // not affect the comparison.
-    //
-    // See: https://github.com/rust-lang/rust/pull/106450
-    let left = Arc::as_ptr(left) as *const ();
-    let right = Arc::as_ptr(right) as *const ();
+    /// Similar to `Arc::ptr_eq`, but only compares the object pointers, not vtables.
+    ///
+    /// Returns `true` if the two `Arc` point to the same layer, false otherwise.
+    #[inline(always)]
+    pub fn compare_arced_layers(left: &Arc<L>, right: &Arc<L>) -> bool {
+        // "dyn Trait" objects are "fat pointers" in that they have two components:
+        // - pointer to the object
+        // - pointer to the vtable
+        //
+        // rust does not provide a guarantee that these vtables are unique, but however
+        // `Arc::ptr_eq` as of writing (at least up to 1.67) uses a comparison where both the
+        // pointer and the vtable need to be equal.
+        //
+        // See: https://github.com/rust-lang/rust/issues/103763
+        //
+        // A future version of rust will most likely use this form below, where we cast each
+        // pointer into a pointer to unit, which drops the inaccessible vtable pointer, making it
+        // not affect the comparison.
+        //
+        // See: https://github.com/rust-lang/rust/pull/106450
+        let left = Arc::as_ptr(left) as *const ();
+        let right = Arc::as_ptr(right) as *const ();
 
-    left == right
+        left == right
+    }
 }
 
 #[cfg(test)]
 mod tests {
-    use super::LayerMap;
-    use crate::tenant::storage_layer::{tests::LayerDescriptor, LayerFileName};
+    use super::{LayerMap, Replacement};
+    use crate::tenant::storage_layer::{Layer, LayerDescriptor, LayerFileName};
     use std::str::FromStr;
     use std::sync::Arc;
 
     mod l0_delta_layers_updated {
 
-        use crate::tenant::storage_layer::{PersistentLayer, PersistentLayerDesc};
-
         use super::*;
 
         #[test]
         fn for_full_range_delta() {
             // l0_delta_layers are used by compaction, and should observe all buffered updates
             l0_delta_layers_updated_scenario(
-                 "000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__0000000053423C21-0000000053424D69",
-                 true
-             )
+                "000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__0000000053423C21-0000000053424D69",
+                true
+            )
         }
 
         #[test]
         fn for_non_full_range_delta() {
             // has minimal uncovered areas compared to l0_delta_layers_updated_on_insert_replace_remove_for_full_range_delta
             l0_delta_layers_updated_scenario(
-                 "000000000000000000000000000000000001-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFE__0000000053423C21-0000000053424D69",
-                 // because not full range
-                 false
-             )
+                "000000000000000000000000000000000001-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFE__0000000053423C21-0000000053424D69",
+                // because not full range
+                false
+            )
         }
 
         #[test]
         fn for_image() {
             l0_delta_layers_updated_scenario(
-                 "000000000000000000000000000000000000-000000000000000000000000000000010000__0000000053424D69",
-                 // code only checks if it is a full range layer, doesn't care about images, which must
-                 // mean we should in practice never have full range images
-                 false
-             )
+                "000000000000000000000000000000000000-000000000000000000000000000000010000__0000000053424D69",
+                // code only checks if it is a full range layer, doesn't care about images, which must
+                // mean we should in practice never have full range images
+                false
+            )
+        }
+
+        #[test]
+        fn replacing_missing_l0_is_notfound() {
+            // original impl had an oversight, and L0 was an anyhow::Error. anyhow::Error should
+            // however only happen for precondition failures.
+
+            let layer = "000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__0000000053423C21-0000000053424D69";
+            let layer = LayerFileName::from_str(layer).unwrap();
+            let layer = LayerDescriptor::from(layer);
+
+            // same skeletan construction; see scenario below
+            let not_found = Arc::new(layer.clone());
+            let new_version = Arc::new(layer);
+
+            let mut map = LayerMap::default();
+
+            let res = map.batch_update().replace_historic(
+                not_found.get_persistent_layer_desc(),
+                &not_found,
+                new_version.get_persistent_layer_desc(),
+                new_version,
+            );
+
+            assert!(matches!(res, Ok(Replacement::NotFound)), "{res:?}");
         }
 
         fn l0_delta_layers_updated_scenario(layer_name: &str, expected_l0: bool) {
@@ -730,31 +906,46 @@ mod tests {
 
             // two disjoint Arcs in different lifecycle phases. even if it seems they must be the
             // same layer, we use LayerMap::compare_arced_layers as the identity of layers.
-            assert_eq!(remote.layer_desc(), downloaded.layer_desc());
+            assert!(!LayerMap::compare_arced_layers(&remote, &downloaded));
 
             let expected_in_counts = (1, usize::from(expected_l0));
 
             map.batch_update()
-                .insert_historic(remote.layer_desc().clone());
-            assert_eq!(
-                count_layer_in(&map, remote.layer_desc()),
-                expected_in_counts
+                .insert_historic(remote.get_persistent_layer_desc(), remote.clone());
+            assert_eq!(count_layer_in(&map, &remote), expected_in_counts);
+
+            let replaced = map
+                .batch_update()
+                .replace_historic(
+                    remote.get_persistent_layer_desc(),
+                    &remote,
+                    downloaded.get_persistent_layer_desc(),
+                    downloaded.clone(),
+                )
+                .expect("name derived attributes are the same");
+            assert!(
+                matches!(replaced, Replacement::Replaced { .. }),
+                "{replaced:?}"
             );
+            assert_eq!(count_layer_in(&map, &downloaded), expected_in_counts);
 
             map.batch_update()
-                .remove_historic(downloaded.layer_desc().clone());
-            assert_eq!(count_layer_in(&map, downloaded.layer_desc()), (0, 0));
+                .remove_historic(downloaded.get_persistent_layer_desc(), downloaded.clone());
+            assert_eq!(count_layer_in(&map, &downloaded), (0, 0));
         }
 
-        fn count_layer_in(map: &LayerMap, layer: &PersistentLayerDesc) -> (usize, usize) {
+        fn count_layer_in<L: Layer + ?Sized>(map: &LayerMap<L>, layer: &Arc<L>) -> (usize, usize) {
             let historic = map
                 .iter_historic_layers()
-                .filter(|x| x.key() == layer.key())
+                .filter(|x| LayerMap::compare_arced_layers(x, layer))
                 .count();
             let l0s = map
                 .get_level0_deltas()
                 .expect("why does this return a result");
-            let l0 = l0s.iter().filter(|x| x.key() == layer.key()).count();
+            let l0 = l0s
+                .iter()
+                .filter(|x| LayerMap::compare_arced_layers(x, layer))
+                .count();
 
             (historic, l0)
         }

pageserver/src/tenant/layer_map/historic_layer_coverage.rs

@@ -3,8 +3,6 @@ use std::ops::Range;
 
 use tracing::info;
 
-use crate::tenant::storage_layer::PersistentLayerDesc;
-
 use super::layer_coverage::LayerCoverageTuple;
 
 /// Layers in this module are identified and indexed by this data.
@@ -55,24 +53,11 @@ impl<'a, L: crate::tenant::storage_layer::Layer + ?Sized> From<&'a L> for LayerK
     }
 }
 
-impl From<&PersistentLayerDesc> for LayerKey {
-    fn from(layer: &PersistentLayerDesc) -> Self {
-        let kr = layer.get_key_range();
-        let lr = layer.get_lsn_range();
-        LayerKey {
-            key: kr.start.to_i128()..kr.end.to_i128(),
-            lsn: lr.start.0..lr.end.0,
-            is_image: !layer.is_incremental(),
-        }
-    }
-}
-
 /// Efficiently queryable layer coverage for each LSN.
 ///
 /// Allows answering layer map queries very efficiently,
 /// but doesn't allow retroactive insertion, which is
 /// sometimes necessary. See BufferedHistoricLayerCoverage.
-#[derive(Clone)]
 pub struct HistoricLayerCoverage<Value> {
     /// The latest state
     head: LayerCoverageTuple<Value>,
@@ -426,7 +411,6 @@ fn test_persistent_overlapping() {
 ///
 /// See this for more on persistent and retroactive techniques:
 /// https://www.youtube.com/watch?v=WqCWghETNDc&t=581s
-#[derive(Clone)]
 pub struct BufferedHistoricLayerCoverage<Value> {
     /// A persistent layer map that we rebuild when we need to retroactively update
     historic_coverage: HistoricLayerCoverage<Value>,
@@ -483,11 +467,6 @@ impl<Value: Clone> BufferedHistoricLayerCoverage<Value> {
     ///
     /// Returns a `Replacement` value describing the outcome; only the case of
     /// `Replacement::Replaced` modifies the map and requires a rebuild.
-    ///
-    /// This function is unlikely to be used in the future because LayerMap now only records the
-    /// layer descriptors. Therefore, anything added to the layer map will only be removed or
-    /// added, and never replaced.
-    #[allow(dead_code)]
     pub fn replace<F>(
         &mut self,
         layer_key: &LayerKey,

pageserver/src/tenant/layer_map/layer_coverage.rs

@@ -15,7 +15,6 @@ use rpds::RedBlackTreeMapSync;
 ///
 /// NOTE The struct is parameterized over Value for easier
 ///      testing, but in practice it's some sort of layer.
-#[derive(Clone)]
 pub struct LayerCoverage<Value> {
     /// For every change in coverage (as we sweep the key space)
     /// we store (lsn.end, value).
@@ -140,7 +139,6 @@ impl<Value: Clone> LayerCoverage<Value> {
 }
 
 /// Image and delta coverage at a specific LSN.
-#[derive(Clone)]
 pub struct LayerCoverageTuple<Value> {
     pub image_coverage: LayerCoverage<Value>,
     pub delta_coverage: LayerCoverage<Value>,

pageserver/src/tenant/layer_map_mgr.rs

@@ -1,146 +0,0 @@
-//! This module implements `LayerMapMgr`, which manages a layer map object and provides lock-free access to the state.
-//!
-//! A common usage pattern is as follows:
-//!
-//! ```ignore
-//! async fn compaction(&self) {
-//!     // Get the current state.
-//!     let state = self.layer_map_mgr.read();
-//!     // No lock held at this point. Do compaction based on the state. This part usually incurs I/O operations and may
-//!     // take a long time.
-//!     let compaction_result = self.do_compaction(&state).await?;
-//!     // Update the state.
-//!     self.layer_map_mgr.update(|mut state| async move {
-//!         // do updates to the state, return it.
-//!         Ok(state)
-//!     }).await?;
-//! }
-//! ```
-use anyhow::Result;
-use arc_swap::ArcSwap;
-use futures::Future;
-use std::sync::Arc;
-
-use super::layer_map::LayerMap;
-
-/// Manages the storage state. Provide utility functions to modify the layer map and get an immutable reference to the
-/// layer map.
-pub struct LayerMapMgr {
-    layer_map: ArcSwap<LayerMap>,
-    state_lock: tokio::sync::Mutex<()>,
-}
-
-impl LayerMapMgr {
-    /// Get the current state of the layer map.
-    pub fn read(&self) -> Arc<LayerMap> {
-        // TODO: it is possible to use `load` to reduce the overhead of cloning the Arc, but read path usually involves
-        // disk reads and layer mapping fetching, and therefore it's not a big deal to use a more optimized version
-        // here.
-        self.layer_map.load_full()
-    }
-
-    /// Clone the layer map for modification.
-    fn clone_for_write(&self, _state_lock_witness: &tokio::sync::MutexGuard<'_, ()>) -> LayerMap {
-        (**self.layer_map.load()).clone()
-    }
-
-    pub fn new(layer_map: LayerMap) -> Self {
-        Self {
-            layer_map: ArcSwap::new(Arc::new(layer_map)),
-            state_lock: tokio::sync::Mutex::new(()),
-        }
-    }
-
-    /// Update the layer map.
-    pub async fn update<O, F>(&self, operation: O) -> Result<()>
-    where
-        O: FnOnce(LayerMap) -> F,
-        F: Future<Output = Result<LayerMap>>,
-    {
-        let state_lock = self.state_lock.lock().await;
-        let state = self.clone_for_write(&state_lock);
-        let new_state = operation(state).await?;
-        self.layer_map.store(Arc::new(new_state));
-        Ok(())
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use utils::{
-        id::{TenantId, TimelineId},
-        lsn::Lsn,
-    };
-
-    use crate::{repository::Key, tenant::storage_layer::PersistentLayerDesc};
-
-    use super::*;
-
-    #[tokio::test]
-    async fn test_layer_map_manage() -> Result<()> {
-        let mgr = LayerMapMgr::new(Default::default());
-        mgr.update(|mut map| async move {
-            let mut updates = map.batch_update();
-            updates.insert_historic(PersistentLayerDesc::new_img(
-                TenantId::generate(),
-                TimelineId::generate(),
-                Key::from_i128(0)..Key::from_i128(1),
-                Lsn(0),
-                false,
-                0,
-            ));
-            updates.flush();
-            Ok(map)
-        })
-        .await?;
-
-        let ref_1 = mgr.read();
-
-        mgr.update(|mut map| async move {
-            let mut updates = map.batch_update();
-            updates.insert_historic(PersistentLayerDesc::new_img(
-                TenantId::generate(),
-                TimelineId::generate(),
-                Key::from_i128(1)..Key::from_i128(2),
-                Lsn(0),
-                false,
-                0,
-            ));
-            updates.flush();
-            Ok(map)
-        })
-        .await?;
-
-        let ref_2 = mgr.read();
-
-        // Modification should not be visible to the old reference.
-        assert_eq!(
-            ref_1
-                .search(Key::from_i128(0), Lsn(1))
-                .unwrap()
-                .layer
-                .key_range,
-            Key::from_i128(0)..Key::from_i128(1)
-        );
-        assert!(ref_1.search(Key::from_i128(1), Lsn(1)).is_none());
-
-        // Modification should be visible to the new reference.
-        assert_eq!(
-            ref_2
-                .search(Key::from_i128(0), Lsn(1))
-                .unwrap()
-                .layer
-                .key_range,
-            Key::from_i128(0)..Key::from_i128(1)
-        );
-        assert_eq!(
-            ref_2
-                .search(Key::from_i128(1), Lsn(1))
-                .unwrap()
-                .layer
-                .key_range,
-            Key::from_i128(1)..Key::from_i128(2)
-        );
-        Ok(())
-    }
-}

pageserver/src/tenant/remote_timeline_client.rs

@@ -1392,12 +1392,7 @@ mod tests {
             let harness = TenantHarness::create(test_name)?;
             let (tenant, ctx) = runtime.block_on(harness.load());
             // create an empty timeline directory
-            let _ = runtime.block_on(tenant.create_test_timeline(
-                TIMELINE_ID,
-                Lsn(8),
-                DEFAULT_PG_VERSION,
-                &ctx,
-            ))?;
+            let _ = tenant.create_test_timeline(TIMELINE_ID, Lsn(8), DEFAULT_PG_VERSION, &ctx)?;
 
             let remote_fs_dir = harness.conf.workdir.join("remote_fs");
             std::fs::create_dir_all(remote_fs_dir)?;

pageserver/src/tenant/storage_layer.rs

@@ -41,6 +41,8 @@ pub use inmemory_layer::InMemoryLayer;
 pub use layer_desc::{PersistentLayerDesc, PersistentLayerKey};
 pub use remote_layer::RemoteLayer;
 
+use super::layer_map::BatchedUpdates;
+
 pub fn range_overlaps<T>(a: &Range<T>, b: &Range<T>) -> bool
 where
     T: PartialOrd<T>,
@@ -174,9 +176,19 @@ impl LayerAccessStats {
     /// Create an empty stats object and record a [`LayerLoad`] event with the given residence status.
     ///
     /// See [`record_residence_event`] for why you need to do this while holding the layer map lock.
-    pub(crate) fn for_loading_layer(status: LayerResidenceStatus) -> Self {
+    pub(crate) fn for_loading_layer<L>(
+        layer_map_lock_held_witness: &BatchedUpdates<'_, L>,
+        status: LayerResidenceStatus,
+    ) -> Self
+    where
+        L: ?Sized + Layer,
+    {
         let new = LayerAccessStats(Mutex::new(LayerAccessStatsLocked::default()));
-        new.record_residence_event(status, LayerResidenceEventReason::LayerLoad);
+        new.record_residence_event(
+            layer_map_lock_held_witness,
+            status,
+            LayerResidenceEventReason::LayerLoad,
+        );
         new
     }
 
@@ -185,16 +197,24 @@ impl LayerAccessStats {
     /// The `new_status` is not recorded in `self`.
     ///
     /// See [`record_residence_event`] for why you need to do this while holding the layer map lock.
-    pub(crate) fn clone_for_residence_change(
+    pub(crate) fn clone_for_residence_change<L>(
         &self,
+        layer_map_lock_held_witness: &BatchedUpdates<'_, L>,
         new_status: LayerResidenceStatus,
-    ) -> LayerAccessStats {
+    ) -> LayerAccessStats
+    where
+        L: ?Sized + Layer,
+    {
         let clone = {
             let inner = self.0.lock().unwrap();
             inner.clone()
         };
         let new = LayerAccessStats(Mutex::new(clone));
-        new.record_residence_event(new_status, LayerResidenceEventReason::ResidenceChange);
+        new.record_residence_event(
+            layer_map_lock_held_witness,
+            new_status,
+            LayerResidenceEventReason::ResidenceChange,
+        );
         new
     }
 
@@ -212,11 +232,14 @@ impl LayerAccessStats {
     /// - Compact: Grab layer map lock, add the new L1 to layer map and remove the L0s, release layer map lock.
     /// - Eviction: observes the new L1 layer whose only activity timestamp is the LayerCreate event.
     ///
-    pub(crate) fn record_residence_event(
+    pub(crate) fn record_residence_event<L>(
         &self,
+        _layer_map_lock_held_witness: &BatchedUpdates<'_, L>,
         status: LayerResidenceStatus,
         reason: LayerResidenceEventReason,
-    ) {
+    ) where
+        L: ?Sized + Layer,
+    {
         let mut locked = self.0.lock().unwrap();
         locked.iter_mut().for_each(|inner| {
             inner
@@ -366,10 +389,10 @@ pub trait Layer: std::fmt::Debug + Send + Sync {
 }
 
 /// Returned by [`Layer::iter`]
-pub type LayerIter<'i> = Box<dyn Iterator<Item = Result<(Key, Lsn, Value)>> + 'i + Send>;
+pub type LayerIter<'i> = Box<dyn Iterator<Item = Result<(Key, Lsn, Value)>> + 'i>;
 
 /// Returned by [`Layer::key_iter`]
-pub type LayerKeyIter<'i> = Box<dyn Iterator<Item = (Key, Lsn, u64)> + 'i + Send>;
+pub type LayerKeyIter<'i> = Box<dyn Iterator<Item = (Key, Lsn, u64)> + 'i>;
 
 /// A Layer contains all data in a "rectangle" consisting of a range of keys and
 /// range of LSNs.
@@ -450,125 +473,94 @@ pub fn downcast_remote_layer(
     }
 }
 
-pub mod tests {
-    use super::*;
+/// Holds metadata about a layer without any content. Used mostly for testing.
+///
+/// To use filenames as fixtures, parse them as [`LayerFileName`] then convert from that to a
+/// LayerDescriptor.
+#[derive(Clone, Debug)]
+pub struct LayerDescriptor {
+    pub key: Range<Key>,
+    pub lsn: Range<Lsn>,
+    pub is_incremental: bool,
+    pub short_id: String,
+}
 
-    /// Holds metadata about a layer without any content. Used mostly for testing.
-    ///
-    /// To use filenames as fixtures, parse them as [`LayerFileName`] then convert from that to a
-    /// LayerDescriptor.
-    #[derive(Clone, Debug)]
-    pub struct LayerDescriptor {
-        base: PersistentLayerDesc,
+impl LayerDescriptor {
+    /// `LayerDescriptor` is only used for testing purpose so it does not matter whether it is image / delta,
+    /// and the tenant / timeline id does not matter.
+    pub fn get_persistent_layer_desc(&self) -> PersistentLayerDesc {
+        PersistentLayerDesc::new_delta(
+            TenantId::from_array([0; 16]),
+            TimelineId::from_array([0; 16]),
+            self.key.clone(),
+            self.lsn.clone(),
+            233,
+        )
+    }
+}
+
+impl Layer for LayerDescriptor {
+    fn get_key_range(&self) -> Range<Key> {
+        self.key.clone()
     }
 
-    impl From<PersistentLayerDesc> for LayerDescriptor {
-        fn from(base: PersistentLayerDesc) -> Self {
-            Self { base }
-        }
+    fn get_lsn_range(&self) -> Range<Lsn> {
+        self.lsn.clone()
     }
 
-    impl Layer for LayerDescriptor {
-        fn get_value_reconstruct_data(
-            &self,
-            _key: Key,
-            _lsn_range: Range<Lsn>,
-            _reconstruct_data: &mut ValueReconstructState,
-            _ctx: &RequestContext,
-        ) -> Result<ValueReconstructResult> {
-            todo!("This method shouldn't be part of the Layer trait")
-        }
-
-        fn dump(&self, _verbose: bool, _ctx: &RequestContext) -> Result<()> {
-            todo!()
-        }
-
-        /// Boilerplate to implement the Layer trait, always use layer_desc for persistent layers.
-        fn get_key_range(&self) -> Range<Key> {
-            self.layer_desc().key_range.clone()
-        }
-
-        /// Boilerplate to implement the Layer trait, always use layer_desc for persistent layers.
-        fn get_lsn_range(&self) -> Range<Lsn> {
-            self.layer_desc().lsn_range.clone()
-        }
-
-        /// Boilerplate to implement the Layer trait, always use layer_desc for persistent layers.
-        fn is_incremental(&self) -> bool {
-            self.layer_desc().is_incremental
-        }
-
-        /// Boilerplate to implement the Layer trait, always use layer_desc for persistent layers.
-        fn short_id(&self) -> String {
-            self.layer_desc().short_id()
-        }
+    fn is_incremental(&self) -> bool {
+        self.is_incremental
     }
 
-    impl PersistentLayer for LayerDescriptor {
-        fn layer_desc(&self) -> &PersistentLayerDesc {
-            &self.base
-        }
-
-        fn local_path(&self) -> Option<PathBuf> {
-            unimplemented!()
-        }
-
-        fn iter(&self, _: &RequestContext) -> Result<LayerIter<'_>> {
-            unimplemented!()
-        }
-
-        fn key_iter(&self, _: &RequestContext) -> Result<LayerKeyIter<'_>> {
-            unimplemented!()
-        }
-
-        fn delete_resident_layer_file(&self) -> Result<()> {
-            unimplemented!()
-        }
-
-        fn info(&self, _: LayerAccessStatsReset) -> HistoricLayerInfo {
-            unimplemented!()
-        }
-
-        fn access_stats(&self) -> &LayerAccessStats {
-            unimplemented!()
-        }
+    fn get_value_reconstruct_data(
+        &self,
+        _key: Key,
+        _lsn_range: Range<Lsn>,
+        _reconstruct_data: &mut ValueReconstructState,
+        _ctx: &RequestContext,
+    ) -> Result<ValueReconstructResult> {
+        todo!("This method shouldn't be part of the Layer trait")
     }
 
-    impl From<DeltaFileName> for LayerDescriptor {
-        fn from(value: DeltaFileName) -> Self {
-            LayerDescriptor {
-                base: PersistentLayerDesc::new_delta(
-                    TenantId::from_array([0; 16]),
-                    TimelineId::from_array([0; 16]),
-                    value.key_range,
-                    value.lsn_range,
-                    233,
-                ),
-            }
-        }
+    fn short_id(&self) -> String {
+        self.short_id.clone()
     }
 
-    impl From<ImageFileName> for LayerDescriptor {
-        fn from(value: ImageFileName) -> Self {
-            LayerDescriptor {
-                base: PersistentLayerDesc::new_img(
-                    TenantId::from_array([0; 16]),
-                    TimelineId::from_array([0; 16]),
-                    value.key_range,
-                    value.lsn,
-                    false,
-                    233,
-                ),
-            }
+    fn dump(&self, _verbose: bool, _ctx: &RequestContext) -> Result<()> {
+        todo!()
+    }
+}
+
+impl From<DeltaFileName> for LayerDescriptor {
+    fn from(value: DeltaFileName) -> Self {
+        let short_id = value.to_string();
+        LayerDescriptor {
+            key: value.key_range,
+            lsn: value.lsn_range,
+            is_incremental: true,
+            short_id,
         }
     }
+}
 
-    impl From<LayerFileName> for LayerDescriptor {
-        fn from(value: LayerFileName) -> Self {
-            match value {
-                LayerFileName::Delta(d) => Self::from(d),
-                LayerFileName::Image(i) => Self::from(i),
-            }
+impl From<ImageFileName> for LayerDescriptor {
+    fn from(value: ImageFileName) -> Self {
+        let short_id = value.to_string();
+        let lsn = value.lsn_as_range();
+        LayerDescriptor {
+            key: value.key_range,
+            lsn,
+            is_incremental: false,
+            short_id,
+        }
+    }
+}
+
+impl From<LayerFileName> for LayerDescriptor {
+    fn from(value: LayerFileName) -> Self {
+        match value {
+            LayerFileName::Delta(d) => Self::from(d),
+            LayerFileName::Image(i) => Self::from(i),
         }
     }
 }

pageserver/src/tenant/storage_layer/delta_layer.rs

@@ -37,7 +37,6 @@ use crate::virtual_file::VirtualFile;
 use crate::{walrecord, TEMP_FILE_SUFFIX};
 use crate::{DELTA_FILE_MAGIC, STORAGE_FORMAT_VERSION};
 use anyhow::{bail, ensure, Context, Result};
-use once_cell::sync::OnceCell;
 use pageserver_api::models::{HistoricLayerInfo, LayerAccessKind};
 use rand::{distributions::Alphanumeric, Rng};
 use serde::{Deserialize, Serialize};
@@ -47,6 +46,7 @@ use std::io::{Seek, SeekFrom};
 use std::ops::Range;
 use std::os::unix::fs::FileExt;
 use std::path::{Path, PathBuf};
+use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 use tracing::*;
 
 use utils::{
@@ -184,7 +184,7 @@ pub struct DeltaLayer {
 
     access_stats: LayerAccessStats,
 
-    inner: OnceCell<DeltaLayerInner>,
+    inner: RwLock<DeltaLayerInner>,
 }
 
 impl std::fmt::Debug for DeltaLayer {
@@ -201,17 +201,21 @@ impl std::fmt::Debug for DeltaLayer {
 }
 
 pub struct DeltaLayerInner {
+    /// If false, the fields below have not been loaded into memory yet.
+    loaded: bool,
+
     // values copied from summary
     index_start_blk: u32,
     index_root_blk: u32,
 
-    /// Reader object for reading blocks from the file.
-    file: FileBlockReader<VirtualFile>,
+    /// Reader object for reading blocks from the file. (None if not loaded yet)
+    file: Option<FileBlockReader<VirtualFile>>,
 }
 
 impl std::fmt::Debug for DeltaLayerInner {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("DeltaLayerInner")
+            .field("loaded", &self.loaded)
             .field("index_start_blk", &self.index_start_blk)
             .field("index_root_blk", &self.index_root_blk)
             .finish()
@@ -242,7 +246,7 @@ impl Layer for DeltaLayer {
             inner.index_start_blk, inner.index_root_blk
         );
 
-        let file = &inner.file;
+        let file = inner.file.as_ref().unwrap();
         let tree_reader = DiskBtreeReader::<_, DELTA_KEY_SIZE>::new(
             inner.index_start_blk,
             inner.index_root_blk,
@@ -311,7 +315,7 @@ impl Layer for DeltaLayer {
             let inner = self.load(LayerAccessKind::GetValueReconstructData, ctx)?;
 
             // Scan the page versions backwards, starting from `lsn`.
-            let file = &inner.file;
+            let file = inner.file.as_ref().unwrap();
             let tree_reader = DiskBtreeReader::<_, DELTA_KEY_SIZE>::new(
                 inner.index_start_blk,
                 inner.index_root_blk,
@@ -496,22 +500,51 @@ impl DeltaLayer {
     /// Open the underlying file and read the metadata into memory, if it's
     /// not loaded already.
     ///
-    fn load(&self, access_kind: LayerAccessKind, ctx: &RequestContext) -> Result<&DeltaLayerInner> {
+    fn load(
+        &self,
+        access_kind: LayerAccessKind,
+        ctx: &RequestContext,
+    ) -> Result<RwLockReadGuard<DeltaLayerInner>> {
         self.access_stats
             .record_access(access_kind, ctx.task_kind());
-        // Quick exit if already loaded
-        self.inner
-            .get_or_try_init(|| self.load_inner())
-            .with_context(|| format!("Failed to load delta layer {}", self.path().display()))
+        loop {
+            // Quick exit if already loaded
+            let inner = self.inner.read().unwrap();
+            if inner.loaded {
+                return Ok(inner);
+            }
+
+            // Need to open the file and load the metadata. Upgrade our lock to
+            // a write lock. (Or rather, release and re-lock in write mode.)
+            drop(inner);
+            let inner = self.inner.write().unwrap();
+            if !inner.loaded {
+                self.load_inner(inner).with_context(|| {
+                    format!("Failed to load delta layer {}", self.path().display())
+                })?;
+            } else {
+                // Another thread loaded it while we were not holding the lock.
+            }
+
+            // We now have the file open and loaded. There's no function to do
+            // that in the std library RwLock, so we have to release and re-lock
+            // in read mode. (To be precise, the lock guard was moved in the
+            // above call to `load_inner`, so it's already been released). And
+            // while we do that, another thread could unload again, so we have
+            // to re-check and retry if that happens.
+        }
     }
 
-    fn load_inner(&self) -> Result<DeltaLayerInner> {
+    fn load_inner(&self, mut inner: RwLockWriteGuard<DeltaLayerInner>) -> Result<()> {
         let path = self.path();
 
-        let file = VirtualFile::open(&path)
-            .with_context(|| format!("Failed to open file '{}'", path.display()))?;
-        let file = FileBlockReader::new(file);
-
+        // Open the file if it's not open already.
+        if inner.file.is_none() {
+            let file = VirtualFile::open(&path)
+                .with_context(|| format!("Failed to open file '{}'", path.display()))?;
+            inner.file = Some(FileBlockReader::new(file));
+        }
+        let file = inner.file.as_mut().unwrap();
         let summary_blk = file.read_blk(0)?;
         let actual_summary = Summary::des_prefix(summary_blk.as_ref())?;
 
@@ -538,13 +571,13 @@ impl DeltaLayer {
             }
         }
 
+        inner.index_start_blk = actual_summary.index_start_blk;
+        inner.index_root_blk = actual_summary.index_root_blk;
+
         debug!("loaded from {}", &path.display());
 
-        Ok(DeltaLayerInner {
-            file,
-            index_start_blk: actual_summary.index_start_blk,
-            index_root_blk: actual_summary.index_root_blk,
-        })
+        inner.loaded = true;
+        Ok(())
     }
 
     /// Create a DeltaLayer struct representing an existing file on disk.
@@ -566,7 +599,12 @@ impl DeltaLayer {
                 file_size,
             ),
             access_stats,
-            inner: once_cell::sync::OnceCell::new(),
+            inner: RwLock::new(DeltaLayerInner {
+                loaded: false,
+                file: None,
+                index_start_blk: 0,
+                index_root_blk: 0,
+            }),
         }
     }
 
@@ -593,7 +631,12 @@ impl DeltaLayer {
                 metadata.len(),
             ),
             access_stats: LayerAccessStats::empty_will_record_residence_event_later(),
-            inner: once_cell::sync::OnceCell::new(),
+            inner: RwLock::new(DeltaLayerInner {
+                loaded: false,
+                file: None,
+                index_start_blk: 0,
+                index_root_blk: 0,
+            }),
         })
     }
 
@@ -757,7 +800,12 @@ impl DeltaLayerWriterInner {
                 metadata.len(),
             ),
             access_stats: LayerAccessStats::empty_will_record_residence_event_later(),
-            inner: once_cell::sync::OnceCell::new(),
+            inner: RwLock::new(DeltaLayerInner {
+                loaded: false,
+                file: None,
+                index_start_blk,
+                index_root_blk,
+            }),
         };
 
         // fsync the file
@@ -892,13 +940,13 @@ struct DeltaValueIter<'a> {
     reader: BlockCursor<Adapter<'a>>,
 }
 
-struct Adapter<'a>(&'a DeltaLayerInner);
+struct Adapter<'a>(RwLockReadGuard<'a, DeltaLayerInner>);
 
 impl<'a> BlockReader for Adapter<'a> {
     type BlockLease = PageReadGuard<'static>;
 
     fn read_blk(&self, blknum: u32) -> Result<Self::BlockLease, std::io::Error> {
-        self.0.file.read_blk(blknum)
+        self.0.file.as_ref().unwrap().read_blk(blknum)
     }
 }
 
@@ -911,8 +959,8 @@ impl<'a> Iterator for DeltaValueIter<'a> {
 }
 
 impl<'a> DeltaValueIter<'a> {
-    fn new(inner: &'a DeltaLayerInner) -> Result<Self> {
-        let file = &inner.file;
+    fn new(inner: RwLockReadGuard<'a, DeltaLayerInner>) -> Result<Self> {
+        let file = inner.file.as_ref().unwrap();
         let tree_reader = DiskBtreeReader::<_, DELTA_KEY_SIZE>::new(
             inner.index_start_blk,
             inner.index_root_blk,
@@ -985,8 +1033,8 @@ impl Iterator for DeltaKeyIter {
 }
 
 impl<'a> DeltaKeyIter {
-    fn new(inner: &'a DeltaLayerInner) -> Result<Self> {
-        let file = &inner.file;
+    fn new(inner: RwLockReadGuard<'a, DeltaLayerInner>) -> Result<Self> {
+        let file = inner.file.as_ref().unwrap();
         let tree_reader = DiskBtreeReader::<_, DELTA_KEY_SIZE>::new(
             inner.index_start_blk,
             inner.index_root_blk,
@@ -1026,21 +1074,3 @@ impl<'a> DeltaKeyIter {
         Ok(iter)
     }
 }
-
-#[cfg(test)]
-mod test {
-    use super::DeltaKeyIter;
-    use super::DeltaLayer;
-    use super::DeltaValueIter;
-
-    // We will soon need the iters to be send in the compaction code.
-    // Cf https://github.com/neondatabase/neon/pull/4462#issuecomment-1587398883
-    // Cf https://github.com/neondatabase/neon/issues/4471
-    #[test]
-    fn is_send() {
-        fn assert_send<T: Send>() {}
-        assert_send::<DeltaLayer>();
-        assert_send::<DeltaValueIter>();
-        assert_send::<DeltaKeyIter>();
-    }
-}

pageserver/src/tenant/storage_layer/inmemory_layer.rs

@@ -304,7 +304,7 @@ impl InMemoryLayer {
         Ok(())
     }
 
-    pub async fn put_tombstone(&self, _key_range: Range<Key>, _lsn: Lsn) -> Result<()> {
+    pub fn put_tombstone(&self, _key_range: Range<Key>, _lsn: Lsn) -> Result<()> {
         // TODO: Currently, we just leak the storage for any deleted keys
 
         Ok(())

pageserver/src/tenant/storage_layer/remote_layer.rs

@@ -4,6 +4,7 @@
 use crate::config::PageServerConf;
 use crate::context::RequestContext;
 use crate::repository::Key;
+use crate::tenant::layer_map::BatchedUpdates;
 use crate::tenant::remote_timeline_client::index::LayerFileMetadata;
 use crate::tenant::storage_layer::{Layer, ValueReconstructResult, ValueReconstructState};
 use anyhow::{bail, Result};
@@ -217,11 +218,15 @@ impl RemoteLayer {
     }
 
     /// Create a Layer struct representing this layer, after it has been downloaded.
-    pub fn create_downloaded_layer(
+    pub fn create_downloaded_layer<L>(
         &self,
+        layer_map_lock_held_witness: &BatchedUpdates<'_, L>,
         conf: &'static PageServerConf,
         file_size: u64,
-    ) -> Arc<dyn PersistentLayer> {
+    ) -> Arc<dyn PersistentLayer>
+    where
+        L: ?Sized + Layer,
+    {
         if self.desc.is_delta {
             let fname = self.desc.delta_file_name();
             Arc::new(DeltaLayer::new(
@@ -230,8 +235,10 @@ impl RemoteLayer {
                 self.desc.tenant_id,
                 &fname,
                 file_size,
-                self.access_stats
-                    .clone_for_residence_change(LayerResidenceStatus::Resident),
+                self.access_stats.clone_for_residence_change(
+                    layer_map_lock_held_witness,
+                    LayerResidenceStatus::Resident,
+                ),
             ))
         } else {
             let fname = self.desc.image_file_name();
@@ -241,8 +248,10 @@ impl RemoteLayer {
                 self.desc.tenant_id,
                 &fname,
                 file_size,
-                self.access_stats
-                    .clone_for_residence_change(LayerResidenceStatus::Resident),
+                self.access_stats.clone_for_residence_change(
+                    layer_map_lock_held_witness,
+                    LayerResidenceStatus::Resident,
+                ),
             ))
         }
     }

pageserver/src/tenant/timeline/eviction_task.rs

@@ -197,11 +197,9 @@ impl Timeline {
         // We don't want to hold the layer map lock during eviction.
         // So, we just need to deal with this.
         let candidates: Vec<Arc<dyn PersistentLayer>> = {
-            let guard = self.lcache.layer_in_use_read().await;
-            let layers = self.layer_mgr.read();
+            let layers = self.layers.read().unwrap();
             let mut candidates = Vec::new();
             for hist_layer in layers.iter_historic_layers() {
-                let hist_layer = self.lcache.get_from_desc(&hist_layer);
                 if hist_layer.is_remote_layer() {
                     continue;
                 }

pageserver/src/tenant/timeline/walreceiver/connection_manager.rs

@@ -1325,7 +1325,6 @@ mod tests {
         let (tenant, ctx) = harness.load().await;
         let timeline = tenant
             .create_test_timeline(TIMELINE_ID, Lsn(0x8), crate::DEFAULT_PG_VERSION, &ctx)
-            .await
             .expect("Failed to create an empty timeline for dummy wal connection manager");
 
         ConnectionManagerState {

pageserver/src/tenant/timeline/walreceiver/walreceiver_connection.rs

@@ -304,15 +304,12 @@ pub(super) async fn handle_walreceiver_connection(
             }
         }
 
-        timeline
-            .check_checkpoint_distance()
-            .await
-            .with_context(|| {
-                format!(
-                    "Failed to check checkpoint distance for timeline {}",
-                    timeline.timeline_id
-                )
-            })?;
+        timeline.check_checkpoint_distance().with_context(|| {
+            format!(
+                "Failed to check checkpoint distance for timeline {}",
+                timeline.timeline_id
+            )
+        })?;
 
         if let Some(last_lsn) = status_update {
             let timeline_remote_consistent_lsn =

pageserver/src/walingest.rs

@@ -333,7 +333,7 @@ impl<'a> WalIngest<'a> {
 
         // Now that this record has been fully handled, including updating the
         // checkpoint data, let the repository know that it is up-to-date to this LSN
-        modification.commit().await?;
+        modification.commit()?;
 
         Ok(())
     }
@@ -1199,7 +1199,7 @@ mod tests {
         let mut m = tline.begin_modification(Lsn(0x10));
         m.put_checkpoint(ZERO_CHECKPOINT.clone())?;
         m.put_relmap_file(0, 111, Bytes::from(""), ctx).await?; // dummy relmapper file
-        m.commit().await?;
+        m.commit()?;
         let walingest = WalIngest::new(tline, Lsn(0x10), ctx).await?;
 
         Ok(walingest)
@@ -1208,9 +1208,7 @@ mod tests {
     #[tokio::test]
     async fn test_relsize() -> Result<()> {
         let (tenant, ctx) = TenantHarness::create("test_relsize")?.load().await;
-        let tline = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(8), DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let tline = tenant.create_test_timeline(TIMELINE_ID, Lsn(8), DEFAULT_PG_VERSION, &ctx)?;
         let mut walingest = init_walingest_test(&tline, &ctx).await?;
 
         let mut m = tline.begin_modification(Lsn(0x20));
@@ -1218,22 +1216,22 @@ mod tests {
         walingest
             .put_rel_page_image(&mut m, TESTREL_A, 0, TEST_IMG("foo blk 0 at 2"), &ctx)
             .await?;
-        m.commit().await?;
+        m.commit()?;
         let mut m = tline.begin_modification(Lsn(0x30));
         walingest
             .put_rel_page_image(&mut m, TESTREL_A, 0, TEST_IMG("foo blk 0 at 3"), &ctx)
             .await?;
-        m.commit().await?;
+        m.commit()?;
         let mut m = tline.begin_modification(Lsn(0x40));
         walingest
             .put_rel_page_image(&mut m, TESTREL_A, 1, TEST_IMG("foo blk 1 at 4"), &ctx)
             .await?;
-        m.commit().await?;
+        m.commit()?;
         let mut m = tline.begin_modification(Lsn(0x50));
         walingest
             .put_rel_page_image(&mut m, TESTREL_A, 2, TEST_IMG("foo blk 2 at 5"), &ctx)
             .await?;
-        m.commit().await?;
+        m.commit()?;
 
         assert_current_logical_size(&tline, Lsn(0x50));
 
@@ -1319,7 +1317,7 @@ mod tests {
         walingest
             .put_rel_truncation(&mut m, TESTREL_A, 2, &ctx)
             .await?;
-        m.commit().await?;
+        m.commit()?;
         assert_current_logical_size(&tline, Lsn(0x60));
 
         // Check reported size and contents after truncation
@@ -1361,7 +1359,7 @@ mod tests {
         walingest
             .put_rel_truncation(&mut m, TESTREL_A, 0, &ctx)
             .await?;
-        m.commit().await?;
+        m.commit()?;
         assert_eq!(
             tline
                 .get_rel_size(TESTREL_A, Lsn(0x68), false, &ctx)
@@ -1374,7 +1372,7 @@ mod tests {
         walingest
             .put_rel_page_image(&mut m, TESTREL_A, 1, TEST_IMG("foo blk 1"), &ctx)
             .await?;
-        m.commit().await?;
+        m.commit()?;
         assert_eq!(
             tline
                 .get_rel_size(TESTREL_A, Lsn(0x70), false, &ctx)
@@ -1399,7 +1397,7 @@ mod tests {
         walingest
             .put_rel_page_image(&mut m, TESTREL_A, 1500, TEST_IMG("foo blk 1500"), &ctx)
             .await?;
-        m.commit().await?;
+        m.commit()?;
         assert_eq!(
             tline
                 .get_rel_size(TESTREL_A, Lsn(0x80), false, &ctx)
@@ -1429,16 +1427,14 @@ mod tests {
     #[tokio::test]
     async fn test_drop_extend() -> Result<()> {
         let (tenant, ctx) = TenantHarness::create("test_drop_extend")?.load().await;
-        let tline = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(8), DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let tline = tenant.create_test_timeline(TIMELINE_ID, Lsn(8), DEFAULT_PG_VERSION, &ctx)?;
         let mut walingest = init_walingest_test(&tline, &ctx).await?;
 
         let mut m = tline.begin_modification(Lsn(0x20));
         walingest
             .put_rel_page_image(&mut m, TESTREL_A, 0, TEST_IMG("foo blk 0 at 2"), &ctx)
             .await?;
-        m.commit().await?;
+        m.commit()?;
 
         // Check that rel exists and size is correct
         assert_eq!(
@@ -1457,7 +1453,7 @@ mod tests {
         // Drop rel
         let mut m = tline.begin_modification(Lsn(0x30));
         walingest.put_rel_drop(&mut m, TESTREL_A, &ctx).await?;
-        m.commit().await?;
+        m.commit()?;
 
         // Check that rel is not visible anymore
         assert_eq!(
@@ -1475,7 +1471,7 @@ mod tests {
         walingest
             .put_rel_page_image(&mut m, TESTREL_A, 0, TEST_IMG("foo blk 0 at 4"), &ctx)
             .await?;
-        m.commit().await?;
+        m.commit()?;
 
         // Check that rel exists and size is correct
         assert_eq!(
@@ -1500,9 +1496,7 @@ mod tests {
     #[tokio::test]
     async fn test_truncate_extend() -> Result<()> {
         let (tenant, ctx) = TenantHarness::create("test_truncate_extend")?.load().await;
-        let tline = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(8), DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let tline = tenant.create_test_timeline(TIMELINE_ID, Lsn(8), DEFAULT_PG_VERSION, &ctx)?;
         let mut walingest = init_walingest_test(&tline, &ctx).await?;
 
         // Create a 20 MB relation (the size is arbitrary)
@@ -1514,7 +1508,7 @@ mod tests {
                 .put_rel_page_image(&mut m, TESTREL_A, blkno, TEST_IMG(&data), &ctx)
                 .await?;
         }
-        m.commit().await?;
+        m.commit()?;
 
         // The relation was created at LSN 20, not visible at LSN 1 yet.
         assert_eq!(
@@ -1559,7 +1553,7 @@ mod tests {
         walingest
             .put_rel_truncation(&mut m, TESTREL_A, 1, &ctx)
             .await?;
-        m.commit().await?;
+        m.commit()?;
 
         // Check reported size and contents after truncation
         assert_eq!(
@@ -1608,7 +1602,7 @@ mod tests {
                 .put_rel_page_image(&mut m, TESTREL_A, blkno, TEST_IMG(&data), &ctx)
                 .await?;
         }
-        m.commit().await?;
+        m.commit()?;
 
         assert_eq!(
             tline
@@ -1642,9 +1636,7 @@ mod tests {
     #[tokio::test]
     async fn test_large_rel() -> Result<()> {
         let (tenant, ctx) = TenantHarness::create("test_large_rel")?.load().await;
-        let tline = tenant
-            .create_test_timeline(TIMELINE_ID, Lsn(8), DEFAULT_PG_VERSION, &ctx)
-            .await?;
+        let tline = tenant.create_test_timeline(TIMELINE_ID, Lsn(8), DEFAULT_PG_VERSION, &ctx)?;
         let mut walingest = init_walingest_test(&tline, &ctx).await?;
 
         let mut lsn = 0x10;
@@ -1655,7 +1647,7 @@ mod tests {
             walingest
                 .put_rel_page_image(&mut m, TESTREL_A, blknum as BlockNumber, img, &ctx)
                 .await?;
-            m.commit().await?;
+            m.commit()?;
         }
 
         assert_current_logical_size(&tline, Lsn(lsn));
@@ -1671,7 +1663,7 @@ mod tests {
         walingest
             .put_rel_truncation(&mut m, TESTREL_A, RELSEG_SIZE, &ctx)
             .await?;
-        m.commit().await?;
+        m.commit()?;
         assert_eq!(
             tline.get_rel_size(TESTREL_A, Lsn(lsn), false, &ctx).await?,
             RELSEG_SIZE
@@ -1684,7 +1676,7 @@ mod tests {
         walingest
             .put_rel_truncation(&mut m, TESTREL_A, RELSEG_SIZE - 1, &ctx)
             .await?;
-        m.commit().await?;
+        m.commit()?;
         assert_eq!(
             tline.get_rel_size(TESTREL_A, Lsn(lsn), false, &ctx).await?,
             RELSEG_SIZE - 1
@@ -1700,7 +1692,7 @@ mod tests {
             walingest
                 .put_rel_truncation(&mut m, TESTREL_A, size as BlockNumber, &ctx)
                 .await?;
-            m.commit().await?;
+            m.commit()?;
             assert_eq!(
                 tline.get_rel_size(TESTREL_A, Lsn(lsn), false, &ctx).await?,
                 size as BlockNumber

safekeeper/src/wal_storage.rs

@@ -98,22 +98,6 @@ pub struct PhysicalStorage {
     /// - points to write_lsn, so no seek is needed for writing
     /// - doesn't point to the end of the segment
     file: Option<File>,
-
-    /// When false, we have just initialized storage using the LSN from find_end_of_wal().
-    /// In this case, [`write_lsn`] can be less than actually written WAL on disk. In particular,
-    /// there can be a case with unexpected .partial file.
-    ///
-    /// Imagine the following:
-    /// - 000000010000000000000001
-    ///   - it was fully written, but the last record is split between 2 segments
-    ///   - after restart, find_end_of_wal() returned 0/1FFFFF0, which is in the end of this segment
-    ///   - write_lsn, write_record_lsn and flush_record_lsn were initialized to 0/1FFFFF0
-    /// - 000000010000000000000002.partial
-    ///   - it has only 1 byte written, which is not enough to make a full WAL record
-    ///
-    /// Partial segment 002 has no WAL records, and it will be removed by the next truncate_wal().
-    /// This flag will be set to true after the first truncate_wal() call.
-    is_truncated_after_restart: bool,
 }
 
 impl PhysicalStorage {
@@ -173,7 +157,6 @@ impl PhysicalStorage {
             flush_record_lsn: flush_lsn,
             decoder: WalStreamDecoder::new(write_lsn, state.server.pg_version / 10000),
             file: None,
-            is_truncated_after_restart: false,
         })
     }
 
@@ -398,10 +381,7 @@ impl Storage for PhysicalStorage {
 
         // Quick exit if nothing to do to avoid writing up to 16 MiB of zeros on
         // disk (this happens on each connect).
-        if self.is_truncated_after_restart
-            && end_pos == self.write_lsn
-            && end_pos == self.flush_record_lsn
-        {
+        if end_pos == self.write_lsn {
             return Ok(());
         }
 
@@ -434,7 +414,6 @@ impl Storage for PhysicalStorage {
         self.write_lsn = end_pos;
         self.write_record_lsn = end_pos;
         self.flush_record_lsn = end_pos;
-        self.is_truncated_after_restart = true;
         Ok(())
     }
 

test_runner/regress/test_proxy.py

@@ -163,6 +163,7 @@ def test_forward_params_to_client(static_proxy: NeonProxy):
                 assert conn.get_parameter_status(name) == value
 
 
+@pytest.mark.timeout(5)
 def test_close_on_connections_exit(static_proxy: NeonProxy):
     # Open two connections, send SIGTERM, then ensure that proxy doesn't exit
     # until after connections close.