From 19f008b48379f0d8f0802da087d055690875e58d Mon Sep 17 00:00:00 2001 From: hoslo <952572463@qq.com> Date: Thu, 23 Feb 2023 18:26:23 +0800 Subject: [PATCH] feat: Add latency to the recover handler context (#2653) * feat: Make recovey log information more abundant * feat: Add latency to the recover handler context * feat: Add latency to the recover handler context * feat: Add latency to the recover handler context --------- Co-authored-by: shuai_yang --- middleware/recovery/recovery.go | 7 ++++++- middleware/recovery/recovery_test.go | 8 +++++++- 2 files changed, 13 insertions(+), 2 deletions(-) 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) }