org.lognet.springboot.grpc.recovery.GRpcExceptionHandlerInterceptor Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of grpc-spring-boot-starter Show documentation
grpc-spring-boot-starter
The newest version!
package org.lognet.springboot.grpc.recovery;

import io.grpc.*;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicBoolean;
import org.lognet.springboot.grpc.FailureHandlingSupport;
import org.lognet.springboot.grpc.MessageBlockingServerCallListener;
import org.lognet.springboot.grpc.autoconfigure.GRpcServerProperties;
import org.springframework.core.Ordered;

public class GRpcExceptionHandlerInterceptor implements ServerInterceptor, Ordered {

  public static final Context.Key EXCEPTION_HANDLED =
      Context.key("org.lognet.springboot.grpc.recovery.EXCEPTION_HANDLED");
  private final GRpcExceptionHandlerMethodResolver methodResolver;
  private final FailureHandlingSupport failureHandlingSupport;

  private Integer order;

  public GRpcExceptionHandlerInterceptor(
      GRpcExceptionHandlerMethodResolver methodResolver,
      FailureHandlingSupport failureHandlingSupport,
      GRpcServerProperties serverProperties) {
    this.methodResolver = methodResolver;
    this.failureHandlingSupport = failureHandlingSupport;
    this.order =
        Optional.ofNullable(serverProperties.getRecovery())
            .map(GRpcServerProperties.RecoveryProperties::getInterceptorOrder)
            .orElse(Ordered.HIGHEST_PRECEDENCE);
  }

  @Override
  public  ServerCall.Listener interceptCall(
      ServerCall call, Metadata headers, ServerCallHandler next) {

    final AtomicBoolean callIsClosed = new AtomicBoolean(false);
    final AtomicBoolean exceptionHandled = new AtomicBoolean(false);

    if (!methodResolver.hasErrorHandlers()) {
      return next.startCall(call, headers);
    }

    final ForwardingServerCall.SimpleForwardingServerCall errorHandlingCall =
        new ForwardingServerCall.SimpleForwardingServerCall(call) {

          @Override
          public void close(Status status, Metadata trailers) {
            // prevent close from being  invoked twice
            //  (like from Reactive service from different thread , close method invoked directly )
            Boolean handled =
                Optional.ofNullable(EXCEPTION_HANDLED.get())
                    .map(AtomicBoolean::get)
                    .orElse(!exceptionHandled.compareAndSet(false, true));
            if (null != status.getCause() && !handled) {
              failureHandlingSupport.closeCall(
                  new GRpcRuntimeExceptionWrapper(status.getCause()), this, trailers);
            }

            if (callIsClosed.compareAndSet(false, true)) {
              super.close(status, trailers);
            }
          }

          @Override
          public void sendMessage(RespT message) {
            try {
              super.sendMessage(message);
            } catch (Exception e) {
              failureHandlingSupport.closeCall(e, this, headers, b -> b.response(message));
            }
          }
        };

    Context context = Context.current().withValue(EXCEPTION_HANDLED, new AtomicBoolean(false));
    return Contexts.interceptCall(
        context,
        errorHandlingCall,
        headers,
        new ServerCallHandler() {
          @Override
          public ServerCall.Listener startCall(
              ServerCall call, Metadata headers) {
            final ServerCall.Listener listener;
            try {

              listener = next.startCall(call, headers);
            } catch (Exception e) {
              failureHandlingSupport.closeCall(e, call, headers);
              return new ServerCall.Listener() {};
            }
            return new MessageBlockingServerCallListener(listener) {
              private ReqT request;

              @Override
              public void onMessage(ReqT message) {
                try {
                  request = message;
                  super.onMessage(message);
                } catch (Exception e) {
                  blockMessage();
                  failureHandlingSupport.closeCall(e, call, headers, b -> b.request(request));
                }
              }

              @Override
              public void onHalfClose() {
                try {
                  if (!callIsClosed.get()) {
                    super.onHalfClose();
                  }
                } catch (Exception e) {
                  failureHandlingSupport.closeCall(e, call, headers, b -> b.request(request));
                }
              }
            };
          }
        });
  }

  @Override
  public int getOrder() {
    return order;
  }
}