Introduce AccessLogWriterUtil and Remove Unnecessary RequestContext Push/Pop

core/src/main/java/com/linecorp/armeria/server/AbstractHttpResponseHandler.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2021 LINE Corporation
+ * Copyright 2024 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
@@ -36,7 +36,6 @@
 import com.linecorp.armeria.common.logging.RequestLogBuilder;
 import com.linecorp.armeria.common.logging.RequestLogProperty;
 import com.linecorp.armeria.common.stream.ClosedStreamException;
-import com.linecorp.armeria.common.util.SafeCloseable;
 import com.linecorp.armeria.internal.common.CancellationScheduler.CancellationTask;
 import com.linecorp.armeria.internal.server.DefaultServiceRequestContext;
 
@@ -221,21 +220,6 @@ final void endLogRequestAndResponse(@Nullable Throwable cause) {
         }
     }
 
-    /**
-     * Writes an access log if the {@link TransientServiceOption#WITH_ACCESS_LOGGING} option is enabled for
-     * the {@link #service()}.
-     */
-    final void maybeWriteAccessLog() {
-        final ServiceConfig config = reqCtx.config();
-        if (config.transientServiceOptions().contains(TransientServiceOption.WITH_ACCESS_LOGGING)) {
-            reqCtx.log().whenComplete().thenAccept(log -> {
-                try (SafeCloseable ignored = reqCtx.push()) {
-                    config.accessLogWriter().log(log);
-                }
-            });
-        }
-    }
-
     /**
      * Schedules a request timeout.
      */

core/src/main/java/com/linecorp/armeria/server/AbstractHttpResponseSubscriber.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2016 LINE Corporation
+ * Copyright 2024 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
@@ -18,6 +18,7 @@
 
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.linecorp.armeria.internal.common.HttpHeadersUtil.mergeTrailers;
+import static com.linecorp.armeria.server.AccessLogWriterUtil.maybeWriteAccessLog;
 
 import java.nio.channels.ClosedChannelException;
 import java.util.concurrent.CompletableFuture;
@@ -339,7 +340,7 @@ private void succeed() {
                 cause = requestLog.responseCause();
             }
             endLogRequestAndResponse(cause);
-            maybeWriteAccessLog();
+            maybeWriteAccessLog(reqCtx);
         }
     }
 
@@ -348,7 +349,7 @@ void fail(Throwable cause) {
         if (tryComplete(cause)) {
             setDone(true);
             endLogRequestAndResponse(cause);
-            maybeWriteAccessLog();
+            maybeWriteAccessLog(reqCtx);
         }
     }
 

core/src/main/java/com/linecorp/armeria/server/AccessLogWriterUtil.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright 2024 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.server;
+
+import com.linecorp.armeria.common.util.SafeCloseable;
+import com.linecorp.armeria.server.logging.AccessLogWriter;
+
+/**
+ * Utility class for {@link AccessLogWriter}.
+ */
+final class AccessLogWriterUtil {
+
+    /**
+     * Writes an access log if the {@link TransientServiceOption#WITH_ACCESS_LOGGING} option is enabled
+     * for the {@link ServiceConfig#transientServiceOptions()} and the {@link ServiceConfig#accessLogWriter()}
+     * is not {@link AccessLogWriter#disabled()} for the given {@link ServiceRequestContext#config()}.
+     */
+    static void maybeWriteAccessLog(ServiceRequestContext reqCtx) {
+        final ServiceConfig config = reqCtx.config();
+        if (shouldWriteAccessLog(config)) {
+            reqCtx.log().whenComplete().thenAccept(log -> {
+                try (SafeCloseable ignored = reqCtx.push()) {
+                    config.accessLogWriter().log(log);
+                }
+            });
+        }
+    }
+
+    /**
+     * Returns whether an access log should be written.
+     *
+     */
+    private static boolean shouldWriteAccessLog(ServiceConfig config) {
+        return config.accessLogWriter() != AccessLogWriter.disabled() &&
+                config.transientServiceOptions().contains(TransientServiceOption.WITH_ACCESS_LOGGING);
+    }
+
+    private AccessLogWriterUtil() {}
+}

core/src/main/java/com/linecorp/armeria/server/AggregatedHttpResponseHandler.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2021 LINE Corporation
+ * Copyright 2024 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
@@ -17,6 +17,7 @@
 package com.linecorp.armeria.server;
 
 import static com.google.common.base.MoreObjects.firstNonNull;
+import static com.linecorp.armeria.server.AccessLogWriterUtil.maybeWriteAccessLog;
 
 import java.nio.channels.ClosedChannelException;
 import java.util.concurrent.CompletableFuture;
@@ -125,7 +126,7 @@ private void recoverAndWrite(Throwable cause) {
     void fail(Throwable cause) {
         if (tryComplete(cause)) {
             endLogRequestAndResponse(cause);
-            maybeWriteAccessLog();
+            maybeWriteAccessLog(reqCtx);
         }
     }
 
@@ -185,7 +186,7 @@ void handleWriteComplete(ChannelFuture future, boolean isSuccess, @Nullable Thro
                     }
                 }
                 endLogRequestAndResponse(cause);
-                maybeWriteAccessLog();
+                maybeWriteAccessLog(reqCtx);
             }
             return;
         }

core/src/main/java/com/linecorp/armeria/server/HttpServerHandler.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2016 LINE Corporation
+ * Copyright 2024 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
@@ -13,6 +13,7 @@
  * License for the specific language governing permissions and limitations
  * under the License.
  */
+
 package com.linecorp.armeria.server;
 
 import static com.google.common.base.MoreObjects.firstNonNull;
@@ -22,6 +23,7 @@
 import static com.linecorp.armeria.common.SessionProtocol.H2C;
 import static com.linecorp.armeria.internal.common.HttpHeadersUtil.CLOSE_STRING;
 import static com.linecorp.armeria.internal.common.RequestContextUtil.NOOP_CONTEXT_HOOK;
+import static com.linecorp.armeria.server.AccessLogWriterUtil.maybeWriteAccessLog;
 import static io.netty.handler.codec.http2.Http2CodecUtil.DEFAULT_WINDOW_SIZE;
 import static java.util.Objects.requireNonNull;
 
@@ -612,11 +614,7 @@ private ChannelFuture respond(ServiceRequestContext reqCtx, ResponseHeadersBuild
                     logBuilder.endResponse(firstNonNull(cause, f.cause()));
                 }
             }
-            reqCtx.log().whenComplete().thenAccept(log -> {
-                try (SafeCloseable ignored = reqCtx.push()) {
-                    reqCtx.config().accessLogWriter().log(log);
-                }
-            });
+            maybeWriteAccessLog(reqCtx);
         });
         return future;
     }

core/src/main/java/com/linecorp/armeria/server/logging/AccessLogWriter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018 LINE Corporation
+ * Copyright 2024 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
@@ -13,6 +13,7 @@
  * License for the specific language governing permissions and limitations
  * under the License.
  */
+
 package com.linecorp.armeria.server.logging;
 
 import static com.google.common.base.Preconditions.checkArgument;

core/src/test/java/com/linecorp/armeria/server/logging/AccessLoggerIntegrationTest.java

@@ -19,48 +19,109 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.awaitility.Awaitility.await;
 
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
 
 import com.linecorp.armeria.common.HttpResponse;
 import com.linecorp.armeria.common.RequestContext;
-import com.linecorp.armeria.common.logging.RequestLog;
+import com.linecorp.armeria.server.HttpService;
 import com.linecorp.armeria.server.ServerBuilder;
 import com.linecorp.armeria.server.ServiceRequestContext;
+import com.linecorp.armeria.server.TransientHttpService;
+import com.linecorp.armeria.server.TransientServiceOption;
 import com.linecorp.armeria.testing.junit5.server.ServerExtension;
 
 class AccessLoggerIntegrationTest {
 
-    private static final AtomicReference<RequestContext> CTX_REF = new AtomicReference<>();
+    private static final AtomicReference<RequestContext> REQUEST_CONTEXT_REFERENCE = new AtomicReference<>();
+    private static final AtomicInteger CONTEXT_HOOK_COUNTER = new AtomicInteger(0);
+
+    private static final AccessLogWriter ACCESS_LOG_WRITER = log ->
+            REQUEST_CONTEXT_REFERENCE.set(RequestContext.currentOrNull());
+
+    private static final Supplier<? extends AutoCloseable> CONTENT_HOOK = () -> {
+        CONTEXT_HOOK_COUNTER.incrementAndGet();
+        return () -> {};
+    };
+
+    private static final HttpService BASE_SERVICE = ((HttpService) (ctx, req) -> HttpResponse.of(200))
+            .decorate((delegate, ctx, req) -> {
+                ctx.hook(CONTENT_HOOK);
+                return delegate.serve(ctx, req);
+            });
 
     @RegisterExtension
-    static ServerExtension server = new ServerExtension() {
+    static final ServerExtension server = new ServerExtension() {
         @Override
-        protected void configure(ServerBuilder sb) throws Exception {
-            sb.service("/", (ctx, req) -> HttpResponse.of(200));
-            sb.accessLogWriter(new AccessLogWriter() {
-                @Override
-                public void log(RequestLog log) {
-                    CTX_REF.set(RequestContext.currentOrNull());
-                }
-            }, false);
+        protected void configure(ServerBuilder sb) {
+            sb.route().path("/default-service")
+                    .build(BASE_SERVICE);
+            sb.route().path("/default-service-with-access-log-writer")
+                    .accessLogWriter(ACCESS_LOG_WRITER, false)
+                    .build(BASE_SERVICE);
+            sb.route().path("/transit-service")
+                    .build(BASE_SERVICE.decorate(TransientHttpService.newDecorator()));
+            sb.route().path("/transit-service-with-access-logger")
+                    .accessLogWriter(ACCESS_LOG_WRITER, false)
+                    .build(BASE_SERVICE.decorate(TransientHttpService.newDecorator()));
+            sb.route().path("/transit-service-with-access-log-option")
+                    .build(BASE_SERVICE.decorate(
+                            TransientHttpService.newDecorator(TransientServiceOption.WITH_ACCESS_LOGGING))
+                    );
+            sb.route().path("/transit-service-with-access-log-option-and-access-logger")
+                    .accessLogWriter(ACCESS_LOG_WRITER, false)
+                    .build(BASE_SERVICE.decorate(
+                            TransientHttpService.newDecorator(TransientServiceOption.WITH_ACCESS_LOGGING))
+                    );
         }
     };
 
     @BeforeEach
-    void beforeEach() {
-        CTX_REF.set(null);
+    void resetState() {
+        REQUEST_CONTEXT_REFERENCE.set(null);
+        CONTEXT_HOOK_COUNTER.set(0);
     }
 
-    @Test
-    void testAccessLogger() throws Exception {
-        assertThat(server.blockingWebClient().get("/").status().code()).isEqualTo(200);
-        assertThat(server.requestContextCaptor().size()).isEqualTo(1);
+    @CsvSource({
+            "/default-service, false",
+            "/default-service-with-access-log-writer, true",
+            "/transit-service, false",
+            "/transit-service-with-access-logger, false",
+            "/transit-service-with-access-log-option, false",
+            "/transit-service-with-access-log-option-and-access-logger, true"
+    })
+    @ParameterizedTest
+    void testAccessLogger(String path, boolean shouldWriteAccessLog) throws Exception {
+        assertThat(server.blockingWebClient().get(path).status().code())
+                .as("Response status for path: %s", path)
+                .isEqualTo(200);
+
+        assertThat(server.requestContextCaptor().size())
+                .as("Expected exactly one captured context for path: %s", path)
+                .isEqualTo(1);
+
         final ServiceRequestContext ctx = server.requestContextCaptor().poll();
-        assertThat(ctx).isNotNull();
-        await().untilAsserted(() -> assertThat(CTX_REF).hasValue(ctx));
+        assertThat(ctx)
+                .as("ServiceRequestContext should not be null for path: %s", path)
+                .isNotNull();
+
+        if (shouldWriteAccessLog) {
+            await().untilAsserted(() ->
+                    assertThat(REQUEST_CONTEXT_REFERENCE)
+                            .as("Expected request context to be set for path: %s", path)
+                            .hasValue(ctx)
+            );
+        }
+
+        final int expectedHookCounter = shouldWriteAccessLog ? 1 : 0;
+        assertThat(CONTEXT_HOOK_COUNTER)
+                .as("Context hook counter mismatch for path: %s", path)
+                .hasValue(expectedHookCounter);
     }
 }