io.quarkus.grpc.runtime.supports.context.GrpcDuplicatedContextGrpcInterceptor Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of quarkus-grpc Show documentation
Serve and consume gRPC services
There is a newer version: 3.15.1
package io.quarkus.grpc.runtime.supports.context;

import static io.quarkus.vertx.core.runtime.context.VertxContextSafetyToggle.setContextSafe;

import java.util.Optional;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;
import java.util.function.Supplier;

import jakarta.enterprise.context.ApplicationScoped;
import jakarta.enterprise.inject.spi.Prioritized;
import jakarta.inject.Inject;

import org.jboss.logging.Logger;

import io.grpc.Metadata;
import io.grpc.ServerCall;
import io.grpc.ServerCallHandler;
import io.grpc.ServerInterceptor;
import io.grpc.StatusException;
import io.quarkus.grpc.ExceptionHandlerProvider;
import io.quarkus.grpc.GlobalInterceptor;
import io.quarkus.grpc.runtime.Interceptors;
import io.smallrye.common.vertx.VertxContext;
import io.vertx.core.Context;
import io.vertx.core.Vertx;

@ApplicationScoped
@GlobalInterceptor
public class GrpcDuplicatedContextGrpcInterceptor implements ServerInterceptor, Prioritized {
    private static final Logger log = Logger.getLogger(GrpcDuplicatedContextGrpcInterceptor.class.getName());

    @Inject
    ExceptionHandlerProvider ehp;

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

        // This interceptor is called first, so, we should be on the event loop.
        Context capturedVertxContext = Vertx.currentContext();

        if (capturedVertxContext != null) {
            // If we are not on a duplicated context, create and switch.
            Context local = VertxContext.getOrCreateDuplicatedContext(capturedVertxContext);
            setContextSafe(local, true);

            // Must be sure to call next.startCall on the right context
            return new ListenedOnDuplicatedContext<>(ehp, call, () -> next.startCall(call, headers), local);
        } else {
            log.warn("Unable to run on a duplicated context - interceptor not called on the Vert.x event loop");
            return next.startCall(call, headers);
        }
    }

    @Override
    public int getPriority() {
        return Interceptors.DUPLICATE_CONTEXT;
    }

    static class ListenedOnDuplicatedContext extends ServerCall.Listener {

        private final Context context;
        private final Supplier> supplier;
        private final ExceptionHandlerProvider ehp;
        private final ServerCall call;
        private ServerCall.Listener delegate;

        private final AtomicBoolean closed = new AtomicBoolean();

        public ListenedOnDuplicatedContext(
                ExceptionHandlerProvider ehp,
                ServerCall call, Supplier> supplier, Context context) {
            this.ehp = ehp;
            this.context = context;
            this.supplier = supplier;
            this.call = call;
        }

        private synchronized ServerCall.Listener getDelegate() {
            if (delegate == null) {
                try {
                    delegate = supplier.get();
                } catch (Throwable t) {
                    // If the interceptor supplier throws an exception, catch it, and close the call.
                    log.warn("Unable to retrieve gRPC Server call listener, see the cause below.");
                    close(t);
                    return null;
                }
            }
            return delegate;
        }

        private void close(Throwable t) {
            // TODO -- "call.isRead" guards against dup calls;
            //  e.g. onComplete, after onError already closed it
            if (closed.compareAndSet(false, true) && call.isReady()) {
                // use EHP so that we're consistent with transforming any user exception
                Throwable nt = ehp.transform(t);
                StatusException sre = (StatusException) ExceptionHandlerProvider.toStatusException(nt, false);
                Optional metadata = ExceptionHandlerProvider.toTrailers(nt);
                log.warn("Closing gRPC call due to an error ...", t);
                call.close(sre.getStatus(), metadata.orElse(new Metadata()));
            }
        }

        private void invoke(Consumer> invocation) {
            if (Vertx.currentContext() == context) {
                ServerCall.Listener listener = getDelegate();
                if (listener == null) {
                    return;
                }
                try {
                    invocation.accept(listener);
                } catch (Throwable t) {
                    close(t);
                }
            } else {
                context.runOnContext(v -> {
                    ServerCall.Listener listener = ListenedOnDuplicatedContext.this.getDelegate();
                    if (listener == null) {
                        return;
                    }
                    try {
                        invocation.accept(listener);
                    } catch (Throwable t) {
                        close(t);
                    }
                });
            }
        }

        @Override
        public void onMessage(ReqT message) {
            invoke(listener -> listener.onMessage(message));
        }

        @Override
        public void onReady() {
            invoke(ServerCall.Listener::onReady);
        }

        @Override
        public void onHalfClose() {
            invoke(ServerCall.Listener::onHalfClose);
        }

        @Override
        public void onCancel() {
            invoke(ServerCall.Listener::onCancel);
        }

        @Override
        public void onComplete() {
            invoke(ServerCall.Listener::onComplete);
        }
    }
}