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

     1  package middleware
     2  
     3  import (
     4  	"fmt"
     5  	"net/http"
     6  	"runtime/debug"
     7  
     8  	"github.com/msales/pkg/v3/log"
     9  	"github.com/msales/pkg/v3/stats"
    10  )
    11  
    12  // RecoveryFunc is used to configure the recovery handler.
    13  type RecoveryFunc func(*Recovery)
    14  
    15  // WithoutStack disables the stack trace dump from the recovery log.
    16  func WithoutStack() RecoveryFunc {
    17  	return func(r *Recovery) {
    18  		r.withStack = false
    19  	}
    20  }
    21  
    22  // Recovery is a middleware that will recover from panics and logs the error.
    23  type Recovery struct {
    24  	handler   http.Handler
    25  	withStack bool
    26  }
    27  
    28  // WithRecovery recovers from panics and log the error.
    29  func WithRecovery(h http.Handler, opts ...RecoveryFunc) http.Handler {
    30  	r := &Recovery{
    31  		handler:   h,
    32  		withStack: true,
    33  	}
    34  
    35  	for _, fn := range opts {
    36  		fn(r)
    37  	}
    38  
    39  	return r
    40  }
    41  
    42  // ServeHTTP serves the request.
    43  func (m *Recovery) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    44  	defer func() {
    45  		if v := recover(); v != nil {
    46  			err := fmt.Errorf("%v", v)
    47  			if v, ok := v.(error); ok {
    48  				err = v
    49  			}
    50  
    51  			logCtx := []interface{}{"url", r.URL.String(),}
    52  			if m.withStack {
    53  				logCtx = append(logCtx, "stack", string(debug.Stack()))
    54  			}
    55  
    56  			log.Error(r.Context(), err.Error(), logCtx...)
    57  			_ = stats.Inc(r.Context(), "panic_recovery", 1, 1)
    58  			w.WriteHeader(500)
    59  		}
    60  	}()
    61  
    62  	m.handler.ServeHTTP(w, r)
    63  }