Merge pull request #307 from laojianzi/perf-gin-recovery-func

perf: allow custom recovery func in egin

server/egin/config.go

@@ -4,6 +4,7 @@ import (
 	"crypto/tls"
 	"embed"
 	"fmt"
+	"net/http"
 	"sync"
 	"time"
 
@@ -52,7 +53,8 @@ type Config struct {
 	blockFallback                 func(*gin.Context)
 	resourceExtract               func(*gin.Context) string
 	aiReqResCelPrg                cel.Program
-	mu                            sync.RWMutex // mutex for EnableAccessInterceptorReq、EnableAccessInterceptorRes、AccessInterceptorReqResFilter、aiReqResCelPrg
+	mu                            sync.RWMutex     // mutex for EnableAccessInterceptorReq、EnableAccessInterceptorRes、AccessInterceptorReqResFilter、aiReqResCelPrg
+	recoveryFunc                  gin.RecoveryFunc // recoveryFunc 处理接口没有被 recover 的 panic，默认返回 500 并且没有任何 response body
 }
 
 // DefaultConfig ...
@@ -68,6 +70,7 @@ func DefaultConfig() *Config {
 		SlowLogThreshold:           xtime.Duration("500ms"),
 		EnableWebsocketCheckOrigin: false,
 		TrustedPlatform:            "",
+		recoveryFunc:               defaultRecoveryFunc,
 	}
 }
 
@@ -93,3 +96,7 @@ func (config *Config) ClientAuthType() tls.ClientAuthType {
 		return tls.NoClientCert
 	}
 }
+
+func defaultRecoveryFunc(ctx *gin.Context, _ interface{}) {
+	ctx.AbortWithStatus(http.StatusInternalServerError)
+}

server/egin/container.go

@@ -98,12 +98,14 @@ func (c *Container) Build(options ...Option) *Component {
 	for _, option := range options {
 		option(c)
 	}
+
 	server := newComponent(c.name, c.config, c.logger)
 	server.Use(healthcheck.Default())
 	server.Use(c.defaultServerInterceptor())
 	if c.config.ContextTimeout > 0 {
 		server.Use(timeoutMiddleware(c.config.ContextTimeout))
 	}
+
 	if c.config.EnableMetricInterceptor {
 		server.Use(metricServerInterceptor())
 	}
@@ -115,6 +117,7 @@ func (c *Container) Build(options ...Option) *Component {
 	if c.config.EnableSentinel {
 		server.Use(c.sentinelMiddleware())
 	}
+
 	econf.OnChange(func(newConf *econf.Configuration) {
 		c.config.mu.Lock()
 		cf := newConf.Sub(c.name)

server/egin/interceptor.go

@@ -199,7 +199,11 @@ func (c *Container) defaultServerInterceptor() gin.HandlerFunc {
 					ctx.Error(rec.(error)) // nolint: errcheck
 					ctx.Abort()
 				} else {
-					ctx.AbortWithStatus(http.StatusInternalServerError)
+					if c.config.recoveryFunc == nil {
+						c.config.recoveryFunc = defaultRecoveryFunc
+					}
+
+					c.config.recoveryFunc(ctx, rec)
 				}
 
 				event = "recover"

server/egin/interceptor_test.go

@@ -58,6 +58,49 @@ func TestPanicInHandler(t *testing.T) {
 	os.Remove(path.Join(logger.ConfigDir(), logger.ConfigName()))
 }
 
+func TestPanicInCustomHandler(t *testing.T) {
+	router := gin.New()
+	// 使用非异步日志
+	logger := elog.DefaultContainer().Build(
+		elog.WithDebug(false),
+		elog.WithEnableAddCaller(true),
+		elog.WithEnableAsync(false),
+	)
+
+	// 自定义 recover
+	var recoverFunc gin.RecoveryFunc = func(ctx *gin.Context, err interface{}) {
+		ctx.String(http.StatusInternalServerError, "%v", err)
+		ctx.Abort()
+	}
+
+	container := DefaultContainer()
+	container.Build(WithLogger(logger), WithRecoveryFunc(recoverFunc))
+
+	// 使用recover组件
+	panicMessage := "we have a panic"
+	router.Use(container.defaultServerInterceptor())
+	router.GET("/recovery", func(_ *gin.Context) {
+		panic(panicMessage)
+	})
+	// 调用触发panic的接口
+	w := performRequest(router, "GET", "/recovery")
+	logged, err := ioutil.ReadFile(path.Join(logger.ConfigDir(), logger.ConfigName()))
+	fmt.Printf("logged--------------->%+v\n", string(logged))
+	assert.Nil(t, err)
+	// TEST
+	assert.Equal(t, http.StatusInternalServerError, w.Code)
+	assert.Equal(t, panicMessage, w.Body.String())
+	var m map[string]interface{}
+	n := strings.Index(string(logged), "{")
+	err = json.Unmarshal(logged[n:], &m)
+	assert.NoError(t, err)
+	assert.Equal(t, m["event"], `recover`)
+	assert.Equal(t, m["error"], panicMessage)
+	assert.Equal(t, m["method"], `GET./recovery`)
+	assert.Contains(t, m["stack"], t.Name())
+	os.Remove(path.Join(logger.ConfigDir(), logger.ConfigName()))
+}
+
 func TestPanicWithBrokenPipe(t *testing.T) {
 	const expectCode = 204
 

server/egin/options.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/gin-gonic/gin"
+
 	"github.com/gotomicro/ego/core/elog"
 )
 
@@ -50,7 +51,7 @@ func WithSentinelBlockFallback(fn func(*gin.Context)) Option {
 	}
 }
 
-// WithTLSSessionCache 限流后的返回数据
+// WithTLSSessionCache TLS Session 缓存
 func WithTLSSessionCache(tsc tls.ClientSessionCache) Option {
 	return func(c *Container) {
 		c.config.TLSSessionCache = tsc
@@ -64,7 +65,7 @@ func WithTrustedPlatform(trustedPlatform string) Option {
 	}
 }
 
-// WithLogger 信任的Header头，获取客户端IP地址
+// WithLogger 设置 logger
 func WithLogger(logger *elog.Component) Option {
 	return func(c *Container) {
 		c.logger = logger
@@ -99,9 +100,16 @@ func WithServerWriteTimeout(timeout time.Duration) Option {
 	}
 }
 
-// WithContextTimeout 设置port
+// WithContextTimeout 设置 context 超时时间
 func WithContextTimeout(timeout time.Duration) Option {
 	return func(c *Container) {
 		c.config.ContextTimeout = timeout
 	}
 }
+
+// WithRecoveryFunc 设置 recovery func
+func WithRecoveryFunc(f gin.RecoveryFunc) Option {
+	return func(c *Container) {
+		c.config.recoveryFunc = f
+	}
+}