From f74e067e6064e3678f1e5ceb176667cf12f213a2 Mon Sep 17 00:00:00 2001 From: Jay Janssen Date: Mon, 17 Aug 2026 07:50:48 -0400 Subject: [PATCH] Bound the SQS handler drain, de-flake the drain test, remove public stop() - Add shutdown_timeout_ms to SqsQueueConfig: an optional per-queue bound on how long doStop() waits for in-progress handlers before cancelling the queue's remaining work. Default null keeps the unbounded join. - De-flake `queues are drained drain`: the Approximate* queue depth attributes are exact in ElasticMQ but eventually consistent on real SQS, so poll until conservation holds instead of asserting on a single read taken right after stop. - Remove the public SqsJobConsumer.stop(): it bypasses the Guava service state machine. Tests use stopAsync().awaitTerminated(). - New integration test for the stuck-handler path, config resolution tests for shutdown_timeout_ms, and the regenerated API dump. Co-Authored-By: Claude Fable 5 --- misk-aws2-sqs/api/misk-aws2-sqs.api | 10 +-- .../misk/aws2/sqs/jobqueue/SqsJobConsumer.kt | 20 +++--- .../aws2/sqs/jobqueue/config/SqsConfig.kt | 15 +++-- .../sqs/jobqueue/config/SqsQueueConfig.kt | 4 +- .../aws2/sqs/jobqueue/SqsJobConsumerTest.kt | 66 +++++++++++++++++-- .../aws2/sqs/jobqueue/config/SqsConfigTest.kt | 65 +++++++++++------- 6 files changed, 134 insertions(+), 46 deletions(-) diff --git a/misk-aws2-sqs/api/misk-aws2-sqs.api b/misk-aws2-sqs/api/misk-aws2-sqs.api index 02ce2ab3454..7d069f70d45 100644 --- a/misk-aws2-sqs/api/misk-aws2-sqs.api +++ b/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 (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 (IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;)V public fun (IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;)V public fun (IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Long;)V - public synthetic fun (IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Long;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Long;Ljava/lang/Long;)V + public synthetic fun (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 diff --git a/misk-aws2-sqs/src/main/kotlin/misk/aws2/sqs/jobqueue/SqsJobConsumer.kt b/misk-aws2-sqs/src/main/kotlin/misk/aws2/sqs/jobqueue/SqsJobConsumer.kt index 88f0e0b1f9b..3a1b422405a 100644 --- a/misk-aws2-sqs/src/main/kotlin/misk/aws2/sqs/jobqueue/SqsJobConsumer.kt +++ b/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,17 +125,21 @@ 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, @@ -142,7 +147,6 @@ constructor( val handlingJobs: List, ) - companion object { private val logger = getLogger() } diff --git a/misk-aws2-sqs/src/main/kotlin/misk/aws2/sqs/jobqueue/config/SqsConfig.kt b/misk-aws2-sqs/src/main/kotlin/misk/aws2/sqs/jobqueue/config/SqsConfig.kt index edb7295d39d..4a6b12f924a 100644 --- a/misk-aws2-sqs/src/main/kotlin/misk/aws2/sqs/jobqueue/config/SqsConfig.kt +++ b/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 = 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 diff --git a/misk-aws2-sqs/src/main/kotlin/misk/aws2/sqs/jobqueue/config/SqsQueueConfig.kt b/misk-aws2-sqs/src/main/kotlin/misk/aws2/sqs/jobqueue/config/SqsQueueConfig.kt index 3144afd0273..d7cdf3fae6d 100644 --- a/misk-aws2-sqs/src/main/kotlin/misk/aws2/sqs/jobqueue/config/SqsQueueConfig.kt +++ b/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 { /** diff --git a/misk-aws2-sqs/src/test/kotlin/misk/aws2/sqs/jobqueue/SqsJobConsumerTest.kt b/misk-aws2-sqs/src/test/kotlin/misk/aws2/sqs/jobqueue/SqsJobConsumerTest.kt index 437fc8cf19a..aed0bfae2e1 100644 --- a/misk-aws2-sqs/src/test/kotlin/misk/aws2/sqs/jobqueue/SqsJobConsumerTest.kt +++ b/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) -> Boolean): Pair { + 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 { val attributes = DockerSqs.client diff --git a/misk-aws2-sqs/src/test/kotlin/misk/aws2/sqs/jobqueue/config/SqsConfigTest.kt b/misk-aws2-sqs/src/test/kotlin/misk/aws2/sqs/jobqueue/config/SqsConfigTest.kt index 4492bf0512d..de9e8b6a793 100644 --- a/misk-aws2-sqs/src/test/kotlin/misk/aws2/sqs/jobqueue/config/SqsConfigTest.kt +++ b/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) + } }