add layer desc guard to ensure it is used within scope

Signed-off-by: Alex Chi Z <chi@neon.tech>

pageserver/src/tenant/layer_map.rs

@@ -53,6 +53,7 @@ use crate::tenant::storage_layer::InMemoryLayer;
 use crate::tenant::storage_layer::Layer;
 use anyhow::Result;
 use std::collections::VecDeque;
+use std::marker::PhantomData;
 use std::ops::Range;
 use std::sync::Arc;
 use utils::lsn::Lsn;
@@ -145,9 +146,32 @@ impl Drop for BatchedUpdates<'_> {
     }
 }
 
+/// Helper struct to keep layer desc being used within the scope of layer map.
+pub struct PersistentLayerDescGuard<'a> {
+    layer_desc: Arc<PersistentLayerDesc>,
+    marker: PhantomData<&'a PersistentLayerDesc>,
+}
+
+impl std::ops::Deref for PersistentLayerDescGuard<'_> {
+    type Target = PersistentLayerDesc;
+
+    fn deref(&self) -> &Self::Target {
+        &self.layer_desc
+    }
+}
+
+impl PersistentLayerDescGuard<'_> {
+    fn new(layer_desc: Arc<PersistentLayerDesc>) -> Self {
+        Self {
+            layer_desc,
+            marker: PhantomData,
+        }
+    }
+}
+
 /// Return value of LayerMap::search
-pub struct SearchResult {
-    pub layer: Arc<PersistentLayerDesc>,
+pub struct SearchResult<'a> {
+    pub layer: PersistentLayerDescGuard<'a>,
     pub lsn_floor: Lsn,
 }
 
@@ -193,14 +217,14 @@ impl LayerMap {
             (None, Some(image)) => {
                 let lsn_floor = image.get_lsn_range().start;
                 Some(SearchResult {
-                    layer: image,
+                    layer: PersistentLayerDescGuard::new(image),
                     lsn_floor,
                 })
             }
             (Some(delta), None) => {
                 let lsn_floor = delta.get_lsn_range().start;
                 Some(SearchResult {
-                    layer: delta,
+                    layer: PersistentLayerDescGuard::new(delta),
                     lsn_floor,
                 })
             }
@@ -210,14 +234,14 @@ impl LayerMap {
                 let image_exact_match = img_lsn + 1 == end_lsn;
                 if image_is_newer || image_exact_match {
                     Some(SearchResult {
-                        layer: image,
+                        layer: PersistentLayerDescGuard::new(image),
                         lsn_floor: img_lsn,
                     })
                 } else {
                     let lsn_floor =
                         std::cmp::max(delta.get_lsn_range().start, image.get_lsn_range().start + 1);
                     Some(SearchResult {
-                        layer: delta,
+                        layer: PersistentLayerDescGuard::new(delta),
                         lsn_floor,
                     })
                 }
@@ -319,8 +343,8 @@ 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 = PersistentLayerDescGuard<'_>> {
+        self.historic.iter().map(PersistentLayerDescGuard::new)
     }
 
     ///

pageserver/src/tenant/timeline.rs

@@ -2379,11 +2379,11 @@ impl Timeline {
         &self,
         // we cannot remove layers otherwise, since gc and compaction will race
         _layer_removal_cs: Arc<tokio::sync::OwnedMutexGuard<()>>,
-        layer: Arc<PersistentLayerDesc>,
+        layer: &PersistentLayerDesc,
         updates: &mut BatchedUpdates<'_>,
         mapping: &mut LayerFileManager,
     ) -> anyhow::Result<()> {
-        let layer = mapping.get_from_desc(&layer);
+        let layer = mapping.get_from_desc(layer);
         if !layer.is_remote_layer() {
             layer.delete_resident_layer_file()?;
             let layer_file_size = layer.file_size();
@@ -4040,7 +4040,7 @@ impl Timeline {
         let mut layer_names_to_delete = Vec::with_capacity(deltas_to_compact.len());
         for l in deltas_to_compact {
             layer_names_to_delete.push(l.filename());
-            self.delete_historic_layer(layer_removal_cs.clone(), l, &mut updates, mapping)?;
+            self.delete_historic_layer(layer_removal_cs.clone(), &l, &mut updates, mapping)?;
         }
         updates.flush();
         drop_wlock(guard);
@@ -4353,7 +4353,8 @@ impl Timeline {
                 l.filename().file_name(),
                 l.is_incremental(),
             );
-            layers_to_remove.push(Arc::clone(&l));
+            let l = mapping.get_from_desc(&l);
+            layers_to_remove.push(l);
         }
         self.wanted_image_layers
             .lock()
@@ -4375,7 +4376,7 @@ impl Timeline {
                     layer_names_to_delete.push(doomed_layer.filename());
                     self.delete_historic_layer(
                         layer_removal_cs.clone(),
-                        doomed_layer,
+                        doomed_layer.layer_desc(),
                         &mut updates,
                         mapping,
                     )?; // FIXME: schedule succeeded deletions before returning?