From ef90a6295eea4cd27a76ae4c1f5143557c861ade Mon Sep 17 00:00:00 2001
From: Andrey Rudenko <me@prepor.dev>
Date: Fri, 7 Jun 2024 13:42:18 +0200
Subject: [PATCH] feat(compute_ctl): add periodic `lease lsn` requests for
 static computes

---
 compute_tools/src/compute.rs  | 55 +++++++++++++++++++++++++++++++++++
 control_plane/src/endpoint.rs |  2 ++
 libs/compute_api/src/spec.rs  |  7 +++++
 3 files changed, 64 insertions(+)

diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs
index 40060f4117..77f641900f 100644
--- a/compute_tools/src/compute.rs
+++ b/compute_tools/src/compute.rs
@@ -1084,6 +1084,8 @@ impl ComputeNode {
         };
         info!(?metrics, "compute start finished");
 
+        start_lsn_lease_for_static(&compute_state);
+
         Ok(pg_process)
     }
 
@@ -1385,3 +1387,56 @@ pub fn forward_termination_signal() {
         kill(pg_pid, Signal::SIGQUIT).ok();
     }
 }
+
+// Start a thread that will periodically renew LSN lease for static compute
+fn start_lsn_lease_for_static(compute_state: &ComputeState) {
+    let spec = compute_state.pspec.as_ref().expect("spec must be set");
+
+    let lsn = match spec.spec.mode {
+        ComputeMode::Static(lsn) => lsn,
+        _ => return,
+    };
+
+    let conn_strings = spec.pageserver_connstr.split(',');
+
+    let configs = conn_strings
+        .map(|connstr| {
+            let mut config = postgres::Config::from_str(connstr).expect("invalid connstr");
+            if let Some(storage_auth_token) = &spec.storage_auth_token {
+                info!("Got storage auth token from spec file");
+                config.password(storage_auth_token.clone());
+            } else {
+                info!("Storage auth token not set");
+            }
+            config
+        })
+        .collect::<Vec<_>>();
+
+    let cmd = format!("lease lsn {} {} {} ", spec.tenant_id, spec.timeline_id, lsn);
+
+    let lsn_lease_interval = spec.spec.lsn_lease_interval;
+    thread::spawn(move || lsn_lease_loop(lsn_lease_interval, configs, cmd));
+}
+
+fn lsn_lease_loop(interval_secs: u64, configs: Vec<postgres::Config>, cmd: String) {
+    loop {
+        match lsn_lease_request(&configs, &cmd) {
+            Ok(_) => {
+                thread::sleep(tokio::time::Duration::from_secs(interval_secs));
+            }
+            Err(e) => {
+                error!("lsn_lease_request failed: {:#}", e);
+                thread::sleep(tokio::time::Duration::from_secs(10));
+            }
+        }
+    }
+}
+
+fn lsn_lease_request(configs: &[postgres::Config], cmd: &str) -> Result<()> {
+    info!("lsn_lease_request: {}", cmd);
+    for config in configs {
+        let mut client = config.connect(NoTls)?;
+        let _ = client.simple_query(&cmd)?;
+    }
+    Ok(())
+}
diff --git a/control_plane/src/endpoint.rs b/control_plane/src/endpoint.rs
index 20371e1cb8..254cf259ff 100644
--- a/control_plane/src/endpoint.rs
+++ b/control_plane/src/endpoint.rs
@@ -46,6 +46,7 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use anyhow::{anyhow, bail, Context, Result};
+use compute_api::spec::default_lsn_lease_interval;
 use compute_api::spec::Database;
 use compute_api::spec::PgIdent;
 use compute_api::spec::RemoteExtSpec;
@@ -593,6 +594,7 @@ impl Endpoint {
             pgbouncer_settings: None,
             shard_stripe_size: Some(shard_stripe_size),
             primary_is_running: None,
+            lsn_lease_interval: default_lsn_lease_interval(),
         };
         let spec_path = self.endpoint_path().join("spec.json");
         std::fs::write(spec_path, serde_json::to_string_pretty(&spec)?)?;
diff --git a/libs/compute_api/src/spec.rs b/libs/compute_api/src/spec.rs
index 1c4ee2089f..0f9bf6a6c6 100644
--- a/libs/compute_api/src/spec.rs
+++ b/libs/compute_api/src/spec.rs
@@ -102,6 +102,13 @@ pub struct ComputeSpec {
     // This is used to determine if replica should wait for
     // RUNNING_XACTS from primary or not.
     pub primary_is_running: Option<bool>,
+
+    #[serde(default = "default_lsn_lease_interval")]
+    pub lsn_lease_interval: u64,
+}
+
+pub fn default_lsn_lease_interval() -> u64 {
+    1800
 }
 
 /// Feature flag to signal `compute_ctl` to enable certain experimental functionality.