layermap: add watermark tracking for lowest used version

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

pageserver/src/tenant.rs

@@ -86,8 +86,8 @@ pub mod block_io;
 pub mod disk_btree;
 pub(crate) mod ephemeral_file;
 pub mod layer_map;
-pub mod manifest;
 pub mod layer_map_mgr;
+pub mod manifest;
 
 pub mod metadata;
 mod par_fsync;

pageserver/src/tenant/layer_map_mgr.rs

@@ -19,20 +19,43 @@
 use anyhow::Result;
 use arc_swap::ArcSwap;
 use futures::Future;
-use std::sync::Arc;
+use std::{cmp::Reverse, collections::BinaryHeap, 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>,
+    layer_map: ArcSwap<LayerMapGuard>,
     state_lock: tokio::sync::Mutex<()>,
+    watermark: Arc<LayerMgrWatermark>,
+}
+
+/// A guard that holds a version of the layer map. When dropped, the version is released and the watermark will be updated.
+#[derive(Clone)]
+pub struct LayerMapGuard {
+    version: u64,
+    layer_map: LayerMap,
+    watermark: Arc<LayerMgrWatermark>,
+}
+
+impl std::ops::Deref for LayerMapGuard {
+    type Target = LayerMap;
+
+    fn deref(&self) -> &Self::Target {
+        &self.layer_map
+    }
+}
+
+impl Drop for LayerMapGuard {
+    fn drop(&mut self) {
+        self.watermark.release(self.version);
+    }
 }
 
 impl LayerMapMgr {
     /// Get the current state of the layer map.
-    pub fn read(&self) -> Arc<LayerMap> {
+    pub fn read(&self) -> Arc<LayerMapGuard> {
         // 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.
@@ -40,13 +63,23 @@ impl LayerMapMgr {
     }
 
     /// Clone the layer map for modification.
-    fn clone_for_write(&self, _state_lock_witness: &tokio::sync::MutexGuard<'_, ()>) -> LayerMap {
+    fn clone_for_write(
+        &self,
+        _state_lock_witness: &tokio::sync::MutexGuard<'_, ()>,
+    ) -> LayerMapGuard {
         (**self.layer_map.load()).clone()
     }
 
     pub fn new(layer_map: LayerMap) -> Self {
+        const INITIAL_VERSION: u64 = 0;
+        let watermark = Arc::new(LayerMgrWatermark::new(INITIAL_VERSION));
         Self {
-            layer_map: ArcSwap::new(Arc::new(layer_map)),
+            layer_map: ArcSwap::new(Arc::new(LayerMapGuard {
+                version: INITIAL_VERSION,
+                layer_map,
+                watermark: watermark.clone(),
+            })),
+            watermark,
             state_lock: tokio::sync::Mutex::new(()),
         }
     }
@@ -58,11 +91,71 @@ impl LayerMapMgr {
         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));
+        let mut guard = self.clone_for_write(&state_lock);
+        guard.version += 1;
+        let layer_map = std::mem::take(&mut guard.layer_map);
+        guard.layer_map = operation(layer_map).await?;
+        self.layer_map.store(Arc::new(guard));
         Ok(())
     }
+
+    pub fn lowest_version_in_use(&self) -> u64 {
+        self.watermark.lowest_version_in_use()
+    }
+}
+
+struct LayerMgrWatermarkCore {
+    lowest_version_in_use: u64,
+    versions_in_use: BinaryHeap<Reverse<u64>>,
+}
+
+/// Computes the lowest version used by any read thread. Once a version is not used any more,
+/// we can remove all layers that are intended to be removed in that version.
+struct LayerMgrWatermark {
+    core: std::sync::Mutex<LayerMgrWatermarkCore>,
+}
+
+impl LayerMgrWatermark {
+    fn new(initial_version: u64) -> Self {
+        Self {
+            core: std::sync::Mutex::new(LayerMgrWatermarkCore {
+                lowest_version_in_use: initial_version,
+                versions_in_use: BinaryHeap::new(),
+            }),
+        }
+    }
+
+    fn lowest_version_in_use(&self) -> u64 {
+        self.core.lock().unwrap().lowest_version_in_use
+    }
+
+    fn release(&self, version: u64) {
+        let mut core = self.core.lock().unwrap();
+        match version.cmp(&core.lowest_version_in_use) {
+            std::cmp::Ordering::Less => {
+                if cfg!(debug_assertions) {
+                    panic!("release a version lower than the lowest version in use.")
+                }
+            }
+            std::cmp::Ordering::Equal => {
+                // Find the next version in use.
+                let mut current_version = version + 1;
+                while let Some(Reverse(next_version)) = core.versions_in_use.peek() {
+                    if *next_version == current_version {
+                        current_version += 1;
+                        core.versions_in_use.pop();
+                    } else {
+                        break;
+                    }
+                }
+                core.lowest_version_in_use = current_version;
+            }
+            std::cmp::Ordering::Greater => {
+                // This version is in use. Add it to the heap.
+                core.versions_in_use.push(Reverse(version));
+            }
+        }
+    }
 }
 
 #[cfg(test)]
@@ -141,6 +234,30 @@ mod tests {
                 .key_range,
             Key::from_i128(1)..Key::from_i128(2)
         );
+
+        assert_eq!(mgr.lowest_version_in_use(), 1);
+        drop(ref_1);
+        drop(ref_2);
+        assert_eq!(mgr.lowest_version_in_use(), 2);
+        mgr.update(|map| async move { Ok(map) }).await?;
+        assert_eq!(mgr.lowest_version_in_use(), 3);
+
         Ok(())
     }
+
+    #[test]
+    fn test_watermark() {
+        let watermark = LayerMgrWatermark::new(0);
+        assert_eq!(watermark.lowest_version_in_use(), 0);
+        watermark.release(0);
+        assert_eq!(watermark.lowest_version_in_use(), 1);
+        watermark.release(1);
+        assert_eq!(watermark.lowest_version_in_use(), 2);
+        watermark.release(3);
+        watermark.release(4);
+        watermark.release(5);
+        assert_eq!(watermark.lowest_version_in_use(), 2);
+        watermark.release(2);
+        assert_eq!(watermark.lowest_version_in_use(), 6);
+    }
 }