refactor: move evict and download to layer cache

Signed-off-by: Alex Chi <iskyzh@gmail.com>

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"

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/src/disk_usage_eviction_task.rs

@@ -379,7 +379,10 @@ pub async fn disk_usage_eviction_task_iteration_impl<U: Usage>(
         debug!(%timeline_id, "evicting batch for timeline");
 
         async {
-            let results = timeline.evict_layers(storage, &batch, cancel.clone()).await;
+            let results = timeline
+                .lcache
+                .evict_layers(storage, &batch, cancel.clone())
+                .await;
 
             match results {
                 Err(e) => {

pageserver/src/http/routes.rs

@@ -719,6 +719,7 @@ async fn layer_download_handler(
 
     let timeline = active_timeline_of_active_tenant(tenant_id, timeline_id).await?;
     let downloaded = timeline
+        .lcache
         .download_layer(layer_file_name)
         .await
         .map_err(ApiError::InternalServerError)?;
@@ -744,6 +745,7 @@ async fn evict_timeline_layer_handler(
 
     let timeline = active_timeline_of_active_tenant(tenant_id, timeline_id).await?;
     let evicted = timeline
+        .lcache
         .evict_layer(layer_file_name)
         .await
         .map_err(ApiError::InternalServerError)?;

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/tenant.rs

@@ -87,7 +87,6 @@ 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;
@@ -557,10 +556,17 @@ 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()
+                    .await
+                    .0
+                    .iter_historic_layers()
+                    .next()
+                    .is_some(),
             "Timeline has no ancestor and no layer files"
         );
 
@@ -1241,7 +1247,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,11 +1259,9 @@ 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
@@ -1277,7 +1281,6 @@ impl Tenant {
             initdb_lsn,
             None,
         )
-        .await
     }
 
     /// Helper for unit tests to create an emtpy timeline.
@@ -1293,9 +1296,7 @@ impl Tenant {
         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));
 
@@ -2755,15 +2756,13 @@ impl Tenant {
             src_timeline.pg_version,
         );
 
-        let uninitialized_timeline = self
-            .prepare_new_timeline(
-                dst_id,
-                &metadata,
-                timeline_uninit_mark,
-                start_lsn + 1,
-                Some(Arc::clone(src_timeline)),
-            )
-            .await?;
+        let uninitialized_timeline = self.prepare_new_timeline(
+            dst_id,
+            &metadata,
+            timeline_uninit_mark,
+            start_lsn + 1,
+            Some(Arc::clone(src_timeline)),
+        )?;
 
         let new_timeline = uninitialized_timeline.finish_creation()?;
 
@@ -2841,15 +2840,13 @@ 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_new_timeline(
+            timeline_id,
+            &new_metadata,
+            timeline_uninit_mark,
+            pgdata_lsn,
+            None,
+        )?;
 
         let tenant_id = raw_timeline.owning_tenant.tenant_id;
         let unfinished_timeline = raw_timeline.raw_timeline()?;
@@ -2902,7 +2899,7 @@ impl Tenant {
     /// 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(
+    fn prepare_new_timeline(
         &self,
         new_timeline_id: TimelineId,
         new_metadata: &TimelineMetadata,
@@ -2929,7 +2926,7 @@ impl Tenant {
             .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?;
+        timeline_struct.init_empty_layer_map(start_lsn);
 
         if let Err(e) =
             self.create_timeline_files(&uninit_mark.timeline_path, new_timeline_id, new_metadata)
@@ -3626,10 +3623,7 @@ mod tests {
             .create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)
             .await?;
 
-        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(),
@@ -4429,9 +4423,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`

pageserver/src/tenant/layer_cache.rs

@@ -1,9 +1,19 @@
+use super::remote_timeline_client::RemoteTimelineClient;
 use super::storage_layer::{PersistentLayer, PersistentLayerDesc, PersistentLayerKey, RemoteLayer};
 use super::Timeline;
 use crate::tenant::layer_map::{self, LayerMap};
-use anyhow::Result;
+use crate::tenant::remote_timeline_client::index::LayerFileMetadata;
+use crate::tenant::storage_layer::LayerFileName;
+use anyhow::{Context, Result};
+use pageserver_api::models::TimelineState;
+use remote_storage::GenericRemoteStorage;
 use std::sync::{Mutex, Weak};
+use std::time::SystemTime;
 use std::{collections::HashMap, sync::Arc};
+use tokio_util::sync::CancellationToken;
+use tracing::instrument;
+use tracing::{error, info, warn};
+use utils::id::{TenantId, TimelineId};
 
 pub struct LayerCache {
     /// Layer removal lock.
@@ -22,18 +32,28 @@ pub struct LayerCache {
     timeline: Weak<Timeline>,
 
     mapping: Mutex<HashMap<PersistentLayerKey, Arc<dyn PersistentLayer>>>,
+
+    tenant_id: TenantId,
+    timeline_id: TimelineId,
 }
 
+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 {
+        let timeline_owned = timeline.upgrade().expect("cannot upgrade");
         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,
+            tenant_id: timeline_owned.tenant_id,
+            timeline_id: timeline_owned.timeline_id,
         }
     }
 
@@ -42,6 +62,18 @@ impl LayerCache {
         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(
@@ -124,4 +156,243 @@ impl LayerCache {
         let mut guard = self.mapping.lock().unwrap();
         guard.remove(&layer.layer_desc().key());
     }
+
+    #[instrument(skip_all, fields(tenant = %self.tenant_id, timeline = %self.timeline_id))]
+    pub async fn download_layer(&self, layer_file_name: &str) -> anyhow::Result<Option<bool>> {
+        let Some(layer) = self.find_layer(layer_file_name) else { return Ok(None) };
+        let Some(remote_layer) = layer.downcast_remote_layer() else { return  Ok(Some(false)) };
+        let Some(timeline) = self.timeline.upgrade() else { return Ok(None) };
+        if timeline.remote_client.is_none() {
+            return Ok(Some(false));
+        }
+
+        timeline.download_remote_layer(remote_layer).await?;
+        Ok(Some(true))
+    }
+
+    fn find_layer(&self, layer_file_name: &str) -> Option<Arc<dyn PersistentLayer>> {
+        let layers = self.mapping.lock().unwrap();
+        for (_, historic_layer) in layers.iter() {
+            let historic_layer_name = historic_layer.filename().file_name();
+            if layer_file_name == historic_layer_name {
+                return Some(historic_layer.clone());
+            }
+        }
+
+        None
+    }
+
+    /// Like [`evict_layer_batch`], but for just one layer.
+    /// Additional case `Ok(None)` covers the case where the layer could not be found by its `layer_file_name`.
+    pub async fn evict_layer(&self, layer_file_name: &str) -> anyhow::Result<Option<bool>> {
+        let Some(local_layer) = self.find_layer(layer_file_name) else { return Ok(None) };
+        let Some(timeline) = self.timeline.upgrade() else { return Ok(None) };
+        let remote_client = timeline
+            .remote_client
+            .as_ref()
+            .ok_or_else(|| anyhow::anyhow!("remote storage not configured; cannot evict"))?;
+
+        let cancel = CancellationToken::new();
+        let results = self
+            .evict_layer_batch(remote_client, &[local_layer], cancel)
+            .await?;
+        assert_eq!(results.len(), 1);
+        let result: Option<anyhow::Result<bool>> = results.into_iter().next().unwrap();
+        match result {
+            None => anyhow::bail!("task_mgr shutdown requested"),
+            Some(Ok(b)) => Ok(Some(b)),
+            Some(Err(e)) => Err(e),
+        }
+    }
+
+    /// Evict a batch of layers.
+    ///
+    /// GenericRemoteStorage reference is required as a witness[^witness_article] for "remote storage is configured."
+    ///
+    /// [^witness_article]: https://willcrichton.net/rust-api-type-patterns/witnesses.html
+    pub async fn evict_layers(
+        &self,
+        _: &GenericRemoteStorage,
+        layers_to_evict: &[Arc<dyn PersistentLayer>],
+        cancel: CancellationToken,
+    ) -> anyhow::Result<Vec<Option<anyhow::Result<bool>>>> {
+        let Some(timeline) = self.timeline.upgrade() else { return Ok(vec![]) };
+        let remote_client = timeline.remote_client.clone().expect(
+            "GenericRemoteStorage is configured, so timeline must have RemoteTimelineClient",
+        );
+
+        self.evict_layer_batch(&remote_client, layers_to_evict, cancel)
+            .await
+    }
+
+    /// Evict multiple layers at once, continuing through errors.
+    ///
+    /// Try to evict the given `layers_to_evict` by
+    ///
+    /// 1. Replacing the given layer object in the layer map with a corresponding [`RemoteLayer`] object.
+    /// 2. Deleting the now unreferenced layer file from disk.
+    ///
+    /// The `remote_client` should be this timeline's `self.remote_client`.
+    /// We make the caller provide it so that they are responsible for handling the case
+    /// where someone wants to evict the layer but no remote storage is configured.
+    ///
+    /// Returns either `Err()` or `Ok(results)` where `results.len() == layers_to_evict.len()`.
+    /// If `Err()` is returned, no eviction was attempted.
+    /// Each position of `Ok(results)` corresponds to the layer in `layers_to_evict`.
+    /// Meaning of each `result[i]`:
+    /// - `Some(Err(...))` if layer replacement failed for an unexpected reason
+    /// - `Some(Ok(true))` if everything went well.
+    /// - `Some(Ok(false))` if there was an expected reason why the layer could not be replaced, e.g.:
+    ///    - evictee was not yet downloaded
+    ///    - replacement failed for an expectable reason (e.g., layer removed by GC before we grabbed all locks)
+    /// - `None` if no eviction attempt was made for the layer because `cancel.is_cancelled() == true`.
+    async fn evict_layer_batch(
+        &self,
+        remote_client: &Arc<RemoteTimelineClient>,
+        layers_to_evict: &[Arc<dyn PersistentLayer>],
+        cancel: CancellationToken,
+    ) -> anyhow::Result<Vec<Option<anyhow::Result<bool>>>> {
+        // ensure that the layers have finished uploading
+        // (don't hold the layer_removal_cs while we do it, we're not removing anything yet)
+        remote_client
+            .wait_completion()
+            .await
+            .context("wait for layer upload ops to complete")?;
+
+        // now lock out layer removal (compaction, gc, timeline deletion)
+        let layer_removal_guard = self.delete_guard().await;
+
+        {
+            // to avoid racing with detach and delete_timeline
+            let Some(timeline) = self.timeline.upgrade() else { return Ok(vec![]) };
+            let state = timeline.current_state();
+            anyhow::ensure!(
+                state == TimelineState::Active,
+                "timeline is not active but {state:?}"
+            );
+        }
+
+        let mut results = Vec::with_capacity(layers_to_evict.len());
+
+        for l in layers_to_evict.iter() {
+            let res = if cancel.is_cancelled() {
+                None
+            } else {
+                Some(self.evict_layer_batch_impl(&layer_removal_guard, l))
+            };
+            results.push(res);
+        }
+
+        // commit the updates & release locks
+        drop(layer_removal_guard);
+
+        assert_eq!(results.len(), layers_to_evict.len());
+        Ok(results)
+    }
+
+    fn evict_layer_batch_impl(
+        &self,
+        _layer_removal_cs: &DeleteGuard,
+        local_layer: &Arc<dyn PersistentLayer>,
+    ) -> anyhow::Result<bool> {
+        if local_layer.is_remote_layer() {
+            // TODO(issue #3851): consider returning an err here instead of false,
+            // which is the same out the match later
+            return Ok(false);
+        }
+
+        let layer_file_size = local_layer.file_size();
+
+        let local_layer_mtime = local_layer
+            .local_path()
+            .expect("local layer should have a local path")
+            .metadata()
+            .context("get local layer file stat")?
+            .modified()
+            .context("get mtime of layer file")?;
+        let local_layer_residence_duration =
+            match SystemTime::now().duration_since(local_layer_mtime) {
+                Err(e) => {
+                    warn!("layer mtime is in the future: {}", e);
+                    None
+                }
+                Ok(delta) => Some(delta),
+            };
+
+        let layer_metadata = LayerFileMetadata::new(layer_file_size);
+
+        let new_remote_layer = Arc::new(match local_layer.filename() {
+            LayerFileName::Image(image_name) => RemoteLayer::new_img(
+                self.tenant_id,
+                self.timeline_id,
+                &image_name,
+                &layer_metadata,
+                local_layer
+                    .access_stats()
+                    .clone_for_residence_change(LayerResidenceStatus::Evicted),
+            ),
+            LayerFileName::Delta(delta_name) => RemoteLayer::new_delta(
+                self.tenant_id,
+                self.timeline_id,
+                &delta_name,
+                &layer_metadata,
+                local_layer
+                    .access_stats()
+                    .clone_for_residence_change(LayerResidenceStatus::Evicted),
+            ),
+        });
+
+        assert_eq!(local_layer.layer_desc(), new_remote_layer.layer_desc());
+
+        let succeed = match self.replace_and_verify(local_layer.clone(), new_remote_layer) {
+            Ok(()) => {
+                if let Err(e) = local_layer.delete_resident_layer_file() {
+                    error!("failed to remove layer file on evict after replacement: {e:#?}");
+                }
+                // Always decrement the physical size gauge, even if we failed to delete the file.
+                // Rationale: we already replaced the layer with a remote layer in the layer map,
+                // and any subsequent download_remote_layer will
+                // 1. overwrite the file on disk and
+                // 2. add the downloaded size to the resident size gauge.
+                //
+                // If there is no re-download, and we restart the pageserver, then load_layer_map
+                // will treat the file as a local layer again, count it towards resident size,
+                // and it'll be like the layer removal never happened.
+                // The bump in resident size is perhaps unexpected but overall a robust behavior.
+
+                let Some(timeline) = self.timeline.upgrade() else { return Ok(false) };
+
+                timeline
+                    .metrics
+                    .resident_physical_size_gauge
+                    .sub(layer_file_size);
+
+                timeline.metrics.evictions.inc();
+
+                if let Some(delta) = local_layer_residence_duration {
+                    timeline
+                        .metrics
+                        .evictions_with_low_residence_duration
+                        .read()
+                        .unwrap()
+                        .observe(delta);
+                    info!(layer=%local_layer.short_id(), residence_millis=delta.as_millis(), "evicted layer after known residence period");
+                } else {
+                    info!(layer=%local_layer.short_id(), "evicted layer after unknown residence period");
+                }
+
+                true
+            }
+            Err(err) => {
+                if cfg!(debug_assertions) {
+                    error!(evicted=?local_layer, "failed to replace: {err}");
+                } else {
+                    error!(evicted=?local_layer, "failed to replace: {err}");
+                }
+                false
+            }
+        };
+
+        Ok(succeed)
+    }
 }

pageserver/src/tenant/layer_map.rs

@@ -66,7 +66,7 @@ use super::storage_layer::PersistentLayerDesc;
 ///
 /// LayerMap tracks what layers exist on a timeline.
 ///
-#[derive(Default, Clone)]
+#[derive(Default)]
 pub struct LayerMap {
     //
     // 'open_layer' holds the current InMemoryLayer that is accepting new

pageserver/src/tenant/layer_map/historic_layer_coverage.rs

@@ -72,7 +72,6 @@ impl From<&PersistentLayerDesc> for LayerKey {
 /// 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 +425,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>,

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/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,16 @@ 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(
+        layer_map_lock_held_witness: &BatchedUpdates<'_>,
+        status: LayerResidenceStatus,
+    ) -> Self {
         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
     }
 
@@ -187,6 +196,7 @@ impl LayerAccessStats {
     /// See [`record_residence_event`] for why you need to do this while holding the layer map lock.
     pub(crate) fn clone_for_residence_change(
         &self,
+        layer_map_lock_held_witness: &BatchedUpdates<'_>,
         new_status: LayerResidenceStatus,
     ) -> LayerAccessStats {
         let clone = {
@@ -194,7 +204,11 @@ impl LayerAccessStats {
             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
     }
 
@@ -214,6 +228,7 @@ impl LayerAccessStats {
     ///
     pub(crate) fn record_residence_event(
         &self,
+        _layer_map_lock_held_witness: &BatchedUpdates<'_>,
         status: LayerResidenceStatus,
         reason: LayerResidenceEventReason,
     ) {

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};
@@ -219,6 +220,7 @@ impl RemoteLayer {
     /// Create a Layer struct representing this layer, after it has been downloaded.
     pub fn create_downloaded_layer(
         &self,
+        layer_map_lock_held_witness: &BatchedUpdates<'_>,
         conf: &'static PageServerConf,
         file_size: u64,
     ) -> Arc<dyn PersistentLayer> {
@@ -230,8 +232,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 +245,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,7 +197,8 @@ 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 layers = self.layer_mgr.read();
+            let guard = self.layers.read().await;
+            let (layers, _) = &*guard;
             let mut candidates = Vec::new();
             for hist_layer in layers.iter_historic_layers() {
                 let hist_layer = self.lcache.get_from_desc(&hist_layer);
@@ -253,6 +254,7 @@ impl Timeline {
         };
 
         let results = match self
+            .lcache
             .evict_layer_batch(remote_client, &candidates[..], cancel.clone())
             .await
         {