From ea4987f941c8ed03be3c084cfd0092617c18903f Mon Sep 17 00:00:00 2001
From: Chris Lu <chrislusf@users.noreply.github.com>
Date: Wed, 21 Oct 2026 12:21:00 -0700
Subject: [PATCH] volume: publish the master only once it is known leader
 (rust)

---
 seaweed-volume/src/server/heartbeat.rs     | 36 ++++++++++++++--------
 seaweed-volume/src/server/volume_server.rs | 10 +++---
 2 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/seaweed-volume/src/server/heartbeat.rs b/seaweed-volume/src/server/heartbeat.rs
index 081ff71a25..b40f33aad5 100644
--- a/seaweed-volume/src/server/heartbeat.rs
+++ b/seaweed-volume/src/server/heartbeat.rs
@@ -258,6 +258,14 @@ fn is_stopping(state: &VolumeServerState) -> bool {
     *state.is_stopping.read().unwrap()
 }
 
+/// Store the master lookups should go to in canonical http `host:port` form.
+/// The heartbeat calls this only for a master known to lead, mirroring Go's
+/// setCurrentMaster.
+async fn publish_current_master(state: &Arc<VolumeServerState>, master: &str) {
+    *state.current_master_url.write().await =
+        super::volume_server::to_http_address(master).into_owned();
+}
+
 fn duplicate_uuid_retry_delay(retry_count: u32) -> Duration {
     Duration::from_secs((1u64 << retry_count) * 2)
 }
@@ -430,14 +438,6 @@ async fn do_heartbeat(
     let mut response_stream = client.send_heartbeat(stream).await?.into_inner();
 
     info!("Heartbeat stream established with {}", grpc_addr);
-    // Publish the master we're now talking to in canonical http host:port
-    // form so Ping admission can recognise it once a leader change moves us
-    // off the seed list. Mirrors Go's vs.setCurrentMaster(masterAddress).
-    {
-        let normalised = super::volume_server::to_http_address(current_master).into_owned();
-        let mut guard = state.current_master_url.write().await;
-        *guard = normalised;
-    }
     if is_stopping(state) {
         state.is_heartbeating.store(false, Ordering::Relaxed);
         send_deregister_heartbeat(config, state, &tx).await;
@@ -451,6 +451,7 @@ async fn do_heartbeat(
     volume_tick.tick().await;
     ec_tick.tick().await;
 
+    let mut confirmed_leader = false;
     loop {
         tokio::select! {
             resp = response_stream.message() => {
@@ -502,10 +503,21 @@ async fn do_heartbeat(
                         if metrics_changed {
                             state.metrics_notify.notify_waiters();
                         }
-                        // Match Go: only redirect if leader is non-empty AND
-                        // different from the current master we're connected to.
-                        if !hb_resp.leader.is_empty() && current_master != hb_resp.leader {
-                            return Ok(Some(hb_resp.leader));
+                        if !hb_resp.leader.is_empty() {
+                            // Only a follower names the leader; publish it now
+                            // so lookups stop landing on this follower while
+                            // the loop reconnects. Mirrors Go.
+                            if super::volume_server::to_http_address(current_master)
+                                != super::volume_server::to_http_address(&hb_resp.leader)
+                            {
+                                publish_current_master(state, &hb_resp.leader).await;
+                                return Ok(Some(hb_resp.leader));
+                            }
+                        } else if !confirmed_leader {
+                            // The leader answers every new stream without naming
+                            // a leader, so this peer is the one lookups go to.
+                            confirmed_leader = true;
+                            publish_current_master(state, current_master).await;
                         }
                     }
                     Ok(None) => return Ok(None),
diff --git a/seaweed-volume/src/server/volume_server.rs b/seaweed-volume/src/server/volume_server.rs
index a440c2a9e1..4dcc4e4d61 100644
--- a/seaweed-volume/src/server/volume_server.rs
+++ b/seaweed-volume/src/server/volume_server.rs
@@ -94,11 +94,11 @@ pub struct VolumeServerState {
     /// Built once at construction so Ping admission stays O(1). Mirrors
     /// Go's `seedMasterSet` on `VolumeServer`.
     pub seed_master_set: std::collections::HashSet<String>,
-    /// Current master this server is heartbeating with, in canonical http
-    /// `host:port` form. Empty when no heartbeat connection is active. The
-    /// heartbeat goroutine writes; admission reads — the lock keeps them
-    /// from racing on a leader change. Mirrors Go's `currentMaster` plus
-    /// `currentMasterLock`.
+    /// Master this server takes as leader, in canonical http `host:port`
+    /// form. Empty until the heartbeat confirms one — a peer that answered a
+    /// stream without naming a leader, or the leader a follower named. The
+    /// heartbeat writes; admission and lookups read. Mirrors Go's
+    /// `currentMaster` plus `currentMasterLock`.
     pub current_master_url: tokio::sync::RwLock<String>,
     /// This server's own address (ip:port) for filtering self from lookup results.
     pub self_url: String,