kpt-grpc-demo/util/grpc/interceptor/grpcsentry/sentry.go

package grpcsentry

import (
	"context"
	"fmt"

	"kpt-grpc-demo/util/xerr"

	"github.com/getsentry/sentry-go"
	grpcMiddleware "github.com/grpc-ecosystem/go-grpc-middleware"
	ctxTag "github.com/grpc-ecosystem/go-grpc-middleware/tags"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

var (
	defaultOptions = &options{
		reportDecider: defaultReportDecider,
		rePanic:       false,
	}
)

type options struct {
	reportDecider func(err error) bool
	rePanic       bool
}

func evaluateServerOpt(opts []Option) *options {
	optCopy := &options{}
	*optCopy = *defaultOptions
	for _, o := range opts {
		o(optCopy)
	}
	return optCopy
}

type Option func(*options)

// WithCodes customizes the function for mapping errors to error codes.
func WithReportDecider(f func(err error) bool) Option {
	return func(o *options) {
		o.reportDecider = f
	}
}

func WithRePanic(v bool) Option {
	return func(o *options) {
		o.rePanic = v
	}
}

func defaultReportDecider(err error) bool {
	_, isCustom := xerr.IsCustomError(err)
	if isCustom {
		return false
	}
	code := status.Code(err)
	switch code {
	case codes.Unknown, codes.Unimplemented, codes.Internal, codes.DeadlineExceeded, codes.DataLoss:
		return true
	default:
		return false
	}
}

// WithUnaryServerHandler intercept unary grpc handler and report error to sentry
func WithUnaryServerHandler(opts ...Option) grpc.UnaryServerInterceptor {
	o := evaluateServerOpt(opts)
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
		var hub *sentry.Hub
		hub, ctx = newHubForCall(ctx, info.FullMethod)
		hub.Scope().SetExtra("request", req)
		// Recover and capture panic
		defer func(ctx context.Context) {
			if rval := recover(); rval != nil {
				capturePanicWithContext(ctx, rval)
				if o.rePanic {
					panic(rval)
				}
				err = status.Error(codes.Internal, fmt.Sprint(rval))
			}
		}(ctx)

		resp, err = handler(ctx, req)
		if o.reportDecider(err) {
			reportError(ctx, err)
		}
		return resp, err
	}
}

// WithStreamServerHandler intercept stream grpc handler and report error to sentry
func WithStreamServerHandler(opts ...Option) grpc.StreamServerInterceptor {
	o := evaluateServerOpt(opts)
	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) (err error) {
		_, newCtx := newHubForCall(stream.Context(), info.FullMethod)
		wrappedStream := grpcMiddleware.WrapServerStream(stream)
		wrappedStream.WrappedContext = newCtx
		stream = wrappedStream
		// Recover and capture panic
		defer func(ctx context.Context) {
			if rval := recover(); rval != nil {
				capturePanicWithContext(ctx, rval)
				if o.rePanic {
					panic(rval)
				}
				err = status.Error(codes.Internal, fmt.Sprint(rval))
			}
		}(stream.Context())

		err = handler(srv, stream)
		if o.reportDecider(err) {
			reportError(stream.Context(), err)
		}

		return err
	}
}

// report if err != nil
func reportError(ctx context.Context, err error) {
	errCode := status.Code(err)
	hub := sentry.GetHubFromContext(ctx)
	if hub == nil {
		return
	}

	sentryExtras := ctxTag.Extract(ctx).Values()
	sentryTags := make(map[string]string)
	sentryTags["grpc.code"] = errCode.String()

	hub.ConfigureScope(func(scope *sentry.Scope) {
		scope.SetTags(sentryTags)
		scope.SetExtras(sentryExtras)
	})
	hub.CaptureException(err)
}

func capturePanicWithContext(ctx context.Context, err interface{}) {
	hub := sentry.GetHubFromContext(ctx)
	if hub == nil {
		return
	}
	hub.ConfigureScope(func(scope *sentry.Scope) {
		scope.SetExtras(ctxTag.Extract(ctx).Values())
	})
	_ = hub.RecoverWithContext(ctx, err)
}

func newHubForCall(ctx context.Context, fullMethodString string) (*sentry.Hub, context.Context) {
	hub := sentry.CurrentHub().Clone()
	hub.ConfigureScope(func(scope *sentry.Scope) {
		scope.SetTag("grpc_method", fullMethodString)
	})
	newCtx := sentry.SetHubOnContext(ctx, hub)
	return hub, newCtx
}