ServerHttpResponse sets the response statusCode correctly

spring/boot2-webflux-autoconfigure/build.gradle

@@ -86,6 +86,8 @@ task generateSources(type: Copy) {
         exclude '**/package-info.java'
         exclude '**/org.springframework.boot.autoconfigure.AutoConfiguration.imports'
         exclude '**/TlsUtil.java'
+        // Micrometer observation is not supported for spring-boot-2
+        exclude '**/ObservationTest.java'
     }
 
     into "${project.ext.genSrcDir}"

spring/boot2-webflux-autoconfigure/src/main/java/com/linecorp/armeria/spring/web/reactive/AbstractServerHttpResponseVersionSpecific.java

@@ -51,7 +51,14 @@ public boolean setStatusCode(@Nullable HttpStatus status) {
     @Override
     @Nullable
     public HttpStatus getStatusCode() {
-        return statusCode != null ? HttpStatus.resolve(statusCode) : null;
+        if (statusCode != null) {
+            return HttpStatus.resolve(statusCode);
+        }
+        final Integer statusCode0 = getStatusCode0();
+        if (statusCode0 != null) {
+            return HttpStatus.resolve(statusCode0);
+        }
+        return null;
     }
 
     @Override
@@ -67,6 +74,7 @@ public boolean setRawStatusCode(@Nullable Integer statusCode) {
     @Override
     @Nullable
     public Integer getRawStatusCode() {
-        return statusCode;
+        final HttpStatus statusCode = getStatusCode();
+        return statusCode != null ? statusCode.value() : null;
     }
 }

spring/boot3-webflux-autoconfigure/src/main/java/com/linecorp/armeria/spring/web/reactive/AbstractServerHttpResponse.java

@@ -303,4 +303,7 @@ protected abstract Mono<Void> writeAndFlushWithInternal(
      */
     protected void touchDataBuffer(DataBuffer buffer) {
     }
+
+    @Nullable
+    abstract Integer getStatusCode0();
 }

spring/boot3-webflux-autoconfigure/src/main/java/com/linecorp/armeria/spring/web/reactive/AbstractServerHttpResponseVersionSpecific.java

@@ -50,7 +50,14 @@ public boolean setStatusCode(@Nullable HttpStatusCode status) {
     @Override
     @Nullable
     public HttpStatusCode getStatusCode() {
-        return statusCode;
+        if (statusCode != null) {
+            return statusCode;
+        }
+        final Integer statusCode0 = getStatusCode0();
+        if (statusCode0 != null) {
+            return HttpStatusCode.valueOf(statusCode0);
+        }
+        return null;
     }
 
     @Override
@@ -62,6 +69,7 @@ public boolean setRawStatusCode(@Nullable Integer statusCode) {
     @Override
     @Nullable
     public Integer getRawStatusCode() {
+        final HttpStatusCode statusCode = getStatusCode();
         return statusCode != null ? statusCode.value() : null;
     }
 }

spring/boot3-webflux-autoconfigure/src/main/java/com/linecorp/armeria/spring/web/reactive/ArmeriaHttpHandlerAdapter.java

@@ -15,6 +15,7 @@
  */
 package com.linecorp.armeria.spring.web.reactive;
 
+import static com.linecorp.armeria.common.logging.RequestLogProperty.RESPONSE_HEADERS;
 import static java.util.Objects.requireNonNull;
 
 import java.util.concurrent.CompletableFuture;
@@ -28,6 +29,7 @@
 import com.linecorp.armeria.common.HttpStatus;
 import com.linecorp.armeria.common.MediaType;
 import com.linecorp.armeria.common.annotation.Nullable;
+import com.linecorp.armeria.common.logging.RequestLog;
 import com.linecorp.armeria.server.ServiceRequestContext;
 import com.linecorp.armeria.spring.internal.common.DataBufferFactoryWrapper;
 
@@ -70,6 +72,15 @@ Mono<Void> handle(ServiceRequestContext ctx, HttpRequest req, CompletableFuture<
                           .doOnError(cause -> {
                               logger.debug("{} Failed to handle a request", ctx, cause);
                               convertedResponse.setComplete(cause).subscribe();
+                          })
+                          .doOnCancel(() -> {
+                              // If Armeria has already returned a response (e.g. RequestTimeout)
+                              // make a best effort to reflect this in the returned response
+                              final RequestLog requestLog = ctx.log().getIfAvailable(RESPONSE_HEADERS);
+                              if (requestLog != null) {
+                                  convertedResponse.setRawStatusCode(requestLog.responseStatus().code());
+                                  convertedResponse.setComplete().subscribe();
+                              }
                           });
     }
 }

spring/boot3-webflux-autoconfigure/src/main/java/com/linecorp/armeria/spring/web/reactive/ArmeriaServerHttpResponse.java

@@ -120,6 +120,15 @@ protected void applyStatusCode() {
         }
     }
 
+    @Nullable
+    @Override
+    Integer getStatusCode0() {
+        if (future.isDone() && !future.isCompletedExceptionally()) {
+            return armeriaHeaders.status().code();
+        }
+        return null;
+    }
+
     @Override
     protected void applyHeaders() {
         getHeaders().forEach((name, values) -> armeriaHeaders.add(HttpHeaderNames.of(name), values));

spring/boot3-webflux-autoconfigure/src/test/java/com/linecorp/armeria/spring/web/reactive/ObservationTest.java

@@ -0,0 +1,137 @@
+/*
+ * 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.spring.web.reactive;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.awaitility.Awaitility.await;
+
+import java.time.Duration;
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.boot.test.context.SpringBootTest.WebEnvironment;
+import org.springframework.boot.test.web.server.LocalServerPort;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.RequestParam;
+import org.springframework.web.bind.annotation.RestController;
+
+import com.linecorp.armeria.client.WebClient;
+import com.linecorp.armeria.common.AggregatedHttpResponse;
+import com.linecorp.armeria.server.logging.LoggingService;
+import com.linecorp.armeria.spring.ArmeriaServerConfigurator;
+
+import io.micrometer.common.KeyValue;
+import io.micrometer.observation.Observation.Context;
+import io.micrometer.observation.ObservationHandler;
+import reactor.core.publisher.Mono;
+
+@SpringBootTest(webEnvironment = WebEnvironment.RANDOM_PORT)
+class ObservationTest {
+
+    private static final AtomicReference<String> ctxStatusRef = new AtomicReference<>();
+
+    @SpringBootApplication
+    @Configuration
+    static class TestConfiguration {
+        @RestController
+        static class TestController {
+            @GetMapping("/hello")
+            Mono<String> hello(@RequestParam String mode) {
+                return switch (mode) {
+                    case "throw" -> throw new RuntimeException("exception thrown");
+                    case "success" -> Mono.just("world");
+                    case "timeout" -> Mono.never();
+                    default -> throw new RuntimeException("Unexpected mode: " + mode);
+                };
+            }
+        }
+
+        @Bean
+        public ArmeriaServerConfigurator serverConfigurator() {
+            return sb -> sb.decorator(LoggingService.newDecorator())
+                           .requestTimeout(Duration.ofSeconds(1));
+        }
+
+        @Bean
+        public ObservationHandler<Context> observationHandler() {
+            return new ObservationHandler<>() {
+                @Override
+                public boolean supportsContext(Context context) {
+                    return true;
+                }
+
+                @Override
+                public void onStop(Context context) {
+                    final KeyValue keyValue = context.getLowCardinalityKeyValue("status");
+                    if (keyValue != null) {
+                        ctxStatusRef.set(keyValue.getValue());
+                    }
+                }
+            };
+        }
+    }
+
+    @LocalServerPort
+    int port;
+
+    @BeforeEach
+    void beforeEach() {
+        ctxStatusRef.set(null);
+    }
+
+    @Test
+    void ok() throws Exception {
+        final WebClient client = WebClient.of("http://127.0.0.1:" + port);
+        final AggregatedHttpResponse response = client.blocking().get("/hello?mode=success");
+        assertThat(response.status().code()).isEqualTo(200);
+        assertThat(response.contentUtf8()).isEqualTo("world");
+
+        await().untilAsserted(() -> assertThat(ctxStatusRef.get()).isNotNull());
+        final String ctxStatus = ctxStatusRef.get();
+        assertThat(ctxStatus).isNotNull();
+        assertThat(ctxStatus).isEqualTo("200");
+    }
+
+    @Test
+    void throwsException() throws Exception {
+        final WebClient client = WebClient.of("http://127.0.0.1:" + port);
+        final AggregatedHttpResponse response = client.blocking().get("/hello?mode=throw");
+        assertThat(response.status().code()).isEqualTo(500);
+
+        await().untilAsserted(() -> assertThat(ctxStatusRef.get()).isNotNull());
+        final String ctxStatus = ctxStatusRef.get();
+        assertThat(ctxStatus).isNotNull();
+        assertThat(ctxStatus).isEqualTo("500");
+    }
+
+    @Test
+    void timesOut() throws Exception {
+        final WebClient client = WebClient.of("http://127.0.0.1:" + port);
+        final AggregatedHttpResponse response = client.blocking().get("/hello?mode=timeout");
+        assertThat(response.status().code()).isEqualTo(503);
+
+        await().untilAsserted(() -> assertThat(ctxStatusRef.get()).isNotNull());
+        final String ctxStatus = ctxStatusRef.get();
+        assertThat(ctxStatus).isNotNull();
+        assertThat(ctxStatus).isEqualTo("503");
+    }
+}