Introduce ContextAwareRunnable|Callable

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

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2023 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.concurrent.Callable;
+
+/**
+ * A delegating {@link Callable} that makes sure an underlying Callable is
+ * executed within the {@link RequestContext}.
+ */
+public interface ContextAwareCallable<T> extends Callable<T>, ContextHolder {
+
+    /**
+     * Returns a new {@link ContextAwareCallable} that sets the specified {@link RequestContext}
+     * before executing an underlying {@link Callable}.
+     */
+    static <T> ContextAwareCallable<T> of(RequestContext context, Callable<T> callable) {
+        return new DefaultContextAwareCallable<T>(context, callable);
+    }
+
+    /**
+     * Returns the {@link RequestContext} that was specified when creating
+     * this {@link ContextAwareCallable}.
+     */
+    @Override
+    RequestContext context();
+
+    /**
+     * Returns the {@link Callable} that's executed without setting
+     * the {@link RequestContext}.
+     */
+    Callable<T> withoutContext();
+}

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

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2023 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;
+
+/**
+ * A delegating {@link Runnable} that makes sure an underlying Runnable is
+ * executed within the {@link RequestContext}.
+ */
+public interface ContextAwareRunnable extends Runnable, ContextHolder {
+
+    /**
+     * Returns a new {@link ContextAwareRunnable} that sets the specified {@link RequestContext}
+     * before executing an underlying {@link Runnable}.
+     */
+    static ContextAwareRunnable of(RequestContext context, Runnable runnable) {
+        return new DefaultContextAwareRunnable(context, runnable);
+    }
+
+    /**
+     * Returns the {@link RequestContext} that was specified when creating
+     * this {@link ContextAwareRunnable}.
+     */
+    @Override
+    RequestContext context();
+
+    /**
+     * Returns the {@link Runnable} that's executed without setting
+     * the {@link RequestContext}.
+     */
+    Runnable withoutContext();
+}

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

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2023 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 java.util.concurrent.Callable;
+
+import com.linecorp.armeria.common.util.SafeCloseable;
+
+final class DefaultContextAwareCallable<T> implements ContextAwareCallable<T> {
+    private final RequestContext context;
+    private final Callable<T> callable;
+
+    DefaultContextAwareCallable(RequestContext context, Callable<T> callable) {
+        requireNonNull(context, "context");
+        requireNonNull(callable, "callable");
+        this.context = context;
+        this.callable = callable;
+    }
+
+    @Override
+    public RequestContext context() {
+        return context;
+    }
+
+    @Override
+    public Callable withoutContext() {
+        return callable;
+    }
+
+    @Override
+    public T call() throws Exception {
+        try (SafeCloseable ignored = context.push()) {
+            return callable.call();
+        }
+    }
+}

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

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2023 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 com.linecorp.armeria.common.util.SafeCloseable;
+
+final class DefaultContextAwareRunnable implements ContextAwareRunnable {
+    private final RequestContext context;
+    private final Runnable runnable;
+
+    DefaultContextAwareRunnable(RequestContext context, Runnable runnable) {
+        requireNonNull(context, "context");
+        requireNonNull(runnable, "runnable");
+        this.context = context;
+        this.runnable = runnable;
+    }
+
+    @Override
+    public RequestContext context() {
+        return context;
+    }
+
+    @Override
+    public Runnable withoutContext() {
+        return runnable;
+    }
+
+    @Override
+    public void run() {
+        try (SafeCloseable ignored = context.push()) {
+            runnable.run();
+        }
+    }
+}

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

@@ -597,26 +597,18 @@ default ContextAwareBlockingTaskExecutor makeContextAware(BlockingTaskExecutor e
      * Returns a {@link Callable} that makes sure the current {@link RequestContext} is set and then invokes
      * the input {@code callable}.
      */
-    default <T> Callable<T> makeContextAware(Callable<T> callable) {
+    default <T> ContextAwareCallable<T> makeContextAware(Callable<T> callable) {
         requireNonNull(callable, "callable");
-        return () -> {
-            try (SafeCloseable ignored = push()) {
-                return callable.call();
-            }
-        };
+        return ContextAwareCallable.of(this, callable);
     }
 
     /**
      * Returns a {@link Runnable} that makes sure the current {@link RequestContext} is set and then invokes
      * the input {@code runnable}.
      */
-    default Runnable makeContextAware(Runnable runnable) {
+    default ContextAwareRunnable makeContextAware(Runnable runnable) {
         requireNonNull(runnable, "runnable");
-        return () -> {
-            try (SafeCloseable ignored = push()) {
-                runnable.run();
-            }
-        };
+        return ContextAwareRunnable.of(this, runnable);
     }
 
     /**