Spring webflux - ServerWebExchangeDecorator code is not executed when an exception is thrown

Question

I have a Webflux application, where I have a ServerWebExchangeDecorator that decorates the request and responses. I have overrides to do some logging and then call the super methods. This is what I have in code:

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.ServerWebExchangeDecorator;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Mono;

@Component
public class LoggingWebFilter implements WebFilter {

  @Override
  public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
    return chain.filter(decorate(exchange));
  }

  private ServerWebExchange decorate(ServerWebExchange exchange) {

    final ServerHttpRequest decoratedRequest = new LoggingServerHttpRequestDecorator(exchange.getRequest());
    final ServerHttpResponse decoratedResponse = new LoggingServerHttpResponseDecorator(exchange.getResponse());

    return new ServerWebExchangeDecorator(exchange) {

      @Override
      public ServerHttpRequest getRequest() {
        return decoratedRequest;
      }

      @Override
      public ServerHttpResponse getResponse() {
        return decoratedResponse;
      }

    };
  }

}

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import reactor.core.publisher.Flux;

public class LoggingServerHttpRequestDecorator extends ServerHttpRequestDecorator {

  private static final Logger logger = LoggerFactory.getLogger(LoggingServerHttpRequestDecorator.class);

  public LoggingServerHttpRequestDecorator(ServerHttpRequest delegate) {
    super(delegate);
  }

  @Override
  public Flux<DataBuffer> getBody() {
    logger.info("getBody method");
    return super.getBody();
  }

}

import org.reactivestreams.Publisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import reactor.core.publisher.Mono;

public class LoggingServerHttpResponseDecorator extends ServerHttpResponseDecorator {

  private static final Logger logger = LoggerFactory.getLogger(LoggingServerHttpResponseDecorator.class);

  public LoggingServerHttpResponseDecorator(ServerHttpResponse delegate) {
    super(delegate);
  }

  @Override
  public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
    logger.info("writeWith method");//THIS LINE IS NOT EXECUTED WHEN AN EXCEPTION IS THROWN
    return super.writeWith(body);
  }

  @Override
  public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
    logger.info("writeAndFlushWith method");
    return super.writeAndFlushWith(body);
  }

}

When I do a happy path with a POST request, this works fine, but when an exception is thrown, the Response Decorator is omitted and my custom code is not being executed.

This is a controller code to replicate the issue:

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

@RestController
@RequestMapping("/decorator-demo")
public class DecoratorDemoController {

  /** The Constant logger. */
  private static final Logger logger = LoggerFactory.getLogger(DecoratorDemoController.class);


  @PostMapping(produces = MediaType.APPLICATION_STREAM_JSON_VALUE, consumes = MediaType.APPLICATION_STREAM_JSON_VALUE)
  public Mono<ResponseEntity<String>> postData(@RequestBody String id) {
    logger.info("attempting to post the data");
    if(id.length() == 1){
      Mono<String> created = Mono.just(id);
      return created.flatMap(vo -> Mono.just(ResponseEntity.status(HttpStatus.CREATED).body(vo)));
    }
    throw new IllegalArgumentException("String length must be 1");
  }

}

When I post a single character, I have the logs I am expecting:

LoggingServerHttpRequestDecorator  : getBody method
DecoratorDemoController            : attempting to post the data
LoggingServerHttpResponseDecorator : writeWith method

But when I post more than one character, this is the logs I am having:

LoggingServerHttpRequestDecorator  : getBody method
DecoratorDemoController            : attempting to post the data
AbstractErrorWebExceptionHandler : [0b933716]  500 Server Error for HTTP POST "/decorator-demo"

Am I doing something wrong, or missing something?

Answer 1

Try this

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import org.reactivestreams.Publisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.filter.NettyWriteResponseFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

@Component
public class ResponseInterceptorFilter implements Ordered, GlobalFilter {

  private static final Logger log = LoggerFactory.getLogger(ResponseInterceptorFilter.class);

  public ResponseInterceptorFilter(SqsPublisher sqsPublisher) {
    this.sqsPublisher = sqsPublisher;
  }

  @Override
  public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {

    ServerHttpResponse response = exchange.getResponse();
    DataBufferFactory dataBufferFactory = response.bufferFactory();
    ServerHttpResponseDecorator decoratedResponse = new ServerHttpResponseDecorator(response) {
      @Override
      public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
        if (!(body instanceof Flux)) {
          return super.writeWith(body);
        }

        Flux<? extends DataBuffer> flux = (Flux<? extends DataBuffer>) body;
        return super.writeWith(flux.buffer().map(dataBuffers -> {
          ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
          dataBuffers.forEach(buffer -> {
            // three byte copies here
            byte[] array = new byte[buffer.readableByteCount()];
            buffer.read(array);
            try {
              outputStream.write(array);
            } catch (IOException e) {
              // TODO: need catch?
            }
            DataBufferUtils.release(buffer);
          });

          byte[] write = outputStream.toByteArray();
          String responseBody = new String(write);
          log.debug("Response ----> {}", responseBody);
          response.getHeaders().setContentLength(write.length);
          return dataBufferFactory.wrap(write);
        }));
      }
    };

    ServerWebExchange serverWebExchange = exchange.mutate().response(decoratedResponse).build();
    return chain.filter(serverWebExchange);
  }

  @Override
  public int getOrder() {
    return NettyWriteResponseFilter.WRITE_RESPONSE_FILTER_ORDER - 1; // this is important
  }
}