Spring cloud gateway response body cache

Reponse body cache

Request body에 이어 Response body도 리턴하기 전에 읽어야 하는 요구사항이 나타났습니다.

최우선으로는 구현된 library나 내부 필터가 존재하면 이용하는 방향으로 가려고 했지만 신기하게도 사용빈도가 떨어지는 이유인지 Response body의 경우 만들어진 캐싱 코드가 없는 것 같아보여 직접 구현을 해야했습니다.

@Component
public class ResponseBodyCheckFilter implements GlobalFilter, Ordered {

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        return chain.filter(
                exchange.mutate().response(new CacheServerHttpResponseDecorator(exchange)).build());
    }

    @Override
    public int getOrder() {
        return FilterOrders.CACHE_RESPONSE_BODY;
    }

    public static final class CacheServerHttpResponseDecorator extends ServerHttpResponseDecorator {

        private final WeakReference<ServerWebExchange> weakRefExchange;

        public CacheServerHttpResponseDecorator(ServerWebExchange exchange) {
            super(exchange.getResponse());
            this.weakRefExchange = new WeakReference<>(exchange);
        }

        @Override
        public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {

            ServerWebExchange exchange = this.weakRefExchange.get();
            if (!Objects.isNull(exchange)) {
                if (body instanceof Flux) {
                    body = ((Flux<DataBuffer>) body).flatMap(
                            dataBuffer -> Flux.just(updateBodyLength(exchange, dataBuffer)));
                } else if (body instanceof Mono) {
                    body = ((Mono<DataBuffer>) body).flatMap(
                            dataBuffer -> Mono.just(updateBodyLength(exchange, dataBuffer)));
                }
            }
            return super.writeWith(body);
        }

        private DataBuffer updateBodyLength(ServerWebExchange exchange, DataBuffer dataBuffer) {
            byte[] bytes = new byte[dataBuffer.readableByteCount()];
            dataBuffer.read(bytes);
            DataBufferUtils.release(dataBuffer);

            long currentLength =
                    (long) exchange.getAttributes().getOrDefault(RESPONSE_BODY_LENGTH, 0L);
            exchange.getAttributes().put(RESPONSE_BODY_LENGTH, currentLength + bytes.length);

            DataBufferFactory bufferFactory = exchange.getResponse().bufferFactory();
            return bufferFactory.wrap(bytes);

        }
    }
}

CacheResponseBodyFilter 필터를 구현했고, Order는 -2로 주었습니다.
Order는 중요합니다.

Global Ordered Filters

RemoveCachedBodyFilter / order = -2147483648

AdaptCachedBodyGlobalFilter / order = -2147482648

NettyWriteResponseFilter / order = -1

ForwardPathFilter / order = 0

GatewayMetricsFilter / order = 0

RouteToRequestUrlFilter / order = 10000

NoLoadBalancerClientFilter / order = 10150

WebsocketRoutingFilter / order = 2147483646

NettyRoutingFilter / order = 2147483647

ForwardRoutingFilter / order = 2147483647

NettyWriteResponseFilter 앞에서 CacheServerHttpResponseDecorator로 ServerHttpResponse를 감싸기 위함입니다.

'writeWith'로 넘어오는 DataBuffer는 실제 Flux의 오브젝트 이므로 형변환을 하여 DataBuffer를 읽어주었고 이를 다시 재포장(?) 하여 새로운 DataBufferFactory를 이용하여 DataBuffer를 재생성하여 기존 flow가 정상적으로 동작할 수 있도록 넘겨주었습니다.

 

여기서 DataBufferUtils.release(dataBuffer); 처음 전달 받았던 DataBuffer는 읽었으므로 release 처리를 해주었습니다. 만약 release를 해주지 않으면 leak이 발생할 것입니다.

 

사실 Cache라는 이름을 가지지만 중간 과정에서 뽑아내는 데이터를 별도로 저장하진 않았지만, 데이터를 concat하여 exchange 내부에 두고 그 후부터 가져다 쓴다면 케시의 역할을 해줄 수 있을 것 같습니다.

 

추후에 body를 두번 사용할 일이 생기면 추가 구현을 해야할 것 입니다.