github.com/msales/pkg/v3@v3.24.0/grpcx/middleware/recovery.go (about)

     1  package middleware
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"runtime/debug"
     7  
     8  	"github.com/msales/pkg/v3/log"
     9  	"google.golang.org/grpc"
    10  )
    11  
    12  // RecoveryFunc is used to configure the recovery interceptors.
    13  type RecoveryFunc func(*recoveryConfig)
    14  
    15  // WithoutStack disables the stack trace dump from the recovery log.
    16  func WithoutStack() RecoveryFunc {
    17  	return func(cfg *recoveryConfig) {
    18  		cfg.withStack = false
    19  	}
    20  }
    21  
    22  // recoveryConfig represents the configuration of the recovery interceptors.
    23  type recoveryConfig struct {
    24  	withStack bool
    25  }
    26  
    27  // newRecoveryConfig returns a new config object with sane defaults.
    28  func newRecoveryConfig(opts ...RecoveryFunc) *recoveryConfig {
    29  	cfg := &recoveryConfig{withStack: true}
    30  	cfg.applyOpts(opts)
    31  
    32  	return cfg
    33  }
    34  
    35  func (cfg *recoveryConfig) applyOpts(opts []RecoveryFunc) {
    36  	for _, fn := range opts {
    37  		fn(cfg)
    38  	}
    39  }
    40  
    41  // WithUnaryServerRecovery returns an interceptor that recovers from panics.
    42  func WithUnaryClientRecovery(recoveryOpts ...RecoveryFunc) grpc.UnaryClientInterceptor {
    43  	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
    44  		cfg := newRecoveryConfig(recoveryOpts...)
    45  
    46  		defer recoveryFunc(ctx, cfg.withStack)
    47  
    48  		return invoker(ctx, method, req, reply, cc, opts...)
    49  	}
    50  }
    51  
    52  // WithUnaryServerRecovery returns an interceptor that recovers from panics.
    53  func WithUnaryServerRecovery(opts ...RecoveryFunc) grpc.UnaryServerInterceptor {
    54  	return func(ctx context.Context, req interface{}, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
    55  		cfg := newRecoveryConfig(opts...)
    56  
    57  		defer recoveryFunc(ctx, cfg.withStack)
    58  
    59  		return handler(ctx, req)
    60  	}
    61  }
    62  
    63  // WithStreamServerRecovery returns an interceptor that recovers from panics.
    64  func WithStreamServerRecovery(opts ...RecoveryFunc) grpc.StreamServerInterceptor {
    65  	return func(srv interface{}, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
    66  		cfg := newRecoveryConfig(opts...)
    67  
    68  		defer recoveryFunc(ss.Context(), cfg.withStack)
    69  
    70  		return handler(srv, ss)
    71  	}
    72  }
    73  
    74  func recoveryFunc(ctx context.Context, withStack bool) {
    75  	if v := recover(); v != nil {
    76  		err := fmt.Errorf("%v", v)
    77  		if v, ok := v.(error); ok {
    78  			err = v
    79  		}
    80  
    81  		var logCtx []interface{}
    82  		if withStack {
    83  			logCtx = append(logCtx, "stack", string(debug.Stack()))
    84  		}
    85  
    86  		log.Error(ctx, err.Error(), logCtx...)
    87  	}
    88  }