Add TraceAbleRequestContextStorage for ease of detect context leaks (#…

…4232)

Motivation:

- Context leaks are hard to find because an exception does not tell where/which context is pushed without poping. By using TraceAbleRequestContextStorage, it helps to report the source of context leaks.
- Details as mentioned in #4100 

By the way, Thanks to @anuraaga for giving a reference to read on [opentelemetry](https://github.com/open-telemetry/opentelemetry-java).

Modifications:

- Add `TraceAbleRequestContextStorage` that stores `RequestContext` stack trace and reports to the user where it happens.
- Add `requestContextLeakDetectionSampler` flag that users can use for enable leak detection. Users can enable it by either system property or SPI flag provider.

Result:

- Closes #4100 
- `TraceAbleRequestContextStorage` is added, so users can use it to report where context leaks happen.

How to enable:
1) By system property
`-Dcom.linecorp.armeria.requestContextLeakDetectionSampler=<sampler-spec>`

2) By providing FlagsProvider SPI
```java
public final class EnableLeakDetectionFlagsProvider implements FlagsProvider {

    @OverRide
    public Sampler<? super RequestContext> requestContextLeakDetectionSampler() {
        return Sampler.always();
    }
    ...
}
```

3) By providing RequestContextStorageProvider SPI (not recommend since RequestContextStorageProvider SPI'll be remove as mentioned in ##4211 )
```java
public final class CustomRequestContextStorageProvider implements RequestContextStorageProvider {

    @OverRide
    public RequestContextStorage newStorage() {
        return new TraceAbleRequestContextStorage(delegate);
    }
}
```
Use case:
Users problematic code
```java
executor.execute(() -> {
    SafeCloseable leaked = fooCtx.push(); //This causes Request context leaks!
    ...
});

executor.execute(() -> {
    try (SafeCloseable ignored = barCtx.push()) { //Exception happen here
        ...
    }
});

```
The above code will produce an error as below. Therefore, users can check the stack trace that which line causes context leaks.
```
java.lang.IllegalStateException: Trying to call object wrapped with context [%New RequestContext%], but context is currently set to TraceableServiceRequestContext[%Previous RequestContext%]
com.linecorp.armeria.internal.common.LeakTracingRequestContextStorage$PendingRequestContextStackTrace: At thread [armeria-testing-eventloop-nio-1-1] previous RequestContext is pushed at the following stacktrace
	at com.linecorp.armeria.internal.common.LeakTracingRequestContextStorage$TraceableServiceRequestContext.<init>(LeakTracingRequestContextStorage.java:111)
	at com.linecorp.armeria.internal.common.LeakTracingRequestContextStorage$TraceableServiceRequestContext.<init>(LeakTracingRequestContextStorage.java:105)
	at com.linecorp.armeria.internal.common.LeakTracingRequestContextStorage.warpRequestContext(LeakTracingRequestContextStorage.java:82)
	at com.linecorp.armeria.internal.common.LeakTracingRequestContextStorage.push(LeakTracingRequestContextStorage.java:62)
	at com.linecorp.armeria.internal.common.RequestContextUtil.getAndSet(RequestContextUtil.java:149)
	at com.linecorp.armeria.server.ServiceRequestContext.push(ServiceRequestContext.java:221)
	at com.linecorp.armeria.internal.common.TraceRequestContextLeakTest.lambda$singleThreadContextLeak$2(TraceRequestContextLeakTest.java:101)  <- This is the line where leaked RequestContext is push 
	...
. This means the callback was called from unexpected thread or forgetting to close previous context.
	at com.linecorp.armeria.internal.common.RequestContextUtil.newIllegalContextPushingException(RequestContextUtil.java:100)
	at com.linecorp.armeria.server.ServiceRequestContext.push(ServiceRequestContext.java:237)
	at com.linecorp.armeria.internal.common.TraceRequestContextLeakTest.lambda$singleThreadContextLeak$3(TraceRequestContextLeakTest.java:107)
	...
```

benchmarks/jmh/src/jmh/java/com/linecorp/armeria/internal/common/LeakTracingRequestContextStorageBenchmark.java

@@ -0,0 +1,95 @@
+/*
+ * Copyright 2022 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 org.openjdk.jmh.annotations.Benchmark;
+
+import com.linecorp.armeria.common.HttpMethod;
+import com.linecorp.armeria.common.HttpRequest;
+import com.linecorp.armeria.common.RequestContext;
+import com.linecorp.armeria.common.RequestContextStorage;
+import com.linecorp.armeria.common.util.Sampler;
+import com.linecorp.armeria.server.ServiceRequestContext;
+
+/**
+ * Microbenchmarks for LeakTracingRequestContextStorage.
+ */
+public class LeakTracingRequestContextStorageBenchmark {
+
+    private static final RequestContextStorage threadLocalReqCtxStorage =
+            RequestContextStorage.threadLocal();
+    private static final RequestContextStorage neverSample =
+            new LeakTracingRequestContextStorage(threadLocalReqCtxStorage, Sampler.never());
+    private static final RequestContextStorage rateLimited1 =
+            new LeakTracingRequestContextStorage(threadLocalReqCtxStorage, Sampler.of("rate-limited=1"));
+    private static final RequestContextStorage rateLimited10 =
+            new LeakTracingRequestContextStorage(threadLocalReqCtxStorage, Sampler.of("rate-limited=10"));
+    private static final RequestContextStorage random1 =
+            new LeakTracingRequestContextStorage(threadLocalReqCtxStorage, Sampler.of("random=0.01"));
+    private static final RequestContextStorage random10 =
+            new LeakTracingRequestContextStorage(threadLocalReqCtxStorage, Sampler.of("random=0.10"));
+    private static final RequestContextStorage alwaysSample =
+            new LeakTracingRequestContextStorage(threadLocalReqCtxStorage, Sampler.always());
+    private static final RequestContext reqCtx = newCtx("/");
+
+    private static ServiceRequestContext newCtx(String path) {
+        return ServiceRequestContext.builder(HttpRequest.of(HttpMethod.GET, path))
+                                    .build();
+    }
+
+    @Benchmark
+    public void baseline_threadLocal() {
+        final RequestContext oldCtx = threadLocalReqCtxStorage.push(reqCtx);
+        threadLocalReqCtxStorage.pop(reqCtx, oldCtx);
+    }
+
+    @Benchmark
+    public void leakTracing_never_sample() {
+        final RequestContext oldCtx = neverSample.push(reqCtx);
+        neverSample.pop(reqCtx, oldCtx);
+    }
+
+    @Benchmark
+    public void leakTracing_rateLimited_1() {
+        final RequestContext oldCtx = rateLimited1.push(reqCtx);
+        rateLimited1.pop(reqCtx, oldCtx);
+    }
+
+    @Benchmark
+    public void leakTracing_rateLimited_10() {
+        final RequestContext oldCtx = rateLimited10.push(reqCtx);
+        rateLimited10.pop(reqCtx, oldCtx);
+    }
+
+    @Benchmark
+    public void leakTracing_random_1() {
+        final RequestContext oldCtx = random1.push(reqCtx);
+        random1.pop(reqCtx, oldCtx);
+    }
+
+    @Benchmark
+    public void leakTracing_random_10() {
+        final RequestContext oldCtx = random10.push(reqCtx);
+        random10.pop(reqCtx, oldCtx);
+    }
+
+    @Benchmark
+    public void leakTracing_always_sample() {
+        final RequestContext oldCtx = alwaysSample.push(reqCtx);
+        alwaysSample.pop(reqCtx, oldCtx);
+    }
+}

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

@@ -216,17 +216,17 @@ static ClientRequestContextBuilder builder(RpcRequest request, URI uri) {
     @MustBeClosed
     default SafeCloseable push() {
         final RequestContext oldCtx = RequestContextUtil.getAndSet(this);
-        if (oldCtx == this) {
-            // Reentrance
-            return noopSafeCloseable();
-        }
-
         if (oldCtx == null) {
             return RequestContextUtil.invokeHookAndPop(this, null);
         }
 
+        if (oldCtx.unwrapAll() == unwrapAll()) {
+            // Reentrance
+            return noopSafeCloseable();
+        }
+
         final ServiceRequestContext root = root();
-        if (oldCtx.root() == root) {
+        if (RequestContextUtil.equalsIgnoreWrapper(oldCtx.root(), root)) {
             return RequestContextUtil.invokeHookAndPop(this, oldCtx);
         }
 

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

@@ -398,4 +398,9 @@ public Path defaultMultipartUploadsLocation() {
                          File.separatorChar + "armeria" +
                          File.separatorChar + "multipart-uploads");
     }
+
+    @Override
+    public Sampler<? super RequestContext> requestContextLeakDetectionSampler() {
+        return Sampler.never();
+    }
 }

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

@@ -49,6 +49,7 @@
 import com.linecorp.armeria.client.retry.RetryingClient;
 import com.linecorp.armeria.client.retry.RetryingRpcClient;
 import com.linecorp.armeria.common.annotation.Nullable;
+import com.linecorp.armeria.common.annotation.UnstableApi;
 import com.linecorp.armeria.common.util.Exceptions;
 import com.linecorp.armeria.common.util.Sampler;
 import com.linecorp.armeria.common.util.SystemInfo;
@@ -113,10 +114,6 @@ public final class Flags {
     static {
         final String strSpec = getNormalized("verboseExceptions",
                                              DefaultFlagsProvider.VERBOSE_EXCEPTION_SAMPLER_SPEC, val -> {
-                    if ("true".equals(val) || "false".equals(val)) {
-                        return true;
-                    }
-
                     try {
                         Sampler.of(val);
                         return true;
@@ -377,6 +374,9 @@ public final class Flags {
     private static final Path DEFAULT_MULTIPART_UPLOADS_LOCATION =
             getValue(FlagsProvider::defaultMultipartUploadsLocation, "defaultMultipartUploadsLocation");
 
+    private static final Sampler<? super RequestContext> REQUEST_CONTEXT_LEAK_DETECTION_SAMPLER =
+            getValue(FlagsProvider::requestContextLeakDetectionSampler, "requestContextLeakDetectionSampler");
+
     /**
      * Returns the specification of the {@link Sampler} that determines whether to retain the stack
      * trace of the exceptions that are thrown frequently by Armeria. A sampled exception will have the stack
@@ -1296,6 +1296,21 @@ public static boolean allowDoubleDotsInQueryString() {
         return ALLOW_DOUBLE_DOTS_IN_QUERY_STRING;
     }
 
+    /**
+     * Returns the {@link Sampler} that determines whether to trace the stack trace of request contexts leaks
+     * and how frequently to keeps stack trace. A sampled exception will have the stack trace while the others
+     * will have an empty stack trace to eliminate the cost of capturing the stack trace.
+     *
+     * <p>The default value of this flag is {@link Sampler#never()}.
+     * Specify the {@code -Dcom.linecorp.armeria.requestContextLeakDetectionSampler=<specification>} JVM option
+     * to override the default. This feature is disabled if {@link Sampler#never()} is specified.
+     * See {@link Sampler#of(String)} for the specification string format.</p>
+     */
+    @UnstableApi
+    public static Sampler<? super RequestContext> requestContextLeakDetectionSampler() {
+        return REQUEST_CONTEXT_LEAK_DETECTION_SAMPLER;
+    }
+
     @Nullable
     private static String nullableCaffeineSpec(Function<FlagsProvider, String> method, String flagName) {
         return caffeineSpec(method, flagName, true);

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

@@ -960,4 +960,20 @@ default Boolean allowDoubleDotsInQueryString() {
     default Path defaultMultipartUploadsLocation() {
         return null;
     }
+
+    /**
+     * Returns the {@link Sampler} that determines whether to trace the stack trace of request contexts leaks
+     * and how frequently to keeps stack trace. A sampled exception will have the stack trace while the others
+     * will have an empty stack trace to eliminate the cost of capturing the stack trace.
+     *
+     * <p>The default value of this flag is {@link Sampler#never()}.
+     * Specify the {@code -Dcom.linecorp.armeria.requestContextLeakDetectionSampler=<specification>} JVM option
+     * to override the default. This feature is disabled if {@link Sampler#never()} is specified.
+     * See {@link Sampler#of(String)} for the specification string format.</p>
+     */
+    @UnstableApi
+    @Nullable
+    default Sampler<? super RequestContext> requestContextLeakDetectionSampler() {
+        return null;
+    }
 }

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

@@ -71,20 +71,12 @@ public Sampler<Class<? extends Throwable>> verboseExceptionSampler() {
         if (spec == null) {
             return null;
         }
-        if ("true".equals(spec) || "always".equals(spec)) {
-            return Sampler.always();
-        }
-        if ("false".equals(spec) || "never".equals(spec)) {
-            return Sampler.never();
-        }
-
         try {
             Sampler.of(spec);
         } catch (Exception e) {
             // Invalid sampler specification
             throw new IllegalArgumentException("invalid sampler spec: " + spec, e);
         }
-
         return new ExceptionSampler(spec);
     }
 
@@ -438,6 +430,20 @@ public Path defaultMultipartUploadsLocation() {
         return getAndParse("defaultMultipartUploadsLocation", Paths::get);
     }
 
+    @Override
+    public Sampler<? super RequestContext> requestContextLeakDetectionSampler() {
+        final String spec = getNormalized("requestContextLeakDetectionSampler");
+        if (spec == null) {
+            return null;
+        }
+        try {
+            return Sampler.of(spec);
+        } catch (Exception e) {
+            // Invalid sampler specification
+            throw new IllegalArgumentException("invalid sampler spec: " + spec, e);
+        }
+    }
+
     @Nullable
     private static Long getLong(String name) {
         return getAndParse(name, Long::parseLong);

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

@@ -46,7 +46,7 @@ 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) {
+        if (contextInThreadLocal == null || current.unwrapAll() != contextInThreadLocal.unwrapAll()) {
             throw newIllegalContextPoppingException(current, contextInThreadLocal);
         }
         context.set(map, toRestore);

core/src/main/java/com/linecorp/armeria/common/util/Samplers.java

@@ -70,8 +70,10 @@ static <T> Sampler<T> of(String specification) {
         requireNonNull(specification, "specification");
         switch (specification.trim()) {
             case "always":
+            case "true":
                 return Sampler.always();
             case "never":
+            case "false":
                 return Sampler.never();
         }
 

core/src/main/java/com/linecorp/armeria/common/util/Unwrappable.java

@@ -129,4 +129,20 @@ default Object unwrapAll() {
             unwrapped = inner;
         }
     }
+
+    /**
+     * Reference checking this {@link Unwrappable} to another {@link Unwrappable}, ignoring wrappers.
+     * Two {@link Unwrappable} are considered equal ignoring wrappers if they are of the same object reference
+     * after {@link #unwrapAll()}.
+     * @param other The {@link Unwrappable} to compare this {@link Unwrappable} against
+     * @return true if the argument is not {@code null}, and it represents a same object reference after
+     *         {@code unwrapAll()}, false otherwise.
+     */
+    @UnstableApi
+    default boolean equalsIgnoreWrapper(@Nullable Unwrappable other) {
+        if (other == null) {
+            return false;
+        }
+        return unwrapAll() == other.unwrapAll();
+    }
 }