Add support for suspending warmup tasks

GitOrigin-RevId: 2b697d9b17e3f8c95f7f2e684e6d094de1ced62d

Author: mateuszmrozewski

misk-warmup/api/misk-warmup.api

@@ -1,3 +1,9 @@
+public abstract class misk/warmup/SuspendingWarmupTask : misk/warmup/WarmupTask {
+	public fun <init> ()V
+	public fun execute ()V
+	public abstract fun executeSuspending (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+}
+
 public final class misk/warmup/WarmupModule : misk/inject/KAbstractModule {
 	public fun <init> (Ljava/lang/String;Lcom/google/inject/Key;)V
 }

misk-warmup/build.gradle.kts

@@ -8,12 +8,14 @@ plugins {
 
 dependencies {
   api(libs.guice)
+  api(libs.kotlinxCoroutinesCore)
   api(project(":misk-inject"))
   implementation(libs.guava)
   implementation(libs.jakartaInject)
   implementation(libs.loggingApi)
-  implementation(project(":wisp:wisp-logging"))
+  implementation(project(":misk-api"))
   implementation(project(":misk-core"))
+  implementation(project(":wisp:wisp-logging"))
 
   testImplementation(libs.assertj)
   testImplementation(libs.junitApi)

misk-warmup/src/main/kotlin/misk/warmup/SuspendingWarmupTask.kt

@@ -0,0 +1,21 @@
+package misk.warmup
+
+import misk.annotation.ExperimentalMiskApi
+
+/**
+ * @see WarmupTask
+ */
+@ExperimentalMiskApi
+abstract class SuspendingWarmupTask : WarmupTask() {
+  override fun execute() {
+    throw UnsupportedOperationException("Suspending warmup task should implement only executeSuspending method")
+  }
+
+  /**
+   * Perform production-like work to cause caches to be seeded, pools to be filled, and hot spots to
+   * be compiled. This should return once warmup is complete.
+   *
+   * This function will be executed on a dedicated warmup thread.
+   */
+  abstract suspend fun executeSuspending()
+}

misk-warmup/src/main/kotlin/misk/warmup/WarmupRunner.kt

@@ -10,6 +10,8 @@ import java.util.concurrent.atomic.AtomicInteger
 import jakarta.inject.Inject
 import com.google.inject.Provider
 import jakarta.inject.Singleton
+import kotlinx.coroutines.runBlocking
+import misk.annotation.ExperimentalMiskApi
 
 /**
  * This class is a health check to defer production traffic until all warmup tasks have completed.
@@ -21,6 +23,7 @@ import jakarta.inject.Singleton
  * Note that if a warmup task fails by throwing an exception, that is not fatal. This will report
  * itself as healthy, and early call latency might not be as low as it should be.
  */
+@OptIn(ExperimentalMiskApi::class)
 @Singleton
 internal class WarmupRunner @Inject constructor(
   private val executorServiceFactory: ExecutorServiceFactory,
@@ -36,8 +39,12 @@ internal class WarmupRunner @Inject constructor(
       executorService.submit {
         val stopwatch = Stopwatch.createStarted()
         try {
-          val task = taskProvider.get()
-          task.execute()
+          when (val task = taskProvider.get()) {
+            is SuspendingWarmupTask -> runBlocking {
+              task.executeSuspending()
+            }
+            else -> task.execute()
+          }
           logger.info { "Warmup task $name completed after $stopwatch" }
         } catch (t: Throwable) {
           logger.error(t) { "Warmup task $name crashed after $stopwatch" }

misk-warmup/src/test/kotlin/misk/warmup/WarmupTest.kt

@@ -22,6 +22,7 @@ import java.util.concurrent.BlockingDeque
 import java.util.concurrent.LinkedBlockingDeque
 import jakarta.inject.Inject
 import jakarta.inject.Singleton
+import misk.annotation.ExperimentalMiskApi
 
 internal class WarmupTest {
   private val events = LinkedBlockingDeque<String>()
@@ -52,6 +53,29 @@ internal class WarmupTest {
       .startsWith("Warmup task LoggingWarmupTask completed")
   }
 
+  @Test
+  fun `suspending warmup task`() {
+    startUpAndShutDown(
+      ServiceModule<LoggingService>(),
+      WarmupModule<SuspendingLoggingWarmupTask>(),
+    )
+
+    assertThat(events).containsExactly(
+      "LoggingService startUp",
+      "SuspendingLoggingWarmupTask created on warmup-0",
+      "SuspendingLoggingWarmupTask warming on warmup-0 @coroutine#1",
+      "HealthChecks all passed",
+      "LoggingService shutDown",
+    )
+
+    assertThat(logCollector.takeMessage(minLevel = Level.INFO))
+      .isEqualTo("Starting ready service")
+    assertThat(logCollector.takeMessage(minLevel = Level.INFO))
+      .isEqualTo("Running warmup tasks: [SuspendingLoggingWarmupTask]")
+    assertThat(logCollector.takeMessage(minLevel = Level.INFO))
+      .startsWith("Warmup task SuspendingLoggingWarmupTask completed")
+  }
+
   @Test
   fun `service is healthy even after warmup task throws`() {
     startUpAndShutDown(
@@ -116,6 +140,20 @@ internal class WarmupTest {
     }
   }
 
+  @OptIn(ExperimentalMiskApi::class)
+  @Singleton
+  class SuspendingLoggingWarmupTask @Inject constructor(
+    private val events: BlockingDeque<String>
+  ) : SuspendingWarmupTask() {
+    init {
+      events += "SuspendingLoggingWarmupTask created on ${Thread.currentThread().name}"
+    }
+
+    override suspend fun executeSuspending() {
+      events += "SuspendingLoggingWarmupTask warming on ${Thread.currentThread().name}"
+    }
+  }
+
   /**
    * This test doesn't use `@MiskTest` so we can start up and shut down the injector in the test.
    * It also customizes the modules per-test.