From bb2e09f680bdd5cd5522d9011bd720f6f5a696e1 Mon Sep 17 00:00:00 2001
From: Cedric Beust <cedric.beust@le.com>
Date: Mon, 24 Apr 2017 09:56:53 -0700
Subject: [PATCH] Kill the Kobalt server after 10 minutes of inactivity.

---
 .../app/remote/GetDependencyGraphHandler.kt   |  1 +
 .../beust/kobalt/app/remote/SparkServer.kt    | 15 +++-
 .../com/beust/kobalt/app/remote/WatchDog.kt   | 68 +++++++++++++++++++
 3 files changed, 81 insertions(+), 3 deletions(-)
 create mode 100644 src/main/kotlin/com/beust/kobalt/app/remote/WatchDog.kt

diff --git a/src/main/kotlin/com/beust/kobalt/app/remote/GetDependencyGraphHandler.kt b/src/main/kotlin/com/beust/kobalt/app/remote/GetDependencyGraphHandler.kt
index 96e8658c..38eb6b0a 100644
--- a/src/main/kotlin/com/beust/kobalt/app/remote/GetDependencyGraphHandler.kt
+++ b/src/main/kotlin/com/beust/kobalt/app/remote/GetDependencyGraphHandler.kt
@@ -35,6 +35,7 @@ class GetDependencyGraphHandler : WebSocketListener {
 
     fun <T> sendWebsocketCommand(endpoint: RemoteEndpoint, commandName: String, payload: T,
             errorMessage: String? = null) {
+        SparkServer.watchDog.rearm()
         val json = Gson().toJson(WebSocketCommand(commandName, payload = Gson().toJson(payload),
                 errorMessage = errorMessage))
         endpoint.sendString(json)
diff --git a/src/main/kotlin/com/beust/kobalt/app/remote/SparkServer.kt b/src/main/kotlin/com/beust/kobalt/app/remote/SparkServer.kt
index 49f3e4a1..f78c8f37 100644
--- a/src/main/kotlin/com/beust/kobalt/app/remote/SparkServer.kt
+++ b/src/main/kotlin/com/beust/kobalt/app/remote/SparkServer.kt
@@ -5,6 +5,7 @@ import com.beust.kobalt.app.Templates
 import com.beust.kobalt.internal.PluginInfo
 import com.google.common.collect.ListMultimap
 import com.google.gson.Gson
+import org.slf4j.Logger
 import spark.ResponseTransformer
 import spark.Route
 import spark.Spark
@@ -14,6 +15,8 @@ class SparkServer(val cleanUpCallback: () -> Unit, val pluginInfo : PluginInfo)
 
     companion object {
         lateinit var cleanUpCallback: () -> Unit
+        val URL_QUIT = "/quit"
+        lateinit var watchDog: WatchDog
     }
 
     init {
@@ -28,19 +31,25 @@ class SparkServer(val cleanUpCallback: () -> Unit, val pluginInfo : PluginInfo)
     private fun jsonRoute(path: String, route: Route)
         = Spark.get(path, "application/json", route, JsonTransformer())
 
-    val log = org.slf4j.LoggerFactory.getLogger("SparkServer")
+    val log: Logger = org.slf4j.LoggerFactory.getLogger("SparkServer")
 
     override fun run(port: Int) {
+        val threadPool = Executors.newFixedThreadPool(2)
+        watchDog = WatchDog(port, 60 * 10 /* 10 minutes */, log)
+        threadPool.submit {
+            watchDog.run()
+        }
         log.debug("Server running")
         Spark.port(port)
         Spark.webSocket("/v1/getDependencyGraph", GetDependencyGraphHandler::class.java)
         Spark.get("/ping") { req, res ->
+            watchDog.rearm()
             log.debug("  Received ping")
             """ { "result" : "ok" } """
         }
-        Spark.get("/quit", { req, res ->
+        Spark.get(URL_QUIT, { req, res ->
             log.debug("  Received quit")
-            Executors.newFixedThreadPool(1).let { executor ->
+            threadPool.let { executor ->
                 executor.submit {
                     Thread.sleep(1000)
                     Spark.stop()
diff --git a/src/main/kotlin/com/beust/kobalt/app/remote/WatchDog.kt b/src/main/kotlin/com/beust/kobalt/app/remote/WatchDog.kt
new file mode 100644
index 00000000..86d5961f
--- /dev/null
+++ b/src/main/kotlin/com/beust/kobalt/app/remote/WatchDog.kt
@@ -0,0 +1,68 @@
+package com.beust.kobalt.app.remote
+
+import com.beust.kobalt.misc.warn
+import org.slf4j.Logger
+import java.net.HttpURLConnection
+import java.net.URL
+import java.time.Duration
+import java.time.LocalDateTime
+import java.time.OffsetDateTime
+import java.time.format.DateTimeFormatter
+
+/**
+ * Wakes up every `WAKE_UP_INTERVAL` and check if a certain period of time (`checkPeriod) has elapsed
+ * without being rearmed. If that time has elapsed, send a QUIT command to the Kobalt server.
+ */
+class WatchDog(val port: Int, val checkPeriodSeconds: Long, val log: Logger) {
+    private val WAKE_UP_INTERVAL: Duration = Duration.ofSeconds(60)
+    private val FORMAT: DateTimeFormatter = DateTimeFormatter.ofPattern("MM/d/y HH:mm:ss")
+
+    private var nextWakeUpMillis: Long = arm()
+    private var stop: Boolean = false
+
+    /**
+     * Rearm for another `checkPeriod`.
+     */
+    fun rearm() {
+        nextWakeUpMillis = arm()
+        log.info("Watchdog rearmed for " + format(nextWakeUpMillis))
+    }
+
+    /**
+     * Start the watch dog.
+     */
+    fun run() {
+        log.info("Next wake up:" + format(nextWakeUpMillis))
+        while (! stop) {
+            Thread.sleep(WAKE_UP_INTERVAL.toMillis())
+            val diffSeconds = (nextWakeUpMillis - System.currentTimeMillis()) / 1000
+            if (diffSeconds <= 0) {
+                log.info("Time to die")
+                stop = true
+            } else {
+                log.info("Dying in $diffSeconds seconds")
+            }
+        }
+
+        try {
+            val connection = (URL("http://localhost:$port" + SparkServer.URL_QUIT)
+                    .openConnection() as HttpURLConnection).apply {
+                requestMethod = "GET"
+            }
+            val code = connection.responseCode
+            if (code == 200) {
+                log.info("Successfully stopped the server")
+            } else {
+                warn("Couldn't stop the server, response: " + code)
+            }
+        } catch(ex: Exception) {
+            warn("Couldn't stop the server: " + ex.message, ex)
+        }
+    }
+
+    private fun arm() = System.currentTimeMillis() + (checkPeriodSeconds * 1000)
+
+    private fun toLocalDate(millis: Long) = LocalDateTime.ofEpochSecond(millis / 1000, 0, OffsetDateTime.now().offset)
+
+    private fun format(millis: Long) = FORMAT.format(toLocalDate(millis))
+}
\ No newline at end of file