From a563231d0438d7de2e4f866844cbe61bbb6633cf Mon Sep 17 00:00:00 2001
From: Cuong Nguyen <ctring23@gmail.com>
Date: Thu, 3 Aug 2023 23:23:06 -0400
Subject: [PATCH] Add new flag for advertising pg address

---
 safekeeper/src/bin/safekeeper.rs | 5 +++++
 safekeeper/src/lib.rs            | 2 ++
 safekeeper/src/timeline.rs       | 5 ++++-
 3 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/safekeeper/src/bin/safekeeper.rs b/safekeeper/src/bin/safekeeper.rs
index 90eebd392c41..acc717294bec 100644
--- a/safekeeper/src/bin/safekeeper.rs
+++ b/safekeeper/src/bin/safekeeper.rs
@@ -79,6 +79,10 @@ struct Args {
     /// Listen http endpoint for management and metrics in the form host:port.
     #[arg(long, default_value = DEFAULT_HTTP_LISTEN_ADDR)]
     listen_http: String,
+    /// Advertised endpoint for receiving/sending WAL in the form host:port. If not
+    /// specified, listen_pg is used to advertise instead.
+    #[arg(long, default_value = None)]
+    advertise_pg: Option<String>,
     /// Availability zone of the safekeeper.
     #[arg(long)]
     availability_zone: Option<String>,
@@ -185,6 +189,7 @@ async fn main() -> anyhow::Result<()> {
         listen_pg_addr: args.listen_pg,
         listen_pg_addr_tenant_only: args.listen_pg_tenant_only,
         listen_http_addr: args.listen_http,
+        advertise_pg_addr: args.advertise_pg,
         availability_zone: args.availability_zone,
         no_sync: args.no_sync,
         broker_endpoint: args.broker_endpoint,
diff --git a/safekeeper/src/lib.rs b/safekeeper/src/lib.rs
index 1a1c0add678d..490af1de9564 100644
--- a/safekeeper/src/lib.rs
+++ b/safekeeper/src/lib.rs
@@ -55,6 +55,7 @@ pub struct SafeKeeperConf {
     pub listen_pg_addr: String,
     pub listen_pg_addr_tenant_only: Option<String>,
     pub listen_http_addr: String,
+    pub advertise_pg_addr: Option<String>,
     pub availability_zone: Option<String>,
     pub no_sync: bool,
     pub broker_endpoint: Uri,
@@ -88,6 +89,7 @@ impl SafeKeeperConf {
             listen_pg_addr: defaults::DEFAULT_PG_LISTEN_ADDR.to_string(),
             listen_pg_addr_tenant_only: None,
             listen_http_addr: defaults::DEFAULT_HTTP_LISTEN_ADDR.to_string(),
+            advertise_pg_addr: None,
             availability_zone: None,
             remote_storage: None,
             my_id: NodeId(0),
diff --git a/safekeeper/src/timeline.rs b/safekeeper/src/timeline.rs
index 30036cc7f2be..9ec89b86d437 100644
--- a/safekeeper/src/timeline.rs
+++ b/safekeeper/src/timeline.rs
@@ -237,7 +237,10 @@ impl SharedState {
             commit_lsn: self.sk.inmem.commit_lsn.0,
             remote_consistent_lsn: remote_consistent_lsn.0,
             peer_horizon_lsn: self.sk.inmem.peer_horizon_lsn.0,
-            safekeeper_connstr: conf.listen_pg_addr.clone(),
+            safekeeper_connstr: conf
+                .advertise_pg_addr
+                .to_owned()
+                .unwrap_or(conf.listen_pg_addr.clone()),
             backup_lsn: self.sk.inmem.backup_lsn.0,
             local_start_lsn: self.sk.state.local_start_lsn.0,
             availability_zone: conf.availability_zone.clone(),