From 1d35639182b256ee51f2ff5201a86d8a05644003 Mon Sep 17 00:00:00 2001
From: Rohit Kulshreshtha <rohit.in.pune@gmail.com>
Date: Fri, 8 Nov 2024 09:49:12 -0800
Subject: [PATCH] feat(Gossip KV): Make member random suffix length
 configurable. (#1556)

---
 datastores/gossip_kv/Cargo.toml           |  2 +-
 datastores/gossip_kv/server/main.rs       |  5 ++-
 datastores/gossip_kv/server/membership.rs | 49 +++++++++++------------
 3 files changed, 28 insertions(+), 28 deletions(-)

diff --git a/datastores/gossip_kv/Cargo.toml b/datastores/gossip_kv/Cargo.toml
index b59765578698..37943072a7d1 100644
--- a/datastores/gossip_kv/Cargo.toml
+++ b/datastores/gossip_kv/Cargo.toml
@@ -7,7 +7,7 @@ license = "Apache-2.0"
 publish = false
 
 [dependencies]
-clap = { version = "4.5.4", features = ["derive"] }
+clap = { version = "4.5.4", features = ["derive", "env"] }
 config = "0.14.0"
 governor = "0.7.0"
 hostname = "0.4.0"
diff --git a/datastores/gossip_kv/server/main.rs b/datastores/gossip_kv/server/main.rs
index 2d2990a7e0d3..bca640e9449b 100644
--- a/datastores/gossip_kv/server/main.rs
+++ b/datastores/gossip_kv/server/main.rs
@@ -39,6 +39,9 @@ struct Opts {
     /// The duration (in seconds) between gossip rounds.
     #[clap(short, long, default_value = "5", value_parser = clap_duration_from_secs)]
     gossip_frequency: Duration,
+
+    #[clap(env = "GOSSIP_MEMBER_SUFFIX_LEN", default_value = "4")]
+    member_suffix_len: usize,
 }
 
 /// Parse duration from float string for clap args.
@@ -133,7 +136,7 @@ async fn main() {
     let gossip_protocol_address =
         SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), opts.gossip_port);
 
-    let member_data = MemberDataBuilder::new(member_name().clone())
+    let member_data = MemberDataBuilder::new(member_name(opts.member_suffix_len).clone())
         .add_protocol(Protocol::new("gossip".into(), gossip_protocol_address))
         .add_protocol(Protocol::new("client".into(), client_protocol_address))
         .build();
diff --git a/datastores/gossip_kv/server/membership.rs b/datastores/gossip_kv/server/membership.rs
index 28a3fcadf137..bdaef84b6c11 100644
--- a/datastores/gossip_kv/server/membership.rs
+++ b/datastores/gossip_kv/server/membership.rs
@@ -1,37 +1,34 @@
 use std::sync::OnceLock;
 
 use gossip_kv::membership::MemberId;
-// use rand::distributions::Distribution;
-// use rand::{Rng};
+use rand::distributions::Distribution;
+use rand::{thread_rng, Rng};
 
-// /// This is a simple distribution that generates a random lower-case alphanumeric
-// struct LowercaseAlphanumeric;
-//
-// impl Distribution<char> for LowercaseAlphanumeric {
-//     fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> char {
-//         let choices = b"abcdefghijklmnopqrstuvwxyz0123456789";
-//         choices[rng.gen_range(0..choices.len())] as char
-//     }
-// }
+/// This is a simple distribution that generates a random lower-case alphanumeric
+struct LowercaseAlphanumeric;
+
+impl Distribution<char> for LowercaseAlphanumeric {
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> char {
+        let choices = b"abcdefghijklmnopqrstuvwxyz0123456789";
+        choices[rng.gen_range(0..choices.len())] as char
+    }
+}
 
 /// Gets a name for the current process.
-pub fn member_name() -> &'static MemberId {
+pub fn member_name(random_suffix_len: usize) -> &'static MemberId {
     static MEMBER_NAME: OnceLock<MemberId> = OnceLock::new();
     MEMBER_NAME.get_or_init(|| {
-        // Generate a lower-case alphanumeric suffix of length 4
-
-        // TODO: Random suffixes are good, but make benchmarking a pain. For now, we'll just use
-        // the hostname as-is.
-
-        // let suffix: String = thread_rng()
-        //     .sample_iter(&LowercaseAlphanumeric)
-        //     .take(4)
-        //     .map(char::from)
-        //     .collect();
-
-        // Retrieve hostname
-        hostname::get().unwrap().to_str().unwrap().to_string()
+        let hostname = hostname::get().unwrap().to_str().unwrap().to_string();
 
-        // format!("{}-{}", hostname, suffix)
+        if random_suffix_len > 0 {
+            let suffix: String = thread_rng()
+                .sample_iter(&LowercaseAlphanumeric)
+                .take(4)
+                .map(char::from)
+                .collect();
+            format!("{}-{}", hostname, suffix)
+        } else {
+            hostname
+        }
     })
 }