You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
57 lines
1.4 KiB
57 lines
1.4 KiB
package recovery
|
|
|
|
import (
|
|
"context"
|
|
"runtime"
|
|
|
|
"github.com/go-kratos/kratos/v2/errors"
|
|
"github.com/go-kratos/kratos/v2/log"
|
|
"github.com/go-kratos/kratos/v2/middleware"
|
|
)
|
|
|
|
// ErrUnknownRequest is unknown request error.
|
|
var ErrUnknownRequest = errors.InternalServer("UNKNOWN", "unknown request error")
|
|
|
|
// HandlerFunc is recovery handler func.
|
|
type HandlerFunc func(ctx context.Context, req, err interface{}) error
|
|
|
|
// Option is recovery option.
|
|
type Option func(*options)
|
|
|
|
type options struct {
|
|
handler HandlerFunc
|
|
}
|
|
|
|
// WithHandler with recovery handler.
|
|
func WithHandler(h HandlerFunc) Option {
|
|
return func(o *options) {
|
|
o.handler = h
|
|
}
|
|
}
|
|
|
|
// Recovery is a server middleware that recovers from any panics.
|
|
func Recovery(opts ...Option) middleware.Middleware {
|
|
op := options{
|
|
handler: func(ctx context.Context, req, err interface{}) error {
|
|
return ErrUnknownRequest
|
|
},
|
|
}
|
|
for _, o := range opts {
|
|
o(&op)
|
|
}
|
|
return func(handler middleware.Handler) middleware.Handler {
|
|
return func(ctx context.Context, req interface{}) (reply interface{}, err error) {
|
|
defer func() {
|
|
if rerr := recover(); rerr != nil {
|
|
buf := make([]byte, 64<<10) //nolint:gomnd
|
|
n := runtime.Stack(buf, false)
|
|
buf = buf[:n]
|
|
log.Context(ctx).Errorf("%v: %+v\n%s\n", rerr, req, buf)
|
|
|
|
err = op.handler(ctx, req, rerr)
|
|
}
|
|
}()
|
|
return handler(ctx, req)
|
|
}
|
|
}
|
|
}
|
|
|