From a09c933de3de7a57b23f814e013dfa0fbcc1356a Mon Sep 17 00:00:00 2001
From: "Alex Chi Z." <4198311+skyzh@users.noreply.github.com>
Date: Tue, 8 Apr 2025 12:08:44 -0400
Subject: [PATCH] test(pageserver): add conditional append test record (#11476)

## Problem

For future gc-compaction tests when we support
https://github.com/neondatabase/neon/issues/10395

## Summary of changes

Add a new type of neon test WAL record that is conditionally applied
(i.e., only when image == the specified value). We can use this to mock
the situation where we lose some records in the middle, firing an error,
and see how gc-compaction reacts to it.

Signed-off-by: Alex Chi Z <chi@neon.tech>
---
 libs/pageserver_api/src/record.rs    | 15 ++++++++++++++
 pageserver/src/tenant.rs             | 29 ++++++++++++++++++++++++++--
 pageserver/src/walredo/apply_neon.rs |  8 ++++++++
 3 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/libs/pageserver_api/src/record.rs b/libs/pageserver_api/src/record.rs
index fda504a26e..73516c5220 100644
--- a/libs/pageserver_api/src/record.rs
+++ b/libs/pageserver_api/src/record.rs
@@ -58,6 +58,8 @@ pub enum NeonWalRecord {
         /// to true. This record does not need the history WALs to reconstruct. See [`NeonWalRecord::will_init`] and
         /// its references in `timeline.rs`.
         will_init: bool,
+        /// Only append the record if the current image is the same as the one specified in this field.
+        only_if: Option<String>,
     },
 }
 
@@ -81,6 +83,17 @@ impl NeonWalRecord {
             append: s.as_ref().to_string(),
             clear: false,
             will_init: false,
+            only_if: None,
+        }
+    }
+
+    #[cfg(feature = "testing")]
+    pub fn wal_append_conditional(s: impl AsRef<str>, only_if: impl AsRef<str>) -> Self {
+        Self::Test {
+            append: s.as_ref().to_string(),
+            clear: false,
+            will_init: false,
+            only_if: Some(only_if.as_ref().to_string()),
         }
     }
 
@@ -90,6 +103,7 @@ impl NeonWalRecord {
             append: s.as_ref().to_string(),
             clear: true,
             will_init: false,
+            only_if: None,
         }
     }
 
@@ -99,6 +113,7 @@ impl NeonWalRecord {
             append: s.as_ref().to_string(),
             clear: true,
             will_init: true,
+            only_if: None,
         }
     }
 }
diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs
index 0c399d4c91..1bfc51d5c8 100644
--- a/pageserver/src/tenant.rs
+++ b/pageserver/src/tenant.rs
@@ -8733,6 +8733,21 @@ mod tests {
                 Lsn(0x20),
                 Value::WalRecord(NeonWalRecord::wal_init("i")),
             ),
+            (
+                get_key(4),
+                Lsn(0x30),
+                Value::WalRecord(NeonWalRecord::wal_append_conditional("j", "i")),
+            ),
+            (
+                get_key(5),
+                Lsn(0x20),
+                Value::WalRecord(NeonWalRecord::wal_init("1")),
+            ),
+            (
+                get_key(5),
+                Lsn(0x30),
+                Value::WalRecord(NeonWalRecord::wal_append_conditional("j", "2")),
+            ),
         ];
         let image1 = vec![(get_key(1), "0x10".into())];
 
@@ -8763,8 +8778,18 @@ mod tests {
 
         // Need to remove the limit of "Neon WAL redo requires base image".
 
-        // assert_eq!(tline.get(get_key(3), Lsn(0x50), &ctx).await?, Bytes::new());
-        // assert_eq!(tline.get(get_key(4), Lsn(0x50), &ctx).await?, Bytes::new());
+        assert_eq!(
+            tline.get(get_key(3), Lsn(0x50), &ctx).await?,
+            Bytes::from_static(b"c")
+        );
+        assert_eq!(
+            tline.get(get_key(4), Lsn(0x50), &ctx).await?,
+            Bytes::from_static(b"ij")
+        );
+
+        // Manual testing required: currently, read errors will panic the process in debug mode. So we
+        // cannot enable this assertion in the unit test.
+        // assert!(tline.get(get_key(5), Lsn(0x50), &ctx).await.is_err());
 
         Ok(())
     }
diff --git a/pageserver/src/walredo/apply_neon.rs b/pageserver/src/walredo/apply_neon.rs
index 61ae1eb970..a3840f1f6f 100644
--- a/pageserver/src/walredo/apply_neon.rs
+++ b/pageserver/src/walredo/apply_neon.rs
@@ -276,6 +276,7 @@ pub(crate) fn apply_in_neon(
             append,
             clear,
             will_init,
+            only_if,
         } => {
             use bytes::BufMut;
             if *will_init {
@@ -288,6 +289,13 @@ pub(crate) fn apply_in_neon(
             if *clear {
                 page.clear();
             }
+            if let Some(only_if) = only_if {
+                if page != only_if.as_bytes() {
+                    return Err(anyhow::anyhow!(
+                        "the current image does not match the expected image, cannot append"
+                    ));
+                }
+            }
             page.put_slice(append.as_bytes());
         }
     }