From 5cf597044d88c6280dd1d4e5c1e8340c4229a755 Mon Sep 17 00:00:00 2001
From: bojanserafimov <bojan.serafimov7@gmail.com>
Date: Mon, 11 Jul 2022 10:31:14 -0400
Subject: [PATCH] Allow prev_lsn hint for fullbackup (#2052)

---
 pageserver/src/basebackup.rs   | 15 +++++++--
 pageserver/src/page_service.rs | 58 +++++++++++++++++++++++++++-------
 2 files changed, 60 insertions(+), 13 deletions(-)

diff --git a/pageserver/src/basebackup.rs b/pageserver/src/basebackup.rs
index ed300b3360..3ec1ec9243 100644
--- a/pageserver/src/basebackup.rs
+++ b/pageserver/src/basebackup.rs
@@ -60,6 +60,7 @@ where
         write: W,
         timeline: &'a Arc<DatadirTimelineImpl>,
         req_lsn: Option<Lsn>,
+        prev_lsn: Option<Lsn>,
         full_backup: bool,
     ) -> Result<Basebackup<'a, W>> {
         // Compute postgres doesn't have any previous WAL files, but the first
@@ -96,16 +97,26 @@ where
             (end_of_timeline.prev, end_of_timeline.last)
         };
 
+        // Consolidate the derived and the provided prev_lsn values
+        let prev_lsn = if let Some(provided_prev_lsn) = prev_lsn {
+            if backup_prev != Lsn(0) {
+                ensure!(backup_prev == provided_prev_lsn)
+            }
+            provided_prev_lsn
+        } else {
+            backup_prev
+        };
+
         info!(
             "taking basebackup lsn={}, prev_lsn={} (full_backup={})",
-            backup_lsn, backup_prev, full_backup
+            backup_lsn, prev_lsn, full_backup
         );
 
         Ok(Basebackup {
             ar: Builder::new(AbortableWrite::new(write)),
             timeline,
             lsn: backup_lsn,
-            prev_record_lsn: backup_prev,
+            prev_record_lsn: prev_lsn,
             full_backup,
             finished: false,
         })
diff --git a/pageserver/src/page_service.rs b/pageserver/src/page_service.rs
index 973a631d23..078edc5c9f 100644
--- a/pageserver/src/page_service.rs
+++ b/pageserver/src/page_service.rs
@@ -772,6 +772,7 @@ impl PageServerHandler {
         pgb: &mut PostgresBackend,
         timelineid: ZTimelineId,
         lsn: Option<Lsn>,
+        prev_lsn: Option<Lsn>,
         tenantid: ZTenantId,
         full_backup: bool,
     ) -> anyhow::Result<()> {
@@ -796,7 +797,8 @@ impl PageServerHandler {
         {
             let mut writer = CopyDataSink { pgb };
 
-            let basebackup = basebackup::Basebackup::new(&mut writer, &timeline, lsn, full_backup)?;
+            let basebackup =
+                basebackup::Basebackup::new(&mut writer, &timeline, lsn, prev_lsn, full_backup)?;
             span.record("lsn", &basebackup.lsn.to_string().as_str());
             basebackup.send_tarball()?;
         }
@@ -899,33 +901,67 @@ impl postgres_backend::Handler for PageServerHandler {
             };
 
             // Check that the timeline exists
-            self.handle_basebackup_request(pgb, timelineid, lsn, tenantid, false)?;
+            self.handle_basebackup_request(pgb, timelineid, lsn, None, tenantid, false)?;
             pgb.write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
         }
+        // return pair of prev_lsn and last_lsn
+        else if query_string.starts_with("get_last_record_rlsn ") {
+            let (_, params_raw) = query_string.split_at("get_last_record_rlsn ".len());
+            let params = params_raw.split_whitespace().collect::<Vec<_>>();
+
+            ensure!(
+                params.len() == 2,
+                "invalid param number for get_last_record_rlsn command"
+            );
+
+            let tenantid = ZTenantId::from_str(params[0])?;
+            let timelineid = ZTimelineId::from_str(params[1])?;
+
+            self.check_permission(Some(tenantid))?;
+            let timeline = tenant_mgr::get_local_timeline_with_load(tenantid, timelineid)
+                .context("Cannot load local timeline")?;
+
+            let end_of_timeline = timeline.tline.get_last_record_rlsn();
+
+            pgb.write_message_noflush(&BeMessage::RowDescription(&[
+                RowDescriptor::text_col(b"prev_lsn"),
+                RowDescriptor::text_col(b"last_lsn"),
+            ]))?
+            .write_message_noflush(&BeMessage::DataRow(&[
+                Some(end_of_timeline.prev.to_string().as_bytes()),
+                Some(end_of_timeline.last.to_string().as_bytes()),
+            ]))?
+            .write_message(&BeMessage::CommandComplete(b"SELECT 1"))?;
+        }
         // same as basebackup, but result includes relational data as well
         else if query_string.starts_with("fullbackup ") {
             let (_, params_raw) = query_string.split_at("fullbackup ".len());
             let params = params_raw.split_whitespace().collect::<Vec<_>>();
 
             ensure!(
-                params.len() == 3,
+                params.len() >= 2,
                 "invalid param number for fullbackup command"
             );
 
             let tenantid = ZTenantId::from_str(params[0])?;
             let timelineid = ZTimelineId::from_str(params[1])?;
 
+            // The caller is responsible for providing correct lsn and prev_lsn.
+            let lsn = if params.len() > 2 {
+                Some(Lsn::from_str(params[2])?)
+            } else {
+                None
+            };
+            let prev_lsn = if params.len() > 3 {
+                Some(Lsn::from_str(params[3])?)
+            } else {
+                None
+            };
+
             self.check_permission(Some(tenantid))?;
 
-            // Lsn is required for fullbackup, because otherwise we would not know
-            // at which lsn to upload this backup.
-            //
-            // The caller is responsible for providing a valid lsn
-            // and using it in the subsequent import.
-            let lsn = Some(Lsn::from_str(params[2])?);
-
             // Check that the timeline exists
-            self.handle_basebackup_request(pgb, timelineid, lsn, tenantid, true)?;
+            self.handle_basebackup_request(pgb, timelineid, lsn, prev_lsn, tenantid, true)?;
             pgb.write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
         } else if query_string.starts_with("import basebackup ") {
             // Import the `base` section (everything but the wal) of a basebackup.