Introduce ContextAwareRunnable|Callable|Function|Consume|BiConsumer|BiFunction

core/src/main/java/com/linecorp/armeria/common/ContextAwareBiConsumer.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 java.util.function.BiConsumer;
+
+import com.linecorp.armeria.common.annotation.UnstableApi;
+
+/**
+ * A delegating {@link BiConsumer} that makes sure an underlying BiConsumer is
+ * executed within the {@link RequestContext}.
+ */
+@UnstableApi
+public interface ContextAwareBiConsumer<T, U> extends BiConsumer<T, U>, ContextHolder {
+
+    /**
+     * Returns a new {@link ContextAwareBiConsumer} that sets the specified {@link RequestContext}
+     * before executing an underlying {@link BiConsumer}.
+     */
+    static <T, U> ContextAwareBiConsumer of(RequestContext context, BiConsumer<T, U> action) {
+        return new DefaultContextAwareBiConsumer(context, action);
+    }
+
+    /**
+     * Returns the {@link RequestContext} that was specified when creating
+     * this {@link ContextAwareBiConsumer}.
+     */
+    @Override
+    RequestContext context();
+
+    /**
+     * Returns the {@link BiConsumer} that's executed without setting
+     * the {@link RequestContext}.
+     */
+    BiConsumer<T, U> withoutContext();
+}

core/src/main/java/com/linecorp/armeria/common/ContextAwareBiFunction.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 java.util.function.BiFunction;
+
+import com.linecorp.armeria.common.annotation.UnstableApi;
+
+/**
+ * A delegating {@link BiFunction} that makes sure an underlying BiFunction is
+ * executed within the {@link RequestContext}.
+ */
+@UnstableApi
+public interface ContextAwareBiFunction<T, U, R> extends BiFunction<T, U, R>, ContextHolder {
+
+    /**
+     * Returns a new {@link ContextAwareBiFunction} that sets the specified {@link RequestContext}
+     * before executing an underlying {@link BiFunction}.
+     */
+    static <T, U, R> ContextAwareBiFunction of(RequestContext context, BiFunction<T, U, R> function) {
+        return new DefaultContextAwareBiFunction(context, function);
+    }
+
+    /**
+     * Returns the {@link RequestContext} that was specified when creating
+     * this {@link ContextAwareBiFunction}.
+     */
+    @Override
+    RequestContext context();
+
+    /**
+     * Returns the {@link BiFunction} that's executed without setting
+     * the {@link RequestContext}.
+     */
+    BiFunction<T, U, R> withoutContext();
+}

core/src/main/java/com/linecorp/armeria/common/ContextAwareCallable.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 java.util.concurrent.Callable;
+
+import com.linecorp.armeria.common.annotation.UnstableApi;
+
+/**
+ * A delegating {@link Callable} that makes sure an underlying Callable is
+ * executed within the {@link RequestContext}.
+ */
+@UnstableApi
+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/ContextAwareConsumer.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 java.util.function.Consumer;
+
+import com.linecorp.armeria.common.annotation.UnstableApi;
+
+/**
+ * A delegating {@link Consumer} that makes sure an underlying Consumer is
+ * executed within the {@link RequestContext}.
+ */
+@UnstableApi
+public interface ContextAwareConsumer<T> extends Consumer<T>, ContextHolder {
+
+    /**
+     * Returns a new {@link ContextAwareConsumer} that sets the specified {@link RequestContext}
+     * before executing an underlying {@link Consumer}.
+     */
+    static <T, R> ContextAwareConsumer of(RequestContext context, Consumer<T> action) {
+        return new DefaultContextAwareConsumer(context, action);
+    }
+
+    /**
+     * Returns the {@link RequestContext} that was specified when creating
+     * this {@link ContextAwareConsumer}.
+     */
+    @Override
+    RequestContext context();
+
+    /**
+     * Returns the {@link Consumer} that's executed without setting
+     * the {@link RequestContext}.
+     */
+    Consumer<T> withoutContext();
+}

core/src/main/java/com/linecorp/armeria/common/ContextAwareFunction.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 java.util.function.Function;
+
+import com.linecorp.armeria.common.annotation.UnstableApi;
+
+/**
+ * A delegating {@link Function} that makes sure an underlying Function is
+ * executed within the {@link RequestContext}.
+ */
+@UnstableApi
+public interface ContextAwareFunction<T, R> extends Function<T, R>, ContextHolder {
+
+    /**
+     * Returns a new {@link ContextAwareFuture} that sets the specified {@link RequestContext}
+     * before executing an underlying {@link Function}.
+     */
+    static <T, R> ContextAwareFunction of(RequestContext context, Function<T, R> function) {
+        return new DefaultContextAwareFunction(context, function);
+    }
+
+    /**
+     * Returns the {@link RequestContext} that was specified when creating
+     * this {@link ContextAwareFunction}.
+     */
+    @Override
+    RequestContext context();
+
+    /**
+     * Returns the {@link Function} that's executed without setting
+     * the {@link RequestContext}.
+     */
+    Function<T, R> withoutContext();
+}

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

@@ -0,0 +1,48 @@
+/*
+ * 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 com.linecorp.armeria.common.annotation.UnstableApi;
+
+/**
+ * A delegating {@link Runnable} that makes sure an underlying Runnable is
+ * executed within the {@link RequestContext}.
+ */
+@UnstableApi
+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/DefaultContextAwareBiConsumer.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 java.util.function.BiConsumer;
+
+import com.linecorp.armeria.common.util.SafeCloseable;
+
+final class DefaultContextAwareBiConsumer<T, U> implements ContextAwareBiConsumer<T, U> {
+    private final RequestContext context;
+    private final BiConsumer<T, U> action;
+
+    DefaultContextAwareBiConsumer(RequestContext context, BiConsumer<T, U> action) {
+        this.context = requireNonNull(context, "context");
+        this.action = requireNonNull(action, "action");
+    }
+
+    @Override
+    public RequestContext context() {
+        return context;
+    }
+
+    @Override
+    public BiConsumer<T, U> withoutContext() {
+        return action;
+    }
+
+    @Override
+    public void accept(T t, U u) {
+        try (SafeCloseable ignored = context.push()) {
+            action.accept(t, u);
+        }
+    }
+}