[LIVY-1071] Add driver.address.use-hostname RSC config flag

conf/livy.conf.template

@@ -93,6 +93,11 @@
 # time of sessions on YARN can be reduced.
 # livy.rsc.jars =
 
+# If true, the RSC driver's address is resolved from the hostname reported by the driver
+# instead of the socket's IP. Useful for Kubernetes + service-mesh (Istio) deployments where
+# pod IPs are not stable/routable but hostnames are.
+# livy.rsc.driver.address.use-hostname = false
+
 # Comma-separated list of Livy REPL jars. By default Livy will upload jars from its installation
 # directory every time a session is started. By caching these files in HDFS, for example, startup
 # time of sessions on YARN can be reduced. Please list all the repl dependencies including

rsc/src/main/java/org/apache/livy/rsc/ContextLauncher.java

@@ -250,6 +250,14 @@ private static void merge(RSCConf conf, String key, String livyConf, String sep)
     conf.set(key, confValue);
   }
 
+  static String resolveDriverAddress(RSCConf conf, BaseProtocol.RemoteDriverAddress msg,
+      InetSocketAddress remoteAddress) {
+    if (conf.getBoolean(DRIVER_ADDRESS_USE_HOSTNAME)) {
+      return msg.host;
+    }
+    return remoteAddress.getAddress().getHostAddress();
+  }
+
   /**
    * Write the configuration to a file readable only by the process's owner. Livy properties
    * are written with an added prefix so that they can be loaded using SparkConf on the driver
@@ -328,12 +336,12 @@ void dispose() {
     //tests fail without it
     public void handle(ChannelHandlerContext ctx, RemoteDriverAddress msg) {
       InetSocketAddress insocket = (InetSocketAddress) ctx.channel().remoteAddress();
-      String ip = insocket.getAddress().getHostAddress();
-      ContextInfo info = new ContextInfo(ip, msg.port, clientId, secret);
+      String driverHost = resolveDriverAddress(conf, msg, insocket);
+      ContextInfo info = new ContextInfo(driverHost, msg.port, clientId, secret);
       if (promise.trySuccess(info)) {
         timeout.cancel(true);
         LOG.debug("Received driver info for client {}: {}/{}.", client.getChannel(),
-                msg.host, msg.port);
+                driverHost, msg.port);
       } else {
         LOG.warn("Connection established but promise is already finalized.");
       }

rsc/src/main/java/org/apache/livy/rsc/RSCConf.java

@@ -58,6 +58,9 @@ public enum Entry implements ConfEntry {
     // How long will the RSC wait for a connection for a Livy server before shutting itself down.
     SERVER_IDLE_TIMEOUT("server.idle-timeout", "10m"),
 
+    // Use driver's reported hostname instead of socket IP. Useful for Kubernetes + Istio mode.
+    DRIVER_ADDRESS_USE_HOSTNAME("driver.address.use-hostname", false),
+
     PROXY_USER("proxy-user", null),
 
     RPC_SERVER_ADDRESS("rpc.server.address", null),

rsc/src/test/java/org/apache/livy/rsc/TestContextLauncher.java

@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.livy.rsc;
+
+import java.net.InetAddress;
+import java.net.InetSocketAddress;
+
+import org.junit.Test;
+import static org.junit.Assert.assertEquals;
+
+import static org.apache.livy.rsc.RSCConf.Entry.DRIVER_ADDRESS_USE_HOSTNAME;
+
+public class TestContextLauncher {
+
+  @Test
+  public void testResolveDriverAddressUsesSocketIpByDefault() throws Exception {
+    RSCConf conf = new RSCConf(null);
+    BaseProtocol.RemoteDriverAddress msg =
+      new BaseProtocol.RemoteDriverAddress("driver.example.internal", 1234);
+    InetSocketAddress remote =
+      new InetSocketAddress(InetAddress.getLoopbackAddress(), 5678);
+
+    String resolved = ContextLauncher.resolveDriverAddress(conf, msg, remote);
+
+    assertEquals(remote.getAddress().getHostAddress(), resolved);
+  }
+
+  @Test
+  public void testResolveDriverAddressUsesHostnameWhenEnabled() throws Exception {
+    RSCConf conf = new RSCConf(null);
+    conf.set(DRIVER_ADDRESS_USE_HOSTNAME, true);
+    BaseProtocol.RemoteDriverAddress msg =
+      new BaseProtocol.RemoteDriverAddress("driver.example.internal", 1234);
+    InetSocketAddress remote =
+      new InetSocketAddress(InetAddress.getLoopbackAddress(), 5678);
+
+    String resolved = ContextLauncher.resolveDriverAddress(conf, msg, remote);
+
+    assertEquals("driver.example.internal", resolved);
+  }
+}