Provide a way to customize context propagation using its own storage

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;
@@ -560,6 +561,8 @@ private ClientFactoryOptions buildOptions() {
      * Returns a newly-created {@link ClientFactory} based on the properties of this builder.
      */
     public ClientFactory build() {
+        // To initialize the context storage when the factory is built not the first request is sent.
+        assert RequestContextUtil.storage() != null;
         return new DefaultClientFactory(new HttpClientFactory(buildOptions()));
     }
 

core/src/main/java/com/linecorp/armeria/client/ClientRequestContext.java

@@ -34,6 +34,7 @@
 
 import com.linecorp.armeria.client.endpoint.EndpointGroup;
 import com.linecorp.armeria.common.ContentTooLargeException;
+import com.linecorp.armeria.common.ContextStorage;
 import com.linecorp.armeria.common.HttpHeaders;
 import com.linecorp.armeria.common.HttpHeadersBuilder;
 import com.linecorp.armeria.common.HttpRequest;
@@ -46,7 +47,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 +323,25 @@ static ClientRequestContextBuilder builder(RpcRequest request, URI uri) {
      */
     @Override
     default SafeCloseable push() {
-        final RequestContext oldCtx = RequestContextThreadLocal.getAndSet(this);
+        final ContextStorage contextStorage = RequestContextUtil.storage();
+        final RequestContext oldCtx = contextStorage.push(this);
         if (oldCtx == this) {
             // Reentrance
             return noopSafeCloseable();
         }
 
         if (oldCtx == null) {
-            return RequestContextThreadLocal::remove;
+            return () -> contextStorage.pop(null);
         }
 
         final ServiceRequestContext root = root();
         if ((oldCtx instanceof ServiceRequestContext && oldCtx == root) ||
             oldCtx instanceof ClientRequestContext && ((ClientRequestContext) oldCtx).root() == root) {
-            return () -> RequestContextThreadLocal.set(oldCtx);
+            return () -> contextStorage.pop(oldCtx);
         }
 
         // Put the oldCtx back before throwing an exception.
-        RequestContextThreadLocal.set(oldCtx);
+        contextStorage.pop(oldCtx);
         throw newIllegalContextPushingException(this, oldCtx);
     }
 

core/src/main/java/com/linecorp/armeria/common/ContextStorage.java

@@ -0,0 +1,45 @@
+/*
+ * 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;
+
+/**
+ * Storage.
+ */
+@UnstableApi
+public interface ContextStorage {
+
+    /**
+     * Push.
+     */
+    @Nullable
+    <T extends RequestContext> T push(RequestContext toPush);
+
+    /**
+     * Pop.
+     */
+    void pop(@Nullable RequestContext toRestore);
+
+    /**
+     * Current.
+     */
+    @Nullable
+    <T extends RequestContext> T currentOrNull();
+}

core/src/main/java/com/linecorp/armeria/common/ContextStorageProvider.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 ContextStorage} dynamically via Java SPI (Service Provider Interface).
+ */
+@UnstableApi
+@FunctionalInterface
+public interface ContextStorageProvider {
+
+    /**
+     * Creates a new {@link ContextStorage}.
+     */
+    ContextStorage newContextStorage();
+}

core/src/main/java/com/linecorp/armeria/common/ContextStorageThreadLocal.java

@@ -0,0 +1,59 @@
+/*
+ * 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 javax.annotation.Nullable;
+
+import com.linecorp.armeria.common.util.UnstableApi;
+
+import io.netty.util.concurrent.FastThreadLocal;
+import io.netty.util.internal.InternalThreadLocalMap;
+
+/**
+ * A {@link ContextStorage} that uses thread-local to store {@link RequestContext}.
+ * Override this.
+ */
+@UnstableApi
+public class ContextStorageThreadLocal implements ContextStorage {
+
+    private static final FastThreadLocal<RequestContext> context = new FastThreadLocal<>();
+
+    @Nullable
+    @Override
+    @SuppressWarnings("unchecked")
+    public <T extends RequestContext> 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(@Nullable RequestContext toRestore) {
+        context.set(toRestore);
+    }
+
+    @Nullable
+    @Override
+    @SuppressWarnings("unchecked")
+    public <T extends RequestContext> T currentOrNull() {
+        return (T) context.get();
+    }
+}

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 extends RequestContext> T current() {
      */
     @Nullable
     static <T extends RequestContext> T currentOrNull() {
-        return RequestContextThreadLocal.get();
+        return RequestContextUtil.storage().currentOrNull();
     }
 
     /**
@@ -328,11 +328,9 @@ 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 ContextStorage contextStorage = RequestContextUtil.storage();
+        final RequestContext oldCtx = contextStorage.push(this);
+        return () -> contextStorage.pop(oldCtx);
     }
 
     /**