Log response body after asynchronous Spring MVC controller method

Question

I've got the following things in my Spring MVC application:

@RestController
public class SomeController {
   @GetMapping(value = "/csv", produces = { "text/csv", MediaType.APPLICATION_JSON_VALUE })
   public Future someAsyncMethod() {
        return CompletableFuture
            .supplyAsync(() -> generateCsvSlowly()))
            .thenApply(csv -> {
                HttpHeaders httpHeaders = new HttpHeaders();
                httpHeaders.add("Content-Disposition", "attachment; filename=" + "Filename_.csv");
                httpHeaders.add("Cookie", "fileDownload=true; path=/");

                return new HttpEntity<>(csv, httpHeaders);
            });
        }
    }
}

So it simply generates csv but so slowly that I have to make this call asynchronous.

I'm trying to log all the response body in the following way:

@Component
public class LoggingFilter extends OncePerRequestFilter {

    private static final Logger LOGGER = LoggerFactory.getLogger(LoggingFilter.class);
    private static final AtomicLong ID = new AtomicLong();

    static final String SOME_FORMAT_STRING = ...;

    @Override
    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain filterChain) throws ServletException, IOException {
        long id = ID.incrementAndGet();

        HttpServletResponse responseToUse = response;
        if (!(response instanceof ContentCachingResponseWrapper)) {
            responseToUse = new ContentCachingResponseWrapper(response);
        }

        try {
            filterChain.doFilter(request, responseToUse);
        }
        finally {
            byte[] responseBodyBytes = ((ContentCachingResponseWrapper) responseToUse).getContentAsByteArray();
            LOGGER.info(SOME_FORMAT_STRING, id, responseToUse.getStatus(), responseToUse.getContentType(),
                new ServletServerHttpResponse(responseToUse).getHeaders(), new String(bodyBytes, UTF_8));
            ((ContentCachingResponseWrapper) responseToUse).copyBodyToResponse();
        }
    }

}

Here is my exception handler:

@ControllerAdvice
public class GlobalExceptionHandler extends ResponseEntityExceptionHandler {

    @ExceptionHandler(ApplicationException.class)
    @ResponseBody
    public ResponseEntity<Status> handleException(ApplicationException exception) {
        Status status = new Status();
        ...

        MultiValueMap<String, String> headers = new LinkedMultiValueMap<>();
        ...

        return new ResponseEntity(status, headers, exception.getHttpCodeMvc());
    }

    @Override
    protected ResponseEntity handleExceptionInternal(Exception ex, Object body, HttpHeaders headers, HttpStatus status, WebRequest request) {
        return handleException(new ApplicationException(ex, ApplicationStatus.GENERIC_ERROR));
    }
}

Here Status is simple POJO and ApplicationException is a custom exception class.

When generateSlowlyCsv throws an exception it gets processed in handleException but nothing get's logged and no body is returned to client. Other non-async controller methods log error (even the same one) just fine and return response body. When csv is generated (I saw it in debugger) without errors the call simply hangs and I can't find where (it returns from completable future). Without LoggingFilter everything works just fine but without logs of course.

How can I not loose response body when an exception occurred and return csv when it is generated? Thank you very much!

P.S. Returning Callable and MvcAsyncTask from the controller method doesn't help too

Answer 1

I had a similar issue.

From the spring.io blog post:

All Spring Framework Servlet filter implementations have been modified as necessary to work in asynchronous request processing. As for any other filters, some will work -- typically those that do pre-processing, and others will need to be modified -- typically those that do post-processing at the end of a request. Such filters will need to recognize when the initial Servlet container thread is being exited, making way for another thread to continue processing, and when they are invoked as part of an asynchronous dispatch to complete processing.

org.springframework.web.filter.ShallowEtagHeaderFilter is a good example.

I got this working by copying methods from ShallowEtagHeaderFilter where required and at the start of doFilter:

    if (!isAsyncDispatch(request) && !(response instanceof ContentCachingResponseWrapper)) {
        response = new ContentCachingResponseWrapper((HttpServletResponse) response);
    }

Then at the end:

    if (!isAsyncStarted(request)) {
        updateResponse(response);
    }

Answer 2

I ended with the following thing:

package com.ololo.filter;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.MethodParameter;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpOutputMessage;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.http.server.ServletServerHttpResponse;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;
import org.springframework.web.util.ContentCachingRequestWrapper;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.concurrent.atomic.AtomicLong;

import javax.servlet.http.HttpServletRequest;

import static java.lang.System.lineSeparator;
import static java.nio.charset.StandardCharsets.UTF_8;

@ControllerAdvice
public class LoggingAdvice implements ResponseBodyAdvice {

    private static final Logger LOGGER = LoggerFactory.getLogger(LoggingAdvice.class);
    private static final AtomicLong ID = new AtomicLong();

    private static final String SOME_RESPONSE_MESSAGE_FORMAT;

    @Override
    public boolean supports(MethodParameter returnType, Class converterType) {
        return true;
    }

    @Override
    public Object beforeBodyWrite(Object body, MethodParameter returnType, MediaType selectedContentType, Class selectedConverterType, ServerHttpRequest request, ServerHttpResponse response) {
        long id = ID.incrementAndGet();

        ServletServerHttpResponse responseToUse = (ServletServerHttpResponse) response;
        HttpMessageConverter httpMessageConverter;
        LoggingHttpOutboundMessageWrapper httpOutputMessage = new LoggingHttpOutboundMessageWrapper();
        try {
            httpMessageConverter = (HttpMessageConverter) selectedConverterType.newInstance();
            httpMessageConverter.write(body, selectedContentType, httpOutputMessage);
            LOGGER.info(SOME_RESPONSE_MESSAGE_FORMAT, id, responseToUse.getServletResponse().getStatus(), responseToUse.getServletResponse().getContentType(),
                    responseToUse.getHeaders(), httpOutputMessage.getResponseBodyInString());
        } catch (InstantiationException | IllegalAccessException | IOException e) {
            e.printStackTrace();
        }

        return body;
    }

    private static final class LoggingHttpOutboundMessageWrapper implements HttpOutputMessage {
        private HttpHeaders httpHeaders = new HttpHeaders();
        private ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();

        @Override
        public OutputStream getBody() throws IOException {
            return byteArrayOutputStream;
        }

        @Override
        public HttpHeaders getHeaders() {
            return httpHeaders;
        }

        public String getResponseBodyInString() {
            return new String(byteArrayOutputStream.toByteArray());
        }
    }

}

And yes, I know, it is terrible as hell but at least it is working for all the @ResponseBody controller methods. I don't think (at least for now) I need something else. I tried to write an asynchronous filter but I couldn't manage to add AsyncListener to AsyncContext (please, see the following question for details). Hope this helps you.