diff --git a/core/src/main/java/com/linecorp/armeria/client/ClientFactoryBuilder.java b/core/src/main/java/com/linecorp/armeria/client/ClientFactoryBuilder.java index 8a3bff0cd93..9b7f374da63 100644 --- a/core/src/main/java/com/linecorp/armeria/client/ClientFactoryBuilder.java +++ b/core/src/main/java/com/linecorp/armeria/client/ClientFactoryBuilder.java @@ -47,6 +47,7 @@ import com.linecorp.armeria.common.CommonPools; import com.linecorp.armeria.common.Flags; import com.linecorp.armeria.common.Request; +import com.linecorp.armeria.internal.common.RequestContextUtil; import io.micrometer.core.instrument.MeterRegistry; import io.netty.channel.ChannelOption; @@ -77,6 +78,10 @@ */ public final class ClientFactoryBuilder { + static { + RequestContextUtil.init(); + } + private final Map, ClientFactoryOptionValue> options = new LinkedHashMap<>(); // Netty-related properties: diff --git a/core/src/main/java/com/linecorp/armeria/client/ClientRequestContext.java b/core/src/main/java/com/linecorp/armeria/client/ClientRequestContext.java index c34a2b35ac6..b45e122d03b 100644 --- a/core/src/main/java/com/linecorp/armeria/client/ClientRequestContext.java +++ b/core/src/main/java/com/linecorp/armeria/client/ClientRequestContext.java @@ -46,7 +46,7 @@ import com.linecorp.armeria.common.logging.RequestLog; import com.linecorp.armeria.common.util.SafeCloseable; import com.linecorp.armeria.common.util.TimeoutMode; -import com.linecorp.armeria.internal.common.RequestContextThreadLocal; +import com.linecorp.armeria.internal.common.RequestContextUtil; import com.linecorp.armeria.server.Service; import com.linecorp.armeria.server.ServiceRequestContext; @@ -322,24 +322,24 @@ static ClientRequestContextBuilder builder(RpcRequest request, URI uri) { */ @Override default SafeCloseable push() { - final RequestContext oldCtx = RequestContextThreadLocal.getAndSet(this); + final RequestContext oldCtx = RequestContextUtil.getAndSet(this); if (oldCtx == this) { // Reentrance return noopSafeCloseable(); } if (oldCtx == null) { - return RequestContextThreadLocal::remove; + return () -> RequestContextUtil.pop(this, null); } final ServiceRequestContext root = root(); if ((oldCtx instanceof ServiceRequestContext && oldCtx == root) || oldCtx instanceof ClientRequestContext && ((ClientRequestContext) oldCtx).root() == root) { - return () -> RequestContextThreadLocal.set(oldCtx); + return () -> RequestContextUtil.pop(this, oldCtx); } // Put the oldCtx back before throwing an exception. - RequestContextThreadLocal.set(oldCtx); + RequestContextUtil.pop(this, oldCtx); throw newIllegalContextPushingException(this, oldCtx); } diff --git a/core/src/main/java/com/linecorp/armeria/common/Flags.java b/core/src/main/java/com/linecorp/armeria/common/Flags.java index 5f75007565f..f6ec8b0a286 100644 --- a/core/src/main/java/com/linecorp/armeria/common/Flags.java +++ b/core/src/main/java/com/linecorp/armeria/common/Flags.java @@ -120,6 +120,10 @@ public final class Flags { private static final boolean VERBOSE_RESPONSES = getBoolean("verboseResponses", false); + @Nullable + private static final String REQUEST_CONTEXT_STORAGE_PROVIDER = + System.getProperty(PREFIX + "requestContextStorageProvider"); + private static final boolean HAS_WSLENV = System.getenv("WSLENV") != null; private static final boolean USE_EPOLL = getBoolean("useEpoll", isEpollAvailable(), value -> isEpollAvailable() || !value); @@ -414,6 +418,20 @@ public static boolean verboseResponses() { return VERBOSE_RESPONSES; } + /** + * Returns the fully qualified class name of {@link RequestContextStorageProvider} that is used to choose + * when multiple {@link RequestContextStorageProvider}s exist. + * + *

The default value of this flag is {@code null}, which means only one + * {@link RequestContextStorageProvider} must be found via Java SPI. If there are more than one, + * you must specify the {@code -Dcom.linecorp.armeria.requestContextStorageProvider=} JVM option to + * choose the {@link RequestContextStorageProvider}. + */ + @Nullable + public static String requestContextStorageProvider() { + return REQUEST_CONTEXT_STORAGE_PROVIDER; + } + /** * Returns whether the JNI-based {@code /dev/epoll} socket I/O is enabled. When enabled on Linux, Armeria * uses {@code /dev/epoll} directly for socket I/O. When disabled, {@code java.nio} socket API is used diff --git a/core/src/main/java/com/linecorp/armeria/common/RequestContext.java b/core/src/main/java/com/linecorp/armeria/common/RequestContext.java index d89f8a4fb02..1f01e1a2a4d 100644 --- a/core/src/main/java/com/linecorp/armeria/common/RequestContext.java +++ b/core/src/main/java/com/linecorp/armeria/common/RequestContext.java @@ -45,7 +45,7 @@ import com.linecorp.armeria.common.logging.RequestLogBuilder; import com.linecorp.armeria.common.util.SafeCloseable; import com.linecorp.armeria.internal.common.JavaVersionSpecific; -import com.linecorp.armeria.internal.common.RequestContextThreadLocal; +import com.linecorp.armeria.internal.common.RequestContextUtil; import com.linecorp.armeria.server.ServiceRequestContext; import io.micrometer.core.instrument.MeterRegistry; @@ -80,7 +80,7 @@ static T current() { */ @Nullable static T currentOrNull() { - return RequestContextThreadLocal.get(); + return RequestContextUtil.get(); } /** @@ -328,11 +328,8 @@ default EventLoop contextAwareEventLoop() { * @see ServiceRequestContext#push() */ default SafeCloseable replace() { - final RequestContext oldCtx = RequestContextThreadLocal.getAndSet(this); - if (oldCtx == null) { - return RequestContextThreadLocal::remove; - } - return () -> RequestContextThreadLocal.set(oldCtx); + final RequestContext oldCtx = RequestContextUtil.getAndSet(this); + return () -> RequestContextUtil.pop(this, oldCtx); } /** diff --git a/core/src/main/java/com/linecorp/armeria/common/RequestContextStorage.java b/core/src/main/java/com/linecorp/armeria/common/RequestContextStorage.java new file mode 100644 index 00000000000..afaae2af45b --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/common/RequestContextStorage.java @@ -0,0 +1,95 @@ +/* + * Copyright 2020 LINE Corporation + * + * LINE Corporation licenses this file to you 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: + * + * https://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 com.linecorp.armeria.common; + +import javax.annotation.Nullable; + +import com.linecorp.armeria.common.util.UnstableApi; + +/** + * The storage for storing {@link RequestContext}. + * + *

If you want to implement your own storage or add some hooks when a {@link RequestContext} is pushed + * and popped, you should use {@link RequestContextStorageProvider}. + * Here's an example that sets MDC before {@link RequestContext} is pushed: + * + * 

{@code
+ * > public class MyStorage implements RequestContextStorageProvider {
+ * >
+ * >     @Override
+ * >     public RequestContextStorage newStorage() {
+ * >         RequestContextStorage threadLocalStorage = RequestContextStorage.threadLocal();
+ * >         return new RequestContextStorage() {
+ * >
+ * >             @Nullable
+ * >             @Override
+ * >             public  T push(RequestContext toPush) {
+ * >                 setMdc(toPush);
+ * >                 return threadLocalStorage.push(toPush);
+ * >             }
+ * >
+ * >             @Override
+ * >             public void pop(RequestContext current, @Nullable RequestContext toRestore) {
+ * >                 clearMdc();
+ * >                 if (toRestore != null) {
+ * >                     setMdc(toRestore);
+ * >                 }
+ * >                 threadLocalStorage.pop(current, toRestore);
+ * >             }
+ * >             ...
+ * >          }
+ * >     }
+ * > }
+ * }
+ */ +@UnstableApi +public interface RequestContextStorage { + + /** + * Returns the default {@link RequestContextStorage} which stores the {@link RequestContext} + * in the thread-local. + */ + static RequestContextStorage threadLocal() { + return ThreadLocalRequestContextStorage.INSTANCE; + } + + /** + * Pushes the specified {@link RequestContext} into the storage. + * + * @return the old {@link RequestContext} which was in the storage before the specified {@code toPush} is + * pushed. {@code null}, if there was no {@link RequestContext}. + */ + @Nullable + T push(RequestContext toPush); + + /** + * Pops the current {@link RequestContext} in the storage and pushes back the specified {@code toRestore}. + * {@code toRestore} is the {@link RequestContext} returned from when + * {@linkplain #push(RequestContext) push(current)} is called, so it can be {@code null}. + * + *

The specified {@code current} must be the {@link RequestContext} in the storage. If it's not, + * it means that {@link RequestContext#push()} is not called using {@code try-with-resources} block, so + * the previous {@link RequestContext} is not popped properly. + */ + void pop(RequestContext current, @Nullable RequestContext toRestore); + + /** + * Returns the {@link RequestContext} in the storage. {@code null} if there is no {@link RequestContext}. + */ + @Nullable + T currentOrNull(); +} diff --git a/core/src/main/java/com/linecorp/armeria/common/RequestContextStorageProvider.java b/core/src/main/java/com/linecorp/armeria/common/RequestContextStorageProvider.java new file mode 100644 index 00000000000..f7c82cab044 --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/common/RequestContextStorageProvider.java @@ -0,0 +1,32 @@ +/* + * Copyright 2020 LINE Corporation + * + * LINE Corporation licenses this file to you 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: + * + * https://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 com.linecorp.armeria.common; + +import com.linecorp.armeria.common.util.UnstableApi; + +/** + * Creates a new {@link RequestContextStorage} dynamically via Java SPI (Service Provider Interface). + */ +@UnstableApi +@FunctionalInterface +public interface RequestContextStorageProvider { + + /** + * Creates a new {@link RequestContextStorage}. + */ + RequestContextStorage newStorage(); +} diff --git a/core/src/main/java/com/linecorp/armeria/common/ThreadLocalRequestContextStorage.java b/core/src/main/java/com/linecorp/armeria/common/ThreadLocalRequestContextStorage.java new file mode 100644 index 00000000000..82967ed006c --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/common/ThreadLocalRequestContextStorage.java @@ -0,0 +1,61 @@ +/* + * Copyright 2020 LINE Corporation + * + * LINE Corporation licenses this file to you 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: + * + * https://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 com.linecorp.armeria.common; + +import static com.linecorp.armeria.internal.common.RequestContextUtil.newIllegalContextPoppingException; +import static java.util.Objects.requireNonNull; + +import javax.annotation.Nullable; + +import io.netty.util.concurrent.FastThreadLocal; +import io.netty.util.internal.InternalThreadLocalMap; + +enum ThreadLocalRequestContextStorage implements RequestContextStorage { + + INSTANCE; + + private static final FastThreadLocal context = new FastThreadLocal<>(); + + @Nullable + @Override + @SuppressWarnings("unchecked") + public T push(RequestContext toPush) { + requireNonNull(toPush, "toPush"); + final InternalThreadLocalMap map = InternalThreadLocalMap.get(); + final RequestContext oldCtx = context.get(map); + context.set(map, toPush); + return (T) oldCtx; + } + + @Override + public void pop(RequestContext current, @Nullable RequestContext toRestore) { + requireNonNull(current, "current"); + final InternalThreadLocalMap map = InternalThreadLocalMap.get(); + final RequestContext contextInThreadLocal = context.get(map); + if (current != contextInThreadLocal) { + throw newIllegalContextPoppingException(current, contextInThreadLocal); + } + context.set(map, toRestore); + } + + @Nullable + @Override + @SuppressWarnings("unchecked") + public T currentOrNull() { + return (T) context.get(); + } +} diff --git a/core/src/main/java/com/linecorp/armeria/internal/common/RequestContextThreadLocal.java b/core/src/main/java/com/linecorp/armeria/internal/common/RequestContextThreadLocal.java deleted file mode 100644 index 1918d717e55..00000000000 --- a/core/src/main/java/com/linecorp/armeria/internal/common/RequestContextThreadLocal.java +++ /dev/null @@ -1,82 +0,0 @@ -/* - * Copyright 2016 LINE Corporation - * - * LINE Corporation licenses this file to you 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: - * - * https://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 com.linecorp.armeria.internal.common; - -import static java.util.Objects.requireNonNull; - -import javax.annotation.Nullable; - -import com.linecorp.armeria.common.RequestContext; - -import io.netty.util.concurrent.FastThreadLocal; -import io.netty.util.internal.InternalThreadLocalMap; - -public final class RequestContextThreadLocal { - - private static final FastThreadLocal context = new FastThreadLocal<>(); - - /** - * Returns the current {@link RequestContext} in the thread-local. - */ - @Nullable - @SuppressWarnings("unchecked") - public static T get() { - return (T) context.get(); - } - - /** - * Sets the specified {@link RequestContext} in the thread-local and returns the old {@link RequestContext}. - */ - @Nullable - @SuppressWarnings("unchecked") - public static T getAndSet(RequestContext ctx) { - requireNonNull(ctx, "ctx"); - final InternalThreadLocalMap map = InternalThreadLocalMap.get(); - final RequestContext oldCtx = context.get(map); - context.set(map, ctx); - return (T) oldCtx; - } - - /** - * Removes the {@link RequestContext} in the thread-local and returns it. - */ - @Nullable - @SuppressWarnings("unchecked") - public static T getAndRemove() { - final InternalThreadLocalMap map = InternalThreadLocalMap.get(); - final RequestContext oldCtx = context.get(map); - context.remove(); - return (T) oldCtx; - } - - /** - * Sets the specified {@link RequestContext} in the thread-local. - */ - public static void set(RequestContext ctx) { - requireNonNull(ctx, "ctx"); - context.set(ctx); - } - - /** - * Removes the current {@link RequestContext} in the thread-local. - */ - public static void remove() { - context.remove(); - } - - private RequestContextThreadLocal() {} -} diff --git a/core/src/main/java/com/linecorp/armeria/internal/common/RequestContextUtil.java b/core/src/main/java/com/linecorp/armeria/internal/common/RequestContextUtil.java index d0ee6ab7475..10e7268b320 100644 --- a/core/src/main/java/com/linecorp/armeria/internal/common/RequestContextUtil.java +++ b/core/src/main/java/com/linecorp/armeria/internal/common/RequestContextUtil.java @@ -19,14 +19,23 @@ import static java.util.Objects.requireNonNull; import java.util.Collections; +import java.util.List; +import java.util.ServiceLoader; import java.util.Set; +import javax.annotation.Nullable; + import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import com.google.common.collect.ImmutableList; import com.google.common.collect.MapMaker; +import com.linecorp.armeria.common.Flags; +import com.linecorp.armeria.common.HttpRequest; import com.linecorp.armeria.common.RequestContext; +import com.linecorp.armeria.common.RequestContextStorage; +import com.linecorp.armeria.common.RequestContextStorageProvider; import com.linecorp.armeria.common.util.SafeCloseable; import io.netty.channel.ChannelFuture; @@ -48,6 +57,64 @@ public final class RequestContextUtil { private static final Set REPORTED_THREADS = Collections.newSetFromMap(new MapMaker().weakKeys().makeMap()); + private static final RequestContextStorage requestContextStorage; + + static { + final List providers = ImmutableList.copyOf( + ServiceLoader.load(RequestContextStorageProvider.class)); + final String providerFqcn = Flags.requestContextStorageProvider(); + if (!providers.isEmpty()) { + + RequestContextStorageProvider provider = null; + if (providers.size() > 1) { + if (providerFqcn == null) { + throw new IllegalStateException( + "Found more than one " + RequestContextStorageProvider.class.getSimpleName() + + ". You must specify -Dcom.linecorp.armeria.requestContextStorageProvider=." + + " providers: " + providers); + } + + for (RequestContextStorageProvider candidate : providers) { + if (candidate.getClass().getName().equals(providerFqcn)) { + if (provider != null) { + throw new IllegalStateException( + providerFqcn + " matches more than one " + + RequestContextStorageProvider.class.getSimpleName() + ". providers: " + + providers); + } else { + provider = candidate; + } + } + } + if (provider == null) { + throw new IllegalStateException( + providerFqcn + " does not match any " + + RequestContextStorageProvider.class.getSimpleName() + ". providers: " + providers); + } + } else { + provider = providers.get(0); + if (logger.isInfoEnabled()) { + logger.info("Using {} as a {}", + provider.getClass().getSimpleName(), + RequestContextStorageProvider.class.getSimpleName()); + } + } + + try { + requestContextStorage = provider.newStorage(); + } catch (Throwable t) { + throw new IllegalStateException("Failed to create context storage. provider: " + provider, t); + } + } else { + requestContextStorage = RequestContextStorage.threadLocal(); + } + } + + /** + * Invoked to initialize this class earlier than when an {@link HttpRequest} is received or sent. + */ + public static void init() { /* no-op */ } + /** * Returns the {@link SafeCloseable} which doesn't do anything. */ @@ -74,21 +141,69 @@ public static IllegalStateException newIllegalContextPushingException( } /** - * Removes the {@link RequestContext} in the thread-local if exists and returns {@link SafeCloseable} which - * pushes the {@link RequestContext} back to the thread-local. + * Returns an {@link IllegalStateException} which is raised when popping a context from + * the unexpected thread or forgetting to close the previous context. + */ + public static IllegalStateException newIllegalContextPoppingException( + RequestContext currentCtx, RequestContext contextInStorage) { + requireNonNull(currentCtx, "currentCtx"); + requireNonNull(contextInStorage, "contextInStorage"); + final IllegalStateException ex = new IllegalStateException( + "The currentCtx " + currentCtx + " is not the same as the context in the storage: " + + contextInStorage + ". This means the callback was called from " + + "unexpected thread or forgetting to close previous context."); + if (REPORTED_THREADS.add(Thread.currentThread())) { + logger.warn("An error occurred while popping a context", ex); + } + return ex; + } + + /** + * Returns the current {@link RequestContext} in the {@link RequestContextStorage}. + */ + @Nullable + @SuppressWarnings("unchecked") + public static T get() { + return (T) requestContextStorage.currentOrNull(); + } + + /** + * Sets the specified {@link RequestContext} in the {@link RequestContextStorage} and + * returns the old {@link RequestContext}. + */ + @Nullable + @SuppressWarnings("unchecked") + public static T getAndSet(RequestContext ctx) { + requireNonNull(ctx, "ctx"); + return (T) requestContextStorage.push(ctx); + } + + /** + * Removes the {@link RequestContext} in the {@link RequestContextStorage} if exists and returns + * {@link SafeCloseable} which pushes the {@link RequestContext} back to the {@link RequestContextStorage}. * *

Because this method pops the {@link RequestContext} arbitrarily, it shouldn't be used in * most cases. One of the examples this can be used is in {@link ChannelFutureListener}. * The {@link ChannelFuture} can be complete when the eventloop handles the different request. The - * eventloop might have the wrong {@link RequestContext} in the thread-local, so we should pop it. + * eventloop might have the wrong {@link RequestContext} in the {@link RequestContextStorage}, + * so we should pop it. */ public static SafeCloseable pop() { - final RequestContext oldCtx = RequestContextThreadLocal.getAndRemove(); + final RequestContext oldCtx = requestContextStorage.currentOrNull(); if (oldCtx == null) { return noopSafeCloseable(); } - return () -> RequestContextThreadLocal.set(oldCtx); + pop(oldCtx, null); + return () -> requestContextStorage.push(oldCtx); + } + + /** + * Pops the current {@link RequestContext} in the storage and pushes back the specified {@code toRestore}. + */ + public static void pop(RequestContext current, @Nullable RequestContext toRestore) { + requireNonNull(current, "current"); + requestContextStorage.pop(current, toRestore); } private RequestContextUtil() {} diff --git a/core/src/main/java/com/linecorp/armeria/server/ServerBuilder.java b/core/src/main/java/com/linecorp/armeria/server/ServerBuilder.java index 20fff98c11b..a2d8ab2af52 100644 --- a/core/src/main/java/com/linecorp/armeria/server/ServerBuilder.java +++ b/core/src/main/java/com/linecorp/armeria/server/ServerBuilder.java @@ -66,6 +66,7 @@ import com.linecorp.armeria.common.RequestId; import com.linecorp.armeria.common.SessionProtocol; import com.linecorp.armeria.common.util.SystemInfo; +import com.linecorp.armeria.internal.common.RequestContextUtil; import com.linecorp.armeria.internal.server.annotation.AnnotatedServiceExtensions; import com.linecorp.armeria.server.annotation.ExceptionHandlerFunction; import com.linecorp.armeria.server.annotation.RequestConverterFunction; @@ -149,6 +150,10 @@ public final class ServerBuilder { ChannelOption.WRITE_BUFFER_HIGH_WATER_MARK, ChannelOption.WRITE_BUFFER_LOW_WATER_MARK, EpollChannelOption.EPOLL_MODE); + static { + RequestContextUtil.init(); + } + private final List ports = new ArrayList<>(); private final List serverListeners = new ArrayList<>(); @VisibleForTesting diff --git a/core/src/main/java/com/linecorp/armeria/server/ServiceRequestContext.java b/core/src/main/java/com/linecorp/armeria/server/ServiceRequestContext.java index 526ddea1131..9d414c5f0c8 100644 --- a/core/src/main/java/com/linecorp/armeria/server/ServiceRequestContext.java +++ b/core/src/main/java/com/linecorp/armeria/server/ServiceRequestContext.java @@ -49,7 +49,7 @@ import com.linecorp.armeria.common.RpcRequest; import com.linecorp.armeria.common.util.SafeCloseable; import com.linecorp.armeria.common.util.TimeoutMode; -import com.linecorp.armeria.internal.common.RequestContextThreadLocal; +import com.linecorp.armeria.internal.common.RequestContextUtil; import com.linecorp.armeria.server.logging.AccessLogWriter; /** @@ -212,22 +212,22 @@ default InetAddress clientAddress() { */ @Override default SafeCloseable push() { - final RequestContext oldCtx = RequestContextThreadLocal.getAndSet(this); + final RequestContext oldCtx = RequestContextUtil.getAndSet(this); if (oldCtx == this) { // Reentrance return noopSafeCloseable(); } - if (oldCtx instanceof ClientRequestContext && ((ClientRequestContext) oldCtx).root() == this) { - return () -> RequestContextThreadLocal.set(oldCtx); + if (oldCtx == null) { + return () -> RequestContextUtil.pop(this, null); } - if (oldCtx == null) { - return RequestContextThreadLocal::remove; + if (oldCtx instanceof ClientRequestContext && ((ClientRequestContext) oldCtx).root() == this) { + return () -> RequestContextUtil.pop(this, oldCtx); } // Put the oldCtx back before throwing an exception. - RequestContextThreadLocal.set(oldCtx); + RequestContextUtil.pop(this, oldCtx); throw newIllegalContextPushingException(this, oldCtx); } diff --git a/it/context-storage/build.gradle b/it/context-storage/build.gradle new file mode 100644 index 00000000000..142ff368292 --- /dev/null +++ b/it/context-storage/build.gradle @@ -0,0 +1,3 @@ +// This module is for testing RequestContextStorage, so we don't need other dependency except `:core`. +// To override RequestContextStorage, we have to use Java SPI, which affects other tests, so we need +// an isolated module. diff --git a/it/context-storage/src/test/java/com/linecorp/armeria/common/CustomRequestContextStorageProvider.java b/it/context-storage/src/test/java/com/linecorp/armeria/common/CustomRequestContextStorageProvider.java new file mode 100644 index 00000000000..c60c188e6da --- /dev/null +++ b/it/context-storage/src/test/java/com/linecorp/armeria/common/CustomRequestContextStorageProvider.java @@ -0,0 +1,79 @@ +/* + * Copyright 2020 LINE Corporation + * + * LINE Corporation licenses this file to you 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: + * + * https://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 com.linecorp.armeria.common; + +import static java.util.Objects.requireNonNull; + +import java.util.concurrent.atomic.AtomicInteger; + +import javax.annotation.Nullable; + +import io.netty.util.concurrent.FastThreadLocal; +import io.netty.util.internal.InternalThreadLocalMap; + +public final class CustomRequestContextStorageProvider implements RequestContextStorageProvider { + + private static final FastThreadLocal context = new FastThreadLocal<>(); + + private static final AtomicInteger pushCalled = new AtomicInteger(); + + private static final AtomicInteger popCalled = new AtomicInteger(); + + @Nullable + static RequestContext current() { + return context.get(); + } + + static int pushCalled() { + return pushCalled.get(); + } + + static int popCalled() { + return popCalled.get(); + } + + @Override + public RequestContextStorage newStorage() { + return new RequestContextStorage() { + + @Nullable + @Override + @SuppressWarnings("unchecked") + public T push(RequestContext toPush) { + requireNonNull(toPush, "toPush"); + pushCalled.incrementAndGet(); + final InternalThreadLocalMap map = InternalThreadLocalMap.get(); + final RequestContext oldCtx = context.get(map); + context.set(map, toPush); + return (T) oldCtx; + } + + @Override + public void pop(RequestContext current, @Nullable RequestContext toRestore) { + popCalled.incrementAndGet(); + context.set(toRestore); + } + + @Nullable + @Override + @SuppressWarnings("unchecked") + public T currentOrNull() { + return (T) context.get(); + } + }; + } +} diff --git a/it/context-storage/src/test/java/com/linecorp/armeria/common/RequestContextStorageCustomizingTest.java b/it/context-storage/src/test/java/com/linecorp/armeria/common/RequestContextStorageCustomizingTest.java new file mode 100644 index 00000000000..2d6eb7b0a21 --- /dev/null +++ b/it/context-storage/src/test/java/com/linecorp/armeria/common/RequestContextStorageCustomizingTest.java @@ -0,0 +1,80 @@ +/* + * Copyright 2020 LINE Corporation + * + * LINE Corporation licenses this file to you 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: + * + * https://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 com.linecorp.armeria.common; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.util.concurrent.CountDownLatch; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +import com.linecorp.armeria.common.util.SafeCloseable; +import com.linecorp.armeria.server.ServiceRequestContext; +import com.linecorp.armeria.testing.junit.common.EventLoopExtension; + +import io.netty.channel.EventLoop; + +class RequestContextStorageCustomizingTest { + + @RegisterExtension + static final EventLoopExtension eventLoopExtension = new EventLoopExtension(); + + @Test + void requestContextStorageDoesNotAffectOtherThread() throws InterruptedException { + final EventLoop eventLoop = eventLoopExtension.get(); + final ServiceRequestContext ctx = newCtx(); + + final CountDownLatch latch1 = new CountDownLatch(1); + final CountDownLatch latch2 = new CountDownLatch(1); + final CountDownLatch latch3 = new CountDownLatch(1); + try (SafeCloseable ignored = ctx.push()) { + assertThat(CustomRequestContextStorageProvider.current()).isEqualTo(ctx); + assertThat(CustomRequestContextStorageProvider.pushCalled()).isOne(); + + eventLoop.execute(() -> { + final ServiceRequestContext ctx1 = newCtx(); + try (SafeCloseable ignored1 = ctx1.push()) { + assertThat(CustomRequestContextStorageProvider.current()).isEqualTo(ctx1); + assertThat(CustomRequestContextStorageProvider.pushCalled()).isEqualTo(2); + latch1.countDown(); + try { + latch2.await(); + } catch (InterruptedException e) { + // ignore + } + } + assertThat(CustomRequestContextStorageProvider.current()).isNull(); + assertThat(CustomRequestContextStorageProvider.popCalled()).isEqualTo(2); + latch3.countDown(); + }); + + latch1.await(); + assertThat(CustomRequestContextStorageProvider.current()).isEqualTo(ctx); + assertThat(CustomRequestContextStorageProvider.pushCalled()).isEqualTo(2); + } + assertThat(CustomRequestContextStorageProvider.current()).isNull(); + assertThat(CustomRequestContextStorageProvider.popCalled()).isOne(); + latch2.countDown(); + latch3.await(); + } + + private static ServiceRequestContext newCtx() { + return ServiceRequestContext.builder(HttpRequest.of(HttpMethod.GET, "/")) + .build(); + } +} diff --git a/it/context-storage/src/test/resources/META-INF/services/com.linecorp.armeria.common.RequestContextStorageProvider b/it/context-storage/src/test/resources/META-INF/services/com.linecorp.armeria.common.RequestContextStorageProvider new file mode 100644 index 00000000000..fa1e8a3258e --- /dev/null +++ b/it/context-storage/src/test/resources/META-INF/services/com.linecorp.armeria.common.RequestContextStorageProvider @@ -0,0 +1 @@ +com.linecorp.armeria.common.CustomRequestContextStorageProvider diff --git a/settings.gradle b/settings.gradle index f73abeb796e..c5780cc856c 100644 --- a/settings.gradle +++ b/settings.gradle @@ -38,6 +38,7 @@ includeWithFlags ':saml', 'java', 'publish', 'relo // Unpublished Java projects includeWithFlags ':benchmarks', 'java' +includeWithFlags ':it:context-storage', 'java' includeWithFlags ':it:server', 'java', 'relocate' includeWithFlags ':it:spring:boot-tomcat', 'java', 'relocate' includeWithFlags ':it:spring:boot-tomcat8.5', 'java', 'relocate'