diff --git a/middleware/recovery/recovery.go b/middleware/recovery/recovery.go index 66dcec182..53e67f81b 100644 --- a/middleware/recovery/recovery.go +++ b/middleware/recovery/recovery.go @@ -3,12 +3,16 @@ package recovery import ( "context" "runtime" + "time" "github.com/go-kratos/kratos/v2/errors" "github.com/go-kratos/kratos/v2/log" "github.com/go-kratos/kratos/v2/middleware" ) +// Latency is recovery latency context key +type Latency struct{} + // ErrUnknownRequest is unknown request error. var ErrUnknownRequest = errors.InternalServer("UNKNOWN", "unknown request error") @@ -41,13 +45,14 @@ func Recovery(opts ...Option) middleware.Middleware { } return func(handler middleware.Handler) middleware.Handler { return func(ctx context.Context, req interface{}) (reply interface{}, err error) { + startTime := time.Now() 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) - + ctx = context.WithValue(ctx, Latency{}, time.Since(startTime).Seconds()) err = op.handler(ctx, req, rerr) } }() diff --git a/middleware/recovery/recovery_test.go b/middleware/recovery/recovery_test.go index cf4a7fbb3..b20b5652c 100644 --- a/middleware/recovery/recovery_test.go +++ b/middleware/recovery/recovery_test.go @@ -18,7 +18,13 @@ func TestOnce(t *testing.T) { next := func(ctx context.Context, req interface{}) (interface{}, error) { panic("panic reason") } - _, e := Recovery()(next)(context.Background(), "panic") + _, e := Recovery(WithHandler(func(ctx context.Context, _, err interface{}) error { + _, ok := ctx.Value(Latency{}).(float64) + if !ok { + t.Errorf("not latency") + } + return errors.InternalServer("RECOVERY", fmt.Sprintf("panic triggered: %v", err)) + }))(next)(context.Background(), "panic") t.Logf("succ and reason is %v", e) }