Bound the SQS handler drain on shutdown + de-flake the drain test

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

@@ -54,7 +54,6 @@ public final class misk/aws2/sqs/jobqueue/SqsJobConsumer : com/google/common/uti
 	public static final field Companion Lmisk/aws2/sqs/jobqueue/SqsJobConsumer$Companion;
 	public fun <init> (Lmisk/aws2/sqs/jobqueue/SqsClientFactory;Lmisk/aws2/sqs/jobqueue/SqsQueueResolver;Lmisk/aws2/sqs/jobqueue/VisibilityTimeoutCalculator;Lcom/squareup/moshi/Moshi;Lmisk/aws2/sqs/jobqueue/DeadLetterQueueProvider;Lmisk/aws2/sqs/jobqueue/SqsMetrics;Ljava/time/Clock;Lio/opentracing/Tracer;Lmisk/inject/AsyncSwitch;)V
 	public fun reset ()V
-	public final fun stop ()V
 	public fun subscribe (Lmisk/jobqueue/QueueName;Lmisk/jobqueue/v2/JobHandler;)V
 	public final fun subscribe (Lmisk/jobqueue/QueueName;Lmisk/jobqueue/v2/JobHandler;Lmisk/aws2/sqs/jobqueue/config/SqsQueueConfig;)V
 	public fun unsubscribe (Lmisk/jobqueue/QueueName;)V
@@ -218,9 +217,11 @@ public final class misk/aws2/sqs/jobqueue/config/SqsQueueConfig {
 	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 fun <init> (IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Long;)V
-	public synthetic fun <init> (IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Long;ILkotlin/jvm/internal/DefaultConstructorMarker;)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
@@ -229,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;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;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
@@ -240,6 +241,7 @@ public final class misk/aws2/sqs/jobqueue/config/SqsQueueConfig {
 	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

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

@@ -7,6 +7,7 @@ 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
@@ -24,7 +25,6 @@ import misk.jobqueue.v2.JobConsumer
 import misk.jobqueue.v2.JobHandler
 import misk.logging.getLogger
 import misk.testing.TestFixture
-import kotlin.time.Duration.Companion.milliseconds
 
 /**
  * Instruments queue consumption.
@@ -112,7 +112,8 @@ constructor(
     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.
+        // 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
@@ -124,25 +125,28 @@ constructor(
           subscription.pollingJob.join()
         }
         // The polling job closes the channel when it completes, which is what lets the handlers finish.
-        subscription.handlingJobs.joinAll()
+        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()
   }
 
-  fun stop() {
-    doStop()
-  }
-
   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/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 {
@@ -38,6 +38,7 @@ constructor(
         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

@@ -13,7 +13,8 @@ package misk.aws2.sqs.jobqueue.config
  * 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].
+ * [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
@@ -28,6 +29,7 @@ constructor(
   val region: String? = null,
   val account_id: String? = null,
   val shutdown_grace_period_ms: Long? = null,
+  val shutdown_timeout_ms: Long? = null,
 ) {
   companion object {
     /**

misk-aws2-sqs/src/test/kotlin/misk/aws2/sqs/jobqueue/SqsJobConsumerTest.kt

@@ -306,16 +306,60 @@ class SqsJobConsumerTest {
     )
 
     repeat(numberOfMessages) { sendMessage(result.queueUrl, "message $it") }
-    jobConsumer.stop()
+    jobConsumer.stopAsync().awaitTerminated()
 
-    val (visible, invisible) = queueDepths(result.queueUrl)
+    // The Approximate* depth attributes are exact in ElasticMQ but eventually consistent on real SQS, so poll
+    // until they settle instead of asserting on a single read taken right after the stop.
+    val (visible, invisible) =
+      awaitSettledQueueDepths(result.queueUrl) { (visible, invisible) ->
+        visible + invisible + (numberOfMessages - latch.count).toInt() == numberOfMessages && invisible == 0
+      }
 
     // Because we shut down consumption early, but gracefully, we expect all received jobs to be processed
     // Some may still be on the queue
-    assertEquals(numberOfMessages, ((visible + invisible + (numberOfMessages - latch.count)).toInt()), "Visible: $visible, invisible: $invisible, latch: ${latch.count} ")
+    assertEquals(
+      numberOfMessages,
+      ((visible + invisible + (numberOfMessages - latch.count)).toInt()),
+      "Visible: $visible, invisible: $invisible, latch: ${latch.count} ",
+    )
     assertEquals(0, invisible)
   }
 
+  @Test
+  fun `handler exceeding the shutdown timeout is cancelled and shutdown terminates predictably`() {
+    val queueName = QueueName("test-queue-1")
+    val result = createQueue(queueName)
+
+    val started = CountDownLatch(1)
+    jobConsumer.subscribe(
+      queueName,
+      object : SuspendingJobHandler {
+        override suspend fun handleJob(job: Job): JobStatus {
+          started.countDown()
+          delay(60_000)
+          return JobStatus.OK
+        }
+      },
+      // Visibility of 30s: the cancelled job's message stays in the not-visible window for the whole test.
+      SqsQueueConfig(region = "us-west-2", visibility_timeout = 30, shutdown_timeout_ms = 1_000),
+    )
+    sendMessage(result.queueUrl, "message")
+    assertTrue(started.await(10, SECONDS), "handler did not start")
+
+    val stopStart = System.nanoTime()
+    jobConsumer.stopAsync()
+    jobConsumer.awaitTerminated(10, SECONDS)
+    val stopMillis = (System.nanoTime() - stopStart) / 1_000_000
+
+    assertTrue(stopMillis >= 1_000, "shutdown returned before the shutdown timeout: ${stopMillis}ms")
+    assertTrue(stopMillis < 8_000, "shutdown did not terminate promptly after the shutdown timeout: ${stopMillis}ms")
+
+    // The handler was cancelled, so its message was never acknowledged and remains in the visibility window.
+    val (visible, invisible) = queueDepths(result.queueUrl)
+    assertEquals(0, visible)
+    assertEquals(1, invisible, "cancelled job's message should remain in the visibility window")
+  }
+
   @Test
   fun `shutdown grace period of zero cancels the in-flight receive`() {
     val queueName = QueueName("test-queue-1")
@@ -330,11 +374,25 @@ class SqsJobConsumerTest {
 
     // Give the poller time to issue its receive before stopping.
     Thread.sleep(1_000)
-    val elapsed = measureTimeMillis { jobConsumer.stop() }
+    val elapsed = measureTimeMillis { jobConsumer.stopAsync().awaitTerminated() }
 
     assertTrue(elapsed < 5_000, "stop() took ${elapsed}ms; the in-flight receive was not canceled")
   }
 
+  /**
+   * Polls the queue depths until [settled] returns true or a ~20s deadline expires, then returns the last read. The
+   * Approximate* depth attributes are exact in ElasticMQ but eventually consistent on real SQS.
+   */
+  private fun awaitSettledQueueDepths(queueUrl: String, settled: (Pair<Int, Int>) -> Boolean): Pair<Int, Int> {
+    val deadline = System.nanoTime() + SECONDS.toNanos(20)
+    var depths = queueDepths(queueUrl)
+    while (!settled(depths) && System.nanoTime() < deadline) {
+      Thread.sleep(100)
+      depths = queueDepths(queueUrl)
+    }
+    return depths
+  }
+
   private fun queueDepths(queueUrl: String): Pair<Int, Int> {
     val attributes =
       DockerSqs.client

misk-aws2-sqs/src/test/kotlin/misk/aws2/sqs/jobqueue/config/SqsConfigTest.kt

@@ -44,9 +44,7 @@ class SqsConfigTest {
 
   @Test
   fun `getQueueConfig returns all_queues when no per_queue_override exists`() {
-    val config = SqsConfig(
-      all_queues = SqsQueueConfig(concurrency = 10, parallelism = 5),
-    )
+    val config = SqsConfig(all_queues = SqsQueueConfig(concurrency = 10, parallelism = 5))
 
     val queueConfig = config.getQueueConfig(misk.jobqueue.QueueName("test-queue"))
 
@@ -56,12 +54,11 @@ class SqsConfigTest {
 
   @Test
   fun `getQueueConfig returns per_queue_override when it exists`() {
-    val config = SqsConfig(
-      all_queues = SqsQueueConfig(concurrency = 1, parallelism = 1),
-      per_queue_overrides = mapOf(
-        "test-queue" to SqsQueueConfig(concurrency = 20, parallelism = 10),
-      ),
-    )
+    val config =
+      SqsConfig(
+        all_queues = SqsQueueConfig(concurrency = 1, parallelism = 1),
+        per_queue_overrides = mapOf("test-queue" to SqsQueueConfig(concurrency = 20, parallelism = 10)),
+      )
 
     val queueConfig = config.getQueueConfig(misk.jobqueue.QueueName("test-queue"))
 
@@ -71,12 +68,11 @@ class SqsConfigTest {
 
   @Test
   fun `getQueueConfig inherits nullable fields from all_queues`() {
-    val config = SqsConfig(
-      all_queues = SqsQueueConfig(region = "us-west-2", wait_timeout = 20),
-      per_queue_overrides = mapOf(
-        "test-queue" to SqsQueueConfig(concurrency = 10),
-      ),
-    )
+    val config =
+      SqsConfig(
+        all_queues = SqsQueueConfig(region = "us-west-2", wait_timeout = 20),
+        per_queue_overrides = mapOf("test-queue" to SqsQueueConfig(concurrency = 10)),
+      )
 
     val queueConfig = config.getQueueConfig(misk.jobqueue.QueueName("test-queue"))
 
@@ -87,16 +83,41 @@ class SqsConfigTest {
 
   @Test
   fun `getQueueConfig resolves shutdown_grace_period_ms`() {
-    val config = SqsConfig(
-      all_queues = SqsQueueConfig(shutdown_grace_period_ms = 5_000),
-      per_queue_overrides = mapOf(
-        "inheriting-queue" to SqsQueueConfig(concurrency = 10),
-        "overriding-queue" to SqsQueueConfig(shutdown_grace_period_ms = 0),
-      ),
-    )
+    val config =
+      SqsConfig(
+        all_queues = SqsQueueConfig(shutdown_grace_period_ms = 5_000),
+        per_queue_overrides =
+          mapOf(
+            "inheriting-queue" to SqsQueueConfig(concurrency = 10),
+            "overriding-queue" to SqsQueueConfig(shutdown_grace_period_ms = 0),
+          ),
+      )
 
     assertEquals(5_000, config.getQueueConfig(misk.jobqueue.QueueName("inheriting-queue")).shutdown_grace_period_ms)
     assertEquals(0, config.getQueueConfig(misk.jobqueue.QueueName("overriding-queue")).shutdown_grace_period_ms)
     assertEquals(5_000, config.getQueueConfig(misk.jobqueue.QueueName("unconfigured-queue")).shutdown_grace_period_ms)
   }
+
+  @Test
+  fun `shutdown_timeout_ms has null default`() {
+    val config = SqsConfig()
+    assertEquals(null, config.all_queues.shutdown_timeout_ms)
+  }
+
+  @Test
+  fun `getQueueConfig resolves shutdown_timeout_ms`() {
+    val config =
+      SqsConfig(
+        all_queues = SqsQueueConfig(shutdown_timeout_ms = 5_000),
+        per_queue_overrides =
+          mapOf(
+            "inheriting-queue" to SqsQueueConfig(concurrency = 10),
+            "overriding-queue" to SqsQueueConfig(shutdown_timeout_ms = 30_000),
+          ),
+      )
+
+    assertEquals(5_000, config.getQueueConfig(misk.jobqueue.QueueName("inheriting-queue")).shutdown_timeout_ms)
+    assertEquals(30_000, config.getQueueConfig(misk.jobqueue.QueueName("overriding-queue")).shutdown_timeout_ms)
+    assertEquals(5_000, config.getQueueConfig(misk.jobqueue.QueueName("unconfigured-queue")).shutdown_timeout_ms)
+  }
 }