fix: propagate TCCL with aws-crt http based client (#1124)

quarkiverse · Feb 14, 2024 · b2534cf · b2534cf
1 parent 76f4a36
commit b2534cf
Show file tree

Hide file tree

Showing 37 changed files with 1,599 additions and 1,504 deletions.
diff --git a/...rc/main/java/io/quarkus/amazon/cognitouserpools/deployment/CognitoUserPoolsProcessor.java b/...rc/main/java/io/quarkus/amazon/cognitouserpools/deployment/CognitoUserPoolsProcessor.java
@@ -34,6 +34,7 @@
 import io.quarkus.deployment.builditem.ExecutorBuildItem;
 import io.quarkus.deployment.builditem.ExtensionSslNativeSupportBuildItem;
 import io.quarkus.deployment.builditem.FeatureBuildItem;
+import io.quarkus.deployment.builditem.LaunchModeBuildItem;
 import software.amazon.awssdk.services.cognitoidentityprovider.CognitoIdentityProviderAsyncClient;
 import software.amazon.awssdk.services.cognitoidentityprovider.CognitoIdentityProviderAsyncClientBuilder;
 import software.amazon.awssdk.services.cognitoidentityprovider.CognitoIdentityProviderClient;
@@ -179,6 +180,7 @@ void createClientBuilders(CognitoUserPoolsRecorder recorder,
             BuildProducer<SyntheticBeanBuildItem> syntheticBeans,
             BuildProducer<AmazonClientSyncResultBuildItem> clientSync,
             BuildProducer<AmazonClientAsyncResultBuildItem> clientAsync,
+            LaunchModeBuildItem launchModeBuildItem,
             ExecutorBuildItem executorBuildItem) {
 
         createClientBuilders(capabilities,
@@ -194,6 +196,7 @@ void createClientBuilders(CognitoUserPoolsRecorder recorder,
                 syntheticBeans,
                 clientSync,
                 clientAsync,
+                launchModeBuildItem,
                 executorBuildItem);
     }
 }
diff --git a/...ent/src/main/java/io/quarkus/amazon/common/deployment/AbstractAmazonServiceProcessor.java b/...ent/src/main/java/io/quarkus/amazon/common/deployment/AbstractAmazonServiceProcessor.java
@@ -40,6 +40,7 @@
 import io.quarkus.deployment.builditem.ExecutorBuildItem;
 import io.quarkus.deployment.builditem.ExtensionSslNativeSupportBuildItem;
 import io.quarkus.deployment.builditem.FeatureBuildItem;
+import io.quarkus.deployment.builditem.LaunchModeBuildItem;
 import io.quarkus.runtime.RuntimeValue;
 import software.amazon.awssdk.awscore.client.builder.AwsClientBuilder;
 import software.amazon.awssdk.awscore.presigner.SdkPresigner;
@@ -261,6 +262,7 @@ protected void createClientBuilders(
             BuildProducer<SyntheticBeanBuildItem> syntheticBeans,
             BuildProducer<AmazonClientSyncResultBuildItem> clientSync,
             BuildProducer<AmazonClientAsyncResultBuildItem> clientAsync,
+            LaunchModeBuildItem launchModeBuildItem,
             ExecutorBuildItem executorBuildItem) {
 
         RuntimeValue<SdkPresigner.Builder> presignerBuilder = null;
@@ -279,12 +281,14 @@ protected void createClientBuilders(
                 syncClientBuilderClass,
                 (syncTransport) -> recorder.createSyncBuilder(syncTransport),
                 asyncClientBuilderClass,
-                (asyncTransport) -> recorder.createAsyncBuilder(asyncTransport, executorBuildItem.getExecutorProxy()),
+                (asyncTransport) -> recorder.createAsyncBuilder(asyncTransport, launchModeBuildItem.getLaunchMode(),
+                        executorBuildItem.getExecutorProxy()),
                 presignerBuilderClass,
                 presignerBuilder,
                 syntheticBeans,
                 clientSync,
                 clientAsync,
+                launchModeBuildItem,
                 executorBuildItem);
     }
 
@@ -306,6 +310,7 @@ private void createClientBuilders(
             BuildProducer<SyntheticBeanBuildItem> syntheticBeans,
             BuildProducer<AmazonClientSyncResultBuildItem> clientSync,
             BuildProducer<AmazonClientAsyncResultBuildItem> clientAsync,
+            LaunchModeBuildItem launchModeBuildItem,
             ExecutorBuildItem executorBuildItem) {
         String configName = configName();
 

diff --git a/common/runtime/src/main/java/io/quarkus/amazon/common/runtime/AmazonClientRecorder.java b/common/runtime/src/main/java/io/quarkus/amazon/common/runtime/AmazonClientRecorder.java
@@ -2,6 +2,7 @@
 
 import java.util.concurrent.Executor;
 
+import io.quarkus.runtime.LaunchMode;
 import io.quarkus.runtime.RuntimeValue;
 import software.amazon.awssdk.awscore.client.builder.AwsAsyncClientBuilder;
 import software.amazon.awssdk.awscore.client.builder.AwsClientBuilder;
@@ -43,24 +44,61 @@ public RuntimeValue<AwsClientBuilder> createSyncBuilder(RuntimeValue<Builder> tr
     }
 
     public RuntimeValue<AwsClientBuilder> createAsyncBuilder(RuntimeValue<SdkAsyncHttpClient.Builder> transport,
+            LaunchMode launchMode,
             Executor executor) {
         AwsAsyncClientBuilder<?, ?> builder = getAsyncClientBuilder();
         AsyncHttpClientConfig config = getAsyncClientConfig();
 
         if (transport != null) {
             builder.httpClientBuilder(transport.getValue());
         }
+
+        Executor configExecutor;
+
         if (!config.advanced().useFutureCompletionThreadPool()) {
-            builder.asyncConfiguration(asyncConfigBuilder -> asyncConfigBuilder
-                    .advancedOption(SdkAdvancedAsyncClientOption.FUTURE_COMPLETION_EXECUTOR, Runnable::run));
+            configExecutor = Runnable::run;
         } else {
-            builder.asyncConfiguration(asyncConfigBuilder -> asyncConfigBuilder
-                    .advancedOption(SdkAdvancedAsyncClientOption.FUTURE_COMPLETION_EXECUTOR, executor));
+            configExecutor = executor;
+        }
+
+        if (launchMode != LaunchMode.NORMAL) {
+            configExecutor = new ClassLoaderExecutorWrapper(executor);
         }
+
+        final Executor futureCompletionExecutor = configExecutor;
+
+        builder.asyncConfiguration(asyncConfigBuilder -> asyncConfigBuilder
+                .advancedOption(SdkAdvancedAsyncClientOption.FUTURE_COMPLETION_EXECUTOR, futureCompletionExecutor));
+
         return new RuntimeValue<>((AwsClientBuilder) builder);
     }
 
     public RuntimeValue<SdkPresigner.Builder> createPresignerBuilder() {
         throw new UnsupportedOperationException();
     }
+
+    /**
+     * Capture the current ClassLoader and restore it to support dev and test mode
+     */
+    private static final class ClassLoaderExecutorWrapper implements Executor {
+
+        private Executor executor;
+        private ClassLoader contextClassLoader;
+
+        private ClassLoaderExecutorWrapper(Executor executor) {
+            this.executor = executor;
+            this.contextClassLoader = Thread.currentThread().getContextClassLoader();
+        }
+
+        @Override
+        public void execute(Runnable command) {
+            final ClassLoader old = Thread.currentThread().getContextClassLoader();
+            try {
+                Thread.currentThread().setContextClassLoader(contextClassLoader);
+                executor.execute(command);
+            } finally {
+                Thread.currentThread().setContextClassLoader(old);
+            }
+        }
+    }
 }