From a778949cbbec17463891955e44bb20db106efd03 Mon Sep 17 00:00:00 2001 From: "Kartikaya Gupta (kats)" Date: Fri, 14 Aug 2026 20:57:42 +0000 Subject: [PATCH 1/5] Update logback to newer patch version to pick up vulnerability fixes (#3896) --- gradle/libs.versions.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index 09e9529ae5d..76b95f76c52 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -192,8 +192,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" } From 4fc004bf388286b40c5bfd0b59321b4854ba5621 Mon Sep 17 00:00:00 2001 From: "Kartikaya Gupta (kats)" Date: Mon, 17 Aug 2026 00:34:32 +0000 Subject: [PATCH 2/5] Migrate to Jackson 3.2.1 (#3872) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bumps the Jackson BOM from 2.21.2 to 3.2.1 and moves the five Jackson-using source files onto the `tools.jackson` packages. `jackson-annotations` is deliberately left on `com.fasterxml` — it is not renamed in 3.x and the 3.2.1 BOM pins it to 2.22, which is what lets Jackson 2 and 3 coexist on a classpath. `jackson-datatype-jsr310` is dropped; java.time support is folded into databind in 3.x. Five changes are behavioural rather than mechanical: - Mappers are immutable in 3.x, but `SecretDeserializer` and `ResourceAwareDeserializer` parse nested documents with the very mapper they are registered on. They now take a `() -> ObjectMapper` supplier that resolves once `builder.build()` returns. `SecretJacksonModule` keeps its `ObjectMapper` constructor and exposes `mapper` as a computed property, so its source shape is unchanged. - `FAIL_ON_UNKNOWN_PROPERTIES` defaults to false in 3.x. Left alone, MiskConfig's "'x' not found in Config, did you mean...?" warning would silently never fire again and config typos would be ignored. It is now explicitly enabled for the first parse attempt; the retry path used to relax the mapper in place and instead rebuilds one. - `SORT_PROPERTIES_ALPHABETICALLY` and `EnumFeature.READ/WRITE_ENUMS_USING_TO_STRING` default on in 3.x and are pinned off. Sorting would reshuffle every service's redacted config dashboard, and config enums are matched by name while `toString()` is frequently overridden for display. - `KotlinFeature.StrictNullChecks` defaults on in 3.x and is pinned off. With it on a null element of a collection nested inside a map is rejected even when that element type is declared nullable, so `Map>` fails where a top-level `Set` is accepted. Config that loads today would stop loading. - A `Map` field with an enum key is deserialized into an `EnumMap` in 3.x (databind#1853), which iterates in enum declaration order rather than in the order the keys appear in the YAML. There is no feature flag for it, but the rewrite only fires when the declared raw type is exactly `Map` and abstract type resolution runs first, so naming `LinkedHashMap` as the implementation settles it beforehand. Iteration order of a config map is observable and this would have changed it silently. This is a breaking change for consumers. The affected ABI, confirmed by the regenerated api dumps, is limited to: the `MiskConfig.load` overloads taking `JsonNode`/`ValueDeserializerModifier`, the three `SimpleModule` subclasses in misk-config, and `BackwardsCompatibleClientsConfigConverter`. Both misk-config and misk expose Jackson via `api(...)`, so consumers relying on the transitive dependency inherit Jackson 3. --- gradle/libs.versions.toml | 13 +- misk-config/api/misk-config.api | 33 +-- misk-config/build.gradle.kts | 1 - .../src/main/kotlin/misk/config/MiskConfig.kt | 242 +++++++++++------- .../test/kotlin/misk/config/MiskConfigTest.kt | 13 +- .../src/test/kotlin/misk/config/TestConfig.kt | 9 + .../test/resources/enum_keyed_map-common.yaml | 5 + .../org/assertj/core/api/AssertExtensions.kt | 8 +- misk/api/misk.api | 6 +- .../kotlin/misk/client/HttpClientsConfig.kt | 2 +- ...HttpClientsConfigBackwardsCompatibility.kt | 6 +- 11 files changed, 215 insertions(+), 123 deletions(-) create mode 100644 misk-config/src/test/resources/enum_keyed_map-common.yaml diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index 76b95f76c52..a8e1d008913 100644 --- a/gradle/libs.versions.toml +++ b/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" } diff --git a/misk-config/api/misk-config.api b/misk-config/api/misk-config.api index 1844ea2dca9..3f3eb619382 100644 --- a/misk-config/api/misk-config.api +++ b/misk-config/api/misk-config.api @@ -20,23 +20,23 @@ public final class misk/config/MiskConfig { public static final field INSTANCE Lmisk/config/MiskConfig; public static final fun filesInDir (Ljava/lang/String;Ljava/io/FilenameFilter;)Ljava/util/List; public static synthetic fun filesInDir$default (Ljava/lang/String;Ljava/io/FilenameFilter;ILjava/lang/Object;)Ljava/util/List; - public static final fun load (Ljava/lang/Class;Ljava/lang/String;Lwisp/deployment/Deployment;Ljava/util/List;Lcom/fasterxml/jackson/databind/JsonNode;Lmisk/resources/ResourceLoader;)Lmisk/config/Config; - public static final fun load (Ljava/lang/Class;Ljava/lang/String;Lwisp/deployment/Deployment;Ljava/util/List;Lcom/fasterxml/jackson/databind/JsonNode;Lmisk/resources/ResourceLoader;Z)Lmisk/config/Config; - public static final fun load (Ljava/lang/Class;Ljava/lang/String;Lwisp/deployment/Deployment;Ljava/util/List;Lcom/fasterxml/jackson/databind/JsonNode;Lmisk/resources/ResourceLoader;ZLcom/fasterxml/jackson/databind/deser/BeanDeserializerModifier;)Lmisk/config/Config; public static final fun load (Ljava/lang/Class;Ljava/lang/String;Lwisp/deployment/Deployment;Ljava/util/List;Lmisk/resources/ResourceLoader;)Lmisk/config/Config; - public static synthetic fun load$default (Ljava/lang/Class;Ljava/lang/String;Lwisp/deployment/Deployment;Ljava/util/List;Lcom/fasterxml/jackson/databind/JsonNode;Lmisk/resources/ResourceLoader;ILjava/lang/Object;)Lmisk/config/Config; - public static synthetic fun load$default (Ljava/lang/Class;Ljava/lang/String;Lwisp/deployment/Deployment;Ljava/util/List;Lcom/fasterxml/jackson/databind/JsonNode;Lmisk/resources/ResourceLoader;ZILjava/lang/Object;)Lmisk/config/Config; - public static synthetic fun load$default (Ljava/lang/Class;Ljava/lang/String;Lwisp/deployment/Deployment;Ljava/util/List;Lcom/fasterxml/jackson/databind/JsonNode;Lmisk/resources/ResourceLoader;ZLcom/fasterxml/jackson/databind/deser/BeanDeserializerModifier;ILjava/lang/Object;)Lmisk/config/Config; + public static final fun load (Ljava/lang/Class;Ljava/lang/String;Lwisp/deployment/Deployment;Ljava/util/List;Ltools/jackson/databind/JsonNode;Lmisk/resources/ResourceLoader;)Lmisk/config/Config; + public static final fun load (Ljava/lang/Class;Ljava/lang/String;Lwisp/deployment/Deployment;Ljava/util/List;Ltools/jackson/databind/JsonNode;Lmisk/resources/ResourceLoader;Z)Lmisk/config/Config; + public static final fun load (Ljava/lang/Class;Ljava/lang/String;Lwisp/deployment/Deployment;Ljava/util/List;Ltools/jackson/databind/JsonNode;Lmisk/resources/ResourceLoader;ZLtools/jackson/databind/deser/ValueDeserializerModifier;)Lmisk/config/Config; public static synthetic fun load$default (Ljava/lang/Class;Ljava/lang/String;Lwisp/deployment/Deployment;Ljava/util/List;Lmisk/resources/ResourceLoader;ILjava/lang/Object;)Lmisk/config/Config; + public static synthetic fun load$default (Ljava/lang/Class;Ljava/lang/String;Lwisp/deployment/Deployment;Ljava/util/List;Ltools/jackson/databind/JsonNode;Lmisk/resources/ResourceLoader;ILjava/lang/Object;)Lmisk/config/Config; + public static synthetic fun load$default (Ljava/lang/Class;Ljava/lang/String;Lwisp/deployment/Deployment;Ljava/util/List;Ltools/jackson/databind/JsonNode;Lmisk/resources/ResourceLoader;ZILjava/lang/Object;)Lmisk/config/Config; + public static synthetic fun load$default (Ljava/lang/Class;Ljava/lang/String;Lwisp/deployment/Deployment;Ljava/util/List;Ltools/jackson/databind/JsonNode;Lmisk/resources/ResourceLoader;ZLtools/jackson/databind/deser/ValueDeserializerModifier;ILjava/lang/Object;)Lmisk/config/Config; public final fun loadConfigYamlMap (Ljava/lang/String;Lwisp/deployment/Deployment;Ljava/util/List;Lmisk/resources/ResourceLoader;)Ljava/util/Map; public static synthetic fun loadConfigYamlMap$default (Lmisk/config/MiskConfig;Ljava/lang/String;Lwisp/deployment/Deployment;Ljava/util/List;Lmisk/resources/ResourceLoader;ILjava/lang/Object;)Ljava/util/Map; public final fun toRedactedYaml (Lmisk/config/Config;Lmisk/resources/ResourceLoader;)Ljava/lang/String; } -public final class misk/config/MiskConfig$DeserializerModifierModule : com/fasterxml/jackson/databind/module/SimpleModule { - public fun (Lcom/fasterxml/jackson/databind/deser/BeanDeserializerModifier;)V - public final fun getDeserializerModifier ()Lcom/fasterxml/jackson/databind/deser/BeanDeserializerModifier; - public fun setupModule (Lcom/fasterxml/jackson/databind/Module$SetupContext;)V +public final class misk/config/MiskConfig$DeserializerModifierModule : tools/jackson/databind/module/SimpleModule { + public fun (Ltools/jackson/databind/deser/ValueDeserializerModifier;)V + public final fun getDeserializerModifier ()Ltools/jackson/databind/deser/ValueDeserializerModifier; + public fun setupModule (Ltools/jackson/databind/JacksonModule$SetupContext;)V } public final class misk/config/MiskConfig$RealSecret : misk/config/Secret { @@ -47,16 +47,17 @@ public final class misk/config/MiskConfig$RealSecret : misk/config/Secret { public fun toString ()Ljava/lang/String; } -public final class misk/config/MiskConfig$RedactSecretJacksonModule : com/fasterxml/jackson/databind/module/SimpleModule { +public final class misk/config/MiskConfig$RedactSecretJacksonModule : tools/jackson/databind/module/SimpleModule { public fun ()V - public fun setupModule (Lcom/fasterxml/jackson/databind/Module$SetupContext;)V + public fun setupModule (Ltools/jackson/databind/JacksonModule$SetupContext;)V } -public final class misk/config/MiskConfig$SecretJacksonModule : com/fasterxml/jackson/databind/module/SimpleModule { - public fun (Lmisk/resources/ResourceLoader;Lcom/fasterxml/jackson/databind/ObjectMapper;)V - public final fun getMapper ()Lcom/fasterxml/jackson/databind/ObjectMapper; +public final class misk/config/MiskConfig$SecretJacksonModule : tools/jackson/databind/module/SimpleModule { + public fun (Lmisk/resources/ResourceLoader;Lkotlin/jvm/functions/Function0;)V + public fun (Lmisk/resources/ResourceLoader;Ltools/jackson/databind/ObjectMapper;)V + public final fun getMapper ()Ltools/jackson/databind/ObjectMapper; public final fun getResourceLoader ()Lmisk/resources/ResourceLoader; - public fun setupModule (Lcom/fasterxml/jackson/databind/Module$SetupContext;)V + public fun setupModule (Ltools/jackson/databind/JacksonModule$SetupContext;)V } public abstract interface annotation class misk/config/Redact : java/lang/annotation/Annotation { diff --git a/misk-config/build.gradle.kts b/misk-config/build.gradle.kts index 864bda7d081..10e8b461fa7 100644 --- a/misk-config/build.gradle.kts +++ b/misk-config/build.gradle.kts @@ -20,7 +20,6 @@ dependencies { implementation(libs.guava) implementation(libs.jacksonCore) implementation(libs.jacksonDataformatYaml) - implementation(libs.jacksonJsr310) implementation(libs.jacksonKotlin) implementation(libs.loggingApi) implementation(libs.okio) diff --git a/misk-config/src/main/kotlin/misk/config/MiskConfig.kt b/misk-config/src/main/kotlin/misk/config/MiskConfig.kt index a116a13d17f..614895feaf0 100644 --- a/misk-config/src/main/kotlin/misk/config/MiskConfig.kt +++ b/misk-config/src/main/kotlin/misk/config/MiskConfig.kt @@ -1,32 +1,6 @@ package misk.config import com.fasterxml.jackson.annotation.JacksonAnnotationsInside -import com.fasterxml.jackson.core.JsonGenerator -import com.fasterxml.jackson.core.JsonParser -import com.fasterxml.jackson.core.JsonToken -import com.fasterxml.jackson.databind.BeanProperty -import com.fasterxml.jackson.databind.DeserializationContext -import com.fasterxml.jackson.databind.DeserializationFeature -import com.fasterxml.jackson.databind.JavaType -import com.fasterxml.jackson.databind.JsonDeserializer -import com.fasterxml.jackson.databind.JsonMappingException -import com.fasterxml.jackson.databind.JsonNode -import com.fasterxml.jackson.databind.JsonSerializer -import com.fasterxml.jackson.databind.ObjectMapper -import com.fasterxml.jackson.databind.SerializerProvider -import com.fasterxml.jackson.databind.annotation.JsonSerialize -import com.fasterxml.jackson.databind.deser.BeanDeserializerModifier -import com.fasterxml.jackson.databind.deser.ContextualDeserializer -import com.fasterxml.jackson.databind.exc.InvalidFormatException -import com.fasterxml.jackson.databind.exc.MismatchedInputException -import com.fasterxml.jackson.databind.exc.UnrecognizedPropertyException -import com.fasterxml.jackson.databind.module.SimpleModule -import com.fasterxml.jackson.databind.node.ObjectNode -import com.fasterxml.jackson.databind.ser.ContextualSerializer -import com.fasterxml.jackson.dataformat.yaml.YAMLFactory -import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule -import com.fasterxml.jackson.module.kotlin.KotlinInvalidNullException -import com.fasterxml.jackson.module.kotlin.KotlinModule import com.google.common.base.Joiner import java.io.File import java.io.FilenameFilter @@ -36,6 +10,32 @@ import kotlin.time.ExperimentalTime import misk.logging.getLogger import misk.resources.ResourceLoader import org.apache.commons.lang3.StringUtils +import tools.jackson.core.JsonGenerator +import tools.jackson.core.JsonParser +import tools.jackson.core.JsonToken +import tools.jackson.databind.BeanProperty +import tools.jackson.databind.DatabindException +import tools.jackson.databind.DeserializationContext +import tools.jackson.databind.DeserializationFeature +import tools.jackson.databind.JavaType +import tools.jackson.databind.JsonNode +import tools.jackson.databind.MapperFeature +import tools.jackson.databind.ObjectMapper +import tools.jackson.databind.SerializationContext +import tools.jackson.databind.ValueDeserializer +import tools.jackson.databind.ValueSerializer +import tools.jackson.databind.annotation.JsonSerialize +import tools.jackson.databind.cfg.EnumFeature +import tools.jackson.databind.deser.ValueDeserializerModifier +import tools.jackson.databind.exc.InvalidFormatException +import tools.jackson.databind.exc.MismatchedInputException +import tools.jackson.databind.exc.UnrecognizedPropertyException +import tools.jackson.databind.module.SimpleModule +import tools.jackson.databind.node.ObjectNode +import tools.jackson.dataformat.yaml.YAMLMapper +import tools.jackson.module.kotlin.KotlinFeature +import tools.jackson.module.kotlin.KotlinInvalidNullException +import tools.jackson.module.kotlin.KotlinModule import wisp.deployment.Deployment object MiskConfig { @@ -125,12 +125,10 @@ object MiskConfig { overrideValues: JsonNode? = null, resourceLoader: ResourceLoader = ResourceLoader.SYSTEM, failOnUnknownProperties: Boolean, - deserializerModifier: BeanDeserializerModifier? = null, + deserializerModifier: ValueDeserializerModifier? = null, ): T { check(!Secret::class.java.isAssignableFrom(configClass)) { "Top level service config cannot be a Secret<*>" } - val mapper = newObjectMapper(resourceLoader, false, deserializerModifier) - val configYamls = loadConfigYamlMap(appName, deployment, overrideResources, resourceLoader) check(configYamls.values.any { it != null }) { "could not find configuration files - checked ${configYamls.keys}" } @@ -139,7 +137,7 @@ object MiskConfig { val configFile = "$appName-${configEnvironmentName.lowercase(Locale.US)}.yaml" return readFlattenedYaml( - mapper, + { failOnUnknown -> newObjectMapper(resourceLoader, false, deserializerModifier, failOnUnknown) }, jsonNode, configClass, configFile, @@ -150,7 +148,7 @@ object MiskConfig { } private fun readFlattenedYaml( - mapper: ObjectMapper, + newMapper: (failOnUnknownProperties: Boolean) -> ObjectMapper, jsonNode: JsonNode, configClass: Class, configFile: String, @@ -160,24 +158,32 @@ object MiskConfig { ): T { try { @Suppress("UNCHECKED_CAST") - return mapper.readValue(jsonNode.toString(), configClass) as T + return newMapper(true).readValue(jsonNode.toString(), configClass) as T } catch (e: UnrecognizedPropertyException) { if (failOnUnknownProperties) { throw IllegalStateException("failed to load configuration for $appName $configEnvironmentName: ${e.message}", e) } - val path = Joiner.on('.').join(e.path.map { it.fieldName ?: it.index }) + val path = Joiner.on('.').join(e.path.map { it.propertyName ?: it.index }) logger.warn(e) { "$configFile: '$path' not found in '${configClass.simpleName}', ignoring " + suggestSpelling(e) } - // Try again, this time ignoring unknown properties. - mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false) - return readFlattenedYaml(mapper, jsonNode, configClass, configFile, appName, configEnvironmentName, false) + // Try again, this time ignoring unknown properties. Mappers are immutable, so build a new one + // rather than reconfiguring this one. + return readFlattenedYaml( + { newMapper(false) }, + jsonNode, + configClass, + configFile, + appName, + configEnvironmentName, + false, + ) } catch (e: KotlinInvalidNullException) { throwMissingPropertyException(e, configClass, configFile, jsonNode) } catch (e: InvalidFormatException) { // The property is present, it just cannot be represented as the declared type. Reporting it as missing (which // the MismatchedInputException branch below would do) sends readers looking for the wrong problem. - val path = Joiner.on('.').join(e.path.map { it.fieldName ?: it.index }) + val path = Joiner.on('.').join(e.path.map { it.propertyName ?: it.index }) throw IllegalStateException( "failed to load configuration for $appName $configEnvironmentName:" + " could not parse '$path' in $configFile: ${e.originalMessage}", @@ -191,12 +197,12 @@ object MiskConfig { } private fun throwMissingPropertyException( - e: JsonMappingException, + e: DatabindException, configClass: Class, configFile: String, jsonNode: JsonNode, ): Nothing { - val path = Joiner.on('.').join(e.path.map { it.fieldName ?: it.index }) + val path = Joiner.on('.').join(e.path.map { it.propertyName ?: it.index }) throw IllegalStateException( "could not find '${path}' of '${configClass.simpleName}'" + " in $configFile or in any of the combined logical config " + @@ -209,11 +215,11 @@ object MiskConfig { if (jsonNode.isObject) { val objectNode = jsonNode as ObjectNode - var seq = objectNode.fieldNames().asSequence().map { Joiner.on('.').join(pathPrefix, it) } + var seq = objectNode.propertyNames().asSequence().map { Joiner.on('.').join(pathPrefix, it) } // Recursively add the field names of any object fields. seq += - objectNode.fields().asSequence().flatMap { + objectNode.properties().asSequence().flatMap { val nextPrefix = Joiner.on('.').join(pathPrefix, it.key) allFieldNames(it.value, nextPrefix) } @@ -240,37 +246,61 @@ object MiskConfig { } fun toRedactedYaml(config: T, resourceLoader: ResourceLoader): String { - val serializingMapper = newObjectMapper(resourceLoader, true, null) + val serializingMapper = newObjectMapper(resourceLoader, true, null, failOnUnknownProperties = true) return serializingMapper.writeValueAsString(config) } private fun newObjectMapper( resourceLoader: ResourceLoader, redactSecrets: Boolean, - deserializerModifier: BeanDeserializerModifier?, + deserializerModifier: ValueDeserializerModifier?, + failOnUnknownProperties: Boolean, ): ObjectMapper { - val mapper = ObjectMapper(YAMLFactory()).registerModules(KotlinModule.Builder().build(), JavaTimeModule()) - - // Fail on null ints/doubles. - mapper.configure(DeserializationFeature.FAIL_ON_NULL_FOR_PRIMITIVES, true) + // The secret and resource deserializers parse nested documents with the very mapper they are + // registered on. Mappers are immutable and built in one shot, so hand the modules a supplier + // that resolves once the build below completes. + lateinit var mapper: ObjectMapper + val mapperProvider = { mapper } + + val builder = + YAMLMapper.builder() + // StrictNullChecks defaults off in Jackson 2 and on in Jackson 3. Leaving it on rejects a null element of a + // collection nested inside a map even when that element type is declared nullable -- Map> + // fails while a top-level Set is accepted -- so existing config stops loading. Keep it off to match + // Jackson 2. + .addModule(KotlinModule.Builder().disable(KotlinFeature.StrictNullChecks).build()) + // Fail on null ints/doubles. + .enable(DeserializationFeature.FAIL_ON_NULL_FOR_PRIMITIVES) + // Jackson 3 defaults this off. Config files are hand-written and a typo'd property should + // still surface as the "did you mean" warning below, so keep the Jackson 2 behaviour. + .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, failOnUnknownProperties) + // Jackson 3 defaults this on. Redacted config is rendered in the dashboard in declaration + // order today; sorting would silently reshuffle every service's config page. + .disable(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY) + // Jackson 3 defaults these on. Config enums are matched by name, and toString() is + // frequently overridden for display, so switching would break existing config files. + .disable(EnumFeature.READ_ENUMS_USING_TO_STRING, EnumFeature.WRITE_ENUMS_USING_TO_STRING) // The SecretDeserializer supports deserializing json, so bind last so it can use previous // mappings. if (redactSecrets) { - mapper.registerModule(RedactSecretJacksonModule()) + builder.addModule(RedactSecretJacksonModule()) } else { - mapper.registerModule(SecretJacksonModule(resourceLoader, mapper)) + builder.addModule(SecretJacksonModule(resourceLoader, mapperProvider)) } // The ResourceAwareDeserializer lets string and other primitive types be loaded by reference using resource loader // paths (classpath, filesystem, environment...) without using the Secret type. // This is useful for non-sensitive data or using environment variables to pass data into non-Secret types in // existing config or framework provided config classes. - mapper.registerModule(ResourceAwareJacksonModule(resourceLoader, mapper)) + builder.addModule(ResourceAwareJacksonModule(resourceLoader, mapperProvider)) + + builder.addModule(LinkedHashMapJacksonModule()) // The deserializerModifier can be null if this mapper is serializing only. - deserializerModifier?.let { mapper.registerModule(DeserializerModifierModule(it)) } + deserializerModifier?.let { builder.addModule(DeserializerModifierModule(it)) } + mapper = builder.build() return mapper } @@ -289,7 +319,7 @@ object MiskConfig { * Returns a JsonNode that combines the YAMLs in `configYamls`. If two nodes define the same value the last one wins. */ private fun flattenYamlMap(configYamls: Map, overrideValues: JsonNode?): JsonNode { - val mapper = ObjectMapper(YAMLFactory()).registerModules(KotlinModule.Builder().build(), JavaTimeModule()) + val mapper = YAMLMapper.builder().addModule(KotlinModule.Builder().build()).build() var result = mapper.createObjectNode() for ((key, value) in configYamls) { @@ -328,35 +358,74 @@ object MiskConfig { private fun embeddedConfigFileNames(appName: String, deployment: Deployment) = listOf("common", deployment.mapToEnvironmentName().lowercase(Locale.US)).map { "$appName-$it.yaml" } - class SecretJacksonModule(val resourceLoader: ResourceLoader, val mapper: ObjectMapper) : SimpleModule() { + class SecretJacksonModule(val resourceLoader: ResourceLoader, private val mapperProvider: () -> ObjectMapper) : + SimpleModule() { + constructor(resourceLoader: ResourceLoader, mapper: ObjectMapper) : this(resourceLoader, { mapper }) + + /** + * The mapper nested secret documents are parsed with. Resolved lazily because a module has to be registered before + * the mapper it belongs to exists. + */ + val mapper: ObjectMapper + get() = mapperProvider() + override fun setupModule(context: SetupContext?) { - addDeserializer(Secret::class.java, SecretDeserializer(resourceLoader, mapper)) + addDeserializer(Secret::class.java, SecretDeserializer(resourceLoader, mapperProvider)) super.setupModule(context) } } - class DeserializerModifierModule(val deserializerModifier: BeanDeserializerModifier) : SimpleModule() { + class DeserializerModifierModule(val deserializerModifier: ValueDeserializerModifier) : SimpleModule() { override fun setupModule(context: SetupContext?) { setDeserializerModifier(deserializerModifier) super.setupModule(context) } } - private class ResourceAwareJacksonModule(val resourceLoader: ResourceLoader, val mapper: ObjectMapper) : + private class ResourceAwareJacksonModule(val resourceLoader: ResourceLoader, val mapperProvider: () -> ObjectMapper) : SimpleModule() { override fun setupModule(context: SetupContext?) { - addDeserializer(String::class.java, ResourceAwareDeserializer(resourceLoader, mapper)) - addDeserializer(Int::class.java, ResourceAwareDeserializer(resourceLoader, mapper)) - addDeserializer(Integer::class.java, ResourceAwareDeserializer(resourceLoader, mapper)) - addDeserializer(Long::class.java, ResourceAwareDeserializer(resourceLoader, mapper)) - addDeserializer(java.lang.Long::class.java, ResourceAwareDeserializer(resourceLoader, mapper)) - addDeserializer(Float::class.java, ResourceAwareDeserializer(resourceLoader, mapper)) - addDeserializer(java.lang.Float::class.java, ResourceAwareDeserializer(resourceLoader, mapper)) - addDeserializer(Boolean::class.java, ResourceAwareDeserializer(resourceLoader, mapper)) + addDeserializer(String::class.java, ResourceAwareDeserializer(resourceLoader, mapperProvider)) + addDeserializer(Int::class.java, ResourceAwareDeserializer(resourceLoader, mapperProvider)) + addDeserializer(Integer::class.java, ResourceAwareDeserializer(resourceLoader, mapperProvider)) + addDeserializer(Long::class.java, ResourceAwareDeserializer(resourceLoader, mapperProvider)) + addDeserializer( + java.lang.Long::class.java, + ResourceAwareDeserializer(resourceLoader, mapperProvider), + ) + addDeserializer(Float::class.java, ResourceAwareDeserializer(resourceLoader, mapperProvider)) + addDeserializer( + java.lang.Float::class.java, + ResourceAwareDeserializer(resourceLoader, mapperProvider), + ) + addDeserializer(Boolean::class.java, ResourceAwareDeserializer(resourceLoader, mapperProvider)) addDeserializer( java.lang.Boolean::class.java, - ResourceAwareDeserializer(resourceLoader, mapper), + ResourceAwareDeserializer(resourceLoader, mapperProvider), + ) + + super.setupModule(context) + } + } + + /** + * Keeps a `Map` field declared with an enum key backed by a [LinkedHashMap], as it was under Jackson 2. + * + * Jackson 3 rewrites `Map` to an [EnumMap] (databind#1853), which iterates in enum declaration order + * rather than in the order the keys appear in the YAML. That is invisible at load time and only shows up wherever + * config iteration order is observable -- a service whose behaviour depends on which region or tier it processes + * first would quietly change, and nothing would fail. `BasicDeserializerFactory.createMapDeserializer` applies the + * rewrite unconditionally with no feature flag, but it only does so when the declared raw type is exactly `Map`, and + * abstract type resolution runs first (`DeserializerCache._createDeserializer`), so naming the implementation here + * settles it before that branch is reached. A field declared as an `EnumMap` still gets one. + */ + private class LinkedHashMapJacksonModule : SimpleModule() { + override fun setupModule(context: SetupContext?) { + @Suppress("UNCHECKED_CAST") + addAbstractTypeMapping( + Map::class.java as Class>, + LinkedHashMap::class.java as Class>, ) super.setupModule(context) @@ -365,19 +434,19 @@ object MiskConfig { private inline fun ResourceAwareDeserializer( resourceLoader: ResourceLoader, - mapper: ObjectMapper, - ): ResourceAwareDeserializer = ResourceAwareDeserializer(T::class, resourceLoader, mapper) + noinline mapperProvider: () -> ObjectMapper, + ): ResourceAwareDeserializer = ResourceAwareDeserializer(T::class, resourceLoader, mapperProvider) private class ResourceAwareDeserializer( val typeClass: KClass, val resourceLoader: ResourceLoader, - val mapper: ObjectMapper, + val mapperProvider: () -> ObjectMapper, val type: JavaType? = null, - ) : JsonDeserializer(), ContextualDeserializer { + ) : ValueDeserializer() { override fun deserialize(jsonParser: JsonParser, deserializationContext: DeserializationContext): T? { if (type == null) { // This only happens if ObjectMapper does not call createContextual for this property. - throw JsonMappingException.from(jsonParser, "Attempting to deserialize an object with no type") + throw DatabindException.from(jsonParser, "Attempting to deserialize an object with no type") } val maybeReferenceWithMarkers = jsonParser.valueAsString @@ -425,7 +494,7 @@ object MiskConfig { val maybeReference = "$scheme:$path" - resourceLoader.loadResource(maybeReference, type, mapper, default) as? T? + resourceLoader.loadResource(maybeReference, type, mapperProvider(), default) as? T? } // Not a resource reference, so convert the scalar itself. This is deliberately evaluated here rather than up // front: a reference like "${environment:PORT}" is not a valid Int, so converting eagerly would reject @@ -433,40 +502,41 @@ object MiskConfig { ?: jsonParser.valueAsTypeOrNull(type) as? T? } - override fun createContextual(ctxt: DeserializationContext?, property: BeanProperty?): JsonDeserializer<*>? { - return ResourceAwareDeserializer(typeClass, resourceLoader, mapper, mapper.constructType(typeClass.java)) + override fun createContextual(ctxt: DeserializationContext, property: BeanProperty?): ValueDeserializer<*> { + val resolved = ctxt.constructType(typeClass.java) + return ResourceAwareDeserializer(typeClass, resourceLoader, mapperProvider, resolved) } } private class SecretDeserializer( val resourceLoader: ResourceLoader, - val mapper: ObjectMapper, + val mapperProvider: () -> ObjectMapper, val type: JavaType? = null, - ) : JsonDeserializer>(), ContextualDeserializer { + ) : ValueDeserializer>() { override fun createContextual( - deserializationContext: DeserializationContext?, + deserializationContext: DeserializationContext, property: BeanProperty, - ): JsonDeserializer<*> { - return SecretDeserializer(resourceLoader, mapper, property.type.bindings.getBoundType(0)) + ): ValueDeserializer<*> { + return SecretDeserializer(resourceLoader, mapperProvider, property.type.bindings.getBoundType(0)) } override fun deserialize(jsonParser: JsonParser, deserializationContext: DeserializationContext): Secret<*>? { if (type == null) { // This only happens if ObjectMapper does not call createContextual for this property. - throw JsonMappingException.from(jsonParser, "Attempting to deserialize an object with no type") + throw DatabindException.from(jsonParser, "Attempting to deserialize an object with no type") } val reference = jsonParser.valueAsString - return RealSecret(resourceLoader.loadResource(reference, type, mapper), reference) + return RealSecret(resourceLoader.loadResource(reference, type, mapperProvider()), reference) } } - internal class RedactSecretJsonSerializer : JsonSerializer(), ContextualSerializer { - override fun serialize(value: Any, gen: JsonGenerator, serializers: SerializerProvider) { + internal class RedactSecretJsonSerializer : ValueSerializer() { + override fun serialize(value: Any, gen: JsonGenerator, ctxt: SerializationContext) { gen.writeString("████████") } - override fun createContextual(prov: SerializerProvider, property: BeanProperty): JsonSerializer<*> { + override fun createContextual(ctxt: SerializationContext, property: BeanProperty): ValueSerializer<*> { return RedactSecretJsonSerializer() } } @@ -478,8 +548,8 @@ object MiskConfig { } } - private class RedactSecretSerializer : JsonSerializer>(), ContextualSerializer { - override fun serialize(value: Secret<*>, gen: JsonGenerator, serializers: SerializerProvider?) { + private class RedactSecretSerializer : ValueSerializer>() { + override fun serialize(value: Secret<*>, gen: JsonGenerator, ctxt: SerializationContext) { if ((value as? RealSecret<*>)?.reference?.isNotBlank() == true) { gen.writeString("${value.reference} -> ████████") } else { @@ -487,7 +557,7 @@ object MiskConfig { } } - override fun createContextual(prov: SerializerProvider?, property: BeanProperty): JsonSerializer<*> { + override fun createContextual(ctxt: SerializationContext, property: BeanProperty): ValueSerializer<*> { return RedactSecretSerializer() } } diff --git a/misk-config/src/test/kotlin/misk/config/MiskConfigTest.kt b/misk-config/src/test/kotlin/misk/config/MiskConfigTest.kt index 17bd1cf05fb..c1e4680105e 100644 --- a/misk-config/src/test/kotlin/misk/config/MiskConfigTest.kt +++ b/misk-config/src/test/kotlin/misk/config/MiskConfigTest.kt @@ -1,6 +1,5 @@ package misk.config -import com.fasterxml.jackson.databind.ObjectMapper import com.google.inject.util.Modules import jakarta.inject.Inject import java.io.File @@ -24,6 +23,7 @@ import org.junit.jupiter.api.BeforeEach import org.junit.jupiter.api.Test import org.junit.jupiter.api.extension.ExtendWith import org.slf4j.event.Level +import tools.jackson.databind.ObjectMapper import uk.org.webcompere.systemstubs.environment.EnvironmentVariables import uk.org.webcompere.systemstubs.jupiter.SystemStub import uk.org.webcompere.systemstubs.jupiter.SystemStubsExtension @@ -182,6 +182,15 @@ class MiskConfigTest { assertEquals(true, actual.boolean_value) } + @Test + fun enumKeyedMapKeepsYamlOrder() { + val actual = MiskConfig.load("enum_keyed_map", TESTING) + + // An EnumMap would answer CHECKING, SAVINGS, BROKERAGE regardless of how the file is written. + assertThat(actual.tiers.keys).containsExactly(AccountTier.BROKERAGE, AccountTier.CHECKING, AccountTier.SAVINGS) + assertThat(actual.tiers).isNotInstanceOf(java.util.EnumMap::class.java) + } + @Test fun configLoadsValuesFromEnvironmentVariables() { // Set environment variables before loading config @@ -269,7 +278,7 @@ class MiskConfigTest { assertFailsWith { MiskConfig.load(TestConfig::class.java, "unknownproperty", TESTING, failOnUnknownProperties = true) } - assertThat(exception).hasMessageContaining("Unrecognized field \"blue_items\"") + assertThat(exception).hasMessageContaining("Unrecognized property \"blue_items\"") } @Test diff --git a/misk-config/src/test/kotlin/misk/config/TestConfig.kt b/misk-config/src/test/kotlin/misk/config/TestConfig.kt index 8a29a5b8c2a..edfa0f29882 100644 --- a/misk-config/src/test/kotlin/misk/config/TestConfig.kt +++ b/misk-config/src/test/kotlin/misk/config/TestConfig.kt @@ -47,3 +47,12 @@ data class EnvironmentTestConfig( val jdbc_url_default: String, val https_url_with_port_default: String, ) : Config + +/** [tiers] is deliberately a plain `Map`: that is the declaration Jackson 3 rewrites to an `EnumMap`. */ +data class EnumKeyedMapConfig(val tiers: Map) : Config + +enum class AccountTier { + CHECKING, + SAVINGS, + BROKERAGE, +} diff --git a/misk-config/src/test/resources/enum_keyed_map-common.yaml b/misk-config/src/test/resources/enum_keyed_map-common.yaml new file mode 100644 index 00000000000..5e7d4593c02 --- /dev/null +++ b/misk-config/src/test/resources/enum_keyed_map-common.yaml @@ -0,0 +1,5 @@ +# Deliberately not in enum declaration order: an EnumMap would iterate CHECKING, SAVINGS, BROKERAGE. +tiers: + BROKERAGE: 3 + CHECKING: 1 + SAVINGS: 2 diff --git a/misk-testing/src/main/kotlin/org/assertj/core/api/AssertExtensions.kt b/misk-testing/src/main/kotlin/org/assertj/core/api/AssertExtensions.kt index cd8a17f04dc..b2e04e927bd 100644 --- a/misk-testing/src/main/kotlin/org/assertj/core/api/AssertExtensions.kt +++ b/misk-testing/src/main/kotlin/org/assertj/core/api/AssertExtensions.kt @@ -1,10 +1,10 @@ package org.assertj.core.api -import com.fasterxml.jackson.core.JacksonException -import com.fasterxml.jackson.core.JsonParser -import com.fasterxml.jackson.databind.ObjectMapper import org.assertj.core.api.Assertions.assertThat import org.assertj.core.api.AssertionsForClassTypes.fail +import tools.jackson.core.JacksonException +import tools.jackson.core.StreamReadFeature +import tools.jackson.databind.json.JsonMapper inline fun MapAssert.containsExactly( vararg p: Pair @@ -12,7 +12,7 @@ inline fun MapAssert.containsExactly( return isEqualTo(mapOf(*p)) } -private val objectMapper = ObjectMapper().configure(JsonParser.Feature.INCLUDE_SOURCE_IN_LOCATION, true) +private val objectMapper = JsonMapper.builder().enable(StreamReadFeature.INCLUDE_SOURCE_IN_LOCATION).build() fun AbstractCharSequenceAssert<*, ACTUAL>.isEqualToAsJson( expected: CharSequence diff --git a/misk/api/misk.api b/misk/api/misk.api index 444d456c0d8..7047dc03c38 100644 --- a/misk/api/misk.api +++ b/misk/api/misk.api @@ -119,12 +119,12 @@ public final class misk/client/BackwardsCompatibleClientsConfig { public fun toString ()Ljava/lang/String; } -public final class misk/client/BackwardsCompatibleClientsConfigConverter : com/fasterxml/jackson/databind/util/Converter { +public final class misk/client/BackwardsCompatibleClientsConfigConverter : tools/jackson/databind/util/StdConverter { public fun ()V public synthetic fun convert (Ljava/lang/Object;)Ljava/lang/Object; public fun convert (Lmisk/client/BackwardsCompatibleClientsConfig;)Lmisk/client/HttpClientsConfig; - public fun getInputType (Lcom/fasterxml/jackson/databind/type/TypeFactory;)Lcom/fasterxml/jackson/databind/JavaType; - public fun getOutputType (Lcom/fasterxml/jackson/databind/type/TypeFactory;)Lcom/fasterxml/jackson/databind/JavaType; + public fun getInputType (Ltools/jackson/databind/type/TypeFactory;)Ltools/jackson/databind/JavaType; + public fun getOutputType (Ltools/jackson/databind/type/TypeFactory;)Ltools/jackson/databind/JavaType; } public final class misk/client/BackwardsCompatibleEndpointConfig { diff --git a/misk/src/main/kotlin/misk/client/HttpClientsConfig.kt b/misk/src/main/kotlin/misk/client/HttpClientsConfig.kt index cca3370e7b2..b171ee12907 100644 --- a/misk/src/main/kotlin/misk/client/HttpClientsConfig.kt +++ b/misk/src/main/kotlin/misk/client/HttpClientsConfig.kt @@ -1,13 +1,13 @@ package misk.client import com.fasterxml.jackson.annotation.JsonAlias -import com.fasterxml.jackson.databind.annotation.JsonDeserialize import java.net.URL import java.time.Duration import misk.config.Config import misk.logging.getLogger import misk.security.ssl.CertStoreConfig import misk.security.ssl.TrustStoreConfig +import tools.jackson.databind.annotation.JsonDeserialize @JsonDeserialize(converter = BackwardsCompatibleClientsConfigConverter::class) data class HttpClientsConfig diff --git a/misk/src/main/kotlin/misk/client/HttpClientsConfigBackwardsCompatibility.kt b/misk/src/main/kotlin/misk/client/HttpClientsConfigBackwardsCompatibility.kt index 94b6126251f..fe14736dee5 100644 --- a/misk/src/main/kotlin/misk/client/HttpClientsConfigBackwardsCompatibility.kt +++ b/misk/src/main/kotlin/misk/client/HttpClientsConfigBackwardsCompatibility.kt @@ -1,9 +1,9 @@ package misk.client import com.fasterxml.jackson.annotation.JsonAlias -import com.fasterxml.jackson.databind.type.TypeFactory -import com.fasterxml.jackson.databind.util.Converter import java.time.Duration +import tools.jackson.databind.type.TypeFactory +import tools.jackson.databind.util.StdConverter data class BackwardsCompatibleEndpointConfig @JvmOverloads @@ -47,7 +47,7 @@ constructor( val logRequests: Boolean = false, ) -class BackwardsCompatibleClientsConfigConverter : Converter { +class BackwardsCompatibleClientsConfigConverter : StdConverter() { override fun getInputType(typeFactory: TypeFactory) = typeFactory.constructType(BackwardsCompatibleClientsConfig::class.java) From 6c28a1603a6a0db6713cc557ccfe911ba57e9a28 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juho=20M=C3=A4kinen?= Date: Mon, 17 Aug 2026 00:50:27 +0000 Subject: [PATCH 3/5] misk-hibernate: resolve Database Query entities by authoritative KClass (#3898) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit HibernateDatabaseQueryDynamicAction authorizes a caller against a query class's metadata, then resolved the entity to actually query by Kotlin simpleName (`transacter.entities().find { it.simpleName == request.entityClass }`). simpleName is not unique: two entities in the same transacter can share one (e.g. a `DbMovie` in two different packages). When a colliding entity that was never exposed to Database Query is registered first, the lookup selects it, and the follow-up binding check added in #3827 also compares simpleName — so both the authorized and the colliding entity pass it — disclosing the unauthorized entity's rows. This is an incomplete-fix bypass of #3827 (VULN-78154), which used simpleName as the identity. Carry the authoritative KClass identities server-side instead of matching by name. A new internal HibernateDatabaseQueryRegistration pairs each entity's DatabaseQueryMetadata with the exact entity/query KClasses it was built from. It is multibound alongside the existing metadata and is never serialized, so the browser-facing JSON API (DatabaseQueryMetadata) is unchanged. - Dynamic action: run against registration.entityClass; still reject a request whose entityClass names a different entity than the authorized query (VULN-78154 contract). - Static action: run registration.queryClass / registration.entityClass directly instead of re-deriving the query by simpleName. - Transacter lookup: match by KClass identity rather than simple name. - Removes the now-unused HibernateQuery multibinding. Adds a regression test that reproduces the collision (two `DbMovie` classes, the protected one mapped to `actors` and registered first) and asserts the authorized `movies` entity is queried. CWE-863 (Incorrect Authorization), CWE-285 (Improper Authorization). Reported via Block's bug bounty program (Bugcrowd), tracked internally as VULN-78266. Co-authored-by: Claude Opus 4.8 --- .../misk/hibernate/HibernateEntityModule.kt | 17 ++- .../HibernateDatabaseQueryDynamicAction.kt | 27 ++-- .../HibernateDatabaseQueryRegistration.kt | 42 ++++++ .../HibernateDatabaseQueryStaticAction.kt | 33 ++--- .../HibernateDatabaseQueryWebActionModule.kt | 20 ++- .../misk/hibernate/actions/HibernateQuery.kt | 11 -- ...ateDatabaseQuerySimpleNameCollisionTest.kt | 138 ++++++++++++++++++ 7 files changed, 226 insertions(+), 62 deletions(-) create mode 100644 misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateDatabaseQueryRegistration.kt delete mode 100644 misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateQuery.kt create mode 100644 misk-hibernate/src/test/kotlin/misk/hibernate/actions/HibernateDatabaseQuerySimpleNameCollisionTest.kt diff --git a/misk-hibernate/src/main/kotlin/misk/hibernate/HibernateEntityModule.kt b/misk-hibernate/src/main/kotlin/misk/hibernate/HibernateEntityModule.kt index 64d46973495..5b0c9736090 100644 --- a/misk-hibernate/src/main/kotlin/misk/hibernate/HibernateEntityModule.kt +++ b/misk-hibernate/src/main/kotlin/misk/hibernate/HibernateEntityModule.kt @@ -6,8 +6,9 @@ import com.google.inject.name.Names import java.util.concurrent.atomic.AtomicInteger import kotlin.reflect.KClass import misk.hibernate.actions.DatabaseQueryMetadataProvider +import misk.hibernate.actions.HibernateDatabaseQueryRegistration +import misk.hibernate.actions.HibernateDatabaseQueryRegistrationProvider import misk.hibernate.actions.HibernateDatabaseQueryWebActionModule -import misk.hibernate.actions.HibernateQuery import misk.inject.KAbstractModule import misk.security.authz.AccessAnnotationEntry import misk.web.metadata.database.DatabaseQueryMetadata @@ -23,8 +24,8 @@ abstract class HibernateEntityModule(private val qualifier: KClass() newMultibinder() + newMultibinder() newMultibinder(qualifier) - newMultibinder() newMultibinder(qualifier) configureHibernate() @@ -60,9 +61,15 @@ abstract class HibernateEntityModule(private val qualifier: KClass().toInstance(HibernateQuery(queryClass as KClass>>)) - } + // Server-side companion to the metadata above that carries the entity/query KClasses the query actions run against. + multibind() + .toProvider( + HibernateDatabaseQueryRegistrationProvider( + dbEntityClass = dbEntityClass, + queryClass = queryClass, + accessAnnotationClass = accessAnnotationClass, + ) + ) } /** Install Entity with a default of no query access from Admin Dashboard */ diff --git a/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateDatabaseQueryDynamicAction.kt b/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateDatabaseQueryDynamicAction.kt index fc7559de60b..da34c4dfeb7 100644 --- a/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateDatabaseQueryDynamicAction.kt +++ b/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateDatabaseQueryDynamicAction.kt @@ -14,7 +14,7 @@ import misk.hibernate.ReflectionQuery import misk.hibernate.Session import misk.hibernate.Transacter import misk.hibernate.actions.HibernateDatabaseQueryWebActionModule.Companion.checkQueryMatchesAction -import misk.hibernate.actions.HibernateDatabaseQueryWebActionModule.Companion.findDatabaseQueryMetadata +import misk.hibernate.actions.HibernateDatabaseQueryWebActionModule.Companion.findDatabaseQueryRegistration import misk.hibernate.actions.HibernateDatabaseQueryWebActionModule.Companion.getTransacterForDatabaseQueryAction import misk.hibernate.actions.HibernateDatabaseQueryWebActionModule.Companion.validateSelectPathsOrDefault import misk.logging.getLogger @@ -26,7 +26,6 @@ import misk.web.ResponseContentType import misk.web.actions.WebAction import misk.web.dashboard.AdminDashboardAccess import misk.web.mediatype.MediaTypes -import misk.web.metadata.database.DatabaseQueryMetadata /** Runs query from Database Query dashboard tab against DB and returns results */ @Singleton @@ -34,7 +33,7 @@ internal class HibernateDatabaseQueryDynamicAction @Inject constructor( @JvmSuppressWildcards private val callerProvider: ActionScoped, - private val databaseQueryMetadata: List, + private val databaseQueryRegistrations: List, private val injector: Injector, private val queryLimitsConfig: ReflectionQuery.QueryLimitsConfig, ) : WebAction { @@ -50,12 +49,13 @@ constructor( checkQueryMatchesAction(queryClass, true) - val metadata = findDatabaseQueryMetadata(databaseQueryMetadata, queryClass) - val transacter = getTransacterForDatabaseQueryAction(injector, metadata) + val registration = findDatabaseQueryRegistration(databaseQueryRegistrations, queryClass) + val metadata = registration.metadata + val transacter = getTransacterForDatabaseQueryAction(injector, registration.entityClass) val results = if (caller.isAllowed(metadata.allowedCapabilities, metadata.allowedServices)) { - runDynamicQuery(transacter, caller.principal, request, metadata) + runDynamicQuery(transacter, caller.principal, request, registration) } else { throw UnauthorizedException("Unauthorized to query [dbEntity=${metadata.entityClass}]") } @@ -67,22 +67,17 @@ constructor( transacter: Transacter, principal: String, request: Request, - metadata: DatabaseQueryMetadata, + registration: HibernateDatabaseQueryRegistration, ) = transacter.transaction { session -> - val dbEntity = - transacter.entities().find { it.simpleName == request.entityClass } - ?: throw BadRequestException("[dbEntity=${request.entityClass}] is not an installed HibernateEntity") - // Authorization is performed against `queryClass` (and its bound `metadata.entityClass`), but the - // entity actually queried is taken from the user-controlled `request.entityClass`. Without the - // binding check below, a caller authorized for one query class could read any registered entity - // by setting `entityClass` to a different entity name. Enforce that the entity executed matches - // the entity the query class is authorized for. - if (dbEntity.simpleName != metadata.entityClass) { + // Query the entity the authorized registration was built from. Reject a request naming a different + // entity than the authorized query class so the caller gets a clear error. + if (request.entityClass != registration.metadata.entityClass) { throw UnauthorizedException( "Requested entity [dbEntity=${request.entityClass}] does not match authorized query [queryClass=${request.queryClass}]" ) } + val dbEntity = registration.entityClass val (selectPaths, rows) = runDynamicQuery(session, principal, dbEntity, request) rows.map { row -> // TODO (adrw) sort the map based on DbEntity order diff --git a/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateDatabaseQueryRegistration.kt b/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateDatabaseQueryRegistration.kt new file mode 100644 index 00000000000..239826ccf03 --- /dev/null +++ b/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateDatabaseQueryRegistration.kt @@ -0,0 +1,42 @@ +package misk.hibernate.actions + +import com.google.inject.Provider +import jakarta.inject.Inject +import kotlin.reflect.KClass +import misk.hibernate.DbEntity +import misk.hibernate.Query +import misk.web.metadata.database.DatabaseQueryMetadata +import misk.web.metadata.database.NoAdminDashboardDatabaseAccess + +/** + * Server-side registration for a single Database Query entity/query. Pairs the browser-facing [DatabaseQueryMetadata] + * with the [KClass] identities it was built from. This is never serialized: the query actions authorize against + * [metadata] and run the query against [entityClass]/[queryClass] directly. + */ +internal data class HibernateDatabaseQueryRegistration( + val metadata: DatabaseQueryMetadata, + val entityClass: KClass>, + /** The static Misk [Query] class, or null for a dynamic query. */ + val queryClass: KClass>>?, +) + +internal class HibernateDatabaseQueryRegistrationProvider>( + private val dbEntityClass: KClass, + private val queryClass: KClass>?, + private val accessAnnotationClass: KClass = NoAdminDashboardDatabaseAccess::class, +) : Provider { + @Inject lateinit var hibernateDatabaseQueryMetadataFactory: HibernateDatabaseQueryMetadataFactory + + @Suppress("UNCHECKED_CAST") + override fun get(): HibernateDatabaseQueryRegistration = + HibernateDatabaseQueryRegistration( + metadata = + hibernateDatabaseQueryMetadataFactory.fromQuery( + dbEntityClass = dbEntityClass, + queryClass = queryClass, + accessAnnotationClass = accessAnnotationClass, + ), + entityClass = dbEntityClass, + queryClass = queryClass as KClass>>?, + ) +} diff --git a/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateDatabaseQueryStaticAction.kt b/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateDatabaseQueryStaticAction.kt index 9c017119694..7849761a0b3 100644 --- a/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateDatabaseQueryStaticAction.kt +++ b/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateDatabaseQueryStaticAction.kt @@ -3,7 +3,6 @@ package misk.hibernate.actions import com.google.inject.Injector import jakarta.inject.Inject import jakarta.inject.Singleton -import java.lang.reflect.ParameterizedType import kotlin.reflect.KClass import misk.MiskCaller import misk.audit.AuditRequestResponse @@ -17,10 +16,9 @@ import misk.hibernate.Session import misk.hibernate.Transacter import misk.hibernate.actions.HibernateDatabaseQueryMetadataFactory.Companion.QUERY_CONFIG_TYPE_NAME import misk.hibernate.actions.HibernateDatabaseQueryWebActionModule.Companion.checkQueryMatchesAction -import misk.hibernate.actions.HibernateDatabaseQueryWebActionModule.Companion.findDatabaseQueryMetadata +import misk.hibernate.actions.HibernateDatabaseQueryWebActionModule.Companion.findDatabaseQueryRegistration import misk.hibernate.actions.HibernateDatabaseQueryWebActionModule.Companion.getTransacterForDatabaseQueryAction import misk.hibernate.actions.HibernateDatabaseQueryWebActionModule.Companion.validateSelectPathsOrDefault -import misk.inject.typeLiteral import misk.logging.getLogger import misk.scope.ActionScoped import misk.web.Post @@ -38,8 +36,7 @@ internal class HibernateDatabaseQueryStaticAction @Inject constructor( @JvmSuppressWildcards private val callerProvider: ActionScoped, - private val databaseQueryMetadata: List, - private val queries: List, + private val databaseQueryRegistrations: List, private val injector: Injector, private val queryLimitsConfig: ReflectionQuery.QueryLimitsConfig, ) : WebAction { @@ -55,12 +52,13 @@ constructor( checkQueryMatchesAction(queryClass, false) - val metadata = findDatabaseQueryMetadata(databaseQueryMetadata, queryClass) - val transacter = getTransacterForDatabaseQueryAction(injector, metadata) + val registration = findDatabaseQueryRegistration(databaseQueryRegistrations, queryClass) + val metadata = registration.metadata + val transacter = getTransacterForDatabaseQueryAction(injector, registration.entityClass) val results = if (caller.isAllowed(metadata.allowedCapabilities, metadata.allowedServices)) { - runStaticQuery(transacter, caller.principal, request, metadata) + runStaticQuery(transacter, caller.principal, request, registration) } else { throw UnauthorizedException("Unauthorized to query [dbEntity=${metadata.entityClass}]") } @@ -72,10 +70,10 @@ constructor( transacter: Transacter, principal: String, request: Request, - metadata: DatabaseQueryMetadata, + registration: HibernateDatabaseQueryRegistration, ) = transacter.transaction { session -> - val (selectPaths, rows) = runStaticQuery(session, principal, request, metadata) + val (selectPaths, rows) = runStaticQuery(session, principal, request, registration) rows.map { row -> // TODO (adrw) sort the map based on DbEntity order // TODO (adrw) Mirror this over to the static path @@ -87,15 +85,12 @@ constructor( session: Session, principal: String, request: Request, - metadata: DatabaseQueryMetadata, + registration: HibernateDatabaseQueryRegistration, ): Pair, List>> { - val query = - queries.map { it.query }.find { it.simpleName == metadata.queryClass } - ?: throw BadRequestException("[query=${metadata.queryClass}] does not exist") - val dbEntity = - ((query.typeLiteral().getSupertype(Query::class.java).type as ParameterizedType).actualTypeArguments.first() - as Class>) - .kotlin + val metadata = registration.metadata + // Use the Query/DbEntity classes from the authorized registration. + val query = registration.queryClass ?: throw BadRequestException("[query=${metadata.queryClass}] does not exist") + val dbEntity = registration.entityClass val maxRows = ((request.query[QUERY_CONFIG_TYPE_NAME] as Map?)?.get("maxRows") as Double?)?.toInt() ?: queryLimitsConfig.maxMaxRows @@ -114,7 +109,7 @@ constructor( private fun getStaticSelectPaths( request: Request, metadata: DatabaseQueryMetadata, - dbEntity: KClass>, + dbEntity: KClass>, ): List { val selectMetadata: DatabaseQueryMetadata.SelectMetadata? = request.query.entries diff --git a/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateDatabaseQueryWebActionModule.kt b/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateDatabaseQueryWebActionModule.kt index f61e1dca819..c0dcdbc8b78 100644 --- a/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateDatabaseQueryWebActionModule.kt +++ b/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateDatabaseQueryWebActionModule.kt @@ -37,22 +37,20 @@ internal class HibernateDatabaseQueryWebActionModule : KAbstractModule() { } } - /** Find the corresponding DatabaseQueryMetadata for a request */ - fun findDatabaseQueryMetadata( - databaseQueryMetadata: List, + /** Find the server-side registration (metadata + authoritative KClasses) for a request's query class */ + fun findDatabaseQueryRegistration( + registrations: List, queryClass: String, - ): DatabaseQueryMetadata = - databaseQueryMetadata.find { it.queryClass == queryClass } ?: throw BadRequestException("Invalid Query Class") + ): HibernateDatabaseQueryRegistration = + registrations.find { it.metadata.queryClass == queryClass } ?: throw BadRequestException("Invalid Query Class") - /** Common to both Dynamic and Static actions, get Transacter for a given request's DbEntity */ - fun getTransacterForDatabaseQueryAction(injector: Injector, metadata: DatabaseQueryMetadata): Transacter = + /** Common to both Dynamic and Static actions, get the Transacter that owns the authorized [entityClass]. */ + fun getTransacterForDatabaseQueryAction(injector: Injector, entityClass: KClass>): Transacter = injector .findBindingsByType(Transacter::class.typeLiteral()) - .find { transacterBinding -> - transacterBinding.provider.get().entities().map { it.simpleName!! }.contains(metadata.entityClass) - } + .find { transacterBinding -> transacterBinding.provider.get().entities().contains(entityClass) } ?.provider - ?.get() ?: throw BadRequestException("[dbEntity=${metadata.entityClass}] has no associated Transacter") + ?.get() ?: throw BadRequestException("[dbEntity=${entityClass.simpleName}] has no associated Transacter") /** Validate provided Select paths or include all (ignoring some that we can't query on like rootId) */ fun validateSelectPathsOrDefault(dbEntity: KClass>, paths: List?): List { diff --git a/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateQuery.kt b/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateQuery.kt deleted file mode 100644 index 767b2a59749..00000000000 --- a/misk-hibernate/src/main/kotlin/misk/hibernate/actions/HibernateQuery.kt +++ /dev/null @@ -1,11 +0,0 @@ -package misk.hibernate.actions - -import kotlin.reflect.KClass -import misk.hibernate.DbEntity -import misk.hibernate.Query - -/** - * [HibernateQuery] is a wrapper class that allows for unqualified binding of Misk.Hibernate.Query classes without - * collision. - */ -internal data class HibernateQuery(val query: KClass>>) diff --git a/misk-hibernate/src/test/kotlin/misk/hibernate/actions/HibernateDatabaseQuerySimpleNameCollisionTest.kt b/misk-hibernate/src/test/kotlin/misk/hibernate/actions/HibernateDatabaseQuerySimpleNameCollisionTest.kt new file mode 100644 index 00000000000..3742e60bfc5 --- /dev/null +++ b/misk-hibernate/src/test/kotlin/misk/hibernate/actions/HibernateDatabaseQuerySimpleNameCollisionTest.kt @@ -0,0 +1,138 @@ +package misk.hibernate.actions + +import jakarta.inject.Inject +import java.time.Instant +import java.time.LocalDate +import javax.persistence.Column +import javax.persistence.Entity +import javax.persistence.GeneratedValue +import javax.persistence.GenerationType +import javax.persistence.Table +import misk.audit.FakeAuditClientModule +import misk.hibernate.DbMovie +import misk.hibernate.DbRoot +import misk.hibernate.DbTimestampedEntity +import misk.hibernate.HibernateEntityModule +import misk.hibernate.Id +import misk.hibernate.Movies +import misk.hibernate.MoviesTestModule +import misk.hibernate.Transacter +import misk.hibernate.actions.HibernateDatabaseQueryTestingModule.Companion.DYNAMIC_MOVIE_QUERY_ACCESS_ENTRY +import misk.hibernate.annotation.Keyspace +import misk.inject.KAbstractModule +import misk.jdbc.DataSourceType +import misk.security.authz.AccessAnnotationEntry +import misk.testing.MiskTest +import misk.testing.MiskTestModule +import org.assertj.core.api.Assertions.assertThat +import org.junit.jupiter.api.BeforeEach +import org.junit.jupiter.api.Test + +/** + * Verifies the Database Query dynamic action queries the entity its authorized query was registered with, even when + * another entity in the same transacter shares its Kotlin [simpleName][kotlin.reflect.KClass.simpleName]. + * + * [SecretModels.DbMovie] (mapped to `actors`, registered only with `addHibernateEntity` and thus never exposed to + * Database Query) sorts ahead of the authorized [misk.hibernate.DbMovie] (mapped to `movies`, exposed via a dynamic + * query). A caller holding only the public movie-query capability must read `movies`, never `actors`. + */ +@MiskTest(startService = true) +class HibernateDatabaseQuerySimpleNameCollisionTest { + @MiskTestModule val module = CollisionTestingModule() + + @Inject + private lateinit var executer: + RealActionRequestExecuter + @Inject @Movies lateinit var transacter: Transacter + + @BeforeEach + fun before() { + executer.requestPath(HibernateDatabaseQueryDynamicAction.HIBERNATE_QUERY_DYNAMIC_WEBACTION_PATH) + + // PUBLIC_CONTROL lives in `movies` (authorized). SECRET_PROOF lives in `actors` (never exposed). + transacter.allowCowrites().transaction { session -> + session.save(DbMovie("PUBLIC_CONTROL")) + session.save(SecretModels.DbMovie("SECRET_PROOF")) + } + } + + @Test + fun `colliding simple name resolves to the authorized entity, not the first-registered one`() { + // Precondition: both entities share the simple name `DbMovie`, and the protected one sorts first. + val collidingOrder = transacter.entities().filter { it.simpleName == "DbMovie" }.map { it.qualifiedName } + assertThat(collidingOrder).hasSize(2) + assertThat(collidingOrder.first()).isEqualTo(SecretModels.DbMovie::class.qualifiedName) + + val response = + executer.executeRequest( + HibernateDatabaseQueryDynamicAction.Request( + entityClass = DbMovie::class.simpleName!!, + queryClass = "DbMovieDynamicQuery", + query = + HibernateDatabaseQueryMetadataFactory.Companion.DynamicQuery( + select = HibernateDatabaseQueryMetadataFactory.Companion.DynamicQuerySelect(paths = listOf("name")) + ), + ), + user = "low-user", + // Authorized only for the public movie query; no capability for the `actors`/SecretModels entity. + capabilities = DYNAMIC_MOVIE_QUERY_ACCESS_ENTRY.capabilities.joinToString() + ",admin_console", + ) + + val names = response.results.map { (it as Map<*, *>)["name"] } + assertThat(names).containsExactly("PUBLIC_CONTROL") + assertThat(names).doesNotContain("SECRET_PROOF") + } +} + +/** + * A second entity whose Kotlin simple name collides with [misk.hibernate.DbMovie] but which maps to a different table + * (`actors`). A distinct JPA `@Entity` name keeps Hibernate happy; the collision is on the Kotlin + * [simpleName][kotlin.reflect.KClass.simpleName]. + */ +object SecretModels { + @Entity(name = "SecretDbMovie") + @Table(name = "actors") + @Keyspace("movies_sharded") + class DbMovie() : DbRoot, DbTimestampedEntity { + @javax.persistence.Id @GeneratedValue(strategy = GenerationType.IDENTITY) override lateinit var id: Id + + @Column override lateinit var updated_at: Instant + + @Column override lateinit var created_at: Instant + + @Column(nullable = false) lateinit var name: String + + @Column var birth_date: LocalDate? = null + + constructor(name: String) : this() { + this.name = name + } + } +} + +class CollisionTestingModule : KAbstractModule() { + override fun configure() { + install(HibernateWebActionTestingModule()) + install( + MoviesTestModule( + type = DataSourceType.MYSQL, + entitiesModule = + object : HibernateEntityModule(Movies::class) { + override fun configureHibernate() { + installHibernateAdminDashboardWebActions() + + // Registered FIRST and WITHOUT a dynamic query, so it is never exposed to Database Query yet + // sorts ahead of the authorized entity in transacter.entities(). + addHibernateEntity(SecretModels.DbMovie::class) + // The authorized, Database-Query-exposed entity, mapped to `movies`. + addEntityWithDynamicQuery() + } + }, + ) + ) + + multibind().toInstance(DYNAMIC_MOVIE_QUERY_ACCESS_ENTRY) + + install(FakeAuditClientModule()) + } +} From 133f1a37819209bb917a0343b3de75973e9735e3 Mon Sep 17 00:00:00 2001 From: Mateusz Mrozewski Date: Mon, 17 Aug 2026 13:33:27 +0000 Subject: [PATCH 4/5] Graceful sqs shutdown (#3895) * Shutdown SQS subscribers gracefully * Update tests to remove long polling * Fix the AsyncSwitch tests * Cancel in-flight SQS receives on shutdown stop() only stopped issuing new receives, so a receive already parked in await() kept the poller alive until its long poll expired. doStop() joins each subscription in turn, so that delay was paid per queue. Subscriber now tracks the in-flight ReceiveMessage futures and exposes cancelInFlightReceives() to abort them. doStop() first gives the poller a grace period to wind down on its own, and only cancels if it overruns: a long poll holding messages returns immediately, and abandoning it would leave those messages invisible until their visibility timeout expired. This also lets the tests go back to long polling. Co-Authored-By: Jay Janssen <691356+jayjanssen@users.noreply.github.com> Co-Authored-By: Claude Opus 5 * API dump --------- Co-authored-by: Jay Janssen <691356+jayjanssen@users.noreply.github.com> Co-authored-by: Claude Opus 5 --- misk-aws2-sqs/api/misk-aws2-sqs.api | 17 ++++- .../misk/aws2/sqs/jobqueue/SqsJobConsumer.kt | 53 ++++++++++--- .../misk/aws2/sqs/jobqueue/Subscriber.kt | 53 +++++++++++-- .../aws2/sqs/jobqueue/config/SqsConfig.kt | 1 + .../sqs/jobqueue/config/SqsQueueConfig.kt | 16 +++- .../aws2/sqs/jobqueue/SqsJobConsumerTest.kt | 76 +++++++++++++++++++ .../misk/aws2/sqs/jobqueue/SubscriberTest.kt | 68 +++++++++++++++++ .../aws2/sqs/jobqueue/config/SqsConfigTest.kt | 15 ++++ 8 files changed, 281 insertions(+), 18 deletions(-) diff --git a/misk-aws2-sqs/api/misk-aws2-sqs.api b/misk-aws2-sqs/api/misk-aws2-sqs.api index 5162a1fc9f5..02ce2ab3454 100644 --- a/misk-aws2-sqs/api/misk-aws2-sqs.api +++ b/misk-aws2-sqs/api/misk-aws2-sqs.api @@ -54,6 +54,7 @@ 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 @@ -137,6 +138,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 (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 +154,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 +205,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 ()V public fun (I)V public fun (II)V @@ -212,8 +217,10 @@ public final class misk/aws2/sqs/jobqueue/config/SqsQueueConfig { public fun (IIIIZLjava/lang/Integer;Ljava/lang/Integer;)V 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 synthetic fun (IIIIZLjava/lang/Integer;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)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 final fun component1 ()I + public final fun component10 ()Ljava/lang/Long; public final fun component2 ()I public final fun component3 ()I public final fun component4 ()I @@ -222,8 +229,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;)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 fun equals (Ljava/lang/Object;)Z public final fun getAccount_id ()Ljava/lang/String; public final fun getChannel_capacity ()I @@ -232,12 +239,16 @@ 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 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 ()V public fun (Ljava/util/Map;)V 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 3a3311706be..88f0e0b1f9b 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 @@ -9,10 +9,14 @@ import java.time.Clock import java.util.concurrent.ConcurrentHashMap 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 @@ -20,6 +24,7 @@ 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. @@ -56,7 +61,7 @@ constructor( ) : JobConsumer, AbstractService(), TestFixture { private val scope = CoroutineScope(Dispatchers.IO.limitedParallelism(1) + SupervisorJob()) - private val handlingScopes = ConcurrentHashMap() + private val subscriptions = ConcurrentHashMap() override fun subscribe(queueName: QueueName, handler: JobHandler) { subscribe(queueName = queueName, handler = handler, queueConfig = SqsQueueConfig()) @@ -83,19 +88,19 @@ 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() { @@ -103,11 +108,41 @@ constructor( } 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. + subscription.handlingJobs.joinAll() + } + } + 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, + ) + + companion object { private val logger = getLogger() } diff --git a/misk-aws2-sqs/src/main/kotlin/misk/aws2/sqs/jobqueue/Subscriber.kt b/misk-aws2-sqs/src/main/kotlin/misk/aws2/sqs/jobqueue/Subscriber.kt index 036ec94f4e7..d0d09a9d779 100644 --- a/misk-aws2-sqs/src/main/kotlin/misk/aws2/sqs/jobqueue/Subscriber.kt +++ b/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>() + + /** 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() 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 ea6dba32749..edb7295d39d 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 @@ -37,6 +37,7 @@ 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, ) } 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 34d8116953f..3144afd0273 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 @@ -11,6 +11,9 @@ 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]. */ data class SqsQueueConfig @JvmOverloads @@ -24,4 +27,15 @@ constructor( val visibility_timeout: Int? = null, val region: String? = null, val account_id: String? = null, -) + val shutdown_grace_period_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 + } +} 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 fd6da375a71..437fc8cf19a 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 @@ -4,7 +4,9 @@ import jakarta.inject.Inject import java.util.concurrent.CountDownLatch import java.util.concurrent.TimeUnit.SECONDS import kotlin.random.Random +import kotlin.system.measureTimeMillis import kotlin.test.assertEquals +import kotlin.test.assertTrue import kotlinx.coroutines.delay import misk.aws2.sqs.jobqueue.config.SqsQueueConfig import misk.jobqueue.QueueName @@ -19,6 +21,7 @@ import misk.testing.MiskTestModule import org.junit.jupiter.api.Disabled import org.junit.jupiter.api.Test import software.amazon.awssdk.services.sqs.model.CreateQueueRequest +import software.amazon.awssdk.services.sqs.model.GetQueueAttributesRequest import software.amazon.awssdk.services.sqs.model.MessageAttributeValue import software.amazon.awssdk.services.sqs.model.QueueAttributeName import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest @@ -279,6 +282,79 @@ class SqsJobConsumerTest { assertEquals(0, latch.count, "Not all messages were consumed") } + @Test + fun `queues are drained drain`() { + // Run the test on a relatively large amount of jobs + val numberOfMessages = 1000 + val queueName = QueueName("test-queue-1") + val result = createQueue(queueName) + + val latch = CountDownLatch(numberOfMessages) + jobConsumer.subscribe( + queueName, + object : SuspendingJobHandler { + override suspend fun handleJob(job: Job): JobStatus { + logger.info { "Handling job: body=${job.body} queue=${job.queueName.value}" } + // Randomize the delay + delay(Random.nextLong(1, 50)) + latch.countDown() + return JobStatus.OK + } + }, + // Use non-standard settings to ensure some randomness of the test + SqsQueueConfig(parallelism = 2, concurrency = 5, channel_capacity = 7, region = "us-west-2"), + ) + + repeat(numberOfMessages) { sendMessage(result.queueUrl, "message $it") } + jobConsumer.stop() + + val (visible, invisible) = queueDepths(result.queueUrl) + + // 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(0, invisible) + } + + @Test + fun `shutdown grace period of zero cancels the in-flight receive`() { + val queueName = QueueName("test-queue-1") + // The queue is created with a 20s receive wait time, so an idle poller sits in a long poll. + createQueue(queueName) + + jobConsumer.subscribe( + queueName, + getHandler(CountDownLatch(1)), + SqsQueueConfig(region = "us-west-2", shutdown_grace_period_ms = 0), + ) + + // Give the poller time to issue its receive before stopping. + Thread.sleep(1_000) + val elapsed = measureTimeMillis { jobConsumer.stop() } + + assertTrue(elapsed < 5_000, "stop() took ${elapsed}ms; the in-flight receive was not canceled") + } + + private fun queueDepths(queueUrl: String): Pair { + val attributes = + DockerSqs.client + .getQueueAttributes( + GetQueueAttributesRequest.builder() + .queueUrl(queueUrl) + .attributeNames( + QueueAttributeName.APPROXIMATE_NUMBER_OF_MESSAGES, + QueueAttributeName.APPROXIMATE_NUMBER_OF_MESSAGES_NOT_VISIBLE, + ) + .build() + ) + .join() + .attributes() + return Pair( + attributes[QueueAttributeName.APPROXIMATE_NUMBER_OF_MESSAGES]!!.toInt(), + attributes[QueueAttributeName.APPROXIMATE_NUMBER_OF_MESSAGES_NOT_VISIBLE]!!.toInt(), + ) + } + private fun getDelayingHandler(latch: CountDownLatch): SuspendingJobHandler { return object : SuspendingJobHandler { override suspend fun handleJob(job: Job): JobStatus { diff --git a/misk-aws2-sqs/src/test/kotlin/misk/aws2/sqs/jobqueue/SubscriberTest.kt b/misk-aws2-sqs/src/test/kotlin/misk/aws2/sqs/jobqueue/SubscriberTest.kt index 120f116ac83..200726de12f 100644 --- a/misk-aws2-sqs/src/test/kotlin/misk/aws2/sqs/jobqueue/SubscriberTest.kt +++ b/misk-aws2-sqs/src/test/kotlin/misk/aws2/sqs/jobqueue/SubscriberTest.kt @@ -5,6 +5,7 @@ import io.prometheus.client.CollectorRegistry import java.time.Clock import java.util.concurrent.CompletableFuture import kotlin.test.assertEquals +import kotlin.test.assertFalse import kotlin.test.assertTrue import kotlinx.coroutines.ExperimentalCoroutinesApi import kotlinx.coroutines.cancelAndJoin @@ -145,6 +146,73 @@ class SubscriberTest { assertTrue(pollingJob.isCancelled) } + @Test + fun `stop lets an in-flight fetch finish`() = runTest { + whenever(sqsQueueResolver.getQueueUrl(queueName)).thenReturn(queueUrl) + val pendingResponse = CompletableFuture() + whenever(client.receiveMessage(any())).thenReturn(pendingResponse) + + val subscriber = subscriber(handler = { JobStatus.OK }) + val pollingJob = backgroundScope.launch { subscriber.poll() } + runCurrent() + + subscriber.stop() + runCurrent() + assertTrue(pollingJob.isActive) + + // The fetch that was already in flight still delivers its messages. + pendingResponse.complete(ReceiveMessageResponse.builder().messages(message("job-1")).build()) + runCurrent() + + assertEquals("job-1", channel.receive().id) + assertTrue(pollingJob.isCompleted) + assertTrue(channel.receiveCatching().isClosed) + assertEquals(0.0, sqsMetrics.sqsReceiveFailures.labels(queueName.value).get()) + } + + @Test + fun `canceling in-flight fetches ends polling and closes the channel`() = runTest { + whenever(sqsQueueResolver.getQueueUrl(queueName)).thenReturn(queueUrl) + val pendingResponse = CompletableFuture() + whenever(client.receiveMessage(any())).thenReturn(pendingResponse) + + val subscriber = subscriber(handler = { JobStatus.OK }) + val pollingJob = backgroundScope.launch { subscriber.poll() } + runCurrent() + assertTrue(pollingJob.isActive) + + subscriber.stop() + subscriber.cancelInFlightReceives() + runCurrent() + + assertTrue(pendingResponse.isCancelled) + assertTrue(pollingJob.isCompleted) + assertFalse(pollingJob.isCancelled) + assertTrue(channel.receiveCatching().isClosed) + // Canceling as part of a stop is expected, so it isn't counted as a receive failure. + assertEquals(0.0, sqsMetrics.sqsReceiveFailures.labels(queueName.value).get()) + } + + @Test + fun `a fetch started after cancellation is canceled immediately`() = runTest { + whenever(sqsQueueResolver.getQueueUrl(queueName)).thenReturn(queueUrl) + val subscriber = subscriber(handler = { JobStatus.OK }) + val pendingResponse = CompletableFuture() + // Stop the subscriber while it is issuing the request, so the set is swept before the future is registered. + whenever(client.receiveMessage(any())).thenAnswer { + subscriber.stop() + subscriber.cancelInFlightReceives() + pendingResponse + } + + val pollingJob = backgroundScope.launch { subscriber.poll() } + runCurrent() + + assertTrue(pendingResponse.isCancelled) + assertTrue(pollingJob.isCompleted) + assertFalse(pollingJob.isCancelled) + } + private fun subscriber( handler: JobHandler, queueConfig: SqsQueueConfig = SqsQueueConfig(install_retry_queue = false), 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 59b8c72ff80..4492bf0512d 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 @@ -84,4 +84,19 @@ class SqsConfigTest { assertEquals("us-west-2", queueConfig.region) // inherited from all_queues assertEquals(20, queueConfig.wait_timeout) // inherited from all_queues } + + @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), + ), + ) + + 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) + } } From f74e067e6064e3678f1e5ceb176667cf12f213a2 Mon Sep 17 00:00:00 2001 From: Jay Janssen Date: Mon, 17 Aug 2026 07:50:48 -0400 Subject: [PATCH 5/5] 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) + } }