@@ -7,27 +7,47 @@ import (
77
88 "github.com/gin-gonic/gin"
99 "github.com/gotomicro/ego/core/eapp"
10+ "github.com/gotomicro/ego/core/econf"
1011 "github.com/gotomicro/ego/core/transport"
1112 "github.com/spf13/cast"
1213)
1314
14- // XResCostTimer wrap gin reponse writer add start time
15+ // XResCostTimer wrap gin response writer add start time
1516type XResCostTimer struct {
1617 gin.ResponseWriter
1718 start time.Time
1819 ginCtx * gin.Context
1920 enableHeaderApp bool
21+ prefixKey string
2022}
2123
22- // 如果写入header,需要这么处理
23- // ctx.Request = ctx.Request.WithContext(sdkCtx.Context )
24+ // WriteHeader 如果写入header,需要这么处理
25+ // ctx.Request = ctx.Request.WithContext(context )
2426func (w * XResCostTimer ) WriteHeader (statusCode int ) {
2527 // header必须在c.json响应。
2628 cost := float64 (time .Since (w .start ).Microseconds ()) / 1000
2729 if w .enableHeaderApp {
28- w .Header ().Set (eapp .EgoHeaderExpose ()+ "time" , strconv .FormatFloat (cost , 'f' , - 1 , 64 )+ "|" + eapp .Name ())
30+ errStr := w .ginCtx .Errors .ByType (gin .ErrorTypePrivate ).String ()
31+ if errStr != "" {
32+ if econf .GetBool (w .prefixKey + ".enableResHeaderError" ) {
33+ w .Header ().Set (eapp .EgoHeaderExpose ()+ "time" , strconv .FormatFloat (cost , 'f' , - 1 , 64 )+ "|" + errStr + "|" + eapp .Name ())
34+ } else {
35+ w .Header ().Set (eapp .EgoHeaderExpose ()+ "time" , strconv .FormatFloat (cost , 'f' , - 1 , 64 )+ "|err|" + eapp .Name ())
36+ }
37+ } else {
38+ w .Header ().Set (eapp .EgoHeaderExpose ()+ "time" , strconv .FormatFloat (cost , 'f' , - 1 , 64 )+ "|ok|" + eapp .Name ())
39+ }
2940 } else {
30- w .Header ().Set (eapp .EgoHeaderExpose ()+ "time" , strconv .FormatFloat (cost , 'f' , - 1 , 64 ))
41+ errStr := w .ginCtx .Errors .ByType (gin .ErrorTypePrivate ).String ()
42+ if errStr != "" {
43+ if econf .GetBool (w .prefixKey + ".enableResHeaderError" ) {
44+ w .Header ().Set (eapp .EgoHeaderExpose ()+ "time" , strconv .FormatFloat (cost , 'f' , - 1 , 64 )+ "|" + errStr )
45+ } else {
46+ w .Header ().Set (eapp .EgoHeaderExpose ()+ "time" , strconv .FormatFloat (cost , 'f' , - 1 , 64 )+ "|err" )
47+ }
48+ } else {
49+ w .Header ().Set (eapp .EgoHeaderExpose ()+ "time" , strconv .FormatFloat (cost , 'f' , - 1 , 64 )+ "|ok" )
50+ }
3151 }
3252 for _ , key := range transport .CustomHeaderKeys () {
3353 if value := cast .ToString (w .ginCtx .Request .Context ().Value (key )); value != "" {
@@ -46,24 +66,14 @@ func (w *XResCostTimer) Write(b []byte) (int, error) {
4666}
4767
4868// NewXResCostTimer middleware to add X-Res-Cost-Time
49- //func NewXResCostTimer(c *gin.Context) {
50- // blw := &XResCostTimer{
51- // ResponseWriter: c.Writer,
52- // start: time.Now(),
53- // ginCtx: c,
54- // }
55- // c.Writer = blw
56- // c.Next()
57- //}
58-
59- // NewXResCostTimer middleware to add X-Res-Cost-Time
60- func NewXResCostTimer (enableHeaderApp bool ) gin.HandlerFunc {
69+ func NewXResCostTimer (prefixKey string , enableHeaderApp bool ) gin.HandlerFunc {
6170 return func (c * gin.Context ) {
6271 blw := & XResCostTimer {
6372 ResponseWriter : c .Writer ,
6473 start : time .Now (),
6574 ginCtx : c ,
6675 enableHeaderApp : enableHeaderApp ,
76+ prefixKey : prefixKey ,
6777 }
6878 c .Writer = blw
6979 c .Next ()
0 commit comments