Sentry gRPC Performance Monitoring

SentryGrpcClientInterceptor.java

/**
The MIT License

Copyright (c) 2021 Casey Duquette, Snap Inc.

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/

import com.google.common.flogger.GoogleLogger;
import io.grpc.CallOptions;
import io.grpc.Channel;
import io.grpc.ClientCall;
import io.grpc.ClientInterceptor;
import io.grpc.ClientInterceptors;
import io.grpc.ForwardingClientCall;
import io.grpc.ForwardingClientCallListener;
import io.grpc.Metadata;
import io.grpc.MethodDescriptor;
import io.grpc.Status;
import io.sentry.ISpan;
import io.sentry.Sentry;

import javax.annotation.Nullable;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * An interceptor that applies tracing via Sentry to all client requests.
 */
public class SentryClientInterceptor implements ClientInterceptor {
    private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();

    public SentryClientInterceptor() {
        logger.atInfo().log("Creating Sentry gRPC interceptor");
    }

    /**
     * Use this interceptor to trace all requests made by this client channel.
     *
     * @param channel to be traced
     * @return intercepted channel
     */
    public Channel intercept(Channel channel) {
        return ClientInterceptors.intercept(channel, this);
    }

    @Override
    public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
            MethodDescriptor<ReqT, RespT> method,
            CallOptions callOptions,
            Channel next
    ) {
        // Name the span after the remote method name being called
        final String operationName = method.getBareMethodName() != null ? method.getBareMethodName() : method.getFullMethodName();
        final String operationMethod = method.getType().toString();

        // Get any transaction/span attached to the scope
        ISpan transaction = Sentry.getSpan();
        boolean txnStarted = false;

        // If there isn't one, start a new transaction
        if (transaction == null) {
            transaction = Sentry.startTransaction(operationName, "gRPC call started", true);
            txnStarted = true;
        }
        final ISpan finalTransaction = transaction;
        final boolean finalTxnStarted = txnStarted;

        // Create the span for this operation
        ISpan span = finalTransaction.startChild(operationName, "gRPC call started");

        // Add breadcrumb for the call
        Sentry.addBreadcrumb(SentryClient.grpcBreadcrumb(operationName, operationMethod, "call started"));

        return new ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT>(next.newCall(method, callOptions)) {
            final private AtomicBoolean finished = new AtomicBoolean(false);

            @Override
            public void start(Listener<RespT> responseListener, final Metadata headers) {
                // Create a response listener to end the span
                Listener<RespT> sentryResponseListener = new ForwardingClientCallListener.SimpleForwardingClientCallListener<RespT>(responseListener) {
                    @Override
                    public void onHeaders(Metadata headers) {
                        // Add breadcrumb for the response headers being received
                        Sentry.addBreadcrumb(SentryClient.grpcBreadcrumb(operationName, operationMethod, "received response headers"));
                        super.onHeaders(headers);
                    }

                    @Override
                    public void onMessage(RespT message) {
                        // Add breadcrumb for the response message being received
                        Sentry.addBreadcrumb(SentryClient.grpcBreadcrumb(operationName, operationMethod, "received response message"));
                        SentryClient.withSpan(span, "onMessage()", "ack received", (span) -> super.onMessage(message));
                    }

                    @Override
                    public void onClose(Status status, Metadata trailers) {
                        if (!finished.compareAndSet(false, true)) {
                            super.onClose(status, trailers);
                            return;
                        }
                        // Add breadcrumb for the response closing
                        Sentry.addBreadcrumb(SentryClient.grpcBreadcrumb(operationName, operationMethod, "call closed", status.getCode().value(), status.getCode().name()));
                        super.onClose(status, trailers);
                        span.finish();
                        if (finalTxnStarted) {
                            finalTransaction.finish();
                        }
                    }
                };

                super.start(sentryResponseListener, headers);
            }

            @Override
            public void sendMessage(ReqT message) {
                // Add breadcrumb for message sent
                Sentry.addBreadcrumb(SentryClient.grpcBreadcrumb(operationName, operationMethod, "sent message"));
                SentryClient.withSpan(span, "sendMessage()", "event sent", (span) -> super.sendMessage(message));
            }

            @Override
            public void halfClose() {
                // Add breadcrumb for all messages flushed
                Sentry.addBreadcrumb(SentryClient.grpcBreadcrumb(operationName, operationMethod, "waiting..."));
                SentryClient.withSpan(span, "halfClose()", "waiting for acks", (span) -> super.halfClose());
            }

            @Override
            public void cancel(@Nullable String message, @Nullable Throwable cause) {
                if (!finished.compareAndSet(false, true)) {
                    super.cancel(message, cause);
                    return;
                }
                Status status = cause == null ? Status.UNKNOWN : Status.fromThrowable(cause);
                // Add breadcrumb for client cancelled
                Sentry.addBreadcrumb(SentryClient.grpcBreadcrumb(operationName, operationMethod, "call cancelled", status.getCode().value(), status.getCode().name()));
                try {
                    SentryClient.withSpan(span, "cancel()", "shutdown", (span) -> super.cancel(message, cause));
                } finally {
                    span.finish();
                    if (finalTxnStarted) {
                        finalTransaction.finish();
                    }
                }
            }
        };
    }
}