context add exit

context.go

@@ -116,6 +116,20 @@ func (c *Context) SetState(sid, key string, v interface{}) {
 	c.Srv.state.Set(sid, key, v)
 }
 
+// Exit 终止后续逻辑执行, code 错误码
+func (c *Context) Exit(code int) {
+	c.Response.Code = code
+	panic(internalExitPanic)
+}
+
+// IfErrExit 如果 err 不为空，则中断执行并直接返回
+func (c *Context) IfErrExit(err error, code int) {
+	if err != nil {
+		c.Err(err, code)
+		c.Exit(code)
+	}
+}
+
 // Err 响应错误，如果错误对象为空则忽略不处理
 func (c *Context) Err(err error, code int) {
 	if err != nil {

middleware.go

@@ -4,3 +4,20 @@ func fillPushResp(c *Context) error {
 	c.Response.fill()
 	return nil
 }
+
+type internalPanic byte
+
+const (
+	internalExitPanic internalPanic = 1 // 应用退出
+)
+
+func internalPanicHandler(c *Context) {
+	defer func() {
+		if data := recover(); data != nil {
+			if _, ok := data.(internalPanic); !ok {
+				panic(data)
+			}
+		}
+	}()
+	c.Next()
+}

srv.go

@@ -11,14 +11,15 @@ import (
 
 // Srv 基于命令的消息处理框架
 type Srv struct {
-	Server         []ServerAdapter // 服务器适配器
-	serverMu       sync.Mutex
-	isRunning      bool                     // 服务是否已经正在运行
-	runErr         chan error               // 服务运行错误通知
-	middleware     []HandlerFunc            // 全局路由中间件
-	pushMiddleware []PushHandlerFunc        // 全局推送中间件
-	routes         map[string][]HandlerFunc // 路由的处理函数
-	state          *State                   // SID 会话的状态数据
+	Server             []ServerAdapter // 服务器适配器
+	serverMu           sync.Mutex
+	isRunning          bool                     // 服务是否已经正在运行
+	runErr             chan error               // 服务运行错误通知
+	middleware         []HandlerFunc            // 全局路由中间件
+	pushMiddleware     []PushHandlerFunc        // 全局推送中间件
+	internalMiddleware []HandlerFunc            // 内部的中间件，执行顺序在洋葱模型的最里层
+	routes             map[string][]HandlerFunc // 路由的处理函数
+	state              *State                   // SID 会话的状态数据
 }
 
 // New 指定服务器实例化一个消息服务
@@ -31,6 +32,9 @@ func New(server ...ServerAdapter) *Srv {
 	}
 	// 推送前填充数据
 	srv.UsePush(fillPushResp)
+
+	// 内部中间件
+	srv.useInternal(internalPanicHandler)
 	return srv
 }
 
@@ -67,6 +71,11 @@ func (s *Srv) Use(handlers ...HandlerFunc) *Srv {
 	return s
 }
 
+func (s *Srv) useInternal(handlers ...HandlerFunc) *Srv {
+	s.internalMiddleware = append(s.internalMiddleware, handlers...)
+	return s
+}
+
 // UsePush 增加推送中间件，该类中间件只会在使用 *Context 服务器主动推送的场景下才会被调用，如 Push, Broadcast, PushSID，在请求-响应模式时不会被调用，使用 ctx.Srv 调用也不会被触发
 func (s *Srv) UsePush(handlers ...PushHandlerFunc) *Srv {
 	s.pushMiddleware = append(s.pushMiddleware, handlers...)
@@ -173,14 +182,15 @@ func (s *Srv) NewContext(server ServerAdapter, sid string, req *Request) *Contex
 	routeHandlers, ok := s.routes[req.Cmd]
 	var handlers []HandlerFunc
 	if ok {
-		handlers = make([]HandlerFunc, 0, len(s.middleware)+len(routeHandlers))
+		handlers = make([]HandlerFunc, 0, len(s.middleware)+len(routeHandlers)+len(s.internalMiddleware))
 		handlers = append(handlers, s.middleware...)
 		handlers = append(handlers, routeHandlers...)
 		ctx.OK() // 匹配到了路由，但是 handler 没有设置响应
 	} else {
 		handlers = make([]HandlerFunc, 0, len(s.middleware)+1)
 		handlers = append(handlers, s.middleware...)
 	}
+	handlers = append(handlers, s.internalMiddleware...)
 	ctx.handlers = handlers
 	ctx.handlerIndex = -1
 	return ctx

srv_test.go

@@ -287,3 +287,53 @@ func TestSrv_State(t *testing.T) {
 		t.Assert(srv.GetState("test.sid", "tk"), nil)
 	})
 }
+
+func TestSrv_Exit(t *testing.T) {
+	srv := cs.New(&testAdapter{
+		request: []*cs.Request{
+			{"a", "1", []byte(`{"x":1}`)},
+			{"b", "2", []byte(`[{"y":2}]`)},
+			{"c", "3", nil},
+			{"d", "4", nil},
+		},
+	})
+	gtest.C(t, func(t *gtest.T) {
+		srv.Use(func(c *cs.Context) {
+			c.Set("t"+c.Cmd, "a")
+			c.Next()
+			c.Set("t"+c.Cmd, c.Get("t"+c.Cmd).(string)+"b")
+		})
+		srv.Use(cs.Recover())
+		srv.Handle("a", func(c *cs.Context) {
+			c.Set("ta", c.Get("ta").(string)+"c")
+			c.Exit(0)
+			c.Set("ta", c.Get("ta").(string)+"d")
+		})
+
+		srv.Handle("b", func(c *cs.Context) {
+			c.Set("tb", c.Get("tb").(string)+"c")
+			panic("test exit")
+		})
+
+		srv.Handle("c", func(c *cs.Context) {
+			c.Set("tc", c.Get("tc").(string)+"c")
+			c.IfErrExit(nil, 0)
+			c.Set("tc", c.Get("tc").(string)+"d")
+		})
+
+		srv.Handle("d", func(c *cs.Context) {
+			c.Set("td", c.Get("td").(string)+"c")
+			c.IfErrExit(errors.New("test error"), 0)
+			c.Set("td", c.Get("td").(string)+"d")
+		})
+
+		srv.Run()
+		time.Sleep(50 * time.Millisecond)
+
+		t.Assert(srv.GetState("test.sid", "ta"), "acb")
+		t.Assert(srv.GetState("test.sid", "tb"), "ac")
+		t.Assert(srv.GetState("test.sid", "tc"), "acdb")
+		t.Assert(srv.GetState("test.sid", "td"), "ac")
+	})
+
+}