Introduce support for suspending SQS Jobqueue

implementation

GitOrigin-RevId: 5981363e7cac52a5f9208645d60f259d585d5dd2

Author: mateuszmrozewski

gradle/libs.versions.toml

@@ -37,6 +37,7 @@ aws2Core = { module = "software.amazon.awssdk:aws-core", version.ref = "aws2" }
 aws2Dynamodb = { module = "software.amazon.awssdk:dynamodb", version.ref = "aws2" }
 aws2DynamodbEnhanced = { module = "software.amazon.awssdk:dynamodb-enhanced", version.ref = "aws2" }
 aws2Regions = { module = "software.amazon.awssdk:regions" }
+aws2Sqs = { module = "software.amazon.awssdk:sqs", version.ref = "aws2" }
 awsCore = { module = "com.amazonaws:aws-java-sdk-core", version.ref = "aws1" }
 awsDynamodb = { module = "com.amazonaws:aws-java-sdk-dynamodb", version.ref = "aws1" }
 awsS3 = { module = "com.amazonaws:aws-java-sdk-s3", version.ref = "aws1" }

misk-aws2-sqs/README.md

@@ -0,0 +1,61 @@
+# Module: AWS SQS Module
+
+**This module is still in experimental state. It's a work in progress towards:**
+* supporting suspending handlers
+* migration to AWS SDK v2
+
+## Differences to the previous implementation
+
+* uses AWS SDK v2
+* exposes suspending API
+* handlers return status and don't make calls to SQS. Acknowledging jobs is done by the framework code
+* no dependency on the lease module. There will be at least one handler per service instance
+* no dependency on the feature flags
+* metrics are updated to v2, names of the metrics have been changed
+
+## Migration
+
+TODO - this section will have detailed steps for migrating from the previous implementation
+
+## Threading model
+
+Receiving and processing messages is handled by separate views on the `Dispatchers.IO`:
+- querying SQS is done on a single thread
+- for each subscribed queue, a receiving coroutine is created
+- received jobs are sent to a dedicated per-queue channel with a default size of 0. This is configurable
+- by default, there is a single thread with a single coroutine running dedicated to processing
+  of jobs from a given queue. This is configurable as well.
+
+How the configuration impacts the processing:
+- increasing the channel size allows to pre-read jobs from SQS. This may be helpful to reduce the
+  latency, but if the handler takes more time to process than the visibility timeout, it may lead
+  to increased duplicated and out of order processing
+- increasing the concurrency spins up more coroutines per queue. This may increase the throughput
+  of processing, if processing mostly uses non-blocking operations
+- increasing the parallelism will increase the thread pool size per-queue. Together with increased
+  concurrency it may increase the throughput of processing, if processing involves heavy computations
+  or blocking operations
+
+It's advised to start with the default settings and adjust based on specific workloads.
+
+![image](concurrency.jpg)
+
+## Outstanding todo items
+
+The module will not be considered beta/GA state until the below items are completed.
+
+Outstanding work that needs to be done:
+* detailed test 
+* tracing
+* test fixtures
+* external queues
+* installing retry queue only on request
+* pass in configuration (make it compatible with previous implementation if needed)
+* detailed documentation
+* batch size configuration
+
+Things that are supported in the old documentation but are questionable:
+* aws queue attribute importer
+
+Outstanding things to document:
+* how batch size plays out with channel size and visibility timeout

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

@@ -0,0 +1,146 @@
+public abstract interface class misk/aws2/sqs/jobqueue/DeadLetterQueueProvider {
+	public abstract fun deadLetterQueueFor (Lmisk/jobqueue/QueueName;)Lmisk/jobqueue/QueueName;
+}
+
+public final class misk/aws2/sqs/jobqueue/DeadLetterQueueProviderKt {
+	public static final fun getParentQueue (Lmisk/jobqueue/QueueName;)Lmisk/jobqueue/QueueName;
+	public static final fun getRetryQueue (Lmisk/jobqueue/QueueName;)Lmisk/jobqueue/QueueName;
+	public static final fun isRetryQueue (Lmisk/jobqueue/QueueName;)Z
+}
+
+public final class misk/aws2/sqs/jobqueue/DefaultDeadLetterQueueProvider : misk/aws2/sqs/jobqueue/DeadLetterQueueProvider {
+	public fun <init> ()V
+	public fun deadLetterQueueFor (Lmisk/jobqueue/QueueName;)Lmisk/jobqueue/QueueName;
+}
+
+public final class misk/aws2/sqs/jobqueue/QueueResolver {
+	public fun <init> (Lsoftware/amazon/awssdk/services/sqs/SqsAsyncClient;)V
+	public final fun getQueueUrl (Lmisk/jobqueue/QueueName;)Ljava/lang/String;
+}
+
+public final class misk/aws2/sqs/jobqueue/SqsJob : misk/jobqueue/v2/Job {
+	public static final field Companion Lmisk/aws2/sqs/jobqueue/SqsJob$Companion;
+	public static final field JOBQUEUE_METADATA_ATTR Ljava/lang/String;
+	public static final field JOBQUEUE_METADATA_IDEMPOTENCE_KEY Ljava/lang/String;
+	public static final field JOBQUEUE_METADATA_ORIGIN_QUEUE Ljava/lang/String;
+	public fun <init> (Lmisk/jobqueue/QueueName;Lcom/squareup/moshi/Moshi;Lsoftware/amazon/awssdk/services/sqs/model/Message;Ljava/lang/String;J)V
+	public fun getAttributes ()Ljava/util/Map;
+	public fun getBody ()Ljava/lang/String;
+	public fun getId ()Ljava/lang/String;
+	public fun getIdempotenceKey ()Ljava/lang/String;
+	public final fun getMessage ()Lsoftware/amazon/awssdk/services/sqs/model/Message;
+	public final fun getPublishToChannelTimestamp ()J
+	public fun getQueueName ()Lmisk/jobqueue/QueueName;
+	public final fun getQueueUrl ()Ljava/lang/String;
+}
+
+public final class misk/aws2/sqs/jobqueue/SqsJob$Companion {
+}
+
+public final class misk/aws2/sqs/jobqueue/SqsJobConsumer : com/google/common/util/concurrent/AbstractService, misk/jobqueue/v2/JobConsumer {
+	public static final field Companion Lmisk/aws2/sqs/jobqueue/SqsJobConsumer$Companion;
+	public fun <init> (Lsoftware/amazon/awssdk/services/sqs/SqsAsyncClient;Lmisk/aws2/sqs/jobqueue/QueueResolver;Lcom/squareup/moshi/Moshi;Lmisk/aws2/sqs/jobqueue/DeadLetterQueueProvider;Lmisk/aws2/sqs/jobqueue/SqsMetrics;Ljava/time/Clock;)V
+	public fun subscribe (Lmisk/jobqueue/QueueName;Lmisk/jobqueue/v2/JobHandler;)V
+	public final fun subscribe (Lmisk/jobqueue/QueueName;Lmisk/jobqueue/v2/JobHandler;III)V
+	public fun unsubscribe (Lmisk/jobqueue/QueueName;)V
+}
+
+public final class misk/aws2/sqs/jobqueue/SqsJobConsumer$Companion {
+	public final fun getLogger ()Lmu/KLogger;
+}
+
+public final class misk/aws2/sqs/jobqueue/SqsJobEnqueuer : misk/jobqueue/v2/JobEnqueuer {
+	public fun <init> (Lsoftware/amazon/awssdk/services/sqs/SqsAsyncClient;Lmisk/aws2/sqs/jobqueue/QueueResolver;Lmisk/aws2/sqs/jobqueue/SqsMetrics;Lcom/squareup/moshi/Moshi;)V
+	public fun enqueue (Lmisk/jobqueue/QueueName;Ljava/lang/String;Ljava/lang/String;Ljava/time/Duration;Ljava/util/Map;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public fun enqueueAsync (Lmisk/jobqueue/QueueName;Ljava/lang/String;Ljava/lang/String;Ljava/time/Duration;Ljava/util/Map;)Ljava/util/concurrent/CompletableFuture;
+	public fun enqueueBlocking (Lmisk/jobqueue/QueueName;Ljava/lang/String;Ljava/lang/String;Ljava/time/Duration;Ljava/util/Map;)V
+}
+
+public final class misk/aws2/sqs/jobqueue/SqsJobHandlerModule : misk/inject/KAbstractModule {
+	public static final field Companion Lmisk/aws2/sqs/jobqueue/SqsJobHandlerModule$Companion;
+	public synthetic fun <init> (Lmisk/jobqueue/QueueName;Lkotlin/reflect/KClass;IIILkotlin/jvm/internal/DefaultConstructorMarker;)V
+}
+
+public final class misk/aws2/sqs/jobqueue/SqsJobHandlerModule$Companion {
+	public final fun create (Lmisk/jobqueue/QueueName;Lkotlin/reflect/KClass;III)Lmisk/aws2/sqs/jobqueue/SqsJobHandlerModule;
+	public static synthetic fun create$default (Lmisk/aws2/sqs/jobqueue/SqsJobHandlerModule$Companion;Lmisk/jobqueue/QueueName;Lkotlin/reflect/KClass;IIIILjava/lang/Object;)Lmisk/aws2/sqs/jobqueue/SqsJobHandlerModule;
+}
+
+public class misk/aws2/sqs/jobqueue/SqsJobQueueModule : misk/inject/KAbstractModule {
+	public fun <init> ()V
+	public fun <init> (Lkotlin/jvm/functions/Function1;)V
+	public synthetic fun <init> (Lkotlin/jvm/functions/Function1;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	protected fun configure ()V
+	public final fun sqsAsyncClient (Lsoftware/amazon/awssdk/auth/credentials/AwsCredentialsProvider;Lmisk/cloud/aws/AwsRegion;)Lsoftware/amazon/awssdk/services/sqs/SqsAsyncClient;
+}
+
+public final class misk/aws2/sqs/jobqueue/SqsMetrics {
+	public final fun getChannelReceiveLag ()Lio/prometheus/client/Histogram;
+	public final fun getHandlerDispatchTime ()Lio/prometheus/client/Histogram;
+	public final fun getHandlerFailures ()Lio/prometheus/client/Counter;
+	public final fun getJobEnqueueFailures ()Lio/prometheus/client/Counter;
+	public final fun getJobsAcknowledged ()Lio/prometheus/client/Counter;
+	public final fun getJobsDeadLettered ()Lio/prometheus/client/Counter;
+	public final fun getJobsEnqueued ()Lio/prometheus/client/Counter;
+	public final fun getJobsReceived ()Lio/prometheus/client/Counter;
+	public final fun getQueueProcessingLag ()Lio/prometheus/client/Histogram;
+	public final fun getSqsApproxNumberOfMessages ()Lio/prometheus/client/Gauge;
+	public final fun getSqsApproxNumberOfMessagesNotVisible ()Lio/prometheus/client/Gauge;
+	public final fun getSqsDeleteTime ()Lio/prometheus/client/Histogram;
+	public final fun getSqsReceiveTime ()Lio/prometheus/client/Histogram;
+	public final fun getSqsSendTime ()Lio/prometheus/client/Histogram;
+}
+
+public final class misk/aws2/sqs/jobqueue/StaticDeadLetterQueueProvider : misk/aws2/sqs/jobqueue/DeadLetterQueueProvider {
+	public fun <init> (Ljava/lang/String;)V
+	public fun deadLetterQueueFor (Lmisk/jobqueue/QueueName;)Lmisk/jobqueue/QueueName;
+}
+
+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;Ljava/lang/String;Lsoftware/amazon/awssdk/services/sqs/SqsAsyncClient;Lmisk/jobqueue/v2/JobHandler;Lkotlinx/coroutines/channels/Channel;Ljava/lang/String;Lmisk/jobqueue/QueueName;Lmisk/aws2/sqs/jobqueue/SqsMetrics;Lmisk/aws2/sqs/jobqueue/QueueResolver;Lcom/squareup/moshi/Moshi;Ljava/time/Clock;)V
+	public final fun getChannel ()Lkotlinx/coroutines/channels/Channel;
+	public final fun getClient ()Lsoftware/amazon/awssdk/services/sqs/SqsAsyncClient;
+	public final fun getClock ()Ljava/time/Clock;
+	public final fun getDeadLetterQueueName ()Lmisk/jobqueue/QueueName;
+	public final fun getHandler ()Lmisk/jobqueue/v2/JobHandler;
+	public final fun getMoshi ()Lcom/squareup/moshi/Moshi;
+	public final fun getQueueName ()Lmisk/jobqueue/QueueName;
+	public final fun getQueueResolver ()Lmisk/aws2/sqs/jobqueue/QueueResolver;
+	public final fun getQueueUrl ()Ljava/lang/String;
+	public final fun getRetryQueueUrl ()Ljava/lang/String;
+	public final fun getSqsMetrics ()Lmisk/aws2/sqs/jobqueue/SqsMetrics;
+	public final fun poll (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public final fun run (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+}
+
+public final class misk/aws2/sqs/jobqueue/Subscriber$Companion {
+}
+
+public final class misk/aws2/sqs/jobqueue/Subscription {
+	public fun <init> (Lmisk/jobqueue/QueueName;Lkotlin/reflect/KClass;)V
+	public fun <init> (Lmisk/jobqueue/QueueName;Lkotlin/reflect/KClass;I)V
+	public fun <init> (Lmisk/jobqueue/QueueName;Lkotlin/reflect/KClass;II)V
+	public fun <init> (Lmisk/jobqueue/QueueName;Lkotlin/reflect/KClass;III)V
+	public synthetic fun <init> (Lmisk/jobqueue/QueueName;Lkotlin/reflect/KClass;IIIILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public final fun component1 ()Lmisk/jobqueue/QueueName;
+	public final fun component2 ()Lkotlin/reflect/KClass;
+	public final fun component3 ()I
+	public final fun component4 ()I
+	public final fun component5 ()I
+	public final fun copy (Lmisk/jobqueue/QueueName;Lkotlin/reflect/KClass;III)Lmisk/aws2/sqs/jobqueue/Subscription;
+	public static synthetic fun copy$default (Lmisk/aws2/sqs/jobqueue/Subscription;Lmisk/jobqueue/QueueName;Lkotlin/reflect/KClass;IIIILjava/lang/Object;)Lmisk/aws2/sqs/jobqueue/Subscription;
+	public fun equals (Ljava/lang/Object;)Z
+	public final fun getChannelCapacity ()I
+	public final fun getConcurrency ()I
+	public final fun getHandler ()Lkotlin/reflect/KClass;
+	public final fun getParallelism ()I
+	public final fun getQueueName ()Lmisk/jobqueue/QueueName;
+	public fun hashCode ()I
+	public fun toString ()Ljava/lang/String;
+}
+
+public final class misk/aws2/sqs/jobqueue/SubscriptionService : com/google/common/util/concurrent/AbstractIdleService {
+	public fun <init> (Lmisk/aws2/sqs/jobqueue/SqsJobConsumer;Ljava/util/Map;Ljava/util/Map;)V
+}
+

misk-aws2-sqs/build.gradle.kts

@@ -0,0 +1,52 @@
+import com.vanniktech.maven.publish.JavadocJar.Dokka
+import com.vanniktech.maven.publish.KotlinJvm
+
+plugins {
+  id("org.jetbrains.kotlin.jvm")
+  id("com.vanniktech.maven.publish.base")
+  id("java-test-fixtures")
+}
+
+dependencies {
+  api(libs.aws2Auth)
+  implementation(libs.aws2Regions)
+  api(libs.aws2Sqs)
+  api(libs.guava)
+  api(libs.guice)
+  api(libs.kotlinxCoroutinesCore)
+  api(libs.jakartaInject)
+  api(libs.moshiCore)
+  api(libs.prometheusClient)
+  implementation(libs.aws2Core)
+  implementation(project(":misk"))
+  implementation(project(":misk-api"))
+  api(project(":misk-aws"))
+  api(project(":misk-inject"))
+  api(project(":misk-jobqueue"))
+  implementation(libs.loggingApi)
+  runtimeOnly(libs.openTracingDatadog)
+  implementation(project(":wisp:wisp-logging"))
+  implementation(project(":misk-metrics"))
+  implementation(project(":misk-service"))
+  testImplementation(libs.assertj)
+  testImplementation(libs.awaitility)
+  testImplementation(libs.dockerApi)
+  testImplementation(libs.junitApi)
+  testImplementation(libs.junitParams)
+  testImplementation(libs.kotlinTest)
+  testImplementation(libs.kotlinxCoroutinesTest)
+  testImplementation(libs.mockitoCore)
+  testImplementation(libs.mockitoKotlin)
+  testImplementation(project(":wisp:wisp-containers-testing"))
+  testImplementation(project(":wisp:wisp-feature-testing"))
+  testImplementation(project(":wisp:wisp-time-testing"))
+  testImplementation(project(":wisp:wisp-logging-testing"))
+  testImplementation(project(":misk-clustering"))
+  testImplementation(project(":misk-testing"))
+}
+
+mavenPublishing {
+  configure(
+    KotlinJvm(javadocJar = Dokka("dokkaGfm"))
+  )
+}

misk-aws2-sqs/concurrency.jpg

@@ -0,0 +1,50 @@
+package misk.aws2.sqs.jobqueue
+
+import com.google.inject.ImplementedBy
+import jakarta.inject.Inject
+import jakarta.inject.Singleton
+import misk.jobqueue.QueueName
+
+/**
+ * Interface for a provider of dead-letter queue names.
+ *
+ * Returns the name of the appropriate dead-letter [QueueName] for a given [QueueName].
+ */
+@ImplementedBy(DefaultDeadLetterQueueProvider::class)
+interface DeadLetterQueueProvider {
+  fun deadLetterQueueFor(queue: QueueName): QueueName
+}
+
+/** Default provider of dead-letter [QueueName]. Returns the name of the main queue suffixed with "_dlq". */
+@Singleton
+class DefaultDeadLetterQueueProvider @Inject constructor() : DeadLetterQueueProvider {
+  override fun deadLetterQueueFor(queue: QueueName): QueueName = queue.deadLetterQueue
+}
+
+/**
+ * Provider of dead-letter [QueueName] that always returns the same value, no matter the supplied queue.
+ *
+ * For apps with queues that share a single dead-letter queue.
+ */
+class StaticDeadLetterQueueProvider(queue: String) : DeadLetterQueueProvider {
+  private val dlq = QueueName(queue)
+  override fun deadLetterQueueFor(queue: QueueName): QueueName = dlq
+}
+
+internal const val deadLetterQueueSuffix = "_dlq"
+
+internal val QueueName.isDeadLetterQueue get() = value.endsWith(deadLetterQueueSuffix)
+internal val QueueName.deadLetterQueue
+  get() = if (isDeadLetterQueue) this else QueueName(parentQueue.value + deadLetterQueueSuffix)
+
+internal const val retryQueueSuffix = "_retryq"
+val QueueName.isRetryQueue get() = value.endsWith(retryQueueSuffix)
+val QueueName.retryQueue
+  get() = if (isRetryQueue) this else QueueName(parentQueue.value + retryQueueSuffix)
+
+val QueueName.parentQueue
+  get() = when {
+    isDeadLetterQueue -> QueueName(value.removeSuffix(deadLetterQueueSuffix))
+    isRetryQueue -> QueueName(value.removeSuffix(retryQueueSuffix))
+    else -> this
+  }

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

@@ -0,0 +1,30 @@
+package misk.aws2.sqs.jobqueue
+
+import com.google.inject.Inject
+import com.google.inject.Singleton
+import misk.jobqueue.QueueName
+import software.amazon.awssdk.services.sqs.SqsAsyncClient
+import software.amazon.awssdk.services.sqs.model.GetQueueUrlRequest
+import java.util.concurrent.ConcurrentHashMap
+
+@Singleton
+class QueueResolver @Inject constructor(
+  private val client: SqsAsyncClient,
+) {
+  private val queueUrlCache = ConcurrentHashMap<QueueName, String>()
+
+  /**
+   * Get sqs queue URL for a given queue.
+   *
+   * Results are cached in-memory, call to SQS is effectively blocking.
+   */
+  fun getQueueUrl(queueName: QueueName): String {
+    return queueUrlCache.getOrPut(queueName) {
+      val retryQueueUrlRequest = GetQueueUrlRequest.builder()
+        .queueName(queueName.value)
+        .build()
+      val response = client.getQueueUrl(retryQueueUrlRequest).join()
+      return response.queueUrl()
+    }
+  }
+}