gRPC(Java)でレスポンスをキャッシュしてみる

諸事情でgRPCのレスポンスをキャッシュする方法を調べたので、メモしておきます。

実現方法

以前以下のエントリーで調べたClientInterceptorを使って実現してみました。

ClientInterceptorの実装例

お試し版なので・・・・Referで始まるメソッドを呼び出した結果をConcurrentMapにキャッシュし、同じパラメータで該当メソッドが呼び出されたらサーバへアクセスせずにキャッシュから復元して呼び出しもとへ結果を返却するようにしています。

package com.example.demo;

import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.MessageOrBuilder;
import com.google.protobuf.util.JsonFormat;
import io.grpc.ForwardingClientCallListener;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import io.grpc.CallOptions;
import io.grpc.Channel;
import io.grpc.ClientCall;
import io.grpc.ClientInterceptor;
import io.grpc.ForwardingClientCall;
import io.grpc.Metadata;
import io.grpc.MethodDescriptor;
import io.grpc.Status;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * Interceptor for caching a response of refer method.
 */
public class CachingClientInterceptor implements ClientInterceptor {
  private static final Logger LOG = LoggerFactory.getLogger(CachingClientInterceptor.class);

  private final ConcurrentMap<String, Object> responseCache = new ConcurrentHashMap<>();

  private final JsonFormat.Printer requestPrinter = JsonFormat.printer().sortingMapKeys()
      .omittingInsignificantWhitespace().printingEnumsAsInts();

  @SuppressWarnings("java:S119")
  @Override
  public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
      MethodDescriptor<ReqT, RespT> method,
      CallOptions callOptions,
      Channel next) {
    return new ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT>(next.newCall(method, callOptions)) {
      private ClientCall.Listener<RespT> responseListener;

      private String cacheKey;

      private boolean restoredFromCache;

      @Override
      public void start(io.grpc.ClientCall.Listener<RespT> responseListener, Metadata headers) {
        LOG.trace("[{}] start(headers={})", method.getFullMethodName(), headers);
        this.responseListener = responseListener;
        super.start(new ForwardingClientCallListener.SimpleForwardingClientCallListener<RespT>(responseListener) {
          @Override
          public void onMessage(RespT message) {
            LOG.trace("[{}] onMessage() response={}", method.getFullMethodName(), message);
            String bareMethodName = method.getBareMethodName();
            if (bareMethodName != null && bareMethodName.startsWith("Refer")) {
              // ★★★★★ キャッシュする
              responseCache.put(cacheKey, message);
            }
            super.onMessage(message);
          }

        }, headers);
      }

      @Override
      public void halfClose() {
        if (!restoredFromCache) {
          super.halfClose();
        }
      }

      @Override
      public void sendMessage(ReqT message) {
        try {
          this.cacheKey = method.getFullMethodName() + "_" + requestPrinter.print((MessageOrBuilder) message);
        }
        catch (InvalidProtocolBufferException e) {
          throw new IllegalStateException(e);
        }
        if (responseCache.containsKey(cacheKey)) {
          // ★★★★★ キャッシュから復元する
          @SuppressWarnings("unchecked")
          RespT response = (RespT) responseCache.get(cacheKey);
          responseListener.onMessage(response);
          responseListener.onClose(Status.OK, null);
          cancel("", null);
          this.restoredFromCache = true;
          LOG.trace("[{}] restored from cache response={}", method.getFullMethodName(), response);
        } else {
          LOG.trace("[{}] sendMessage() request={}", method.getFullMethodName(), message);
          super.sendMessage(message);
        }
      }
    };

  }

}

まとめ

実アプリでキャッシュする場合は、キャッシュの削除などもっと多くのことを考慮する必要があると思いますが、CachingClientInterceptorを使ってレスポンスのキャッシュができそうです。