Use responseCause instead of CapturedServiceException (#4914)

line · Jun 12, 2023 · 0b93c98 · 0b93c98
1 parent 988aee6
commit 0b93c98
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 78 deletions.
diff --git a/core/src/main/java/com/linecorp/armeria/common/logging/DefaultRequestLog.java b/core/src/main/java/com/linecorp/armeria/common/logging/DefaultRequestLog.java
@@ -1186,8 +1186,7 @@ public void responseCause(Throwable cause) {
         }
 
         requireNonNull(cause, "cause");
-        responseCause = cause;
-        updateFlags(RequestLogProperty.RESPONSE_CAUSE);
+        setResponseCause(cause, true);
     }
 
     @Override
@@ -1285,11 +1284,8 @@ public void responseContent(@Nullable Object responseContent, @Nullable Object r
                 throw new IllegalArgumentException("responseContent must be complete: " + responseContent);
             }
 
-            if (rpcResponse.cause() != null && responseCause == null) {
-                // Don't use 'responseCause(Throwable)' to set 'RpcResponse.cause()'.
-                // 'responseCause(Throwable)' performs nothing if the HTTP response was ended successfully.
-                responseCause = rpcResponse.cause();
-                updateFlags(RequestLogProperty.RESPONSE_CAUSE);
+            if (rpcResponse.cause() != null) {
+                setResponseCause(rpcResponse.cause(), true);
             }
         }
 
@@ -1388,16 +1384,25 @@ private void endResponse0(@Nullable Throwable responseCause, long responseEndTim
                 responseHeaders = DUMMY_RESPONSE_HEADERS;
             }
         }
-        if (this.responseCause == null) {
-            if (responseCause instanceof HttpStatusException ||
-                responseCause instanceof HttpResponseException) {
-                // Log the responseCause only when an Http{Status,Response}Exception was created with a cause.
-                this.responseCause = responseCause.getCause();
-            } else {
-                this.responseCause = responseCause;
+        setResponseCause(responseCause, false);
+        updateFlags(flags);
+    }
+
+    private void setResponseCause(@Nullable Throwable responseCause, boolean updateFlag) {
+        if (this.responseCause != null) {
+            return;
+        }
+        if (responseCause instanceof HttpStatusException ||
+            responseCause instanceof HttpResponseException) {
+            // Log the responseCause only when an Http{Status,Response}Exception was created with a cause.
+            responseCause = responseCause.getCause();
+        }
+        if (responseCause != null) {
+            this.responseCause = responseCause;
+            if (updateFlag) {
+                updateFlags(RequestLogProperty.RESPONSE_CAUSE);
             }
         }
-        updateFlags(flags);
     }
 
     @Override

diff --git a/core/src/main/java/com/linecorp/armeria/server/AbstractHttpResponseHandler.java b/core/src/main/java/com/linecorp/armeria/server/AbstractHttpResponseHandler.java
@@ -33,6 +33,7 @@
 import com.linecorp.armeria.common.ResponseHeaders;
 import com.linecorp.armeria.common.annotation.Nullable;
 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.internal.common.CancellationScheduler.CancellationTask;
 import com.linecorp.armeria.internal.server.DefaultServiceRequestContext;
@@ -111,7 +112,10 @@ final boolean failIfStreamOrSessionClosed() {
         // 3. After successfully flushing, the listener requests next data and
         //    the subscriber attempts to write the next data to the stream closed at 2).
         if (!isWritable()) {
-            Throwable cause = CapturedServiceException.get(reqCtx);
+            Throwable cause = null;
+            if (reqCtx.log().isAvailable(RequestLogProperty.RESPONSE_CAUSE)) {
+                cause = reqCtx.log().ensureAvailable(RequestLogProperty.RESPONSE_CAUSE).responseCause();
+            }
             if (cause == null) {
                 if (reqCtx.sessionProtocol().isMultiplex()) {
                     cause = ClosedStreamException.get();
@@ -209,14 +213,14 @@ final AggregatedHttpResponse toAggregatedHttpResponse(HttpStatusException cause)
         return response;
     }
 
-    final void endLogRequestAndResponse(Throwable cause) {
-        logBuilder().endRequest(cause);
-        logBuilder().endResponse(cause);
-    }
-
-    final void endLogRequestAndResponse() {
-        logBuilder().endRequest();
-        logBuilder().endResponse();
+    final void endLogRequestAndResponse(@Nullable Throwable cause) {
+        if (cause != null) {
+            logBuilder().endRequest(cause);
+            logBuilder().endResponse(cause);
+        } else {
+            logBuilder().endRequest();
+            logBuilder().endResponse();
+        }
     }
 
     /**

diff --git a/core/src/main/java/com/linecorp/armeria/server/AbstractHttpResponseSubscriber.java b/core/src/main/java/com/linecorp/armeria/server/AbstractHttpResponseSubscriber.java
@@ -37,6 +37,7 @@
 import com.linecorp.armeria.common.HttpObject;
 import com.linecorp.armeria.common.ResponseHeaders;
 import com.linecorp.armeria.common.annotation.Nullable;
+import com.linecorp.armeria.common.logging.RequestLogProperty;
 import com.linecorp.armeria.common.stream.CancelledSubscriptionException;
 import com.linecorp.armeria.common.stream.ClosedStreamException;
 import com.linecorp.armeria.common.util.Exceptions;
@@ -329,12 +330,11 @@ public void onComplete() {
 
     private void succeed() {
         if (tryComplete(null)) {
-            final Throwable capturedException = CapturedServiceException.get(reqCtx);
-            if (capturedException != null) {
-                endLogRequestAndResponse(capturedException);
-            } else {
-                endLogRequestAndResponse();
+            Throwable cause = null;
+            if (reqCtx.log().isAvailable(RequestLogProperty.RESPONSE_CAUSE)) {
+                cause = reqCtx.log().ensureAvailable(RequestLogProperty.RESPONSE_CAUSE).responseCause();
             }
+            endLogRequestAndResponse(cause);
             maybeWriteAccessLog();
         }
     }

diff --git a/core/src/main/java/com/linecorp/armeria/server/AggregatedHttpResponseHandler.java b/core/src/main/java/com/linecorp/armeria/server/AggregatedHttpResponseHandler.java
@@ -29,6 +29,7 @@
 import com.linecorp.armeria.common.CancellationException;
 import com.linecorp.armeria.common.EmptyHttpResponseException;
 import com.linecorp.armeria.common.annotation.Nullable;
+import com.linecorp.armeria.common.logging.RequestLogProperty;
 import com.linecorp.armeria.common.util.Exceptions;
 import com.linecorp.armeria.common.util.SafeCloseable;
 import com.linecorp.armeria.internal.common.Http1ObjectEncoder;
@@ -172,15 +173,10 @@ void handleWriteComplete(ChannelFuture future, boolean isSuccess, @Nullable Thro
         if (isSuccess) {
             logBuilder().responseFirstBytesTransferred();
             if (tryComplete(cause)) {
-                if (cause == null) {
-                    cause = CapturedServiceException.get(reqCtx);
-                }
-
-                if (cause == null) {
-                    endLogRequestAndResponse();
-                } else {
-                    endLogRequestAndResponse(cause);
+                if (cause == null && reqCtx.log().isAvailable(RequestLogProperty.RESPONSE_CAUSE)) {
+                    cause = reqCtx.log().ensureAvailable(RequestLogProperty.RESPONSE_CAUSE).responseCause();
                 }
+                endLogRequestAndResponse(cause);
                 maybeWriteAccessLog();
             }
             return;

diff --git a/core/src/main/java/com/linecorp/armeria/server/CapturedServiceException.java b/core/src/main/java/com/linecorp/armeria/server/CapturedServiceException.java
diff --git a/core/src/main/java/com/linecorp/armeria/server/HttpServerHandler.java b/core/src/main/java/com/linecorp/armeria/server/HttpServerHandler.java
@@ -371,8 +371,7 @@ private void handleRequest(ChannelHandlerContext ctx, DecodedHttpRequest req) th
             }
 
             serviceResponse = serviceResponse.recover(cause -> {
-                // Store the cause to set as the log.responseCause().
-                CapturedServiceException.set(reqCtx, cause);
+                reqCtx.logBuilder().responseCause(cause);
                 // Recover the failed response with the error handler.
                 return serviceCfg.errorHandler().onServiceException(reqCtx, cause);
             });