org.zalando.logbook.spring.webflux.BufferingServerHttpResponse Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of logbook-spring-webflux Show documentation

There is a newer version: 3.10.0

Show newest version

package org.zalando.logbook.spring.webflux;

import org.reactivestreams.Publisher;
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;


@SuppressWarnings({"NullableProblems"})
class BufferingServerHttpResponse extends ServerHttpResponseDecorator {
    private final ServerResponse serverResponse;

    BufferingServerHttpResponse(ServerHttpResponse delegate, ServerResponse serverResponse, Runnable writeHook) {
        super(delegate);
        this.serverResponse = serverResponse;
        beforeCommit(() -> {
            writeHook.run();
            return Mono.empty();
        });
    }

    @Override
    public Mono writeWith(Publisher body) {
        return super.writeWith(bufferingWrap(body));
    }

    private Publisher bufferingWrap(Publisher body) {
        if (serverResponse.shouldBuffer()) {
            return DataBufferCopyUtils.wrapAndBuffer(body, serverResponse::buffer);
        } else {
            return body;
        }
    }
}