basebackup import: pre-lock the layer map for the flush() calls

The checkpointer loop isn't running anyway, so, there's no risk of
blocking it through the pre-lock.

(cherry picked from commit 1b2663350c)

pageserver/src/pgdatadir_mapping.rs

@@ -1118,11 +1118,13 @@ impl<'a> DatadirModification<'a> {
 
         let writer = self.tline.writer();
 
+        let mut layer_map = self.tline.layers.write().unwrap();
+
         // Flush relation and  SLRU data blocks, keep metadata.
         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);
+                result = writer.put_locked(key, self.lsn, value, &mut layer_map);
                 false
             } else {
                 true

pageserver/src/tenant/timeline.rs

@@ -119,7 +119,7 @@ pub struct Timeline {
 
     pub pg_version: u32,
 
-    pub(super) layers: RwLock<LayerMap<dyn PersistentLayer>>,
+    pub(crate) layers: RwLock<LayerMap<dyn PersistentLayer>>,
 
     /// Set of key ranges which should be covered by image layers to
     /// allow GC to remove old layers. This set is created by GC and its cutoff LSN is also stored.
@@ -238,6 +238,8 @@ pub struct Timeline {
     eviction_task_timeline_state: tokio::sync::Mutex<EvictionTaskTimelineState>,
 }
 
+type LayerMapWriteLockGuard<'t> = std::sync::RwLockWriteGuard<'t, LayerMap<dyn PersistentLayer>>;
+
 /// Internal structure to hold all data needed for logical size calculation.
 ///
 /// Calculation consists of two stages:
@@ -2619,7 +2621,14 @@ impl Timeline {
     ///
     fn get_layer_for_write(&self, lsn: Lsn) -> anyhow::Result<Arc<InMemoryLayer>> {
         let mut layers = self.layers.write().unwrap();
+        self.get_layer_for_write_locked(lsn, &mut layers)
+    }
 
+    fn get_layer_for_write_locked(
+        &self,
+        lsn: Lsn,
+        layers: &mut LayerMapWriteLockGuard,
+    ) -> anyhow::Result<Arc<InMemoryLayer>> {
         ensure!(lsn.is_aligned());
 
         let last_record_lsn = self.get_last_record_lsn();
@@ -2669,6 +2678,19 @@ impl Timeline {
         Ok(())
     }
 
+    fn put_value_locked(
+        &self,
+        key: Key,
+        lsn: Lsn,
+        val: &Value,
+        pre_locked_layer_map: &mut LayerMapWriteLockGuard,
+    ) -> anyhow::Result<()> {
+        //info!("PUT: key {} at {}", key, lsn);
+        let layer = self.get_layer_for_write_locked(lsn, pre_locked_layer_map)?;
+        layer.put_value(key, lsn, val)?;
+        Ok(())
+    }
+
     fn put_tombstone(&self, key_range: Range<Key>, lsn: Lsn) -> anyhow::Result<()> {
         let layer = self.get_layer_for_write(lsn)?;
         layer.put_tombstone(key_range, lsn)?;
@@ -4515,6 +4537,17 @@ impl<'a> TimelineWriter<'a> {
         self.tl.put_value(key, lsn, value)
     }
 
+    pub fn put_locked(
+        &self,
+        key: Key,
+        lsn: Lsn,
+        value: &Value,
+        pre_locked_layer_map: &mut LayerMapWriteLockGuard,
+    ) -> anyhow::Result<()> {
+        self.tl
+            .put_value_locked(key, lsn, value, pre_locked_layer_map)
+    }
+
     pub fn delete(&self, key_range: Range<Key>, lsn: Lsn) -> anyhow::Result<()> {
         self.tl.put_tombstone(key_range, lsn)
     }