diff --git a/core/eapp/env.go b/core/eapp/env.go index f5b2183..5a724b0 100644 --- a/core/eapp/env.go +++ b/core/eapp/env.go @@ -44,7 +44,7 @@ func initEnv() { egoLogTimeType = "%Y-%m-%d %H:%M:%S" } egoLogEnableAddCaller = ienv.EnvOrBool(constant.EgoLogEnableAddCaller, false) - egoHeaderExpose = ienv.EnvOrStr(constant.EgoHeaderExpose, "x-expose") + egoHeaderExpose = ienv.EnvOrStr(constant.EgoHeaderExpose, "x-res-cost-") } // AppMode returns application running mode. diff --git a/examples/cost/main.go b/examples/cost/main.go new file mode 100644 index 0000000..2a46715 --- /dev/null +++ b/examples/cost/main.go @@ -0,0 +1,49 @@ +package main + +import ( + "strconv" + "time" + + "github.com/gin-gonic/gin" +) + +// XResponseTimer wrap gin reponse writer add start time +type XResponseTimer struct { + gin.ResponseWriter + start time.Time +} + +func (w *XResponseTimer) WriteHeader(statusCode int) { + duration := time.Since(w.start) + lan := strconv.FormatInt(duration.Milliseconds(), 10) + "ms" + w.Header().Set("X-Response-Time", lan) + w.ResponseWriter.WriteHeader(statusCode) +} + +func (w *XResponseTimer) Write(b []byte) (int, error) { + return w.ResponseWriter.Write(b) +} + +// NewXResponseTimer middleware to add X-Response-Time +func NewXResponseTimer(c *gin.Context) { + blw := &XResponseTimer{ResponseWriter: c.Writer, start: time.Now()} + c.Writer = blw + c.Next() +} +func main() { + r := gin.Default() + + // 全局使用耗时中间件 + r.Use(NewXResponseTimer) + + r.GET("/fast", func(c *gin.Context) { + c.JSON(200, gin.H{"message": "fast response"}) + }) + + r.GET("/slow", func(c *gin.Context) { + time.Sleep(500 * time.Millisecond) + c.JSON(200, gin.H{"message": "slow response"}) + }) + + r.Run(":8082") +} diff --git a/examples/http/server/main.go b/examples/http/server/main.go index 0590b19..413b177 100644 --- a/examples/http/server/main.go +++ b/examples/http/server/main.go @@ -1,11 +1,13 @@ package main import ( + "context" "fmt" "github.com/gin-gonic/gin" "github.com/gotomicro/ego" "github.com/gotomicro/ego/core/elog" + "github.com/gotomicro/ego/core/transport" "github.com/gotomicro/ego/server/egin" ) @@ -13,6 +15,11 @@ import ( func main() { if err := ego.New().Serve(func() *egin.Component { server := egin.Load("server.http").Build() + server.Use(func(c *gin.Context) { + c.Header("haha1", "haha") + c.Next() + c.Header("haha2", "haha") + }) server.GET("/panic", func(ctx *gin.Context) { <-ctx.Request.Context().Done() @@ -26,6 +33,8 @@ func main() { }) server.GET("/hello", func(ctx *gin.Context) { + transport.SetHeaderKeys([]string{"x-expose-1"}) + ctx.Request = ctx.Request.WithContext(context.WithValue(ctx.Request.Context(), "x-expose-1", "expose-1")) ctx.JSON(200, "Hello client: "+ctx.GetHeader("app")) }) server.POST("/hello", func(ctx *gin.Context) { diff --git a/server/egin/container.go b/server/egin/container.go index 286ccf7..4462a5c 100644 --- a/server/egin/container.go +++ b/server/egin/container.go @@ -104,6 +104,7 @@ func (c *Container) Build(options ...Option) *Component { server := newComponent(c.name, c.config, c.logger) server.Use(healthcheck.Default()) server.Use(c.defaultServerInterceptor()) + server.Use(NewXResCostTimer) if c.config.ContextTimeout > 0 { server.Use(timeoutMiddleware(c.config.ContextTimeout)) } diff --git a/server/egin/interceptor.go b/server/egin/interceptor.go index 5fa5b40..7aa690d 100644 --- a/server/egin/interceptor.go +++ b/server/egin/interceptor.go @@ -160,16 +160,6 @@ func (c *Container) defaultServerInterceptor() gin.HandlerFunc { } } - for _, key := range transport.CustomHeaderKeys() { - if value := tools.ContextValue(ctx.Request.Context(), key); value != "" { - // x-expose 需要在这里获取 - if strings.HasPrefix(key, eapp.EgoHeaderExpose()) { - // 设置到ctx response header - ctx.Writer.Header().Set(key, value) - } - } - } - if etrace.IsGlobalTracerRegistered() { fields = append(fields, elog.FieldTid(etrace.ExtractTraceID(ctx.Request.Context()))) } diff --git a/server/egin/interceptor_cost.go b/server/egin/interceptor_cost.go new file mode 100644 index 0000000..f4a565e --- /dev/null +++ b/server/egin/interceptor_cost.go @@ -0,0 +1,52 @@ +package egin + +import ( + "strconv" + "strings" + "time" + + "github.com/gin-gonic/gin" + "github.com/gotomicro/ego/core/eapp" + "github.com/gotomicro/ego/core/transport" + "github.com/spf13/cast" +) + +// XResCostTimer wrap gin reponse writer add start time +type XResCostTimer struct { + gin.ResponseWriter + start time.Time + ginCtx *gin.Context +} + +// 如果写入header,需要这么处理 +// ctx.Request = ctx.Request.WithContext(sdkCtx.Context) +func (w *XResCostTimer) WriteHeader(statusCode int) { + // header必须在c.json响应。 + cost := float64(time.Since(w.start).Microseconds()) / 1000 + w.Header().Set(eapp.EgoHeaderExpose()+"time", strconv.FormatFloat(cost, 'f', -1, 64)) + for _, key := range transport.CustomHeaderKeys() { + if value := cast.ToString(w.ginCtx.Request.Context().Value(key)); value != "" { + // x-expose 需要在这里获取 + if strings.HasPrefix(key, eapp.EgoHeaderExpose()) { + // 设置到ctx response header + w.Header().Set(key, value) + } + } + } + w.ResponseWriter.WriteHeader(statusCode) +} + +func (w *XResCostTimer) Write(b []byte) (int, error) { + return w.ResponseWriter.Write(b) +} + +// NewXResCostTimer middleware to add X-Res-Cost-Time +func NewXResCostTimer(c *gin.Context) { + blw := &XResCostTimer{ + ResponseWriter: c.Writer, + start: time.Now(), + ginCtx: c, + } + c.Writer = blw + c.Next() +} diff --git a/server/egin/options_test.go b/server/egin/options_test.go index 47d33d0..2875388 100644 --- a/server/egin/options_test.go +++ b/server/egin/options_test.go @@ -16,7 +16,7 @@ import ( func TestInterceptor(t *testing.T) { comp := DefaultContainer().Build() // healthcheck,默认中间件,监控中间件,限流中间件 - assert.Equal(t, 3, len(comp.Handlers)) + assert.Equal(t, 4, len(comp.Handlers)) } func TestWithTrustedPlatform(t *testing.T) {