From 49673c43144abe774e1d9c5923cc9110c95bfb96 Mon Sep 17 00:00:00 2001
From: Kevin Kim <95660882+kevink-sq@users.noreply.github.com>
Date: Tue, 24 Sep 2024 11:38:01 -0700
Subject: [PATCH] Undo `shutdown_sleep_ms` workaround for Jetty shutdown
 behavior

GitOrigin-RevId: 59b25694ddd68a7157b42a89d6ce7faa19d854bb
---
 .../kotlin/misk/testing/MiskTestExtension.kt  | 20 -------
 misk/src/main/kotlin/misk/web/WebConfig.kt    |  6 +-
 .../kotlin/misk/web/jetty/JettyService.kt     |  4 +-
 .../test/kotlin/misk/web/JettyShutdownTest.kt |  1 -
 .../web/actions/LivenessCheckActionTest.kt    | 57 +++++--------------
 5 files changed, 16 insertions(+), 72 deletions(-)

diff --git a/misk-testing/src/main/kotlin/misk/testing/MiskTestExtension.kt b/misk-testing/src/main/kotlin/misk/testing/MiskTestExtension.kt
index dbd2ec29a8b..8be1a30772f 100644
--- a/misk-testing/src/main/kotlin/misk/testing/MiskTestExtension.kt
+++ b/misk-testing/src/main/kotlin/misk/testing/MiskTestExtension.kt
@@ -196,19 +196,6 @@ private fun ExtensionContext.startService(): Boolean {
   }
 }
 
-
-private fun ServiceManager.jettyIsRunning() : Boolean {
-  return this
-    .servicesByState()
-    .values()
-    .toList()
-    .any { it.toString().startsWith("JettyService") }
-}
-
-private fun ExtensionContext.stopJetty() {
-  this.retrieve<Injector>("injector").getInstance<JettyService>().stop()
-}
-
 // The injector is reused across tests if
 //   1. The tests module(s) used in the test extend ReusableTestModules, AND
 //   2. The environment variable MISK_TEST_REUSE_INJECTOR is set to true
@@ -232,11 +219,4 @@ private fun ExtensionContext.getExternalDependencies(): Iterable<ExternalDepende
 
 private fun ServiceManager.stop(context: ExtensionContext) {
   this.stopAsync().awaitStopped(30, TimeUnit.SECONDS)
-  /**
-   * By default, Jetty shutdown is not managed by Guava so needs to be
-   * done explicitly. This is being done in MiskApplication.
-   */
-  if (this.jettyIsRunning()) {
-    context.stopJetty()
-  }
 }
diff --git a/misk/src/main/kotlin/misk/web/WebConfig.kt b/misk/src/main/kotlin/misk/web/WebConfig.kt
index 99b3936da81..bfd6a99999e 100644
--- a/misk/src/main/kotlin/misk/web/WebConfig.kt
+++ b/misk/src/main/kotlin/misk/web/WebConfig.kt
@@ -105,11 +105,7 @@ data class WebConfig @JvmOverloads constructor(
     log_level = concurrency_limiter_log_level,
   ),
 
-  /**
-   * The number of milliseconds to sleep before commencing service shutdown. 0 or
-   * greater will defer the actual shutdown of Jetty to after all Guava managed
-   * services are shutdown.
-   * */
+  /** The number of milliseconds to sleep before commencing service shutdown. */
   val shutdown_sleep_ms: Int = 0,
 
   /** The maximum allowed size in bytes for the HTTP request line and HTTP request headers. */
diff --git a/misk/src/main/kotlin/misk/web/jetty/JettyService.kt b/misk/src/main/kotlin/misk/web/jetty/JettyService.kt
index 9ca4305c065..dc47319facf 100644
--- a/misk/src/main/kotlin/misk/web/jetty/JettyService.kt
+++ b/misk/src/main/kotlin/misk/web/jetty/JettyService.kt
@@ -400,9 +400,7 @@ class JettyService @Inject internal constructor(
     //
     // Ideally we could call jetty.awaitInflightRequests() but that's not available
     // for us.
-    //
-    // Default is to shutdown jetty after all guava managed services are shutdown.
-    if (webConfig.shutdown_sleep_ms >= 0) {
+    if (webConfig.shutdown_sleep_ms > 0) {
       sleep(webConfig.shutdown_sleep_ms.toLong())
     } else {
       stop()
diff --git a/misk/src/test/kotlin/misk/web/JettyShutdownTest.kt b/misk/src/test/kotlin/misk/web/JettyShutdownTest.kt
index 6211a57f1e5..87e03ac5974 100644
--- a/misk/src/test/kotlin/misk/web/JettyShutdownTest.kt
+++ b/misk/src/test/kotlin/misk/web/JettyShutdownTest.kt
@@ -80,7 +80,6 @@ internal abstract class AbstractJettyShutdownTest {
 
     jetty.stopAsync()
     try {
-      jetty.stop()
       jetty.awaitTerminated(timeoutMs, TimeUnit.MILLISECONDS)
     } catch (e: TimeoutException) {
       assertThat(timeoutExpected).isTrue()
diff --git a/misk/src/test/kotlin/misk/web/actions/LivenessCheckActionTest.kt b/misk/src/test/kotlin/misk/web/actions/LivenessCheckActionTest.kt
index bc37624f9b6..56f64797882 100644
--- a/misk/src/test/kotlin/misk/web/actions/LivenessCheckActionTest.kt
+++ b/misk/src/test/kotlin/misk/web/actions/LivenessCheckActionTest.kt
@@ -4,63 +4,34 @@ import com.google.common.util.concurrent.ServiceManager
 import com.google.inject.util.Modules
 import jakarta.inject.Inject
 import misk.MiskTestingServiceModule
-import misk.ReadyService
-import misk.ServiceModule
-import misk.inject.KAbstractModule
+import misk.services.FakeServiceModule
 import misk.testing.MiskTest
 import misk.testing.MiskTestModule
-import misk.time.ClockModule
-import misk.web.FakeService
 import misk.web.WebActionModule
-import misk.web.WebServerTestingModule
-import misk.web.jetty.JettyService
-import okhttp3.HttpUrl
-import okhttp3.OkHttpClient
-import okhttp3.Request
-import okhttp3.Response
 import org.assertj.core.api.Assertions.assertThat
-import org.assertj.core.api.Assertions.assertThatThrownBy
 import org.junit.jupiter.api.Test
-import java.time.Duration
 
-@MiskTest(startService = true)
+@MiskTest
 class LivenessCheckActionTest {
   @MiskTestModule
-  val module = TestModule()
-  class TestModule : KAbstractModule() {
-    override fun configure() {
-      install(Modules.override(MiskTestingServiceModule()).with(ClockModule()))
-      install(WebServerTestingModule())
-      install(WebActionModule.create<LivenessCheckAction>())
-      install(ServiceModule<FakeService>().enhancedBy<ReadyService>())
-    }
-  }
+  val module = Modules.combine(
+    MiskTestingServiceModule(),
+    FakeServiceModule(),
+    WebActionModule.create<LivenessCheckAction>()
+  )
 
-  @Inject lateinit var jettyService: JettyService
+  @Inject lateinit var livenessCheckAction: LivenessCheckAction
   @Inject lateinit var serviceManager: ServiceManager
 
-  private val client = OkHttpClient().newBuilder()
-    .readTimeout(Duration.ofSeconds(30))
-    .connectTimeout(Duration.ofSeconds(30))
-    .writeTimeout(Duration.ofSeconds(30))
-    .build()
-
-  /**
-   * Liveness should only fail when Jetty is shut down.
-   */
   @Test
-  fun liveness() {
-    val url = jettyService.httpServerUrl.resolve("/_liveness")!!
-    assertThat(get(url).code).isEqualTo(200)
+  fun livenessDependsOnServiceState() {
+    serviceManager.startAsync()
+    serviceManager.awaitHealthy()
+    assertThat(livenessCheckAction.livenessCheck().statusCode).isEqualTo(200)
     serviceManager.stopAsync()
     serviceManager.awaitStopped()
-    assertThat(get(url).code).isEqualTo(200)
-    jettyService.stop()
-    assertThatThrownBy { get(url)}
-  }
 
-  private fun get(url: HttpUrl) : Response{
-    val req = Request.Builder().url(url).build();
-    return client.newCall(req).execute();
+    // liveness should not fail for terminated services (only failed)
+    assertThat(livenessCheckAction.livenessCheck().statusCode).isEqualTo(200)
   }
 }