From 0e49d748b828eb5ec0ac8323f263a533b6a8d308 Mon Sep 17 00:00:00 2001
From: Konstantin Knizhnik <knizhnik@garret.ru>
Date: Sat, 30 Apr 2022 00:41:04 +0300
Subject: [PATCH] Fix bug in dictinary creation

---
 pageserver/src/layered_repository.rs                | 12 ++++++++----
 pageserver/src/layered_repository/inmemory_layer.rs |  7 ++++---
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/pageserver/src/layered_repository.rs b/pageserver/src/layered_repository.rs
index f7731afd8d..0d1d91d139 100644
--- a/pageserver/src/layered_repository.rs
+++ b/pageserver/src/layered_repository.rs
@@ -1989,13 +1989,17 @@ impl LayeredTimeline {
                     prefetched.push((key, lsn, value));
 
                     while let Some(y) = all_values_iter.next() {
-                        let (key, lsn, value) = y?;
-                        samples.push(Value::ser(&value)?);
-                        prefetched.push((key, lsn, value));
-                        if samples.len() == config::ZSTD_MAX_SAMPLES {
+                        if let Ok((key, lsn, value)) = y {
+                            samples.push(Value::ser(&value)?);
+                            prefetched.push((key, lsn, value));
+                            if samples.len() == config::ZSTD_MAX_SAMPLES {
+                                break;
+                            }
+                        } else {
                             break;
                         }
                     }
+
                     let dictionary = if samples.len() >= config::ZSTD_MIN_SAMPLES {
                         zstd::dict::from_samples(&samples, config::ZSTD_MAX_DICTIONARY_SIZE)?
                     } else {
diff --git a/pageserver/src/layered_repository/inmemory_layer.rs b/pageserver/src/layered_repository/inmemory_layer.rs
index 5026f89c1b..a65c249fbd 100644
--- a/pageserver/src/layered_repository/inmemory_layer.rs
+++ b/pageserver/src/layered_repository/inmemory_layer.rs
@@ -327,17 +327,18 @@ impl InMemoryLayer {
 
         let mut cursor = inner.file.block_cursor();
 
-        // First learn dictionary */
-        for (_key, vec_map) in keys.iter() {
+        // First learn dictionary
+        'train: for (_key, vec_map) in keys.iter() {
             // Write all page versions
             for (_lsn, pos) in vec_map.as_slice() {
                 cursor.read_blob_into_buf(*pos, &mut buf)?;
                 samples.push(buf.clone());
                 if samples.len() == config::ZSTD_MAX_SAMPLES {
-                    break;
+                    break 'train;
                 }
             }
         }
+
         let dictionary = if samples.len() >= config::ZSTD_MIN_SAMPLES {
             zstd::dict::from_samples(&samples, config::ZSTD_MAX_DICTIONARY_SIZE)?
         } else {