fix: overhaul endpoint discovery (#1221)

Author: ianbotsf

.changes/49af01b8-6fed-4add-ace0-9f027e83425a.json

@@ -0,0 +1,9 @@
+{
+    "id": "49af01b8-6fed-4add-ace0-9f027e83425a",
+    "type": "feature",
+    "description": "⚠️ **IMPORTANT**: Refactor endpoint discoverer classes into interfaces so custom implementations may be provided",
+    "issues": [
+        "awslabs/aws-sdk-kotlin#1413"
+    ],
+    "requiresMinorVersionBump": true
+}

.changes/e6515649-dab5-4be9-b4b4-b289369960d5.json

@@ -0,0 +1,8 @@
+{
+    "id": "e6515649-dab5-4be9-b4b4-b289369960d5",
+    "type": "bugfix",
+    "description": "Favor `endpointUrl` instead of endpoint discovery if both are provided",
+    "issues": [
+        "awslabs/aws-sdk-kotlin#1413"
+    ]
+}

codegen/smithy-kotlin-codegen-testutils/src/main/kotlin/software/amazon/smithy/kotlin/codegen/test/ModelTestUtils.kt

@@ -8,6 +8,7 @@ import software.amazon.smithy.build.MockManifest
 import software.amazon.smithy.codegen.core.SymbolProvider
 import software.amazon.smithy.kotlin.codegen.*
 import software.amazon.smithy.kotlin.codegen.core.CodegenContext
+import software.amazon.smithy.kotlin.codegen.core.GenerationContext
 import software.amazon.smithy.kotlin.codegen.core.KotlinDelegator
 import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
 import software.amazon.smithy.kotlin.codegen.model.OperationNormalizer
@@ -122,9 +123,11 @@ fun Model.newTestContext(
     val manifest = MockManifest()
     val provider: SymbolProvider = KotlinCodegenPlugin.createSymbolProvider(model = this, rootNamespace = packageName, serviceName = serviceName, settings = settings)
     val service = this.getShape(ShapeId.from("$packageName#$serviceName")).get().asServiceShape().get()
-    val delegator = KotlinDelegator(settings, this, manifest, provider, integrations)
 
-    val ctx = ProtocolGenerator.GenerationContext(
+    val codegenCtx = GenerationContext(this, provider, settings, generator, integrations)
+    val delegator = KotlinDelegator(codegenCtx, manifest, integrations)
+
+    val generationCtx = ProtocolGenerator.GenerationContext(
         settings,
         this,
         service,
@@ -133,7 +136,8 @@ fun Model.newTestContext(
         generator.protocol,
         delegator,
     )
-    return TestContext(ctx, manifest, generator)
+
+    return TestContext(generationCtx, manifest, generator)
 }
 
 fun TestContext.toCodegenContext() = object : CodegenContext {

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/CodegenVisitor.kt

@@ -87,12 +87,12 @@ class CodegenVisitor(context: PluginContext) : ShapeVisitor.Default<Unit>() {
             integration.decorateSymbolProvider(settings, model, provider)
         }
 
-        writers = KotlinDelegator(settings, model, fileManifest, symbolProvider, integrations)
-
         protocolGenerator = resolveProtocolGenerator(integrations, model, service, settings)
-        applicationProtocol = protocolGenerator?.applicationProtocol ?: ApplicationProtocol.createDefaultHttpApplicationProtocol()
-
         baseGenerationContext = GenerationContext(model, symbolProvider, settings, protocolGenerator, integrations)
+
+        writers = KotlinDelegator(baseGenerationContext, fileManifest, integrations)
+
+        applicationProtocol = protocolGenerator?.applicationProtocol ?: ApplicationProtocol.createDefaultHttpApplicationProtocol()
     }
 
     private fun resolveProtocolGenerator(

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/core/AbstractCodeWriterExt.kt

@@ -150,6 +150,15 @@ inline fun <W : AbstractCodeWriter<W>, reified V> AbstractCodeWriter<W>.getConte
  */
 inline fun <W : AbstractCodeWriter<W>, reified V> AbstractCodeWriter<W>.getContextValue(key: SectionKey<V>): V = getContextValue(key.name)
 
+/**
+ * Convenience function to set a typed value in the context
+ * @param key
+ */
+inline fun <W : AbstractCodeWriter<W>, reified V> AbstractCodeWriter<W>.putContextValue(
+    key: SectionKey<V>,
+    value: V,
+): W = putContext(key.name, value)
+
 /**
  * Convenience function to set context only if there is no value already associated with the given [key]
  */

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/core/CodegenContext.kt

@@ -8,6 +8,7 @@ package software.amazon.smithy.kotlin.codegen.core
 import software.amazon.smithy.codegen.core.SymbolProvider
 import software.amazon.smithy.kotlin.codegen.KotlinSettings
 import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+import software.amazon.smithy.kotlin.codegen.integration.SectionKey
 import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolGenerator
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.Shape
@@ -16,6 +17,10 @@ import software.amazon.smithy.model.shapes.Shape
  * Common codegen properties required across different codegen contexts
  */
 interface CodegenContext {
+    companion object {
+        val Key = SectionKey<CodegenContext>("CodegenContext")
+    }
+
     val model: Model
     val symbolProvider: SymbolProvider
     val settings: KotlinSettings

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/core/KotlinDelegator.kt

@@ -6,11 +6,9 @@ package software.amazon.smithy.kotlin.codegen.core
 
 import software.amazon.smithy.build.FileManifest
 import software.amazon.smithy.codegen.core.*
-import software.amazon.smithy.kotlin.codegen.KotlinSettings
 import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
 import software.amazon.smithy.kotlin.codegen.model.SymbolProperty
 import software.amazon.smithy.kotlin.codegen.utils.namespaceToPath
-import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.Shape
 import java.nio.file.Paths
 
@@ -21,13 +19,10 @@ const val DEFAULT_TEST_SOURCE_SET_ROOT = "./src/test/kotlin/"
  * Manages writers for Kotlin files.
  */
 class KotlinDelegator(
-    private val settings: KotlinSettings,
-    private val model: Model,
+    private val ctx: CodegenContext,
     val fileManifest: FileManifest,
-    private val symbolProvider: SymbolProvider,
     private val integrations: List<KotlinIntegration> = listOf(),
 ) {
-
     private val writers: MutableMap<String, KotlinWriter> = mutableMapOf()
 
     // Tracks dependencies for source not provided by codegen that may reside in the service source tree.
@@ -91,7 +86,7 @@ class KotlinDelegator(
         shape: Shape,
         block: (KotlinWriter) -> Unit,
     ) {
-        val symbol = symbolProvider.toSymbol(shape)
+        val symbol = ctx.symbolProvider.toSymbol(shape)
         useSymbolWriter(symbol, block)
     }
 
@@ -151,7 +146,9 @@ class KotlinDelegator(
         val needsNewline = writers.containsKey(formattedFilename)
         val writer = writers.getOrPut(formattedFilename) {
             val kotlinWriter = KotlinWriter(namespace)
-            if (settings.debug) kotlinWriter.enableStackTraceComments(true)
+            kotlinWriter.putContextValue(CodegenContext.Key, ctx)
+
+            if (ctx.settings.debug) kotlinWriter.enableStackTraceComments(true)
 
             // Register all integrations [SectionWriterBindings] on the writer.
             integrations.forEach { integration ->

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/core/RuntimeTypes.kt

@@ -45,11 +45,11 @@ object RuntimeTypes {
     object HttpClient : RuntimeTypePackage(KotlinDependency.HTTP_CLIENT) {
         val SdkHttpClient = symbol("SdkHttpClient")
 
-        object Middleware : RuntimeTypePackage(KotlinDependency.HTTP, "middleware") {
+        object Middleware : RuntimeTypePackage(KotlinDependency.HTTP_CLIENT, "middleware") {
             val MutateHeadersMiddleware = symbol("MutateHeaders")
         }
 
-        object Operation : RuntimeTypePackage(KotlinDependency.HTTP, "operation") {
+        object Operation : RuntimeTypePackage(KotlinDependency.HTTP_CLIENT, "operation") {
             val AuthSchemeResolver = symbol("AuthSchemeResolver")
             val context = symbol("context")
             val EndpointResolver = symbol("EndpointResolver")
@@ -68,18 +68,19 @@ object RuntimeTypes {
             val setResolvedEndpoint = symbol("setResolvedEndpoint")
         }
 
-        object Config : RuntimeTypePackage(KotlinDependency.HTTP, "config") {
+        object Config : RuntimeTypePackage(KotlinDependency.HTTP_CLIENT, "config") {
             val HttpClientConfig = symbol("HttpClientConfig")
             val HttpEngineConfig = symbol("HttpEngineConfig")
         }
 
-        object Engine : RuntimeTypePackage(KotlinDependency.HTTP, "engine") {
+        object Engine : RuntimeTypePackage(KotlinDependency.HTTP_CLIENT, "engine") {
             val HttpClientEngine = symbol("HttpClientEngine")
             val manage = symbol("manage", "engine.internal", isExtension = true)
         }
 
-        object Interceptors : RuntimeTypePackage(KotlinDependency.HTTP, "interceptors") {
+        object Interceptors : RuntimeTypePackage(KotlinDependency.HTTP_CLIENT, "interceptors") {
             val ContinueInterceptor = symbol("ContinueInterceptor")
+            val DiscoveredEndpointErrorInterceptor = symbol("DiscoveredEndpointErrorInterceptor")
             val HttpInterceptor = symbol("HttpInterceptor")
             val HttpChecksumRequiredInterceptor = symbol("HttpChecksumRequiredInterceptor")
             val FlexibleChecksumsRequestInterceptor = symbol("FlexibleChecksumsRequestInterceptor")
@@ -97,7 +98,6 @@ object RuntimeTypes {
     }
 
     object Core : RuntimeTypePackage(KotlinDependency.CORE) {
-        val Clock = symbol("Clock", "time")
         val ExecutionContext = symbol("ExecutionContext", "operation")
         val ErrorMetadata = symbol("ErrorMetadata")
         val ServiceErrorMetadata = symbol("ServiceErrorMetadata")
@@ -125,11 +125,12 @@ object RuntimeTypes {
             val attributesOf = symbol("attributesOf")
             val AttributeKey = symbol("AttributeKey")
             val createOrAppend = symbol("createOrAppend")
+            val ExpiringKeyedCache = symbol("ExpiringKeyedCache")
             val get = symbol("get")
             val mutableMultiMapOf = symbol("mutableMultiMapOf")
+            val PeriodicSweepCache = symbol("PeriodicSweepCache")
             val putIfAbsent = symbol("putIfAbsent")
             val putIfAbsentNotNull = symbol("putIfAbsentNotNull")
-            val ReadThroughCache = symbol("ReadThroughCache")
             val toMutableAttributes = symbol("toMutableAttributes")
             val emptyAttributes = symbol("emptyAttributes")
         }

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/ExceptionBaseClassGenerator.kt

@@ -10,7 +10,6 @@ import software.amazon.smithy.codegen.core.Symbol
 import software.amazon.smithy.kotlin.codegen.KotlinSettings
 import software.amazon.smithy.kotlin.codegen.core.*
 import software.amazon.smithy.kotlin.codegen.integration.SectionId
-import software.amazon.smithy.kotlin.codegen.integration.SectionKey
 import software.amazon.smithy.kotlin.codegen.model.buildSymbol
 import software.amazon.smithy.kotlin.codegen.model.namespace
 import software.amazon.smithy.model.knowledge.TopDownIndex
@@ -31,12 +30,10 @@ object ExceptionBaseClassGenerator {
     /**
      * Defines a section in which code can be added to the body of the base exception type.
      */
-    object ExceptionBaseClassSection : SectionId {
-        val CodegenContext: SectionKey<CodegenContext> = SectionKey("CodegenContext")
-    }
+    object ExceptionBaseClassSection : SectionId
 
     fun render(ctx: CodegenContext, writer: KotlinWriter) {
-        writer.declareSection(ExceptionBaseClassSection, mapOf(ExceptionBaseClassSection.CodegenContext to ctx)) {
+        writer.declareSection(ExceptionBaseClassSection) {
             ServiceExceptionBaseClassGenerator().render(ctx, writer)
         }
     }

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/endpoints/discovery/DefaultEndpointDiscovererGenerator.kt

@@ -0,0 +1,103 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.kotlin.codegen.rendering.endpoints.discovery
+
+import software.amazon.smithy.aws.traits.clientendpointdiscovery.ClientEndpointDiscoveryTrait
+import software.amazon.smithy.codegen.core.Symbol
+import software.amazon.smithy.kotlin.codegen.KotlinSettings
+import software.amazon.smithy.kotlin.codegen.core.*
+import software.amazon.smithy.kotlin.codegen.lang.KotlinTypes
+import software.amazon.smithy.kotlin.codegen.model.buildSymbol
+import software.amazon.smithy.kotlin.codegen.model.expectShape
+import software.amazon.smithy.kotlin.codegen.model.expectTrait
+import software.amazon.smithy.kotlin.codegen.rendering.endpoints.SdkEndpointBuiltinIntegration
+import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.shapes.ServiceShape
+
+class DefaultEndpointDiscovererGenerator(private val ctx: CodegenContext, private val delegator: KotlinDelegator) {
+    private val symbol = symbolFor(ctx.settings)
+    private val service = ctx.model.expectShape<ServiceShape>(ctx.settings.service)
+    private val clientSymbol = ctx.symbolProvider.toSymbol(service)
+    private val operationName = run {
+        val epDiscoveryTrait = service.expectTrait<ClientEndpointDiscoveryTrait>()
+        val operation = ctx.model.expectShape<OperationShape>(epDiscoveryTrait.operation)
+        operation.defaultName()
+    }
+
+    companion object {
+        fun symbolFor(settings: KotlinSettings): Symbol = buildSymbol {
+            val clientName = clientName(settings.sdkId)
+            name = "Default${clientName}EndpointDiscoverer"
+            namespace = "${settings.pkg.name}.endpoints"
+        }
+    }
+
+    fun render() {
+        delegator.applyFileWriter(symbol) {
+            val service = clientName(ctx.settings.sdkId)
+            dokka(
+                """
+                    A class which looks up specific endpoints for $service calls via the `$operationName` API. These
+                    unique endpoints are cached as appropriate to avoid unnecessary latency in subsequent calls.
+                    @param cache An [ExpiringKeyedCache] implementation used to cache discovered hosts
+                """.trimIndent(),
+            )
+
+            withBlock(
+                "#1L class #2T(#1L val cache: #3T<DiscoveryParams, #4T> = #5T(10.#6T)) : #7T {",
+                "}",
+                ctx.settings.api.visibility,
+                symbol,
+                RuntimeTypes.Core.Collections.ExpiringKeyedCache,
+                RuntimeTypes.Core.Net.Host,
+                RuntimeTypes.Core.Collections.PeriodicSweepCache,
+                KotlinTypes.Time.minutes,
+                EndpointDiscovererInterfaceGenerator.symbolFor(ctx.settings),
+            ) {
+                renderAsEndpointResolver()
+                write("")
+                renderInvalidate()
+            }
+        }
+    }
+
+    private fun KotlinWriter.renderAsEndpointResolver() {
+        withBlock(
+            "override fun asEndpointResolver(client: #1T, delegate: #2T): #2T = #2T { request ->",
+            "}",
+            clientSymbol,
+            RuntimeTypes.HttpClient.Operation.EndpointResolver,
+        ) {
+            withBlock("if (client.config.#L == null) {", "}", SdkEndpointBuiltinIntegration.EndpointUrlProp.propertyName) {
+                write("val identity = request.identity")
+                write(
+                    """require(identity is #T) { "Endpoint discovery requires AWS credentials" }""",
+                    RuntimeTypes.Auth.Credentials.AwsCredentials.Credentials,
+                )
+                write("")
+                write("val cacheKey = DiscoveryParams(client.config.region, identity.accessKeyId)")
+                write("request.context[DiscoveryParamsKey] = cacheKey")
+                write("val discoveredHost = cache.get(cacheKey) { discoverHost(client) }")
+                write("")
+                write("val originalEndpoint = delegate.resolve(request)")
+                withBlock("#T(", ")", RuntimeTypes.SmithyClient.Endpoints.Endpoint) {
+                    write("originalEndpoint.uri.copy { host = discoveredHost },")
+                    write("originalEndpoint.headers,")
+                    write("originalEndpoint.attributes,")
+                }
+
+                // If user manually specifies endpointUrl, skip endpoint discovery
+                closeAndOpenBlock("} else {")
+                write("delegate.resolve(request)")
+            }
+        }
+    }
+
+    private fun KotlinWriter.renderInvalidate() {
+        withBlock("override public suspend fun invalidate(context: #T) {", "}", RuntimeTypes.Core.ExecutionContext) {
+            write("context.getOrNull(DiscoveryParamsKey)?.let { cache.invalidate(it) }")
+        }
+    }
+}