github.com/grafana/pyroscope@v1.18.0/pkg/util/recovery.go (about)

     1  package util
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net/http"
     7  	"os"
     8  	"runtime"
     9  
    10  	"connectrpc.com/connect"
    11  	"github.com/grafana/dskit/httpgrpc"
    12  	"github.com/grafana/dskit/middleware"
    13  	grpc_recovery "github.com/grpc-ecosystem/go-grpc-middleware/recovery"
    14  	"github.com/prometheus/client_golang/prometheus"
    15  	"github.com/prometheus/client_golang/prometheus/promauto"
    16  
    17  	httputil "github.com/grafana/pyroscope/pkg/util/http"
    18  )
    19  
    20  const maxStacksize = 8 * 1024
    21  
    22  var (
    23  	panicTotal = promauto.NewCounter(prometheus.CounterOpts{
    24  		Namespace: "pyroscope",
    25  		Name:      "panic_total",
    26  		Help:      "The total number of panic triggered",
    27  	})
    28  
    29  	RecoveryHTTPMiddleware = middleware.Func(func(next http.Handler) http.Handler {
    30  		return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
    31  			defer func() {
    32  				if p := recover(); p != nil {
    33  					httputil.Error(w, httpgrpc.Errorf(http.StatusInternalServerError, "error while processing request: %v", PanicError(p)))
    34  				}
    35  			}()
    36  			next.ServeHTTP(w, req)
    37  		})
    38  	})
    39  
    40  	RecoveryInterceptor     recoveryInterceptor
    41  	RecoveryInterceptorGRPC = grpc_recovery.UnaryServerInterceptor(grpc_recovery.WithRecoveryHandler(PanicError))
    42  )
    43  
    44  func PanicError(p interface{}) error {
    45  	stack := make([]byte, maxStacksize)
    46  	stack = stack[:runtime.Stack(stack, true)]
    47  	// keep a multiline stack
    48  	fmt.Fprintf(os.Stderr, "panic: %v\n%s", p, stack)
    49  	panicTotal.Inc()
    50  	return fmt.Errorf("%v", p)
    51  }
    52  
    53  // RecoverPanic is a helper function to recover from panic and return an error.
    54  func RecoverPanic(f func() error) func() error {
    55  	return func() (err error) {
    56  		defer func() {
    57  			if p := recover(); p != nil {
    58  				err = PanicError(p)
    59  			}
    60  		}()
    61  		return f()
    62  	}
    63  }
    64  
    65  func Recover(f func()) {
    66  	defer func() {
    67  		if p := recover(); p != nil {
    68  			_ = PanicError(p)
    69  		}
    70  	}()
    71  	f()
    72  }
    73  
    74  type recoveryInterceptor struct{}
    75  
    76  func (recoveryInterceptor) WrapUnary(next connect.UnaryFunc) connect.UnaryFunc {
    77  	return func(ctx context.Context, req connect.AnyRequest) (resp connect.AnyResponse, err error) {
    78  		defer func() {
    79  			if p := recover(); p != nil {
    80  				err = connect.NewError(connect.CodeInternal, PanicError(p))
    81  			}
    82  		}()
    83  		return next(ctx, req)
    84  	}
    85  }
    86  
    87  func (recoveryInterceptor) WrapStreamingHandler(next connect.StreamingHandlerFunc) connect.StreamingHandlerFunc {
    88  	return func(ctx context.Context, conn connect.StreamingHandlerConn) (err error) {
    89  		defer func() {
    90  			if p := recover(); p != nil {
    91  				err = connect.NewError(connect.CodeInternal, PanicError(p))
    92  			}
    93  		}()
    94  		return next(ctx, conn)
    95  	}
    96  }
    97  
    98  func (recoveryInterceptor) WrapStreamingClient(next connect.StreamingClientFunc) connect.StreamingClientFunc {
    99  	return next
   100  }