From 6901bf465e83e7d191e1cf07627e9cbe53b458ad Mon Sep 17 00:00:00 2001
From: John Spray <john@neon.tech>
Date: Fri, 10 Nov 2023 11:05:12 +0000
Subject: [PATCH] Use secondary mode bits in neon_local migration

---
 control_plane/src/pageserver.rs       | 29 ++++++++++++++++++++------
 control_plane/src/tenant_migration.rs | 30 +++++++++++++++++++--------
 2 files changed, 44 insertions(+), 15 deletions(-)

diff --git a/control_plane/src/pageserver.rs b/control_plane/src/pageserver.rs
index 3037665692..ebef32a4e7 100644
--- a/control_plane/src/pageserver.rs
+++ b/control_plane/src/pageserver.rs
@@ -11,6 +11,7 @@ use std::io::{BufReader, Write};
 use std::num::NonZeroU64;
 use std::path::PathBuf;
 use std::process::{Child, Command};
+use std::time::Duration;
 use std::{io, result};
 
 use anyhow::{bail, Context};
@@ -531,17 +532,33 @@ impl PageServerNode {
         &self,
         tenant_id: TenantId,
         config: LocationConfig,
+        flush_ms: Option<Duration>,
     ) -> anyhow::Result<()> {
         let req_body = TenantLocationConfigRequest { tenant_id, config };
 
+        let path = format!(
+            "{}/tenant/{}/location_config",
+            self.http_base_url, tenant_id
+        );
+        let path = if let Some(flush_ms) = flush_ms {
+            format!("{}?flush_ms={}", path, flush_ms.as_millis())
+        } else {
+            path
+        };
+
+        self.http_request(Method::PUT, path)?
+            .json(&req_body)
+            .send()?
+            .error_from_body()?;
+
+        Ok(())
+    }
+
+    pub fn secondary_download(&self, tenant_id: TenantId) -> anyhow::Result<()> {
         self.http_request(
-            Method::PUT,
-            format!(
-                "{}/tenant/{}/location_config",
-                self.http_base_url, tenant_id
-            ),
+            Method::POST,
+            format!("{}/secondary/{}/download", self.http_base_url, tenant_id),
         )?
-        .json(&req_body)
         .send()?
         .error_from_body()?;
 
diff --git a/control_plane/src/tenant_migration.rs b/control_plane/src/tenant_migration.rs
index d28d1f9fe8..a13b7fb0b4 100644
--- a/control_plane/src/tenant_migration.rs
+++ b/control_plane/src/tenant_migration.rs
@@ -41,9 +41,9 @@ fn await_lsn(
     loop {
         let latest = match get_lsns(tenant_id, pageserver) {
             Ok(l) => l,
-            Err(e) => {
+            Err(_e) => {
                 println!(
-                    "🕑 Can't get LSNs on pageserver {} yet, waiting ({e})",
+                    "🕑 Waiting for pageserver {} to activate...",
                     pageserver.conf.id
                 );
                 std::thread::sleep(Duration::from_millis(500));
@@ -90,10 +90,10 @@ pub fn migrate_tenant(
     tenant_id: TenantId,
     dest_ps: PageServerNode,
 ) -> anyhow::Result<()> {
-    // Get a new generation
+    println!("🤔 Checking existing status...");
     let attachment_service = AttachmentService::from_env(env);
-
     let previous = attachment_service.inspect(tenant_id)?;
+
     let mut baseline_lsns = None;
     if let Some((generation, origin_ps_id)) = &previous {
         let origin_ps = PageServerNode::from_env(env, env.get_pageserver_conf(*origin_ps_id)?);
@@ -107,7 +107,7 @@ pub fn migrate_tenant(
                 secondary_conf: None,
                 tenant_conf: TenantConfig::default(),
             };
-            dest_ps.location_config(tenant_id, dest_conf)?;
+            dest_ps.location_config(tenant_id, dest_conf, None)?;
             println!("✅ Migration complete");
             return Ok(());
         }
@@ -120,11 +120,23 @@ pub fn migrate_tenant(
             secondary_conf: None,
             tenant_conf: TenantConfig::default(),
         };
-        origin_ps.location_config(tenant_id, stale_conf)?;
+        origin_ps.location_config(tenant_id, stale_conf, Some(Duration::from_secs(10)))?;
 
         baseline_lsns = Some(get_lsns(tenant_id, &origin_ps)?);
     }
 
+    println!(
+        "🔁 Downloading latest layers to destination pageserver {}",
+        dest_ps.conf.id
+    );
+    match dest_ps.secondary_download(tenant_id) {
+        Ok(()) => {}
+        Err(_) => {
+            println!("  (skipping, destination wasn't in secondary mode)")
+        }
+    }
+
+    // Get a new generation
     let gen = attachment_service.attach_hook(tenant_id, dest_ps.conf.id)?;
     let dest_conf = LocationConfig {
         mode: LocationConfigMode::AttachedMulti,
@@ -134,7 +146,7 @@ pub fn migrate_tenant(
     };
 
     println!("🔁 Attaching to pageserver {}", dest_ps.conf.id);
-    dest_ps.location_config(tenant_id, dest_conf)?;
+    dest_ps.location_config(tenant_id, dest_conf, None)?;
 
     if let Some(baseline) = baseline_lsns {
         println!("🕑 Waiting for LSN to catch up...");
@@ -181,7 +193,7 @@ pub fn migrate_tenant(
             "💤 Switching to secondary mode on pageserver {}",
             other_ps.conf.id
         );
-        other_ps.location_config(tenant_id, secondary_conf)?;
+        other_ps.location_config(tenant_id, secondary_conf, None)?;
     }
 
     println!(
@@ -194,7 +206,7 @@ pub fn migrate_tenant(
         secondary_conf: None,
         tenant_conf: TenantConfig::default(),
     };
-    dest_ps.location_config(tenant_id, dest_conf)?;
+    dest_ps.location_config(tenant_id, dest_conf, None)?;
 
     println!("✅ Migration complete");