Bound the handler drain + de-flake the drain test

gradle/libs.versions.toml

@@ -12,7 +12,7 @@ googleAuth = "1.39.1"
 googleHttp = "2.0.0"
 guava = "33.5.0-jre"
 hoplite = "2.7.5"
-jackson = "2.21.2"
+jackson = "3.2.1"
 jooq = "3.19.29"
 junit = "5.14.2"
 kotest = "6.0.7"
@@ -123,12 +123,11 @@ hopliteYaml = { module = "com.sksamuel.hoplite:hoplite-yaml", version.ref = "hop
 hsqldb = { module = "org.hsqldb:hsqldb", version = "2.7.4" }
 jCommander = { module = "com.beust:jcommander", version = "1.82" }
 jacksonAnnotations = { module = "com.fasterxml.jackson.core:jackson-annotations" }
-jacksonBom = { module = "com.fasterxml.jackson:jackson-bom", version.ref = "jackson" }
-jacksonCore = { module = "com.fasterxml.jackson.core:jackson-core" }
-jacksonDatabind = { module = "com.fasterxml.jackson.core:jackson-databind" }
-jacksonDataformatYaml = { module = "com.fasterxml.jackson.dataformat:jackson-dataformat-yaml" }
-jacksonJsr310 = { module = "com.fasterxml.jackson.datatype:jackson-datatype-jsr310" }
-jacksonKotlin = { module = "com.fasterxml.jackson.module:jackson-module-kotlin" }
+jacksonBom = { module = "tools.jackson:jackson-bom", version.ref = "jackson" }
+jacksonCore = { module = "tools.jackson.core:jackson-core" }
+jacksonDatabind = { module = "tools.jackson.core:jackson-databind" }
+jacksonDataformatYaml = { module = "tools.jackson.dataformat:jackson-dataformat-yaml" }
+jacksonKotlin = { module = "tools.jackson.module:jackson-module-kotlin" }
 jakartaInject = { module = "jakarta.inject:jakarta.inject-api", version = "2.0.1" }
 javaxAnnotation = { module = "javax.annotation:javax.annotation-api", version = "1.3.2" }
 javaxInject = { module = "javax.inject:javax.inject", version = "1" }
@@ -192,8 +191,8 @@ kubernetesClient = { module = "io.kubernetes:client-java", version = "27.0.0" }
 kubernetesClientApi = { module = "io.kubernetes:client-java-api", version = "27.0.0" }
 launchDarkly = { module = "com.launchdarkly:launchdarkly-java-server-sdk", version = "7.10.2" }
 lettuceCore = { module = "io.lettuce:lettuce-core", version = "6.8.1.RELEASE" }
-logbackClassic = { module = "ch.qos.logback:logback-classic", version = "1.5.33" }
-logbackCore = { module = "ch.qos.logback:logback-core", version = "1.5.33" }
+logbackClassic = { module = "ch.qos.logback:logback-classic", version = "1.5.36" }
+logbackCore = { module = "ch.qos.logback:logback-core", version = "1.5.36" }
 loggingApi = { module = "io.github.microutils:kotlin-logging", version = "3.0.5" }
 mcpKotlinSdkClient = { module = "io.modelcontextprotocol:kotlin-sdk-client", version.ref = "mcp" }
 mcpKotlinSdkCore = { module = "io.modelcontextprotocol:kotlin-sdk-core", version.ref = "mcp" }

misk-aws2-sqs/api/misk-aws2-sqs.api

@@ -137,6 +137,7 @@ public final class misk/aws2/sqs/jobqueue/StaticDeadLetterQueueProvider : misk/a
 public final class misk/aws2/sqs/jobqueue/Subscriber {
 	public static final field Companion Lmisk/aws2/sqs/jobqueue/Subscriber$Companion;
 	public fun <init> (Lmisk/jobqueue/QueueName;Lmisk/aws2/sqs/jobqueue/config/SqsQueueConfig;Lmisk/jobqueue/QueueName;Lmisk/jobqueue/v2/JobHandler;Lkotlinx/coroutines/channels/Channel;Lsoftware/amazon/awssdk/services/sqs/SqsAsyncClient;Lmisk/aws2/sqs/jobqueue/SqsQueueResolver;Lmisk/aws2/sqs/jobqueue/SqsMetrics;Lcom/squareup/moshi/Moshi;Ljava/time/Clock;Lio/opentracing/Tracer;Lmisk/aws2/sqs/jobqueue/VisibilityTimeoutCalculator;Lmisk/inject/AsyncSwitch;)V
+	public final fun cancelInFlightReceives ()V
 	public final fun getAsyncSwitch ()Lmisk/inject/AsyncSwitch;
 	public final fun getChannel ()Lkotlinx/coroutines/channels/Channel;
 	public final fun getClient ()Lsoftware/amazon/awssdk/services/sqs/SqsAsyncClient;
@@ -152,6 +153,7 @@ public final class misk/aws2/sqs/jobqueue/Subscriber {
 	public final fun getVisibilityTimeoutCalculator ()Lmisk/aws2/sqs/jobqueue/VisibilityTimeoutCalculator;
 	public final fun poll (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public final fun run (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public final fun stop ()V
 }
 
 public final class misk/aws2/sqs/jobqueue/Subscriber$Companion {
@@ -202,6 +204,8 @@ public final class misk/aws2/sqs/jobqueue/config/SqsConfig : misk/config/Config
 }
 
 public final class misk/aws2/sqs/jobqueue/config/SqsQueueConfig {
+	public static final field Companion Lmisk/aws2/sqs/jobqueue/config/SqsQueueConfig$Companion;
+	public static final field DEFAULT_SHUTDOWN_GRACE_PERIOD_MS J
 	public fun <init> ()V
 	public fun <init> (I)V
 	public fun <init> (II)V
@@ -212,8 +216,12 @@ public final class misk/aws2/sqs/jobqueue/config/SqsQueueConfig {
 	public fun <init> (IIIIZLjava/lang/Integer;Ljava/lang/Integer;)V
 	public fun <init> (IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;)V
 	public fun <init> (IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;)V
-	public synthetic fun <init> (IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun <init> (IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Long;)V
+	public fun <init> (IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Long;Ljava/lang/Long;)V
+	public synthetic fun <init> (IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Long;Ljava/lang/Long;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun component1 ()I
+	public final fun component10 ()Ljava/lang/Long;
+	public final fun component11 ()Ljava/lang/Long;
 	public final fun component2 ()I
 	public final fun component3 ()I
 	public final fun component4 ()I
@@ -222,8 +230,8 @@ public final class misk/aws2/sqs/jobqueue/config/SqsQueueConfig {
 	public final fun component7 ()Ljava/lang/Integer;
 	public final fun component8 ()Ljava/lang/String;
 	public final fun component9 ()Ljava/lang/String;
-	public final fun copy (IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;)Lmisk/aws2/sqs/jobqueue/config/SqsQueueConfig;
-	public static synthetic fun copy$default (Lmisk/aws2/sqs/jobqueue/config/SqsQueueConfig;IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;ILjava/lang/Object;)Lmisk/aws2/sqs/jobqueue/config/SqsQueueConfig;
+	public final fun copy (IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Long;Ljava/lang/Long;)Lmisk/aws2/sqs/jobqueue/config/SqsQueueConfig;
+	public static synthetic fun copy$default (Lmisk/aws2/sqs/jobqueue/config/SqsQueueConfig;IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Long;Ljava/lang/Long;ILjava/lang/Object;)Lmisk/aws2/sqs/jobqueue/config/SqsQueueConfig;
 	public fun equals (Ljava/lang/Object;)Z
 	public final fun getAccount_id ()Ljava/lang/String;
 	public final fun getChannel_capacity ()I
@@ -232,12 +240,17 @@ public final class misk/aws2/sqs/jobqueue/config/SqsQueueConfig {
 	public final fun getMax_number_of_messages ()I
 	public final fun getParallelism ()I
 	public final fun getRegion ()Ljava/lang/String;
+	public final fun getShutdown_grace_period_ms ()Ljava/lang/Long;
+	public final fun getShutdown_timeout_ms ()Ljava/lang/Long;
 	public final fun getVisibility_timeout ()Ljava/lang/Integer;
 	public final fun getWait_timeout ()Ljava/lang/Integer;
 	public fun hashCode ()I
 	public fun toString ()Ljava/lang/String;
 }
 
+public final class misk/aws2/sqs/jobqueue/config/SqsQueueConfig$Companion {
+}
+
 public final class misk/aws2/sqs/jobqueue/coordinated/AwsSqsJobQueueConfig : misk/config/Config {
 	public fun <init> ()V
 	public fun <init> (Ljava/util/Map;)V

misk-aws2-sqs/src/main/kotlin/misk/aws2/sqs/jobqueue/SqsJobConsumer.kt

@@ -7,12 +7,17 @@ import io.opentracing.Tracer
 import jakarta.inject.Inject
 import java.time.Clock
 import java.util.concurrent.ConcurrentHashMap
+import kotlin.time.Duration.Companion.milliseconds
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.joinAll
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.withTimeoutOrNull
 import misk.aws2.sqs.jobqueue.config.SqsQueueConfig
 import misk.inject.AsyncSwitch
 import misk.jobqueue.QueueName
@@ -56,7 +61,7 @@ constructor(
 ) : JobConsumer, AbstractService(), TestFixture {
   private val scope = CoroutineScope(Dispatchers.IO.limitedParallelism(1) + SupervisorJob())
 
-  private val handlingScopes = ConcurrentHashMap<QueueName, CoroutineScope>()
+  private val subscriptions = ConcurrentHashMap<QueueName, Subscription>()
 
   override fun subscribe(queueName: QueueName, handler: JobHandler) {
     subscribe(queueName = queueName, handler = handler, queueConfig = SqsQueueConfig())
@@ -83,31 +88,65 @@ constructor(
         asyncSwitch = asyncSwitch,
       )
 
-    scope.launch { subscriber.poll() }
-    handlingScopes[queueName] =
-      CoroutineScope(Dispatchers.IO.limitedParallelism(queueConfig.parallelism) + SupervisorJob())
-    repeat(queueConfig.concurrency) { handlingScopes[queueName]?.launch { subscriber.run() } }
+    val pollingJob = scope.launch { subscriber.poll() }
+    val handlingScope = CoroutineScope(Dispatchers.IO.limitedParallelism(queueConfig.parallelism) + SupervisorJob())
+    val handlingJobs = List(queueConfig.concurrency) { handlingScope.launch { subscriber.run() } }
+    subscriptions[queueName] = Subscription(subscriber, pollingJob, handlingScope, handlingJobs)
   }
 
   override fun unsubscribe(queueName: QueueName) {
-    handlingScopes[queueName]?.cancel()
+    subscriptions[queueName]?.handlingScope?.cancel()
   }
 
   /** Called automatically between every test to prevent long-running scopes or test timeouts. */
   override fun reset() {
-    handlingScopes.forEach { _, scope -> scope.cancel() }
+    subscriptions.forEach { _, subscription -> subscription.handlingScope.cancel() }
   }
 
   override fun doStart() {
     notifyStarted()
   }
 
   override fun doStop() {
-    scope.cancel()
-    handlingScopes.values.forEach { it.cancel() }
+    logger.info("Stopping job consumer")
+    runBlocking(scope.coroutineContext) {
+      subscriptions.forEach { (queueName, subscription) ->
+        // Stop issuing new receives, and give the in-flight one a chance to finish its long poll. Messages it already
+        // fetched are handled normally; abandoning it would leave them invisible until their visibility timeout
+        // expires.
+        subscription.subscriber.stop()
+        val gracePeriod =
+          (subscription.subscriber.queueConfig.shutdown_grace_period_ms
+              ?: SqsQueueConfig.DEFAULT_SHUTDOWN_GRACE_PERIOD_MS)
+            .milliseconds
+        if (withTimeoutOrNull(gracePeriod) { subscription.pollingJob.join() } == null) {
+          logger.info { "Polling for queue ${queueName.value} did not stop within $gracePeriod; canceling it" }
+          subscription.subscriber.cancelInFlightReceives()
+          subscription.pollingJob.join()
+        }
+        // The polling job closes the channel when it completes, which is what lets the handlers finish.
+        val shutdownTimeout = subscription.subscriber.queueConfig.shutdown_timeout_ms
+        if (shutdownTimeout == null) {
+          subscription.handlingJobs.joinAll()
+        } else if (withTimeoutOrNull(shutdownTimeout.milliseconds) { subscription.handlingJobs.joinAll() } == null) {
+          logger.warn {
+            "Handlers for queue ${queueName.value} did not finish within ${shutdownTimeout}ms; cancelling remaining work"
+          }
+          subscription.handlingScope.cancel()
+        }
+      }
+    }
+    logger.info("Stopped job consumer")
     notifyStopped()
   }
 
+  private data class Subscription(
+    val subscriber: Subscriber,
+    val pollingJob: Job,
+    val handlingScope: CoroutineScope,
+    val handlingJobs: List<Job>,
+  )
+
   companion object {
     private val logger = getLogger<SqsJobConsumer>()
   }

misk-aws2-sqs/src/main/kotlin/misk/aws2/sqs/jobqueue/Subscriber.kt

@@ -5,6 +5,8 @@ import io.opentracing.Tracer
 import io.opentracing.tag.Tags
 import java.time.Clock
 import java.util.concurrent.CompletableFuture
+import java.util.concurrent.ConcurrentHashMap
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.delay
@@ -28,6 +30,7 @@ import software.amazon.awssdk.services.sqs.model.MessageSystemAttributeName
 import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest
 import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse
 import software.amazon.awssdk.services.sqs.model.SendMessageRequest
+import kotlin.math.log
 
 /**
  * Subscriber reads jobs from the channel and passes them to handler.
@@ -50,10 +53,28 @@ class Subscriber(
   val asyncSwitch: AsyncSwitch,
 ) {
   private var wasDisabled = false
+  @Volatile private var isRunning = true
+
+  /** In-flight `ReceiveMessage` requests, so they can be aborted when the subscriber is stopped. */
+  private val inFlightReceives = ConcurrentHashMap.newKeySet<CompletableFuture<*>>()
+
+  /** Stops issuing new receives. In-flight ones are left to complete on their own. */
+  fun stop() {
+    isRunning = false
+  }
+
+  /**
+   * Aborts the in-flight receives, so shutdown doesn't have to wait out the long poll.
+   *
+   * Messages that SQS already handed to a canceled request are not lost, but they stay invisible until their visibility
+   * timeout expires, so prefer giving [stop] a chance to wind down on its own first.
+   */
+  fun cancelInFlightReceives() {
+    inFlightReceives.forEach { it.cancel(false) }
+  }
 
   suspend fun run() {
-    while (true) {
-      val job = tracer.withSpan("channel-receive-queue-${queueName.value}") { channel.receive() }
+    for (job in channel) {
       tracer.withSpan("process-queue-${queueName.value}") {
         val receiveFromChannelTimestamp = clock.millis()
         sqsMetrics.channelReceiveLag
@@ -192,12 +213,15 @@ class Subscriber(
       } else {
         messageFlow(queueName)
       }
-      .collect { received -> channel.send(received) }
+      .collect { received ->
+        channel.send(received)
+      }
+    channel.close()
   }
 
   private fun messageFlow(queueName: QueueName) = flow {
     val queueUrl = sqsQueueResolver.getQueueUrl(queueName)
-    while (true) {
+    while (isRunning) {
       if (!asyncSwitch.isEnabled("sqs")) {
         if (!wasDisabled) {
           logger.info { "Async SQS tasks disabled. Polling paused for queue ${queueName.value}." }
@@ -213,7 +237,26 @@ class Subscriber(
       val startTime = clock.millis()
       val response =
         try {
-          fetchMessages(queueUrl).await()
+          val future = fetchMessages(queueUrl)
+          inFlightReceives.add(future)
+          // Closes the race with a cancelInFlightReceives() that already swept the set.
+          if (!isRunning) {
+            future.cancel(false)
+          }
+          try {
+            future.await()
+          } finally {
+            inFlightReceives.remove(future)
+          }
+        } catch (e: CancellationException) {
+          // Propagate cancellation of this subscriber, but recover if only the receive was canceled.
+          currentCoroutineContext().ensureActive()
+          if (isRunning) {
+            logger.warn(e) { "Receive was canceled for queue ${queueName.value}; retrying" }
+            sqsMetrics.sqsReceiveFailures.labels(queueName.value).inc()
+            continue
+          }
+          break
         } catch (e: Exception) {
           // Propagate cancellation of this subscriber, but recover if only the failed operation was canceled.
           currentCoroutineContext().ensureActive()

misk-aws2-sqs/src/main/kotlin/misk/aws2/sqs/jobqueue/config/SqsConfig.kt

@@ -10,10 +10,10 @@ import misk.jobqueue.QueueName
  * overriding configuration for a given queue `buffered_batch_flush_frequency_ms` controls how often buffered messages
  * are flushed to SQS when using enqueueBuffered
  *
- * `config_feature_flag` allows specifying a dynamic config name that returns a JSON object matching the
- * structure of SqsConfig. When set, the dynamic config is evaluated at service startup and **completely replaces**
- * the YAML configuration. This allows dynamic configuration changes with a service restart (without requiring a code
- * deploy). If not set, or if the dynamic config returns null/empty, the YAML configuration is used.
+ * `config_feature_flag` allows specifying a dynamic config name that returns a JSON object matching the structure of
+ * SqsConfig. When set, the dynamic config is evaluated at service startup and **completely replaces** the YAML
+ * configuration. This allows dynamic configuration changes with a service restart (without requiring a code deploy). If
+ * not set, or if the dynamic config returns null/empty, the YAML configuration is used.
  */
 data class SqsConfig
 @JvmOverloads
@@ -22,9 +22,9 @@ constructor(
   val per_queue_overrides: Map<String, SqsQueueConfig> = emptyMap(),
   val buffered_batch_flush_frequency_ms: Long = 50,
   /**
-   * Dynamic config name that returns a JSON object matching SqsConfig structure.
-   * When set and returns a valid config, it completely replaces the YAML config.
-   * Example value: {"all_queues": {"concurrency": 10}, "per_queue_overrides": {"my_queue": {"concurrency": 20}}}
+   * Dynamic config name that returns a JSON object matching SqsConfig structure. When set and returns a valid config,
+   * it completely replaces the YAML config. Example value: {"all_queues": {"concurrency": 10}, "per_queue_overrides":
+   * {"my_queue": {"concurrency": 20}}}
    */
   val config_feature_flag: String? = null,
 ) : Config {
@@ -37,6 +37,8 @@ constructor(
         visibility_timeout = override.visibility_timeout ?: all_queues.visibility_timeout,
         region = override.region ?: all_queues.region,
         account_id = override.account_id ?: all_queues.account_id,
+        shutdown_grace_period_ms = override.shutdown_grace_period_ms ?: all_queues.shutdown_grace_period_ms,
+        shutdown_timeout_ms = override.shutdown_timeout_ms ?: all_queues.shutdown_timeout_ms,
       )
     } else {
       all_queues

misk-aws2-sqs/src/main/kotlin/misk/aws2/sqs/jobqueue/config/SqsQueueConfig.kt

@@ -11,6 +11,10 @@ package misk.aws2.sqs.jobqueue.config
  * will be invisible for subsequent requests. If configured to null, the queue settings will be used. `region` AWS
  * Region of the consumed queue, defaults to the current region. `account_id` AWS Account ID of the consumed queue,
  * defaults to the current account. `queue_name` AWS Queue Name, defaults to the application provided name of the queue.
+ * `shutdown_grace_period_ms` defines how long shutdown waits for an in-flight receive to complete before aborting it.
+ * Set it to 0 to abort in-flight receives immediately. Defaults to null, which uses
+ * [SqsQueueConfig.DEFAULT_SHUTDOWN_GRACE_PERIOD_MS]. `shutdown_timeout_ms` defines how long shutdown waits for
+ * in-progress handlers to finish before cancelling them. Defaults to null, which waits indefinitely.
  */
 data class SqsQueueConfig
 @JvmOverloads
@@ -24,4 +28,16 @@ constructor(
   val visibility_timeout: Int? = null,
   val region: String? = null,
   val account_id: String? = null,
-)
+  val shutdown_grace_period_ms: Long? = null,
+  val shutdown_timeout_ms: Long? = null,
+) {
+  companion object {
+    /**
+     * How long shutdown waits for an in-flight receive to complete before aborting it.
+     *
+     * A long poll that has messages returns immediately, so this only needs to cover a receive that is about to
+     * deliver. One that is still waiting out its `wait_timeout` has nothing to lose by being canceled.
+     */
+    const val DEFAULT_SHUTDOWN_GRACE_PERIOD_MS = 1000L
+  }
+}