[LIVY-11] Enable HA support

conf/livy.conf.template

@@ -100,6 +100,16 @@
 # on user request and then livy server classpath automatically.
 # livy.repl.enable-hive-context =
 
+# High Availability mode of Livy. Possible values:
+# off: Default. Turn off High Availability. 
+# on: Livy uses Zookeeper as a state store to ensure a livy server is always available with the
+#     correct state.
+# Must set livy.server.ha.zookeeper-url to configure HA
+# livy.server.ha.mode = off
+
+# For zookeeper, the address to the Zookeeper servers. e.g. host1:port1,host2:port2
+# livy.server.ha.zookeeper-url =
+
 # Recovery mode of Livy. Possible values:
 # 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

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

@@ -184,6 +184,19 @@ object LivyConf {
    * zookeeper: Store state in a Zookeeper instance.
    */
   val RECOVERY_STATE_STORE = Entry("livy.server.recovery.state-store", null)
+
+  /**
+  * High Availability mode of Livy. Possible values:
+  * off: Default. Turn off High Availability.
+  * on: Livy uses Zookeeper as a state store to ensure a livy server is always available with the
+  *     correct state.
+  * Must set livy.server.ha.zookeeper-url to configure HA
+  */
+  val HA_MODE = Entry("livy.server.ha.mode", "off")
+
+  // For zookeeper, the address to the Zookeeper servers. e.g. host1:port1,host2:port2
+  val HA_ZOOKEEPER_URL = Entry("livy.server.ha.zookeeper-url", "")
+
   /**
    * 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:///.

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

@@ -0,0 +1,123 @@
+/*
+ * 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
+
+import java.io.Closeable
+import java.io.IOException
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.AtomicInteger
+
+import org.apache.curator.framework.CuratorFramework
+import org.apache.curator.framework.CuratorFrameworkFactory
+import org.apache.curator.framework.recipes.leader.LeaderLatch
+import org.apache.curator.framework.recipes.leader.LeaderLatchListener
+import org.apache.curator.retry.RetryNTimes
+
+import org.apache.livy.{LivyConf, Logging}
+import org.apache.livy.LivyConf.Entry
+
+object CuratorElectorService {
+  val HA_KEY_PREFIX_CONF = Entry("livy.server.ha.key-prefix", "livy_ha")
+  val HA_RETRY_CONF = Entry("livy.server.ha.retry-policy", "5,100")
+}
+
+class CuratorElectorService(livyConf : LivyConf, livyServer : LivyServer)
+  extends LeaderLatchListener
+  with Logging
+{
+
+  import CuratorElectorService._
+
+  val haAddress = livyConf.get(LivyConf.HA_ZOOKEEPER_URL)
+  require(!haAddress.isEmpty, s"Please configure ${LivyConf.HA_ZOOKEEPER_URL.key}.")
+  val haKeyPrefix = livyConf.get(HA_KEY_PREFIX_CONF)
+  val retryValue = livyConf.get(HA_RETRY_CONF)
+  // a regex to match patterns like "m, n" where m and n both are integer values
+  val retryPattern = """\s*(\d+)\s*,\s*(\d+)\s*""".r
+  val retryPolicy = retryValue match {
+    case retryPattern(n, sleepMs) => new RetryNTimes(n.toInt, sleepMs.toInt)
+    case _ => throw new IllegalArgumentException(
+      s"$HA_KEY_PREFIX_CONF contains bad value: $retryValue. " +
+        "Correct format is <max retry count>,<sleep ms between retry>. e.g. 5,100")
+  }
+
+  val client: CuratorFramework = CuratorFrameworkFactory.newClient(haAddress, retryPolicy)
+  val leaderKey = "/$haKeyPrefix/leader"
+
+  var server : LivyServer = livyServer
+
+  var leaderLatch = new LeaderLatch(client, leaderKey)
+  leaderLatch.addListener(this)
+
+  object HAState extends Enumeration{
+    type HAState = Value
+    val Active, Standby = Value
+  }
+  var currentState = HAState.Standby
+
+  def isLeader() {
+    transitionToActive();
+  }
+
+  def notLeader(){
+    transitionToStandby();
+  }
+
+  def start() : Unit = {
+    transitionToStandby()
+
+    client.start()
+    leaderLatch.start()
+    leaderLatch.await()
+
+    // This instance is now the leader. Joining the webserver to the main thread
+    info("starting join")
+    server.join()
+    info("join completed?")
+    close()
+  }
+
+  def close() : Unit = {
+    transitionToStandby();
+    leaderLatch.close();
+  }
+
+  def transitionToActive() : Unit = {
+    info("Transitioning to Active state")
+    if(currentState == HAState.Active){
+      info("Already in Active State")
+    }
+    else {
+      server.start()
+      currentState = HAState.Active
+      info("Transition complete")
+    }
+  }
+
+  def transitionToStandby() : Unit = {
+    info("Transitioning to Standby state")
+    if(currentState == HAState.Standby){
+      info("Already in Standby State");
+    }
+    else {
+      server.stop();
+      currentState = HAState.Standby
+      info("Transition complete");
+    }
+  }
+}

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

@@ -27,6 +27,7 @@ import scala.collection.JavaConverters._
 import scala.concurrent.ExecutionContext.Implicits.global
 import scala.concurrent.Future
 
+import org.apache.curator.utils.CloseableUtils
 import org.apache.hadoop.security.{SecurityUtil, UserGroupInformation}
 import org.apache.hadoop.security.authentication.server._
 import org.eclipse.jetty.servlet.FilterHolder
@@ -38,7 +39,7 @@ import org.scalatra.servlet.{MultipartConfig, ServletApiImplicits}
 import org.apache.livy._
 import org.apache.livy.server.batch.BatchSessionServlet
 import org.apache.livy.server.interactive.InteractiveSessionServlet
-import org.apache.livy.server.recovery.{SessionStore, StateStore}
+import org.apache.livy.server.recovery.{SessionStore, StateStore, ZooKeeperStateStore}
 import org.apache.livy.server.ui.UIServlet
 import org.apache.livy.sessions.{BatchSessionManager, InteractiveSessionManager}
 import org.apache.livy.sessions.SessionManager.SESSION_RECOVERY_MODE_OFF
@@ -394,16 +395,29 @@ class LivyServer extends Logging {
   }
 }
 
+object HighAvailabilitySettings {
+  val HA_ON = "on"
+  val HA_OFF = "off"
+}
+
 object LivyServer {
 
   def main(args: Array[String]): Unit = {
     val server = new LivyServer()
-    try {
-      server.start()
-      server.join()
-    } finally {
-      server.stop()
+    val livyConf = new LivyConf().loadFromFile("livy.conf")
+
+    if(livyConf.get(LivyConf.HA_MODE) == HighAvailabilitySettings.HA_ON){
+      info("Starting HA connection")
+      val electorService: CuratorElectorService = new CuratorElectorService(livyConf, server)
+      electorService.start()
+    }
+    else {
+      try {
+        server.start()
+        server.join()
+      } finally {
+        server.stop()
+      }
     }
   }
-
 }

server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperStateStore.scala

@@ -46,7 +46,7 @@ class ZooKeeperStateStore(
   }
 
   private val zkAddress = livyConf.get(LivyConf.RECOVERY_STATE_STORE_URL)
-  require(!zkAddress.isEmpty, s"Please config ${LivyConf.RECOVERY_STATE_STORE_URL.key}.")
+  require(!zkAddress.isEmpty, s"Please configure ${LivyConf.RECOVERY_STATE_STORE_URL.key}.")
   private val zkKeyPrefix = livyConf.get(ZK_KEY_PREFIX_CONF)
   private val retryValue = livyConf.get(ZK_RETRY_CONF)
   // a regex to match patterns like "m, n" where m and n both are integer values