From 51aa53ab90be60450d0a78f2a80c3169e7ad24cb Mon Sep 17 00:00:00 2001
From: Konstantin Knizhnik <knizhnik@garret.ru>
Date: Thu, 6 Oct 2022 16:06:39 +0300
Subject: [PATCH] Check that image contains search key in LayerMap::search

---
 pageserver/src/tenant/delta_layer.rs    | 47 +++++++++++++++++++------
 pageserver/src/tenant/image_layer.rs    |  4 +++
 pageserver/src/tenant/inmemory_layer.rs |  5 +++
 pageserver/src/tenant/layer_map.rs      |  3 ++
 pageserver/src/tenant/storage_layer.rs  |  3 ++
 5 files changed, 51 insertions(+), 11 deletions(-)

diff --git a/pageserver/src/tenant/delta_layer.rs b/pageserver/src/tenant/delta_layer.rs
index 155f58c560..8dc8ca7743 100644
--- a/pageserver/src/tenant/delta_layer.rs
+++ b/pageserver/src/tenant/delta_layer.rs
@@ -424,6 +424,30 @@ impl Layer for DeltaLayer {
 
         Ok(())
     }
+
+    fn contains(&self, key: &Key) -> Result<bool> {
+        // Open the file and lock the metadata in memory
+        let inner = self.load()?;
+
+        // Scan the page versions backwards, starting from `lsn`.
+        let file = inner.file.as_ref().unwrap();
+        let reader = DiskBtreeReader::<_, DELTA_KEY_SIZE>::new(
+            inner.index_start_blk,
+            inner.index_root_blk,
+            file,
+        );
+        let search_key = DeltaKey::from_key_lsn(&key, Lsn(0));
+        let mut found = false;
+        reader.visit(
+            &search_key.0,
+            VisitDirection::Forwards,
+            |delta_key, _val| {
+                found = DeltaKey::extract_key_from_buf(delta_key) == *key;
+                false
+            },
+        )?;
+        Ok(found)
+    }
 }
 
 impl DeltaLayer {
@@ -899,32 +923,33 @@ impl<'a> DeltaKeyIter {
         );
 
         let mut all_keys: Vec<(DeltaKey, u64)> = Vec::new();
+        let mut last_pos = 0u64;
+        let mut last_delta: Option<DeltaKey> = None;
         tree_reader.visit(
             &[0u8; DELTA_KEY_SIZE],
             VisitDirection::Forwards,
             |key, value| {
                 let blob_ref = BlobRef(value);
                 if !blob_ref.is_image() || !skip_images {
-                    let delta_key = DeltaKey::from_slice(key);
+                    let next_delta = DeltaKey::from_slice(key);
                     let pos = blob_ref.pos();
-                    if let Some(last) = all_keys.last_mut() {
-                        if last.0.key() == delta_key.key() {
+                    if let Some(prev_delta) = last_delta.take() {
+                        if prev_delta.key() == next_delta.key() {
+                            last_delta = Some(next_delta);
                             return true;
-                        } else {
-                            // subtract offset of new key BLOB and first blob of this key
-                            // to get total size if values associated with this key
-                            let first_pos = last.1;
-                            last.1 = pos - first_pos;
                         }
+                        all_keys.push((prev_delta, pos - last_pos));
                     }
-                    all_keys.push((delta_key, pos));
+                    last_delta = Some(next_delta);
+                    last_pos = pos;
                 }
                 true
             },
         )?;
-        if let Some(last) = all_keys.last_mut() {
+        if let Some(prev_delta) = last_delta.take() {
             // Last key occupies all space till end of layer
-            last.1 = std::fs::metadata(&file.file.path)?.len() - last.1;
+            let file_size = std::fs::metadata(&file.file.path)?.len();
+            all_keys.push((prev_delta, file_size - last_pos));
         }
         let iter = DeltaKeyIter {
             all_keys,
diff --git a/pageserver/src/tenant/image_layer.rs b/pageserver/src/tenant/image_layer.rs
index 92bf022fee..b68de939ee 100644
--- a/pageserver/src/tenant/image_layer.rs
+++ b/pageserver/src/tenant/image_layer.rs
@@ -223,6 +223,10 @@ impl Layer for ImageLayer {
 
         Ok(())
     }
+
+    fn contains(&self, key: &Key) -> Result<bool> {
+        Ok(self.get_key_range().contains(key))
+    }
 }
 
 impl ImageLayer {
diff --git a/pageserver/src/tenant/inmemory_layer.rs b/pageserver/src/tenant/inmemory_layer.rs
index d645fc4710..356dcc5f5c 100644
--- a/pageserver/src/tenant/inmemory_layer.rs
+++ b/pageserver/src/tenant/inmemory_layer.rs
@@ -235,6 +235,11 @@ impl Layer for InMemoryLayer {
 
         Ok(())
     }
+
+    fn contains(&self, key: &Key) -> Result<bool> {
+        let inner = self.inner.read().unwrap();
+        Ok(inner.index.get(&key).is_some())
+    }
 }
 
 impl InMemoryLayer {
diff --git a/pageserver/src/tenant/layer_map.rs b/pageserver/src/tenant/layer_map.rs
index 369d02fd10..c271fb0b3f 100644
--- a/pageserver/src/tenant/layer_map.rs
+++ b/pageserver/src/tenant/layer_map.rs
@@ -276,6 +276,9 @@ impl LayerMap {
                 continue;
             }
             assert!(l.get_key_range().contains(&key));
+            if !l.contains(&key)? {
+                continue;
+            }
             if l.get_lsn_range().start >= end_lsn {
                 info!(
                     "Candidate delta layer {}..{} is too new for lsn {}",
diff --git a/pageserver/src/tenant/storage_layer.rs b/pageserver/src/tenant/storage_layer.rs
index 7f025a5aba..c01796cbfd 100644
--- a/pageserver/src/tenant/storage_layer.rs
+++ b/pageserver/src/tenant/storage_layer.rs
@@ -150,4 +150,7 @@ pub trait Layer: Send + Sync {
 
     /// Dump summary of the contents of the layer to stdout
     fn dump(&self, verbose: bool) -> Result<()>;
+
+    // Check if ayer contains particular key
+    fn contains(&self, key: &Key) -> Result<bool>;
 }