twitter-server: Add c.t.server.Hook#onExit lifecycle callback

Problem

There may be cases where a Hook needs to run in the `App#onExit`
lifecycle.

Solution

Add an `onExit` callback function which can be implemented by a
`c.t.server.Hook` and implement calling it in the Hooks trait
which loads and fires service-loaded hooks.

Result

A `c.t.server.Hook` can define logic which will be triggered during
the `onExit` lifecycle phase of an App.

Differential Revision: https://phabricator.twitter.biz/D198379

CHANGES

@@ -7,6 +7,14 @@ Note that ``PHAB_ID=#`` and ``RB_ID=#`` correspond to associated messages in com
 Unreleased
 ----------
 
+New Features
+~~~~~~~~~~~~
+
+  * Add `onExit` lifecycle callback to `c.t.server.Hook` (which is now an abstract class) to allow
+    implemented hooks to execute functions in the `App#onExit` lifecycle phase. Note:
+    `c.t.server.Hook#premain` now has a default implementation and requires the `override` modifier.
+    ``PHAB_ID=D198379``
+
 18.7.0
 -------
 

server/src/main/scala/com/twitter/server/Hook.scala

@@ -0,0 +1,18 @@
+package com.twitter.server
+
+/**
+ * Defines a hook into an [[com.twitter.app.App]].
+ */
+abstract class Hook {
+
+  /**
+   * Executed in `App#premain` which is right before the user's main is invoked.
+   */
+  def premain(): Unit = ()
+
+  /**
+   * Executed in `App#onExit` which is when shutdown is requested. All registered exit hooks
+   * run in parallel and are executed after all `App#postmain` functions complete.
+   */
+  def onExit(): Unit = ()
+}

server/src/main/scala/com/twitter/server/Hooks.scala

@@ -3,29 +3,19 @@ package com.twitter.server
 import com.twitter.app.App
 import com.twitter.finagle.util.LoadService
 
-/**
- * Defines a hook into an App.
- */
-trait Hook {
-  def premain(): Unit
-}
-
-/**
- * Create a new hook for the given App. NewHooks are
- * service-loaded.
- */
-trait NewHook extends (App => Hook)
-
 /**
  * Mix-in to include service-loaded hooks.
  */
 trait Hooks { self: App =>
-  private val hooks = LoadService[NewHook]() map { newHook =>
+  private val hooks = LoadService[NewHook]().map { newHook =>
     newHook(self)
   }
 
   premain {
     for (hook <- hooks)
       hook.premain()
   }
-}
+
+  for (hook <- hooks)
+    onExit { hook.onExit() }
+}

server/src/main/scala/com/twitter/server/NewHook.scala

@@ -0,0 +1,28 @@
+package com.twitter.server
+
+import com.twitter.app.App
+
+/**
+ * Create a new hook for the given App. NewHooks are service-loaded.
+ *
+ * To use, extend the [[NewHook]] trait and implement an apply method which
+ * returns a [[Hook]] implementation, e.g.,
+ *
+ * Add the Hook as a service-loaded class in
+ * /META-INF/services/com.twitter.server.NewHook
+ *
+ * {{{
+ *   class MyHook extends NewHook {
+ *     def apply(app: App) = new Hook {
+ *
+ *       override def premain(): Unit = ???
+ *
+ *       override def onExit(): Unit = ???
+ *     }
+ *   }
+ * }}}
+ *
+ * @see [[com.twitter.server.Hook]]
+ * @see [[com.twitter.finagle.util.LoadService]]
+ */
+trait NewHook extends (App => Hook)

server/src/test/java/BUILD

@@ -1,18 +1,20 @@
 junit_tests(
-  dependencies=[
-    '3rdparty/jvm/junit:junit',
-    '3rdparty/jvm/org/mockito:mockito-all',
-    '3rdparty/jvm/org/scala-lang:scala-library',
-    '3rdparty/jvm/org/scalacheck:scalacheck',
-    'finagle/finagle-core/src/main/scala:scala',
-    'finagle/finagle-http/src/main/scala',
-    'twitter-server/server/src/main/scala',
-    'util/util-app/src/main/scala:scala',
-    'util/util-core/src/main/scala:scala',
-    'util/util-lint/src/main/scala:scala',
-    'util/util-slf4j-api/src/main/scala:scala',
-    'util/util-stats/src/main/scala:scala',
-  ],
-  fatal_warnings=False,
-  sources=rglobs('*.java')
-)
+    sources = rglobs("*.java"),
+    fatal_warnings = False,
+    dependencies = [
+        "3rdparty/jvm/junit",
+        "3rdparty/jvm/org/mockito:mockito-all",
+        "3rdparty/jvm/org/scala-lang:scala-library",
+        "3rdparty/jvm/org/scalacheck",
+        "3rdparty/jvm/org/slf4j:slf4j-api",
+        "3rdparty/jvm/org/slf4j:slf4j-nop",
+        "finagle/finagle-core/src/main/scala",
+        "finagle/finagle-http/src/main/scala",
+        "twitter-server/server/src/main/scala",
+        "util/util-app/src/main/scala",
+        "util/util-core/src/main/scala",
+        "util/util-lint/src/main/scala",
+        "util/util-slf4j-api/src/main/scala",
+        "util/util-stats/src/main/scala",
+    ],
+)

server/src/test/scala/BUILD

@@ -8,6 +8,8 @@ junit_tests(
         "3rdparty/jvm/org/mockito:mockito-all",
         "3rdparty/jvm/org/scalacheck",
         "3rdparty/jvm/org/scalatest",
+        "3rdparty/jvm/org/slf4j:slf4j-api",
+        "3rdparty/jvm/org/slf4j:slf4j-nop",
         "finagle/finagle-core/src/main/scala",
         "finagle/finagle-http/src/main/scala",
         "finagle/finagle-toggle",
@@ -21,4 +23,4 @@ junit_tests(
         "util/util-stats/src/main/scala",
         "util/util-tunable/src/main/scala",
     ],
-)
+)

server/src/test/scala/com/twitter/server/HookTest.scala

@@ -0,0 +1,70 @@
+package com.twitter.server
+
+import org.scalatest.FunSuite
+import scala.collection.mutable
+
+class OnExitHook1 extends Hook {
+  override def onExit(): Unit = {
+    throw new Exception("onExit 1")
+  }
+}
+
+class OnExitHook2 extends Hook {
+  override def onExit(): Unit = {
+    throw new Exception("onExit 2")
+  }
+}
+
+class HookTestTwitterServer(hooks: Seq[Hook]) extends TwitterServer {
+  override val defaultAdminPort = 0
+
+  val bootstrapSeq: mutable.MutableList[Symbol] = mutable.MutableList.empty[Symbol]
+
+  def main(): Unit = {
+    bootstrapSeq += 'Main
+  }
+
+  override def exitOnError(throwable: Throwable): Unit = {
+    throw throwable
+  }
+
+  init {
+    bootstrapSeq += 'Init
+  }
+
+  premain {
+    bootstrapSeq += 'PreMain
+    for (hook <- hooks)
+      hook.premain()
+  }
+
+  onExit {
+    bootstrapSeq += 'Exit
+  }
+
+  for (hook <- hooks)
+    onExit { hook.onExit() }
+
+  postmain {
+    bootstrapSeq += 'PostMain
+  }
+}
+
+class HookTest extends FunSuite {
+
+  val hooks: Seq[Hook] = Seq(new OnExitHook1(), new OnExitHook2())
+
+  test("Hooks which error in OnExit are properly captured in onExit blocks") {
+    val twitterServer = new HookTestTwitterServer(hooks)
+
+    val e = intercept[Exception] {
+      twitterServer.main(args = Array.empty[String])
+    }
+
+    assert(e.getSuppressed.length == 2)
+    assert(
+      twitterServer.bootstrapSeq ==
+        Seq('Init, 'PreMain, 'Main, 'PostMain, 'Exit)
+    )
+  }
+}

server/src/test/scala/com/twitter/server/TwitterServerTest.scala

@@ -3,17 +3,15 @@ package com.twitter.server
 import com.twitter.finagle.Service
 import com.twitter.finagle.http.{Request, Response}
 import com.twitter.util._
-import java.net.{InetAddress, InetSocketAddress}
 import org.junit.runner.RunWith
 import org.scalatest.FunSuite
 import org.scalatest.junit.JUnitRunner
 import scala.collection.mutable
 
 class TestTwitterServer extends TwitterServer {
-  override val adminPort =
-    flag("admin.port", new InetSocketAddress(InetAddress.getLoopbackAddress, 0), "")
+  override val defaultAdminPort = 0
 
-  val bootstrapSeq = mutable.MutableList.empty[Symbol]
+  val bootstrapSeq: mutable.MutableList[Symbol] = mutable.MutableList.empty[Symbol]
 
   def main(): Unit = {
     bootstrapSeq += 'Main
@@ -70,4 +68,4 @@ class TwitterServerTest extends FunSuite {
     twitterServer.main(args = Array.empty[String])
     assert(twitterServer.bootstrapSeq == Seq('Init, 'PreMain, 'Main, 'Exit, 'PostMain))
   }
-}
+}