[LIVY-11] Enable HA support

conf/livy.conf.template

@@ -100,6 +100,25 @@
 # 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 =
+
+# The list of ids for all livy servers used for HA
+# livy.server.ha.server-ids =
+
+# The list of hostnames for all livy servers used for HA
+# livy.server.ha.server-hostnames =
+
+# The list of server addresses for all livy servers used for HA
+# livy.server.ha.server-addresses =
+
 # 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

integration-test/src/main/scala/org/apache/livy/test/framework/MiniCluster.scala

@@ -154,6 +154,7 @@ object MiniLivyMain extends MiniClusterBase {
     saveProperties(livyConf, new File(configPath + "/livy.conf"))
 
     val server = new LivyServer()
+    server.init()
     server.start()
     server.livyConf.set(LivyConf.ENABLE_HIVE_CONTEXT, false)
     // Write a serverUrl.conf file to the conf directory with the location of the Livy

server/pom.xml

@@ -251,6 +251,12 @@
       <scope>test</scope>
     </dependency>
 
+    <dependency>
+      <groupId>org.springframework</groupId>
+      <artifactId>spring-web</artifactId>
+      <version>5.2.0.RELEASE</version>
+    </dependency>
+
     <dependency>
       <groupId>org.scalatra</groupId>
       <artifactId>scalatra-test_${scala.binary.version}</artifactId>

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

@@ -196,6 +196,25 @@ 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", "")
+
+  // The ids of all servers used in HA
+  val HA_SERVER_IDS = Entry("livy.server.ha.server-ids", "")
+
+  // The hostnames of all servers used in HA
+  val HA_SERVER_HOSTNAMES = Entry("livy.server.ha.server-hostnames", "")
+
   /**
    * 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:///.
@@ -406,3 +425,4 @@ class LivyConf(loadDefaults: Boolean) extends ClientConf[LivyConf](null) {
   }
 
 }
+

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

@@ -0,0 +1,150 @@
+/*
+ * 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.net.InetAddress
+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")
+}
+
+object HAState extends Enumeration{
+  type HAState = Value
+  val Active, Standby = Value
+}
+
+
+class CuratorElectorService(livyConf : LivyConf, livyServer : LivyServer,
+    mockCuratorClient: Option[CuratorFramework] = None,
+    mockLeaderLatch: Option[LeaderLatch] = None)
+  extends LeaderLatchListener
+  with Logging
+{
+  import CuratorElectorService._
+  import HAState._
+
+  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")
+  }
+
+  var server : LivyServer = livyServer
+
+  val client: CuratorFramework = mockCuratorClient.getOrElse {
+    CuratorFrameworkFactory.newClient(haAddress, retryPolicy)
+  }
+  val leaderKey = s"/$haKeyPrefix/leader"
+
+  val leaderIds = livyConf.configToSeq(LivyConf.HA_SERVER_IDS)
+  val leaderHostnames = livyConf.configToSeq(LivyConf.HA_SERVER_HOSTNAMES)
+
+  var leaderLatch = mockLeaderLatch.getOrElse {
+    new LeaderLatch(client, leaderKey, getCurrentId())
+  }
+  leaderLatch.addListener(this)
+
+  var currentState = HAState.Standby
+  override def isLeader() {
+    transitionToActive()
+  }
+
+  override def notLeader(){
+    transitionToStandby()
+  }
+
+  def getCurrentId(): String = {
+    val currentHostname = java.net.InetAddress.getLocalHost().getCanonicalHostName();
+    debug("This server's current hostname is: " + currentHostname)
+    debug("This hostnames for valid leaders are: " + leaderHostnames)
+    val currentId = leaderIds(leaderHostnames indexOf currentHostname)
+    debug("This server's designated ID is: " + currentId)
+    currentId
+  }
+
+  def getActiveHostname(): String = {
+    val activeLeaderId = leaderLatch.getLeader().getId()
+    val activeAddress = leaderHostnames(leaderIds indexOf activeLeaderId)
+    activeAddress
+  }
+
+  def start(): Unit = {
+    server.initHa(this)
+    transitionToStandby()
+
+    server.start()
+    client.start()
+    leaderLatch.start()
+
+    try {
+      Thread.currentThread.join()
+    } finally {
+      transitionToStandby()
+    }
+  }
+
+  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.restart()
+      currentState = HAState.Active
+      info("Transition complete")
+    }
+  }
+
+  def transitionToStandby(): Unit = {
+    info("Transitioning to Standby state")
+    if(currentState == HAState.Standby) {
+      info("Already in Standby State")
+    }
+    else {
+      currentState = HAState.Standby
+      info("Transition complete")
+    }
+  }
+}
+

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

@@ -0,0 +1,84 @@
+/*
+ * 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.IOException
+import java.net.URL
+import javax.servlet.Filter
+import javax.servlet.FilterChain
+import javax.servlet.FilterConfig
+import javax.servlet.ServletContext
+import javax.servlet.ServletException
+import javax.servlet.ServletRequest
+import javax.servlet.ServletRequestWrapper
+import javax.servlet.ServletResponse
+import javax.servlet.http.HttpServletRequest
+import javax.servlet.http.HttpServletRequestWrapper
+import javax.servlet.http.HttpServletResponse
+
+import org.springframework.web.util.UriComponentsBuilder
+
+import org.apache.livy.{LivyConf, Logging}
+
+class DomainRedirectionFilter(haService: CuratorElectorService) extends Filter
+  with Logging
+{
+
+  val METHODS_TO_IGNORE = Set("GET", "OPTIONS", "HEAD")
+
+  val HEADER_NAME = "X-Requested-By"
+
+  def isLeader(): Boolean = {
+    haService.currentState == HAState.Active
+  }
+
+  override def init(filterConfig: FilterConfig): Unit = {}
+
+  override def doFilter(request: ServletRequest,
+                        response: ServletResponse,
+                        chain: FilterChain): Unit = {
+    if (!isLeader()) {
+        debug("active leader's hostnames are:" + haService.getActiveHostname())
+        debug("current id:" + haService.getCurrentId())
+        val httpRequest = request.asInstanceOf[HttpServletRequest]
+        val queryOpt: Option[String] = Option(httpRequest.getQueryString())
+        val requestURL = httpRequest.getRequestURL().toString()
+        debug("requested url: " + requestURL)
+
+        val builder = UriComponentsBuilder.fromHttpUrl(requestURL)
+        val activeURL = builder.host(haService.getActiveHostname()).toUriString()
+        val redirectURL = if (queryOpt.isEmpty) activeURL else activeURL + "?" + queryOpt.get
+        debug("redirected url:" + redirectURL)
+
+        val httpServletResponse = response.asInstanceOf[HttpServletResponse];
+        val redirectMsg = "This is a standby Livy Instance. The redirect url is: " + redirectURL
+        val out = httpServletResponse.getWriter()
+        // scalastyle:off println
+        out.println(redirectMsg)
+        // scalastyle:on println
+
+        httpServletResponse.setHeader("Location", redirectURL)
+        httpServletResponse.setStatus(HttpServletResponse.SC_TEMPORARY_REDIRECT)
+    } else {
+        chain.doFilter(request, response);
+    }
+  }
+
+  override def destroy(): Unit = {}
+}
+