From ce160df29fd2da9c37e95bf0fad1f1f4865977cf Mon Sep 17 00:00:00 2001
From: Daniel Kec <daniel.kec@oracle.com>
Date: Thu, 3 Feb 2022 15:17:30 +0100
Subject: [PATCH] Jersey JAX-RS client context propagation #3762

Signed-off-by: Daniel Kec <daniel.kec@oracle.com>
---
 jersey/client/pom.xml                         |  19 +-
 .../jersey/client/ClientBuilderListener.java  |  49 +++++
 .../jersey/client/ExecutorProvider.java       |  85 +++++++++
 .../client/ScheduledExecutorProvider.java     |  49 +++++
 .../helidon/jersey/client/package-info.java   |  21 +++
 jersey/client/src/main/java/module-info.java  |   9 +-
 .../io/helidon/jersey/client/ContextTest.java | 172 ++++++++++++++++++
 7 files changed, 402 insertions(+), 2 deletions(-)
 create mode 100644 jersey/client/src/main/java/io/helidon/jersey/client/ClientBuilderListener.java
 create mode 100644 jersey/client/src/main/java/io/helidon/jersey/client/ExecutorProvider.java
 create mode 100644 jersey/client/src/main/java/io/helidon/jersey/client/ScheduledExecutorProvider.java
 create mode 100644 jersey/client/src/main/java/io/helidon/jersey/client/package-info.java
 create mode 100644 jersey/client/src/test/java/io/helidon/jersey/client/ContextTest.java

diff --git a/jersey/client/pom.xml b/jersey/client/pom.xml
index 3785c0d2b92..e6d3546c7a9 100644
--- a/jersey/client/pom.xml
+++ b/jersey/client/pom.xml
@@ -1,6 +1,6 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <!--
-    Copyright (c) 2019, 2021 Oracle and/or its affiliates.
+    Copyright (c) 2019, 2022 Oracle and/or its affiliates.
 
     Licensed under the Apache License, Version 2.0 (the "License");
     you may not use this file except in compliance with the License.
@@ -50,11 +50,28 @@
             <groupId>jakarta.inject</groupId>
             <artifactId>jakarta.inject-api</artifactId>
         </dependency>
+        <dependency>
+            <groupId>jakarta.annotation</groupId>
+            <artifactId>jakarta.annotation-api</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>jakarta.enterprise</groupId>
+            <artifactId>jakarta.enterprise.cdi-api</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>io.helidon.common</groupId>
+            <artifactId>helidon-common-context</artifactId>
+        </dependency>
         <dependency>
             <groupId>org.junit.jupiter</groupId>
             <artifactId>junit-jupiter-params</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>io.helidon.webserver</groupId>
+            <artifactId>helidon-webserver</artifactId>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>org.hamcrest</groupId>
             <artifactId>hamcrest-all</artifactId>
diff --git a/jersey/client/src/main/java/io/helidon/jersey/client/ClientBuilderListener.java b/jersey/client/src/main/java/io/helidon/jersey/client/ClientBuilderListener.java
new file mode 100644
index 00000000000..1dd6b602914
--- /dev/null
+++ b/jersey/client/src/main/java/io/helidon/jersey/client/ClientBuilderListener.java
@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) 2022 Oracle and/or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package io.helidon.jersey.client;
+
+import jakarta.annotation.Priority;
+import jakarta.ws.rs.Priorities;
+import jakarta.ws.rs.client.ClientBuilder;
+
+/**
+ * Listen to the Jersey's client builder being invoked and register Helidon executor provider.
+ */
+// Lowest priority gets invoked first and first registered provider wins,
+// user's custom ClientBuilderListeners needs to be able to override this
+@Priority(Priorities.USER + 100)
+public class ClientBuilderListener implements org.glassfish.jersey.client.spi.ClientBuilderListener {
+
+    private static final String USE_CONTEXT_AWARE_EXEC_PROVIDER = "io.helidon.jersey.client.useContextAwareExecutorProvider";
+    private static Boolean useProvider;
+
+    @Override
+    public void onNewBuilder(ClientBuilder builder) {
+        if (useContextAwareProvider()) {
+            builder.register(ExecutorProvider.class);
+            builder.register(ScheduledExecutorProvider.class);
+        }
+    }
+
+    private boolean useContextAwareProvider() {
+        if (useProvider == null) {
+            useProvider = Boolean.parseBoolean(System.getProperty(USE_CONTEXT_AWARE_EXEC_PROVIDER, "true"));
+        }
+        return useProvider;
+    }
+}
diff --git a/jersey/client/src/main/java/io/helidon/jersey/client/ExecutorProvider.java b/jersey/client/src/main/java/io/helidon/jersey/client/ExecutorProvider.java
new file mode 100644
index 00000000000..80ed7cd630a
--- /dev/null
+++ b/jersey/client/src/main/java/io/helidon/jersey/client/ExecutorProvider.java
@@ -0,0 +1,85 @@
+/*
+ * Copyright (c) 2022 Oracle and/or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package io.helidon.jersey.client;
+
+import java.util.Optional;
+import java.util.concurrent.ExecutorService;
+import java.util.logging.Logger;
+
+import io.helidon.common.context.Contexts;
+
+import jakarta.inject.Inject;
+import jakarta.inject.Named;
+import org.glassfish.jersey.client.ClientAsyncExecutor;
+import org.glassfish.jersey.client.internal.LocalizationMessages;
+import org.glassfish.jersey.internal.util.collection.LazyValue;
+import org.glassfish.jersey.internal.util.collection.Value;
+import org.glassfish.jersey.internal.util.collection.Values;
+import org.glassfish.jersey.spi.ThreadPoolExecutorProvider;
+
+/**
+ * Wraps default executor to enable Helidon context propagation for Jersey async calls.
+ */
+@ClientAsyncExecutor
+public class ExecutorProvider extends ThreadPoolExecutorProvider {
+
+    static final String THREAD_NAME_PREFIX = "helidon-client-async-executor";
+    private static final Logger LOGGER = Logger.getLogger(ExecutorProvider.class.getName());
+    private final LazyValue<Integer> asyncThreadPoolSize;
+
+    /**
+     * Create new instance of the context aware executor provider.
+     * @param poolSize Maximum pool size
+     */
+    @Inject
+    public ExecutorProvider(@Named("ClientAsyncThreadPoolSize") final Optional<Integer> poolSize) {
+        super(THREAD_NAME_PREFIX);
+
+        this.asyncThreadPoolSize = Values.lazy((Value<Integer>) () -> {
+            if (poolSize.filter(i -> i > 0).isEmpty()) {
+                LOGGER.config(LocalizationMessages.IGNORED_ASYNC_THREADPOOL_SIZE(poolSize.orElse(0)));
+                // using default
+                return Integer.MAX_VALUE;
+            } else {
+                LOGGER.config(LocalizationMessages.USING_FIXED_ASYNC_THREADPOOL(poolSize.get()));
+                return poolSize.get();
+            }
+        });
+    }
+
+    @Override
+    protected int getMaximumPoolSize() {
+        return asyncThreadPoolSize.get();
+    }
+
+    @Override
+    protected int getCorePoolSize() {
+        // Mimicking the Executors.newCachedThreadPool and newFixedThreadPool configuration values.
+        final int maximumPoolSize = getMaximumPoolSize();
+        if (maximumPoolSize != Integer.MAX_VALUE) {
+            return maximumPoolSize;
+        } else {
+            return 0;
+        }
+    }
+
+    @Override
+    public ExecutorService getExecutorService() {
+        return Contexts.wrap(super.getExecutorService());
+    }
+}
diff --git a/jersey/client/src/main/java/io/helidon/jersey/client/ScheduledExecutorProvider.java b/jersey/client/src/main/java/io/helidon/jersey/client/ScheduledExecutorProvider.java
new file mode 100644
index 00000000000..b2853d2eccb
--- /dev/null
+++ b/jersey/client/src/main/java/io/helidon/jersey/client/ScheduledExecutorProvider.java
@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) 2022 Oracle and/or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package io.helidon.jersey.client;
+
+import java.util.concurrent.ScheduledExecutorService;
+
+import io.helidon.common.context.Contexts;
+
+import org.glassfish.jersey.client.ClientBackgroundScheduler;
+import org.glassfish.jersey.spi.ScheduledThreadPoolExecutorProvider;
+
+/**
+ * Wraps default executor to enable Helidon context propagation for Jersey async calls.
+ */
+@ClientBackgroundScheduler
+public class ScheduledExecutorProvider extends ScheduledThreadPoolExecutorProvider {
+
+    /**
+     * Creates a new instance.
+     */
+    public ScheduledExecutorProvider() {
+        super("helidon-client-background-scheduler");
+    }
+
+    @Override
+    protected int getCorePoolSize() {
+        return 1;
+    }
+
+    @Override
+    public ScheduledExecutorService getExecutorService() {
+        return Contexts.wrap(super.getExecutorService());
+    }
+}
diff --git a/jersey/client/src/main/java/io/helidon/jersey/client/package-info.java b/jersey/client/src/main/java/io/helidon/jersey/client/package-info.java
new file mode 100644
index 00000000000..7984c695a3a
--- /dev/null
+++ b/jersey/client/src/main/java/io/helidon/jersey/client/package-info.java
@@ -0,0 +1,21 @@
+/*
+ * Copyright (c) 2022 Oracle and/or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+/**
+ * Jersey JAX-RS client with context propagation support.
+ */
+package io.helidon.jersey.client;
diff --git a/jersey/client/src/main/java/module-info.java b/jersey/client/src/main/java/module-info.java
index cb696d9667e..99174ae3e15 100644
--- a/jersey/client/src/main/java/module-info.java
+++ b/jersey/client/src/main/java/module-info.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018, 2021 Oracle and/or its affiliates.
+ * Copyright (c) 2018, 2022 Oracle and/or its affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -21,4 +21,11 @@
     requires transitive jakarta.ws.rs;
     requires transitive jersey.common;
     requires transitive jersey.client;
+    requires java.logging;
+    requires jakarta.inject;
+    requires jakarta.annotation;
+    requires io.helidon.common.context;
+
+    provides org.glassfish.jersey.client.spi.ClientBuilderListener
+            with io.helidon.jersey.client.ClientBuilderListener;
 }
diff --git a/jersey/client/src/test/java/io/helidon/jersey/client/ContextTest.java b/jersey/client/src/test/java/io/helidon/jersey/client/ContextTest.java
new file mode 100644
index 00000000000..95946050cc6
--- /dev/null
+++ b/jersey/client/src/test/java/io/helidon/jersey/client/ContextTest.java
@@ -0,0 +1,172 @@
+/*
+ * Copyright (c) 2022 Oracle and/or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package io.helidon.jersey.client;
+
+import java.time.Duration;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+import java.util.function.Function;
+
+import io.helidon.common.context.Context;
+import io.helidon.common.context.Contexts;
+import io.helidon.common.reactive.Single;
+import io.helidon.webserver.Routing;
+import io.helidon.webserver.WebServer;
+
+import org.glassfish.jersey.client.ClientAsyncExecutor;
+import org.glassfish.jersey.spi.ThreadPoolExecutorProvider;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Test;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.core.StringStartsWith.startsWith;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import jakarta.ws.rs.client.ClientBuilder;
+import jakarta.ws.rs.client.Entity;
+import jakarta.ws.rs.client.InvocationCallback;
+
+public class ContextTest {
+
+    private static final Duration TIME_OUT = Duration.ofSeconds(5);
+    private static final String TEST_KEY = "test-key";
+    private static final String TEST_VAL = "test-val";
+    private static final String TEST_THREAD_NAME_PREFIX = "test-exec";
+    private static WebServer server;
+
+    @BeforeAll
+    static void beforeAll() {
+        server = WebServer.builder()
+                .host("localhost")
+                .routing(Routing.builder()
+                        .put((req, res) -> res.send("I'm Frank!"))
+                        .build())
+                .build()
+                .start()
+                .await(TIME_OUT);
+    }
+
+    @AfterAll
+    static void afterAll() {
+        if (server != null) {
+            server.shutdown().await(TIME_OUT);
+        }
+    }
+
+    @Test
+    void defaultThreadPool() {
+        test(Function.identity(), ExecutorProvider.THREAD_NAME_PREFIX, TEST_VAL);
+    }
+
+    @Test
+    void overrideDefaultThreadPoolWithBuilder() throws InterruptedException {
+        ExecutorService executorService = Executors.newSingleThreadExecutor(r -> new Thread(r, TEST_THREAD_NAME_PREFIX));
+        try {
+            test(cb -> cb.executorService(executorService), TEST_THREAD_NAME_PREFIX, null);
+        } finally {
+            executorService.shutdown();
+            assertTrue(executorService.awaitTermination(TIME_OUT.toSeconds(), TimeUnit.SECONDS));
+        }
+    }
+
+    @Test
+    @Disabled("Disabled until better defaulting spi is available in Jersey")
+    void overrideDefaultThreadPoolWithExecutorProvider() throws InterruptedException {
+        test(cb -> cb.register(TestNoCtxExecutorProvider.class), TEST_THREAD_NAME_PREFIX, null);
+    }
+
+    void test(Function<ClientBuilder, ClientBuilder> builderTweaker, String expectedThreadPrefix, String expectedContextValue) {
+        Context ctx = Context.create();
+        ctx.register(TEST_KEY, TEST_VAL);
+
+        CompletableFuture<String> contextValFuture = new CompletableFuture<>();
+        CompletableFuture<String> threadNameFuture = new CompletableFuture<>();
+
+        Contexts.runInContext(ctx, () ->
+                builderTweaker.apply(ClientBuilder.newBuilder())
+                        .build()
+                        .target("http://localhost:" + server.port())
+                        .request()
+                        .async()
+                        .put(Entity.text(""), ClientCallback.create(s -> {
+                            threadNameFuture.complete(Thread.currentThread().getName());
+                            contextValFuture.complete(Contexts.context()
+                                    .flatMap(c -> c.get(TEST_KEY, String.class))
+                                    .orElse(null));
+                        }, threadNameFuture::completeExceptionally)));
+
+        String actThreadName = Single.create(threadNameFuture, true).await(TIME_OUT);
+        assertThat(actThreadName, startsWith(expectedThreadPrefix));
+        String actContextVal = Single.create(contextValFuture, true).await(TIME_OUT);
+        assertThat(actContextVal, is(expectedContextValue));
+    }
+
+    @ClientAsyncExecutor
+    public static class TestNoCtxExecutorProvider extends ThreadPoolExecutorProvider {
+
+        public TestNoCtxExecutorProvider() {
+            super(TEST_THREAD_NAME_PREFIX);
+        }
+
+        @Override
+        public ExecutorService getExecutorService() {
+            return Executors.newSingleThreadExecutor();
+        }
+
+        @Override
+        public void dispose(ExecutorService executorService) {
+            executorService.shutdown();
+            try {
+                assertTrue(executorService.awaitTermination(TIME_OUT.toSeconds(), TimeUnit.SECONDS));
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+        }
+    }
+
+    static class ClientCallback implements InvocationCallback<String> {
+        Consumer<String> completed;
+        Consumer<Throwable> exceptionally;
+
+        private ClientCallback() {
+        }
+
+        static ClientCallback create(Consumer<String> completed, Consumer<Throwable> exceptionally) {
+            ClientCallback c = new ClientCallback();
+            c.completed = completed;
+            c.exceptionally = exceptionally;
+            return c;
+        }
+
+        @Override
+        public void completed(String s) {
+            completed.accept(s);
+        }
+
+        @Override
+        public void failed(Throwable t) {
+            exceptionally.accept(t);
+        }
+    }
+}