[incubator-uniffle] 13/17: [Improvement] Add RSS_IP environment variable support for K8S (#204)

[jshao]

This is an automated email from the ASF dual-hosted git repository.

jshao pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-uniffle.git

commit 6937631876052425b8d808d26caf78c79b24536a
Author: roryqi <jerqi1242949...@gmail.com>
AuthorDate: Wed Jun 29 10:06:31 2022 +0800

    [Improvement] Add RSS_IP environment variable support for K8S (#204)
    
    ### What changes were proposed in this pull request?
    Method `getHostIp` can acquire IP by environment variable.
    
    ### Why are the changes needed?
    For K8S, there are too many IPs, it's hard to decide which we should use. 
So we use the environment variable to tell RSS to use which one.
    
    ### Does this PR introduce _any_ user-facing change?
    NO
    
    ### How was this patch tested?
    UT
---
 .../java/com/tencent/rss/common/util/RssUtils.java | 10 +++++++++
 .../com/tencent/rss/common/util/RssUtilsTest.java  | 26 ++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/common/src/main/java/com/tencent/rss/common/util/RssUtils.java 
b/common/src/main/java/com/tencent/rss/common/util/RssUtils.java
index 1b7200e..7ecae6b 100644
--- a/common/src/main/java/com/tencent/rss/common/util/RssUtils.java
+++ b/common/src/main/java/com/tencent/rss/common/util/RssUtils.java
@@ -41,6 +41,7 @@ import java.util.Map;
 import java.util.Properties;
 
 import com.google.common.collect.Lists;
+import com.google.common.net.InetAddresses;
 import org.roaringbitmap.longlong.Roaring64NavigableMap;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -102,6 +103,15 @@ public class RssUtils {
   // loop back, etc.). If the network interface in the machine is more than 
one, we
   // will choose the first IP.
   public static String getHostIp() throws Exception {
+    // For K8S, there are too many IPs, it's hard to decide which we should 
use.
+    // So we use the environment variable to tell RSS to use which one.
+    String ip = System.getenv("RSS_IP");
+    if (ip != null) {
+      if (!InetAddresses.isInetAddress(ip)) {
+        throw new RuntimeException("Environment RSS_IP: " + ip + " is wrong 
format");
+      }
+      return ip;
+    }
     Enumeration<NetworkInterface> nif = 
NetworkInterface.getNetworkInterfaces();
     String siteLocalAddress = null;
     while (nif.hasMoreElements()) {
diff --git a/common/src/test/java/com/tencent/rss/common/util/RssUtilsTest.java 
b/common/src/test/java/com/tencent/rss/common/util/RssUtilsTest.java
index 95fd55f..220cb5c 100644
--- a/common/src/test/java/com/tencent/rss/common/util/RssUtilsTest.java
+++ b/common/src/test/java/com/tencent/rss/common/util/RssUtilsTest.java
@@ -18,6 +18,7 @@
 
 package com.tencent.rss.common.util;
 
+import java.lang.reflect.Field;
 import java.net.InetAddress;
 import java.nio.ByteBuffer;
 import java.util.Arrays;
@@ -62,6 +63,18 @@ public class RssUtilsTest {
       if (!address.equals("127.0.0.1")) {
         assertEquals(address, realIp);
       }
+      setEnv("RSS_IP", "8.8.8.8");
+      assertEquals("8.8.8.8", RssUtils.getHostIp());
+      setEnv("RSS_IP", "xxxx");
+      boolean isException = false;
+      try {
+        RssUtils.getHostIp();
+      } catch (Exception e) {
+        isException = true;
+      }
+      setEnv("RSS_IP", realIp);
+      RssUtils.getHostIp();
+      assertTrue(isException);
     } catch (Exception e) {
       fail(e.getMessage());
     }
@@ -185,6 +198,19 @@ public class RssUtilsTest {
     }
   }
 
+  public static void setEnv(String key, String value) {
+    try {
+      Map<String, String> env = System.getenv();
+      Class<?> cl = env.getClass();
+      Field field = cl.getDeclaredField("m");
+      field.setAccessible(true);
+      Map<String, String> writableEnv = (Map<String, String>) field.get(env);
+      writableEnv.put(key, value);
+    } catch (Exception e) {
+      throw new IllegalStateException("Failed to set environment variable", e);
+    }
+  }
+
   public static class RssUtilTestDummySuccess implements RssUtilTestDummy {
     private final String s;