Add TraceAbleRequestContextStorage for ease of detect context leaks

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,15 +216,15 @@ 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) {
             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 (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/internal/common/LeakTracingRequestContextStorage.java

@@ -0,0 +1,134 @@
+/*
+ * 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 static java.lang.Thread.currentThread;
+import static java.util.Objects.requireNonNull;
+
+import java.io.PrintWriter;
+import java.io.StringWriter;
+
+import com.linecorp.armeria.client.ClientRequestContext;
+import com.linecorp.armeria.client.ClientRequestContextWrapper;
+import com.linecorp.armeria.common.RequestContext;
+import com.linecorp.armeria.common.RequestContextStorage;
+import com.linecorp.armeria.common.RequestContextWrapper;
+import com.linecorp.armeria.common.annotation.Nullable;
+import com.linecorp.armeria.common.annotation.UnstableApi;
+import com.linecorp.armeria.common.util.Sampler;
+import com.linecorp.armeria.server.ServiceRequestContext;
+import com.linecorp.armeria.server.ServiceRequestContextWrapper;
+
+/**
+ * A {@link RequestContextStorage} which keeps track of {@link RequestContext}s, reporting pushed thread
+ * information if a {@link RequestContext} is leaked.
+ */
+@UnstableApi
+final class LeakTracingRequestContextStorage implements RequestContextStorage {
+
+    private final RequestContextStorage delegate;
+    private final Sampler<? super RequestContext> sampler;
+
+    /**
+     * Creates a new instance.
+     * @param delegate the underlying {@link RequestContextStorage} that stores {@link RequestContext}
+     * @param sampler the {@link Sampler} that determines whether to retain the stacktrace of the context leaks
+     */
+    LeakTracingRequestContextStorage(RequestContextStorage delegate,
+                                            Sampler<? super RequestContext> sampler) {
+        this.delegate = requireNonNull(delegate, "delegate");
+        this.sampler = requireNonNull(sampler, "sampler");
+    }
+
+    @Nullable
+    @Override
+    public <T extends RequestContext> T push(RequestContext toPush) {
+        requireNonNull(toPush, "toPush");
+        if (sampler.isSampled(toPush)) {
+            return delegate.push(warpRequestContext(toPush));
+        }
+        return delegate.push(toPush);
+    }
+
+    @Override
+    public void pop(RequestContext current, @Nullable RequestContext toRestore) {
+        requireNonNull(current, "current");
+        delegate.pop(current, toRestore);
+    }
+
+    @Nullable
+    @Override
+    public <T extends RequestContext> T currentOrNull() {
+        return delegate.currentOrNull();
+    }
+
+    private static RequestContextWrapper<?> warpRequestContext(RequestContext ctx) {
+        return ctx instanceof ClientRequestContext ?
+               new TraceableClientRequestContext((ClientRequestContext) ctx)
+               : new TraceableServiceRequestContext((ServiceRequestContext) ctx);
+    }
+
+    private static final class TraceableClientRequestContext extends ClientRequestContextWrapper {
+
+        private final PendingRequestContextStackTrace stackTrace;
+
+        private TraceableClientRequestContext(ClientRequestContext delegate) {
+            super(delegate);
+            stackTrace = new PendingRequestContextStackTrace();
+        }
+
+        @Override
+        public String toString() {
+            final StringWriter sw = new StringWriter();
+            stackTrace.printStackTrace(new PrintWriter(sw));
+            return new StringBuilder().append(getClass().getSimpleName())
+                                      .append(super.toString())
+                                      .append(System.lineSeparator())
+                                      .append(sw).toString();
+        }
+    }
+
+    private static final class TraceableServiceRequestContext extends ServiceRequestContextWrapper {
+
+        private final PendingRequestContextStackTrace stackTrace;
+
+        private TraceableServiceRequestContext(ServiceRequestContext delegate) {
+            super(delegate);
+            stackTrace = new PendingRequestContextStackTrace();
+        }
+
+        @Override
+        public String toString() {
+            final StringWriter sw = new StringWriter();
+            stackTrace.printStackTrace(new PrintWriter(sw));
+            return new StringBuilder().append(getClass().getSimpleName())
+                           .append(super.toString())
+                           .append(System.lineSeparator())
+                           .append(sw).toString();
+        }
+    }
+
+    private static final class PendingRequestContextStackTrace extends RuntimeException {
+
+        private static final long serialVersionUID = -689451606253441556L;
+
+        private PendingRequestContextStackTrace() {
+            super("At thread [" + currentThread().getName() + "] previous RequestContext is pushed at " +
+                  "the following stacktrace");
+        }
+    }
+}