Not getting response bodies logged for async 'chunked' responses

[aukevanleeuwen]

In an application in our (integration) test landscape we're using using Logbook to log requests and responses including bodies for all requests. This helps us analyze other applications. For this it is important the the complete request and response body are logged, but I'm having some issues with a Spring Boot application that has a controller method somewhat like below:

    @GetMapping(value = "/{width}x{height}.png", produces = MediaType.IMAGE_PNG_VALUE)
    public ResponseEntity<StreamingResponseBody> png(
            HttpServletRequest request,
            @PathVariable("width") @Max(5_000) int width,
            @PathVariable("height") @Max(5_000) int height) {
        ShallowEtagHeaderFilter.disableContentCaching(request);
        BufferedImage image = ImageGenerator.generate(width, height);

        StreamingResponseBody stream = outputStream -> ImageIO.write(image, "png", outputStream);

        return ResponseEntity.ok().contentType(MediaType.IMAGE_PNG).body(stream);
    }

It's a bit of an odd controller perhaps, but as I said it's used in our testing landscape to see how other applications deal with specific HTTP behaviour. So this one will return Transfer-Encoding: chunked responses because it's returning a StreamingResponseBody and it will also trigger Request#startAsync while handling it.

My response bodies however are not visible. I've debugged the code a bit and I don't quite understand some of the code in the

logbook/logbook-servlet/src/main/java/org/zalando/logbook/servlet/LogbookFilter.java

Lines 57 to 83 in 9b168df

	public void doFilter(final HttpServletRequest httpRequest, final HttpServletResponse httpResponse,
	final FilterChain chain) throws ServletException, IOException {
	final RemoteRequest request = new RemoteRequest(httpRequest, formRequestMode);
	final LocalResponse response = new LocalResponse(httpResponse, request.getProtocolVersion());
	final ResponseProcessingStage processing;
	if (request.getDispatcherType() == ASYNC) {
	processing = (ResponseProcessingStage) request.getAttribute(responseProcessingStageName);
	} else {
	processing = process(request).write();
	request.setAttribute(responseProcessingStageName, processing);
	}
	final ResponseWritingStage writing = processing.process(response);
	request.setAsyncListener(Optional.of(new LogbookAsyncListener(event -> write(request, response, writing))));
	request.setAttribute(responseWritingStageSynchronizationName, new AtomicBoolean(false));
	chain.doFilter(request, response);
	if (request.isAsyncStarted()) {
	return;
	}
	write(request, response, writing);
	}

.

To me it reads like it should roughly do the following:

1. Do the 'logbook request' write if this is NOT an async dispatch, if it IS an async dispatch: retrieve the previously stored ResponseProcessingStage.
2. An async listener is set to run 'on complete' of an async request to do the 'logbook response' write.
3. Return early if request.isAsyncStarted is true, otherwise do the 'logbook response' write.

However, I don't know if it works correctly, because

logbook/logbook-servlet/src/main/java/org/zalando/logbook/servlet/LogbookFilter.java

Line 82 in 9b168df

write(request, response, writing);

is called also for ASYNC dispatched requests, whereas I think that should only be done for the original (non ASYNC) request, right? Now it's writing early and the response stream is not consumed yet.
What happens in the end is that the write(...) from the LogbookAsyncListener#onComplete handler is actually called, but the implementation of

logbook/logbook-servlet/src/main/java/org/zalando/logbook/servlet/LogbookFilter.java

Line 85 in 9b168df

private void write(RemoteRequest request, LocalResponse response, ResponseWritingStage writing) throws IOException {

is in such a way that it will only actually write to the 'logbook' once. It's guarded by an AtomicBoolean:

logbook/logbook-servlet/src/main/java/org/zalando/logbook/servlet/LogbookFilter.java

Line 87 in 9b168df

if (attribute != null && !attribute.getAndSet(true)) {

.

I can give more debug output if you want to, but maybe it's easier to give an idea of how I think it makes more sense to me:

@Override
public void doFilter(final HttpServletRequest httpRequest, final HttpServletResponse httpResponse,
                     final FilterChain chain) throws ServletException, IOException {

    final RemoteRequest request = new RemoteRequest(httpRequest, formRequestMode);
    final LocalResponse response = new LocalResponse(httpResponse, request.getProtocolVersion());

    final ResponseProcessingStage processing;

    if (request.getDispatcherType() == ASYNC) {
        processing = (ResponseProcessingStage) request.getAttribute(responseProcessingStageName);
    } else {
        processing = process(request).write();
        request.setAttribute(responseProcessingStageName, processing);
    }
	
    // ↓↓ I'm not quite sure if this is intended to run on both regular and async dispatch (?)
    final ResponseWritingStage writing = processing.process(response);
    // ↓↓ Not sure if this is needed now to be honest
    request.setAttribute(responseWritingStageSynchronizationName, new AtomicBoolean(false));

    chain.doFilter(request, response);

    // ↓↓ Only add the listener if async was started during this request invocation?
    if (request.isAsyncStarted()) {
        request.getAsyncContext().addListener(new LogbookAsyncListener(event -> write(request, response, writing)));

        // ↓↓ return early because we've started async so we expect the write to be done async
        return;
    }

    // ↓↓ The ASYNC writing is handled by the listener attached above, return early.
    if (request.getDispatcherType() != ASYNC) {
        write(request, response, writing);
    }
}

Feedback would be much appreciated.

[aukevanleeuwen]

I've opened a PR here just to let you know where it's failing for me: #2249. However the build is probably going to break because of code coverage, but I'm not sure if I can remove that AtomicBoolean altogether. I don't think it should normally be writing responses twice anyway with the way that it's set up now?

At the very least the PR has an added test for StreamingResponseBody responses that was failing in the original implementation.