From ff55d6f4df9c002329e3aed1c93e265f0150cd44 Mon Sep 17 00:00:00 2001
From: Andrey Rudenko <me@prepor.dev>
Date: Wed, 12 Jun 2024 11:45:26 +0200
Subject: [PATCH] feedback: using valid_until (it doesn't work currently)

---
 compute_tools/src/compute.rs  | 46 ++++++++++++++++++++++++++---------
 control_plane/src/endpoint.rs |  2 --
 libs/compute_api/src/spec.rs  |  7 ------
 3 files changed, 34 insertions(+), 21 deletions(-)

diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs
index d14092fc2b..54e51d0557 100644
--- a/compute_tools/src/compute.rs
+++ b/compute_tools/src/compute.rs
@@ -10,7 +10,9 @@ use std::sync::atomic::AtomicU32;
 use std::sync::atomic::Ordering;
 use std::sync::{Condvar, Mutex, RwLock};
 use std::thread;
+use std::time;
 use std::time::Instant;
+use std::time::SystemTime;
 
 use anyhow::{Context, Result};
 use chrono::{DateTime, Utc};
@@ -1414,29 +1416,49 @@ fn start_lsn_lease_for_static(compute_state: &ComputeState) {
 
     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));
+    thread::spawn(move || lsn_lease_loop(configs, cmd));
 }
 
-fn lsn_lease_loop(interval_secs: u64, configs: Vec<postgres::Config>, cmd: String) {
+fn lsn_lease_loop(configs: Vec<postgres::Config>, cmd: String) {
     loop {
         match lsn_lease_request(&configs, &cmd) {
-            Ok(_) => {
-                thread::sleep(tokio::time::Duration::from_secs(interval_secs));
+            Ok(valid_until) => {
+                let now = SystemTime::now()
+                    .duration_since(SystemTime::UNIX_EPOCH)
+                    .expect("Time went backwards")
+                    .as_millis();
+
+                let sleep_duration_millis = valid_until - now - 60_000;
+
+                let sleep_duration = time::Duration::from_millis(sleep_duration_millis as u64);
+
+                info!(
+                    "lsn_lease_request succeeded, sleeping for {} seconds",
+                    sleep_duration.as_secs()
+                );
+                thread::sleep(sleep_duration);
             }
             Err(e) => {
                 error!("lsn_lease_request failed: {:#}", e);
-                thread::sleep(tokio::time::Duration::from_secs(10));
+                thread::sleep(time::Duration::from_secs(10));
             }
         }
     }
 }
 
-fn lsn_lease_request(configs: &[postgres::Config], cmd: &str) -> Result<()> {
+fn lsn_lease_request(configs: &[postgres::Config], cmd: &str) -> Result<u128> {
     info!("lsn_lease_request: {}", cmd);
-    for config in configs {
-        let mut client = config.connect(NoTls)?;
-        let _ = client.simple_query(cmd)?;
-    }
-    Ok(())
+    let valid_until = configs
+        .iter()
+        .map(|config| {
+            let mut client = config.connect(NoTls)?;
+            let msg = client.query_one(cmd, &[])?;
+            let bytes: &[u8] = msg.try_get("valid_until")?;
+            Ok(u128::from_be_bytes(bytes.try_into()?))
+        })
+        .collect::<Result<Vec<u128>>>()?
+        .into_iter()
+        .min()
+        .unwrap();
+    Ok(valid_until)
 }
diff --git a/control_plane/src/endpoint.rs b/control_plane/src/endpoint.rs
index 254cf259ff..20371e1cb8 100644
--- a/control_plane/src/endpoint.rs
+++ b/control_plane/src/endpoint.rs
@@ -46,7 +46,6 @@ 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;
@@ -594,7 +593,6 @@ 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 0f9bf6a6c6..1c4ee2089f 100644
--- a/libs/compute_api/src/spec.rs
+++ b/libs/compute_api/src/spec.rs
@@ -102,13 +102,6 @@ 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.