Add RequestContextStorage.hook() (line#2723)

Motivation:

There's currently no way to perform an additional operation on top of
the `RequestContextStorage` implementation discovered via `ServiceLoader`.

Modifications:

- Added `RequestContextStorage.hook(Function)` so that a user is allowed
  to modify the current `RequestContextStorage`.
- `RequestContextStorage` is now `Unwrappable`.
  - Added `RequestContextStorageWrapper`

Result:

- A user can install a custom hook in runtime, usually at startup time,
  like other frameworks.

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

@@ -16,9 +16,15 @@
 
 package com.linecorp.armeria.common;
 
+import static java.util.Objects.requireNonNull;
+
+import java.util.function.Function;
+
 import javax.annotation.Nullable;
 
 import com.linecorp.armeria.common.util.UnstableApi;
+import com.linecorp.armeria.common.util.Unwrappable;
+import com.linecorp.armeria.internal.common.RequestContextUtil;
 
 /**
  * The storage for storing {@link RequestContext}.
@@ -57,7 +63,22 @@
  * }</pre>
  */
 @UnstableApi
-public interface RequestContextStorage {
+public interface RequestContextStorage extends Unwrappable {
+
+    /**
+     * Customizes the current {@link RequestContextStorage} by applying the specified {@link Function} to it.
+     * This method is useful when you need to perform an additional operation when a {@link RequestContext}
+     * is pushed or popped. Note:
+     * <ul>
+     *   <li>All {@link RequestContextStorage} operations are highly performance-sensitive operation and thus
+     *       it's not a good idea to run a time-consuming task.</li>
+     *   <li>This method must be invoked at the beginning of application startup, e.g.
+     *       Never call this method in the middle of request processing.</li>
+     * </ul>
+     */
+    static void hook(Function<? super RequestContextStorage, ? extends RequestContextStorage> function) {
+        RequestContextUtil.hook(requireNonNull(function, "function"));
+    }
 
     /**
      * Returns the default {@link RequestContextStorage} which stores the {@link RequestContext}

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

@@ -0,0 +1,55 @@
+/*
+ * 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 java.util.function.Function;
+
+import javax.annotation.Nullable;
+
+import com.linecorp.armeria.common.util.AbstractUnwrappable;
+
+/**
+ * Wraps an existing {@link RequestContextStorage}.
+ *
+ * @see RequestContextStorage#hook(Function)
+ */
+public class RequestContextStorageWrapper
+        extends AbstractUnwrappable<RequestContextStorage> implements RequestContextStorage {
+
+    /**
+     * Creates a new instance that wraps the specified {@link RequestContextStorage}.
+     */
+    protected RequestContextStorageWrapper(RequestContextStorage delegate) {
+        super(delegate);
+    }
+
+    @Nullable
+    @Override
+    public <T extends RequestContext> T push(RequestContext toPush) {
+        return delegate().push(toPush);
+    }
+
+    @Override
+    public void pop(RequestContext current, @Nullable RequestContext toRestore) {
+        delegate().pop(current, toRestore);
+    }
+
+    @Nullable
+    @Override
+    public <T extends RequestContext> T currentOrNull() {
+        return delegate().currentOrNull();
+    }
+}

core/src/main/java/com/linecorp/armeria/internal/common/RequestContextUtil.java

@@ -16,12 +16,14 @@
 
 package com.linecorp.armeria.internal.common;
 
+import static com.google.common.base.Preconditions.checkState;
 import static java.util.Objects.requireNonNull;
 
 import java.util.Collections;
 import java.util.List;
 import java.util.ServiceLoader;
 import java.util.Set;
+import java.util.function.Function;
 
 import javax.annotation.Nullable;
 
@@ -57,7 +59,7 @@ public final class RequestContextUtil {
     private static final Set<Thread> REPORTED_THREADS =
             Collections.newSetFromMap(new MapMaker().weakKeys().makeMap());
 
-    private static final RequestContextStorage requestContextStorage;
+    private static RequestContextStorage requestContextStorage;
 
     static {
         final List<RequestContextStorageProvider> providers = ImmutableList.copyOf(
@@ -158,6 +160,12 @@ public static IllegalStateException newIllegalContextPoppingException(
         return ex;
     }
 
+    public static void hook(Function<? super RequestContextStorage, ? extends RequestContextStorage> function) {
+        final RequestContextStorage newStorage = function.apply(requestContextStorage);
+        checkState(newStorage != null, "function.apply() returned null: %s", function);
+        requestContextStorage = newStorage;
+    }
+
     /**
      * Returns the current {@link RequestContext} in the {@link RequestContextStorage}.
      */

it/context-storage/src/test/java/com/linecorp/armeria/common/CustomRequestContextStorageProvider.java

@@ -48,32 +48,33 @@ static int popCalled() {
 
     @Override
     public RequestContextStorage newStorage() {
-        return new RequestContextStorage() {
-
-            @Nullable
-            @Override
-            @SuppressWarnings("unchecked")
-            public <T extends RequestContext> 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 extends RequestContext> T currentOrNull() {
-                return (T) context.get();
-            }
-        };
+        return new CustomRequestContextStorage();
+    }
+
+    static final class CustomRequestContextStorage implements RequestContextStorage {
+        @Nullable
+        @Override
+        @SuppressWarnings("unchecked")
+        public <T extends RequestContext> 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 extends RequestContext> T currentOrNull() {
+            return (T) context.get();
+        }
     }
 }

it/context-storage/src/test/java/com/linecorp/armeria/common/RequestContextStorageCustomizingTest.java

@@ -19,10 +19,14 @@
 import static org.assertj.core.api.Assertions.assertThat;
 
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import javax.annotation.Nullable;
 
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
+import com.linecorp.armeria.common.CustomRequestContextStorageProvider.CustomRequestContextStorage;
 import com.linecorp.armeria.common.util.SafeCloseable;
 import com.linecorp.armeria.server.ServiceRequestContext;
 import com.linecorp.armeria.testing.junit.common.EventLoopExtension;
@@ -73,6 +77,57 @@ void requestContextStorageDoesNotAffectOtherThread() throws InterruptedException
         latch3.await();
     }
 
+    @Test
+    void hook() {
+        final AtomicBoolean pushed = new AtomicBoolean();
+        final AtomicBoolean popped = new AtomicBoolean();
+        final AtomicBoolean got = new AtomicBoolean();
+
+        RequestContextStorage.hook(delegate -> new RequestContextStorageWrapper(delegate) {
+            @Nullable
+            @Override
+            public <T extends RequestContext> T push(RequestContext toPush) {
+                pushed.set(true);
+                return super.push(toPush);
+            }
+
+            @Override
+            public void pop(RequestContext current, @Nullable RequestContext toRestore) {
+                popped.set(true);
+                super.pop(current, toRestore);
+            }
+
+            @Nullable
+            @Override
+            public <T extends RequestContext> T currentOrNull() {
+                got.set(true);
+                return super.currentOrNull();
+            }
+        });
+
+        try {
+            final ServiceRequestContext ctx = newCtx();
+
+            assertThat(pushed).isFalse();
+            assertThat(popped).isFalse();
+            assertThat(got).isFalse();
+
+            try (SafeCloseable ignored = ctx.push()) {
+                assertThat(pushed).isTrue();
+                assertThat(popped).isFalse();
+                assertThat(got).isFalse();
+
+                assertThat(ServiceRequestContext.current()).isSameAs(ctx);
+                assertThat(popped).isFalse();
+                assertThat(got).isTrue();
+            }
+
+            assertThat(popped).isTrue();
+        } finally {
+            RequestContextStorage.hook(storage -> storage.as(CustomRequestContextStorage.class));
+        }
+    }
+
     private static ServiceRequestContext newCtx() {
         return ServiceRequestContext.builder(HttpRequest.of(HttpMethod.GET, "/"))
                                     .build();