From 8c933a96864d7dd9d186646bc88b950aa19da403 Mon Sep 17 00:00:00 2001
From: Christian Schwarz <christian@neon.tech>
Date: Mon, 28 Aug 2023 18:35:32 +0200
Subject: [PATCH] page_cache: two-level immutable_page_map

This change alone makes things strictly worse, but on top of this, we'll
be able to rewrite drop_buffers_for_immutable to remove pages based on
what's in the `drop_file_id`'s immutable_page_map second-level hash map.
That's better than what we do currently, i.e., iterating over all the
slots.
---
 pageserver/src/page_cache.rs | 39 +++++++++++++++++++++++++++++-------
 1 file changed, 32 insertions(+), 7 deletions(-)

diff --git a/pageserver/src/page_cache.rs b/pageserver/src/page_cache.rs
index e1e696ddad..c9873fbd45 100644
--- a/pageserver/src/page_cache.rs
+++ b/pageserver/src/page_cache.rs
@@ -205,6 +205,8 @@ impl Slot {
     }
 }
 
+type ImmutablePageMap = HashMap<FileId, RwLock<HashMap<u32, usize>>>;
+
 pub struct PageCache {
     /// This contains the mapping from the cache key to buffer slot that currently
     /// contains the page, if any.
@@ -217,7 +219,7 @@ pub struct PageCache {
     /// can have a separate mapping map, next to this field.
     materialized_page_map: RwLock<HashMap<MaterializedPageHashKey, Vec<Version>>>,
 
-    immutable_page_map: RwLock<HashMap<(FileId, u32), usize>>,
+    immutable_page_map: RwLock<ImmutablePageMap>,
 
     /// The actual buffers with their metadata.
     slots: Box<[Slot]>,
@@ -657,7 +659,9 @@ impl PageCache {
             }
             CacheKey::ImmutableFilePage { file_id, blkno } => {
                 let map = self.immutable_page_map.read().unwrap();
-                Some(*map.get(&(*file_id, *blkno))?)
+                let block_nos = map.get(file_id)?;
+                let block_nos = block_nos.read().unwrap();
+                block_nos.get(blkno).copied()
             }
         }
     }
@@ -680,7 +684,9 @@ impl PageCache {
             }
             CacheKey::ImmutableFilePage { file_id, blkno } => {
                 let map = self.immutable_page_map.read().unwrap();
-                Some(*map.get(&(*file_id, *blkno))?)
+                let block_nos = map.get(file_id)?;
+                let block_nos = block_nos.read().unwrap();
+                block_nos.get(blkno).copied()
             }
         }
     }
@@ -712,10 +718,27 @@ impl PageCache {
                 }
             }
             CacheKey::ImmutableFilePage { file_id, blkno } => {
-                let mut map = self.immutable_page_map.write().unwrap();
-                map.remove(&(*file_id, *blkno))
-                    .expect("could not find old key in mapping");
+                let map = self.immutable_page_map.read().unwrap();
+                let block_nos = map.get(file_id).expect("could not find file_id in mapping");
+                let mut block_nos = block_nos.write().unwrap();
+                block_nos
+                    .remove(blkno)
+                    .expect("could not find blkno in mapping");
                 self.size_metrics.current_bytes_immutable.sub_page_sz(1);
+                if block_nos.is_empty() {
+                    drop(block_nos);
+                    // re-lock map in write mode
+                    drop(map);
+                    let mut map = self.immutable_page_map.write().unwrap();
+                    let Some(block_nos_rwl) = map.get(file_id) else {
+                        return;
+                    };
+                    let block_nos = block_nos_rwl.read().unwrap();
+                    if block_nos.is_empty() {
+                        drop(block_nos);
+                        map.remove(file_id);
+                    }
+                }
             }
         }
     }
@@ -753,7 +776,9 @@ impl PageCache {
 
             CacheKey::ImmutableFilePage { file_id, blkno } => {
                 let mut map = self.immutable_page_map.write().unwrap();
-                match map.entry((*file_id, *blkno)) {
+                let block_nos = map.entry(*file_id).or_default();
+                let mut block_nos = block_nos.write().unwrap();
+                match block_nos.entry(*blkno) {
                     Entry::Occupied(entry) => Some(*entry.get()),
                     Entry::Vacant(entry) => {
                         entry.insert(slot_idx);