[LIVY-616] Livy Server discovery

api/src/main/java/org/apache/livy/LivyClient.java

@@ -107,4 +107,11 @@ public interface LivyClient {
    */
   Future<?> addFile(URI uri);
 
+  /**
+   * Get Livy Server URI.
+   * URI will be propagated from LivyClientBuilder during creating LivyClient.
+   *
+   * @return A future with Livy Server URI
+   */
+  Future<URI> getServerUri();
 }

api/src/main/java/org/apache/livy/LivyClientBuilder.java

@@ -81,6 +81,14 @@ public LivyClientBuilder(boolean loadDefaults) throws IOException {
     }
   }
 
+  /**
+   * Set Livy Server URI.
+   * This is possible to set it manually or get URI from LivyDiscoveryManager
+   * ({@code livy.zookeeper.url} should be configured).
+   *
+   * @param uri Livy Server URI
+   * @return this builder
+   */
   public LivyClientBuilder setURI(URI uri) {
     config.setProperty(LIVY_URI_KEY, uri.toString());
     return this;

api/src/test/java/org/apache/livy/TestClientFactory.java

@@ -81,6 +81,11 @@ public Future<?> addFile(URI uri) {
       throw new UnsupportedOperationException();
     }
 
+    @Override
+    public Future<URI> getServerUri() {
+      throw new UnsupportedOperationException();
+    }
+
 }
 
 }

client-http/pom.xml

@@ -78,6 +78,11 @@
       <artifactId>httpmime</artifactId>
       <version>4.5.1</version>
     </dependency>
+    <dependency>
+      <groupId>io.netty</groupId>
+      <artifactId>netty-all</artifactId>
+      <version>${netty.version}</version>
+    </dependency>
 
     <dependency>
       <groupId>org.apache.livy</groupId>

client-http/src/main/java/org/apache/livy/client/http/HttpClient.java

@@ -30,6 +30,9 @@
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
+import io.netty.util.concurrent.ImmediateEventExecutor;
+import io.netty.util.concurrent.Promise;
+
 import org.apache.livy.Job;
 import org.apache.livy.JobHandle;
 import org.apache.livy.LivyClient;
@@ -47,12 +50,14 @@ class HttpClient implements LivyClient {
   private final int sessionId;
   private final ScheduledExecutorService executor;
   private final Serializer serializer;
+  private final Promise<URI> serverUriPromise;
 
   private boolean stopped;
 
   HttpClient(URI uri, HttpConf httpConf) {
     this.config = httpConf;
     this.stopped = false;
+    this.serverUriPromise = ImmediateEventExecutor.INSTANCE.newPromise();
 
     // If the given URI looks like it refers to an existing session, then try to connect to
     // an existing session. Note this means that any Spark configuration in httpConf will be
@@ -77,6 +82,7 @@ class HttpClient implements LivyClient {
         ClientMessage create = new CreateClientRequest(sessionConf);
         this.conn = new LivyConnection(uri, httpConf);
         this.sessionId = conn.post(create, SessionInfo.class, "/").id;
+        serverUriPromise.setSuccess(uri);
       }
     } catch (Exception e) {
       throw propagate(e);
@@ -145,6 +151,11 @@ public Future<?> addFile(URI uri) {
     return addResource("add-file", uri);
   }
 
+  @Override
+  public Future<URI> getServerUri() {
+    return serverUriPromise;
+  }
+
   private Future<?> uploadResource(final File file, final String command, final String paramName) {
     Callable<Void> task = new Callable<Void>() {
       @Override

client-http/src/test/scala/org/apache/livy/client/http/HttpClientSpec.scala

@@ -98,6 +98,10 @@ class HttpClientSpec extends FunSpecLike with BeforeAndAfterAll with LivyBaseUni
       testJob(false)
     }
 
+    withClient("should get Livy Server URI") {
+      assume(Option(client.getServerUri.get()).isDefined)
+    }
+
     withClient("should propagate errors from jobs") {
       val errorMessage = "This job throws an error."
       val (jobId, handle) = runJob(false, { id => Seq(

conf/livy.conf.template

@@ -100,8 +100,9 @@
 # off: Default. Turn off recovery. Every time Livy shuts down, it stops and forgets all sessions.
 # recovery: Livy persists session info to the state store. When Livy restarts, it recovers
 #           previous sessions from the state store.
-# Must set livy.server.recovery.state-store and livy.server.recovery.state-store.url to
-# configure the state store.
+# Must set livy.server.recovery.state-store to needed state store (filesystem or zookeeper)
+# Set livy.server.recovery.state-store.url for filesystem state store
+# or livy.zookeeper.url for zookeeper state store.
 # livy.server.recovery.mode = off
 
 # Where Livy should store state to for recovery. Possible values:
@@ -112,8 +113,22 @@
 
 # For filesystem state store, the path of the state store directory. Please don't use a filesystem
 # that doesn't support atomic rename (e.g. S3). e.g. file:///tmp/livy or hdfs:///.
-# For zookeeper, the address to the Zookeeper servers. e.g. host1:port1,host2:port2
 # livy.server.recovery.state-store.url =
+# For zookeeper, the address to the Zookeeper servers. e.g. host1:port1,host2:port2
+# livy.zookeeper.url =
+
+# Livy Server discovery
+# ZooKeeper quorum URLs, e.g. host1:port1,host2:port2
+# livy.zookeeper.url =
+# Name of base Livy znode. Default livy
+# livy.zookeeper.namespace = livy
+# Name of Livy Server znode. Uses livy.zookeeper.namespace as parent.
+# By default, the full path to znode is /livy/server.uri
+# livy.server.zookeeper.namespace = server.uri
+# Number of trials to establish the connection to ZooKeeper quorum
+# livy.server.zookeeper.connection.max.retries = 3
+# Sleep time between connection retries to ZooKeeper quorum
+# livy.server.zookeeper.connection.retry.interval.ms = 500
 
 # If Livy can't find the yarn app within this time, consider it lost.
 # livy.server.yarn.app-lookup-timeout = 120s

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

@@ -192,6 +192,7 @@ public void onFailure(Throwable error) throws Exception {
     return promise;
   }
 
+  @Override
   public Future<URI> getServerUri() {
     return serverUriPromise;
   }

scala-api/src/main/scala/org/apache/livy/scalaapi/LivyScalaClient.scala

@@ -142,6 +142,13 @@ class LivyScalaClient(livyJavaClient: LivyClient) {
    */
   def addFile(uri: URI): Future[_] = new PollingContainer(livyJavaClient.addFile(uri)).poll()
 
+  /**
+    * Get Livy Server URI.
+    *
+    * @return A future with Livy Server URI
+    */
+  def getServerUri: Future[URI] = new PollingContainer(livyJavaClient.getServerUri).poll()
+
   private class PollingContainer[T] private[livy] (jFuture: JFuture[T]) extends Runnable {
 
     private val initialDelay = 1

scala-api/src/test/scala/org/apache/livy/scalaapi/ScalaClientTest.scala

@@ -117,6 +117,13 @@ class ScalaClientTest extends FunSuite
     ScalaClientTestUtils.assertTestPassed(sFuture, "test file")
   }
 
+  test("test get uri") {
+    configureClient(true)
+    val getUriFuture = client.getServerUri
+    val uri = Await.result(getUriFuture, ScalaClientTestUtils.Timeout second)
+    assert(Option(uri).isDefined)
+  }
+
   test("test add jar") {
     configureClient(true)
     val jar = File.createTempFile("test", ".resource")

server/src/main/scala/org/apache/livy/LivyConf.scala

@@ -173,7 +173,9 @@ object LivyConf {
    * off: Default. Turn off recovery. Every time Livy shuts down, it stops and forgets all sessions.
    * recovery: Livy persists session info to the state store. When Livy restarts, it recovers
    *   previous sessions from the state store.
-   * Must set livy.server.recovery.state-store and livy.server.recovery.state-store.url to
+   * Must set livy.server.recovery.state-store to needed state store (filesystem or zookeeper)
+   * Set livy.server.recovery.state-store.url for filesystem state store
+   * or livy.zookeeper.url for zookeeper state store.
    * configure the state store.
    */
   val RECOVERY_MODE = Entry("livy.server.recovery.mode", "off")
@@ -187,10 +189,31 @@ object LivyConf {
   /**
    * For filesystem state store, the path of the state store directory. Please don't use a
    * filesystem that doesn't support atomic rename (e.g. S3). e.g. file:///tmp/livy or hdfs:///.
-   * For zookeeper, the address to the Zookeeper servers. e.g. host1:port1,host2:port2
+   * For zookeeper, use livy.zookeeper.url.
    */
   val RECOVERY_STATE_STORE_URL = Entry("livy.server.recovery.state-store.url", "")
 
+  /**
+    * ZooKeeper address witch will be used for Livy Server discovery
+    * and Zookeeper state store. e.g. host1:port1,host2:port2
+    */
+  val LIVY_ZOOKEEPER_URL = Entry("livy.zookeeper.url", null)
+
+  // Name of base Livy znode.
+  val LIVY_ZOOKEEPER_NAMESPACE = Entry("livy.zookeeper.namespace", "livy")
+
+  // Number of trials to establish the connection to ZooKeeper quorum.
+  val LIVY_ZOOKEEPER_CONNECTION_MAX_RETRIES =
+    Entry("livy.server.zookeeper.connection.max.retries", 3)
+
+  // Sleep time between connection retries to ZooKeeper quorum.
+  val LIVY_ZOOKEEPER_CONNECTION_RETRY_INTERVAL =
+    Entry("livy.server.zookeeper.connection.retry.interval.ms", 500)
+
+  // Name of Livy Server znode. Uses LIVY_ZOOKEEPER_NAMESPACE as parent.
+  // By default, the full path to znode is /livy/server.uri
+  val LIVY_SERVER_ZOOKEEPER_NAMESPACE = Entry("livy.server.zookeeper.namespace", "server.uri")
+
   // Livy will cache the max no of logs specified. 0 means don't cache the logs.
   val SPARK_LOGS_SIZE = Entry("livy.cache-log.size", 200)
 

server/src/main/scala/org/apache/livy/server/LivyServer.scala

@@ -18,15 +18,16 @@
 package org.apache.livy.server
 
 import java.io.{BufferedInputStream, InputStream}
-import java.net.InetAddress
-import java.util.concurrent._
+import java.net.{InetAddress, URI}
 import java.util.EnumSet
+import java.util.concurrent._
 import javax.servlet._
 
 import scala.collection.JavaConverters._
 import scala.concurrent.ExecutionContext.Implicits.global
 import scala.concurrent.Future
 
+import org.apache.curator.framework.CuratorFramework
 import org.apache.hadoop.security.{SecurityUtil, UserGroupInformation}
 import org.apache.hadoop.security.authentication.server._
 import org.eclipse.jetty.servlet.FilterHolder
@@ -37,6 +38,7 @@ import org.scalatra.servlet.{MultipartConfig, ServletApiImplicits}
 
 import org.apache.livy._
 import org.apache.livy.server.batch.BatchSessionServlet
+import org.apache.livy.server.discovery.LivyDiscoveryManager
 import org.apache.livy.server.interactive.InteractiveSessionServlet
 import org.apache.livy.server.recovery.{SessionStore, StateStore}
 import org.apache.livy.server.ui.UIServlet
@@ -72,6 +74,8 @@ class LivyServer extends Logging {
         maxFileSize = Some(livyConf.getLong(LivyConf.FILE_UPLOAD_MAX_SIZE))
       ).toMultipartConfigElement
 
+    setServerUri(livyConf)
+
     // Make sure the `spark-submit` program exists, otherwise much of livy won't work.
     testSparkHome(livyConf)
 
@@ -385,6 +389,25 @@ class LivyServer extends Logging {
     }
   }
 
+  private[livy] def setServerUri(livyConf: LivyConf,
+                                 mockCuratorClient: Option[CuratorFramework] = None): Unit = {
+    if (Option(livyConf.get(LIVY_ZOOKEEPER_URL)).isDefined) {
+      val discoveryManager = LivyDiscoveryManager(livyConf, mockCuratorClient)
+      val host = resolvedSeverHost(livyConf)
+      val uri = new URI(s"http://$host:${livyConf.getInt(LivyConf.SERVER_PORT)}")
+      discoveryManager.setServerUri(uri)
+    }
+  }
+
+  private def resolvedSeverHost(livyConf: LivyConf) = {
+    val host = livyConf.get(LivyConf.SERVER_HOST)
+    if (host.equals(livyConf.get(LivyConf.SERVER_HOST.dflt.toString))) {
+      InetAddress.getLocalHost.getHostName
+    } else {
+      host
+    }
+  }
+
   private[livy] def testRecovery(livyConf: LivyConf): Unit = {
     if (!livyConf.isRunningOnYarn()) {
       // If recovery is turned on but we are not running on YARN, quit.

server/src/main/scala/org/apache/livy/server/discovery/JsonMapper.scala

@@ -0,0 +1,35 @@
+/*
+ * 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.server.discovery
+
+import scala.reflect.{classTag, ClassTag}
+
+import com.fasterxml.jackson.databind.ObjectMapper
+import com.fasterxml.jackson.module.scala.DefaultScalaModule
+
+import org.apache.livy.sessions.SessionKindModule
+
+protected[server] trait JsonMapper {
+  protected val mapper = new ObjectMapper()
+    .registerModule(DefaultScalaModule)
+    .registerModule(new SessionKindModule())
+
+  def serializeToBytes(value: Object): Array[Byte] = mapper.writeValueAsBytes(value)
+
+  def deserialize[T: ClassTag](json: Array[Byte]): T =
+    mapper.readValue(json, classTag[T].runtimeClass.asInstanceOf[Class[T]])
+}

server/src/main/scala/org/apache/livy/server/discovery/LivyDiscoveryManager.scala

@@ -0,0 +1,65 @@
+/*
+ * 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.server.discovery
+
+import java.net.URI
+
+import org.apache.curator.framework.CuratorFramework
+
+import org.apache.livy.LivyConf
+
+/**
+  * Livy Server Discovery manager.
+  * Stores information about Livy Server location in ZooKeeper.
+  * The address will be stored in
+  * "/{@code LIVY_ZOOKEEPER_NAMESPACE}/{@code LIVY_SERVER_ZOOKEEPER_NAMESPACE}" znode
+  * By default, the full path to znode is /livy/server.uri.
+  * Need to set {@code livy.zookeeper.url} to be able to get information from ZooKeeper.
+  *
+  * @param livyConf - Livy configurations
+  * @param mockCuratorClient - used for testing
+  */
+class LivyDiscoveryManager(val livyConf: LivyConf,
+                           val mockCuratorClient: Option[CuratorFramework] = None)
+  extends ZooKeeperManager {
+
+  private val LIVY_SERVER_URI_KEY = livyConf.get(LivyConf.LIVY_SERVER_ZOOKEEPER_NAMESPACE)
+
+  /**
+    * Save Livy Server URI to ZooKeeper.
+    * @param address - URI address of Livy Server
+    */
+  def setServerUri(address: URI): Unit = {
+    setData(LIVY_SERVER_URI_KEY, address)
+  }
+
+  /**
+    * Get Livy Server URI from ZooKeeper.
+    * @return Livy Server URI
+    */
+  def getServerUri(): URI = {
+    getData[URI](LIVY_SERVER_URI_KEY).getOrElse(URI.create(""))
+  }
+}
+
+object LivyDiscoveryManager {
+
+  def apply(livyConf: LivyConf,
+            mockCuratorClient: Option[CuratorFramework] = None): LivyDiscoveryManager = {
+    new LivyDiscoveryManager(livyConf, mockCuratorClient)
+  }
+}