Add TraceAbleRequestContextStorage for ease of detect context leaks

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

@@ -0,0 +1,91 @@
+/*
+ * 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.common;
+
+import org.openjdk.jmh.annotations.Benchmark;
+
+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/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

@@ -376,6 +376,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
@@ -1295,6 +1298,20 @@ 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 users provide the {@link Sampler#never()}.
+     * See {@link Sampler#of(String)} for the specification string format.</p>
+     */
+    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,19 @@ 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 users provide the {@link Sampler#never()}.
+     * See {@link Sampler#of(String)} for the specification string format.</p>
+     */
+    @Nullable
+    default Sampler<? super RequestContext> requestContextLeakDetectionSampler() {
+        return null;
+    }
 }

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

@@ -0,0 +1,116 @@
+/*
+ * 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.common;
+
+import static com.linecorp.armeria.internal.common.RequestContextUtil.newIllegalContextPushingException;
+import static java.lang.Thread.currentThread;
+import static java.util.Objects.requireNonNull;
+
+import com.linecorp.armeria.client.ClientRequestContext;
+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 io.netty.util.concurrent.FastThreadLocal;
+
+/**
+ * A {@link RequestContextStorage} which keeps track of {@link RequestContext}s, reporting pushed thread
+ * information if a {@link RequestContext} is leaked.
+ */
+@UnstableApi
+public final class LeakTracingRequestContextStorage implements RequestContextStorage {
+
+    private final RequestContextStorage delegate;
+    private final FastThreadLocal<PendingRequestContextStackTrace> pendingRequestCtx;
+    private final Sampler<Object> sampler;
+
+    /**
+     * Creates a new instance.
+     * @param delegate the underlying {@link RequestContextStorage} that stores {@link RequestContext}
+     */
+    public LeakTracingRequestContextStorage(RequestContextStorage delegate) {
+        this(delegate, Flags.verboseExceptionSampler());
+    }
+
+    /**
+     * 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
+     */
+    public LeakTracingRequestContextStorage(RequestContextStorage delegate,
+                                            Sampler<?> sampler) {
+        this.delegate = requireNonNull(delegate, "delegate");
+        pendingRequestCtx = new FastThreadLocal<>();
+        this.sampler = (Sampler<Object>) requireNonNull(sampler, "sampler");
+    }
+
+    @Nullable
+    @Override
+    public <T extends RequestContext> T push(RequestContext toPush) {
+        requireNonNull(toPush, "toPush");
+
+        final RequestContext prevContext = delegate.currentOrNull();
+        if (prevContext != null) {
+            if (prevContext == toPush) {
+                // Re-entrance
+            } else if (toPush instanceof ServiceRequestContext &&
+                       prevContext.root() == toPush) {
+                // The delegate has the ServiceRequestContext whose root() is toPush
+            } else if (toPush instanceof ClientRequestContext &&
+                       prevContext.root() == toPush.root()) {
+                // The delegate has the ClientRequestContext whose root() is the same as toPush.root()
+            } else {
+                throw newIllegalContextPushingException(prevContext, toPush, pendingRequestCtx.get());
+            }
+        }
+
+        if (sampler.isSampled(PendingRequestContextStackTrace.class)) {
+            pendingRequestCtx.set(new PendingRequestContextStackTrace(toPush, true));
+        } else {
+            pendingRequestCtx.set(new PendingRequestContextStackTrace(toPush, false));
+        }
+
+        return delegate.push(toPush);
+    }
+
+    @Override
+    public void pop(RequestContext current, @Nullable RequestContext toRestore) {
+        try {
+            delegate.pop(current, toRestore);
+        } finally {
+            pendingRequestCtx.remove();
+        }
+    }
+
+    @Nullable
+    @Override
+    public <T extends RequestContext> T currentOrNull() {
+        return delegate.currentOrNull();
+    }
+
+    static class PendingRequestContextStackTrace extends RuntimeException {
+
+        private static final long serialVersionUID = -689451606253441556L;
+
+        PendingRequestContextStackTrace(RequestContext context, boolean isSample) {
+            super("At thread [" + currentThread().getName() + "], previous RequestContext didn't popped : " +
+                  context + (isSample ? ", It is pushed at the following stacktrace" : ""), null,
+                  true, isSample);
+        }
+    }
+}

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

@@ -438,6 +438,26 @@ public Path defaultMultipartUploadsLocation() {
         return getAndParse("defaultMultipartUploadsLocation", Paths::get);
     }
 
+    @Override
+    public Sampler<? super RequestContext> requestContextLeakDetectionSampler() {
+        final String spec = getNormalized("requestContextLeakDetectionSampler");
+        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 {
+            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/internal/common/RequestContextUtil.java

@@ -33,11 +33,13 @@
 import com.linecorp.armeria.client.DefaultClientRequestContext;
 import com.linecorp.armeria.common.Flags;
 import com.linecorp.armeria.common.HttpRequest;
+import com.linecorp.armeria.common.LeakTracingRequestContextStorage;
 import com.linecorp.armeria.common.RequestContext;
 import com.linecorp.armeria.common.RequestContextStorage;
 import com.linecorp.armeria.common.RequestContextStorageProvider;
 import com.linecorp.armeria.common.annotation.Nullable;
 import com.linecorp.armeria.common.util.SafeCloseable;
+import com.linecorp.armeria.common.util.Sampler;
 import com.linecorp.armeria.server.DefaultServiceRequestContext;
 
 import io.netty.channel.ChannelFuture;
@@ -63,8 +65,14 @@ public final class RequestContextUtil {
 
     static {
         final RequestContextStorageProvider provider = Flags.requestContextStorageProvider();
+        final Sampler<? super RequestContext> sampler = Flags.requestContextLeakDetectionSampler();
         try {
-            requestContextStorage = provider.newStorage();
+            if (sampler.equals(Sampler. never())) {
+                requestContextStorage = provider.newStorage();
+            } else {
+                requestContextStorage =
+                        new LeakTracingRequestContextStorage(provider.newStorage(), sampler);
+            }
         } catch (Throwable t) {
             throw new IllegalStateException("Failed to create context storage. provider: " + provider, t);
         }
@@ -89,12 +97,22 @@ public static SafeCloseable noopSafeCloseable() {
      */
     public static IllegalStateException newIllegalContextPushingException(
             RequestContext newCtx, RequestContext oldCtx) {
+        return newIllegalContextPushingException(newCtx, oldCtx, null);
+    }
+
+    /**
+     * Returns an {@link IllegalStateException} which is raised when pushing a context from
+     * the unexpected thread or forgetting to close the previous context.
+     */
+    public static IllegalStateException newIllegalContextPushingException(
+            RequestContext newCtx, RequestContext oldCtx, @Nullable Throwable cause) {
         requireNonNull(newCtx, "newCtx");
         requireNonNull(oldCtx, "oldCtx");
-        final IllegalStateException ex = new IllegalStateException(
-                "Trying to call object wrapped with context " + newCtx + ", but context is currently " +
-                "set to " + oldCtx + ". This means the callback was called from " +
-                "unexpected thread or forgetting to close previous context.");
+        final String message = "Trying to call object wrapped with context " + newCtx + ", but context is " +
+                               "currently set to " + oldCtx + ". This means the callback was called from " +
+                               "unexpected thread or forgetting to close previous context.";
+        final IllegalStateException ex = cause == null ? new IllegalStateException(message)
+                                                       : new IllegalStateException(message, cause);
         if (REPORTED_THREADS.add(Thread.currentThread())) {
             logger.warn("An error occurred while pushing a context", ex);
         }

core/src/test/java/com/linecorp/armeria/common/FlagsTest.java

@@ -222,6 +222,32 @@ void invalidSystemPropertyVerboseExceptionSampler() throws Throwable {
                 .isEqualTo(new ExceptionSampler("rate-limit=10"));
     }
 
+    @Test
+    void defaultRequestContextLeakDetectionSampler() throws Exception {
+        final Method method = flags.getDeclaredMethod("requestContextLeakDetectionSampler");
+        assertThat(method.invoke(flags))
+                .usingRecursiveComparison()
+                .isEqualTo(Sampler.never());
+    }
+
+    @Test
+    @SetSystemProperty(key = "com.linecorp.armeria.requestContextLeakDetectionSampler", value = "always")
+    void systemPropertyRequestContextLeakDetectionSampler() throws Exception {
+        final Method method = flags.getDeclaredMethod("requestContextLeakDetectionSampler");
+        assertThat(method.invoke(flags))
+                .usingRecursiveComparison()
+                .isEqualTo(Sampler.always());
+    }
+
+    @Test
+    @SetSystemProperty(key = "com.linecorp.armeria.requestContextLeakDetectionSampler", value = "invalid-spec")
+    void invalidSystemPropertyRequestContextLeakDetectionSampler() throws Exception {
+        final Method method = flags.getDeclaredMethod("requestContextLeakDetectionSampler");
+        assertThat(method.invoke(flags))
+                .usingRecursiveComparison()
+                .isEqualTo(Sampler.never());
+    }
+
     @Test
     void testApiConsistencyBetweenFlagsAndFlagsProvider() {
         //Check method consistency between Flags and FlagsProvider excluding deprecated methods