github.com/blend/go-sdk@v1.20220411.3/grpcutil/server_recovery.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package grpcutil
     9  
    10  import (
    11  	"context"
    12  
    13  	"google.golang.org/grpc"
    14  	"google.golang.org/grpc/codes"
    15  	"google.golang.org/grpc/status"
    16  
    17  	"github.com/blend/go-sdk/ex"
    18  	"github.com/blend/go-sdk/logger"
    19  )
    20  
    21  // LoggedRecoveryHandler is a recovery handler shim.
    22  func LoggedRecoveryHandler(log logger.Log) RecoveryHandlerFunc {
    23  	return func(p interface{}) error {
    24  		logger.MaybeError(log, ex.New(p))
    25  		return status.Errorf(codes.Internal, "%+v", p)
    26  	}
    27  }
    28  
    29  type serverRecoveryOptions struct {
    30  	recoveryHandlerFunc RecoveryHandlerFunc
    31  }
    32  
    33  // ServerRecoveryOption is a type that provides a recovery option.
    34  type ServerRecoveryOption func(*serverRecoveryOptions)
    35  
    36  // WithServerRecoveryHandler customizes the function for recovering from a panic.
    37  func WithServerRecoveryHandler(f RecoveryHandlerFunc) ServerRecoveryOption {
    38  	return func(o *serverRecoveryOptions) {
    39  		o.recoveryHandlerFunc = f
    40  	}
    41  }
    42  
    43  // RecoveryHandlerFunc is a function that recovers from the panic `p` by returning an `error`.
    44  type RecoveryHandlerFunc func(p interface{}) (err error)
    45  
    46  // RecoverServerUnary returns a new unary server interceptor for panic recovery.
    47  func RecoverServerUnary(opts ...ServerRecoveryOption) grpc.UnaryServerInterceptor {
    48  	o := evaluateOptions(opts)
    49  	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (_ interface{}, err error) {
    50  		defer func() {
    51  			if r := recover(); r != nil {
    52  				err = recoverFrom(r, o.recoveryHandlerFunc)
    53  			}
    54  		}()
    55  		return handler(ctx, req)
    56  	}
    57  }
    58  
    59  // RecoverServerStream returns a new streaming server interceptor for panic recovery.
    60  func RecoverServerStream(opts ...ServerRecoveryOption) grpc.StreamServerInterceptor {
    61  	o := evaluateOptions(opts)
    62  	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) (err error) {
    63  		defer func() {
    64  			if r := recover(); r != nil {
    65  				err = recoverFrom(r, o.recoveryHandlerFunc)
    66  			}
    67  		}()
    68  
    69  		return handler(srv, stream)
    70  	}
    71  }
    72  
    73  func recoverFrom(p interface{}, r RecoveryHandlerFunc) error {
    74  	if r == nil {
    75  		return status.Errorf(codes.Internal, "%s", p)
    76  	}
    77  	return r(p)
    78  }
    79  
    80  var (
    81  	defaultOptions = &serverRecoveryOptions{
    82  		recoveryHandlerFunc: nil,
    83  	}
    84  )
    85  
    86  func evaluateOptions(opts []ServerRecoveryOption) *serverRecoveryOptions {
    87  	optCopy := new(serverRecoveryOptions)
    88  	*optCopy = *defaultOptions
    89  	for _, o := range opts {
    90  		o(optCopy)
    91  	}
    92  	return optCopy
    93  }