exception recovery

examples/echo/main.go

@@ -9,6 +9,7 @@ func main() {
 	upgrader := gws.NewUpgrader(&Handler{}, &gws.ServerOption{
 		CompressEnabled:  true,
 		CheckUtf8Enabled: true,
+		Caller:           gws.Recovery(gws.StdLogger),
 	})
 	http.HandleFunc("/connect", func(writer http.ResponseWriter, request *http.Request) {
 		socket, err := upgrader.Upgrade(writer, request)

init.go

@@ -5,4 +5,6 @@ import "github.com/lxzan/gws/internal"
 var (
 	myPadding  = frameHeader{}            // 帧头填充物
 	binaryPool = internal.NewBufferPool() // 缓冲池
+
+	StdLogger = new(stdLogger) // 标准日志输出
 )

option.go

@@ -74,6 +74,10 @@ type (
 		// 是否检查文本utf8编码, 关闭性能会好点
 		// Whether to check the text utf8 encoding, turn off the performance will be better
 		CheckUtf8Enabled bool
+
+		// OnMessage调用器, 可用于异常恢复
+		// OnMessage caller, can be used for exception recovery
+		Caller Caller
 	}
 
 	ServerOption struct {
@@ -93,6 +97,7 @@ type (
 		CompressThreshold   int
 		CompressorNum       int
 		CheckUtf8Enabled    bool
+		Caller              Caller
 
 		// TLS设置
 		TlsConfig *tls.Config
@@ -168,6 +173,9 @@ func initServerOption(c *ServerOption) *ServerOption {
 	if c.HandshakeTimeout <= 0 {
 		c.HandshakeTimeout = defaultHandshakeTimeout
 	}
+	if c.Caller == nil {
+		c.Caller = func(f func()) { f() }
+	}
 	c.CompressorNum = internal.ToBinaryNumber(c.CompressorNum)
 	c.deleteProtectedHeaders()
 
@@ -184,6 +192,7 @@ func initServerOption(c *ServerOption) *ServerOption {
 		CompressThreshold:   c.CompressThreshold,
 		CheckUtf8Enabled:    c.CheckUtf8Enabled,
 		CompressorNum:       c.CompressorNum,
+		Caller:              c.Caller,
 	}
 	if c.config.CompressEnabled {
 		c.config.compressors = new(compressors).initialize(c.CompressorNum, c.config.CompressLevel)
@@ -210,6 +219,7 @@ type ClientOption struct {
 	CompressLevel       int
 	CompressThreshold   int
 	CheckUtf8Enabled    bool
+	Caller              Caller
 
 	// 连接地址, 例如 wss://example.com/connect
 	// server address, eg: wss://example.com/connect
@@ -277,6 +287,9 @@ func initClientOption(c *ClientOption) *ClientOption {
 	if c.NewSessionStorage == nil {
 		c.NewSessionStorage = func() SessionStorage { return new(sliceMap) }
 	}
+	if c.Caller == nil {
+		c.Caller = func(f func()) { f() }
+	}
 	return c
 }
 
@@ -292,6 +305,7 @@ func (c *ClientOption) getConfig() *Config {
 		CompressLevel:       c.CompressLevel,
 		CompressThreshold:   c.CompressThreshold,
 		CheckUtf8Enabled:    c.CheckUtf8Enabled,
+		Caller:              c.Caller,
 		CompressorNum:       1,
 	}
 	if config.CompressEnabled {

option_test.go

@@ -24,6 +24,7 @@ func validateServerOption(as *assert.Assertions, u *Upgrader) {
 	as.Equal(config.WriteBufferSize, option.WriteBufferSize)
 	as.Equal(config.CompressorNum, option.CompressorNum)
 	as.NotNil(config.readerPool)
+	as.NotNil(config.Caller)
 
 	_, ok := u.option.NewSessionStorage().(*sliceMap)
 	as.True(ok)
@@ -42,6 +43,7 @@ func validateClientOption(as *assert.Assertions, option *ClientOption) {
 	as.Equal(config.ReadBufferSize, option.ReadBufferSize)
 	as.Equal(config.WriteBufferSize, option.WriteBufferSize)
 	as.Nil(config.readerPool)
+	as.NotNil(config.Caller)
 
 	_, ok := option.NewSessionStorage().(*sliceMap)
 	as.True(ok)

reader.go

@@ -153,9 +153,15 @@ func (c *Conn) emitMessage(msg *Message) (err error) {
 	}
 
 	if c.config.ReadAsyncEnabled {
-		c.readQueue.Go(func() { c.handler.OnMessage(c, msg) })
+		c.readQueue.Go(func() {
+			c.config.Caller(func() {
+				c.handler.OnMessage(c, msg)
+			})
+		})
 	} else {
-		c.handler.OnMessage(c, msg)
+		c.config.Caller(func() {
+			c.handler.OnMessage(c, msg)
+		})
 	}
 	return nil
 }

recovery.go

@@ -0,0 +1,35 @@
+package gws
+
+import (
+	"log"
+	"runtime"
+	"unsafe"
+)
+
+type Logger interface {
+	Error(v ...any)
+}
+
+type stdLogger struct{}
+
+func (c *stdLogger) Error(v ...any) {
+	log.Println(v...)
+}
+
+type Caller func(f func())
+
+func Recovery(logger Logger) Caller {
+	return func(f func()) {
+		defer func() {
+			if e := recover(); e != nil {
+				const size = 64 << 10
+				buf := make([]byte, size)
+				buf = buf[:runtime.Stack(buf, false)]
+				msg := *(*string)(unsafe.Pointer(&buf))
+				logger.Error(e, msg)
+			}
+		}()
+
+		f()
+	}
+}

writer_test.go

@@ -269,3 +269,20 @@ func (b broadcastHandler) OnMessage(socket *Conn, message *Message) {
 	defer message.Close()
 	b.wg.Done()
 }
+
+func TestRecovery(t *testing.T) {
+	var as = assert.New(t)
+	var serverHandler = new(webSocketMocker)
+	var clientHandler = new(webSocketMocker)
+	var serverOption = &ServerOption{Caller: Recovery(StdLogger)}
+	var clientOption = &ClientOption{}
+	serverHandler.onMessage = func(socket *Conn, message *Message) {
+		var m map[string]uint8
+		m[""] = 1
+	}
+	server, client := newPeer(serverHandler, serverOption, clientHandler, clientOption)
+	go server.ReadLoop()
+	go client.ReadLoop()
+	as.NoError(client.WriteString("hi"))
+	time.Sleep(100 * time.Millisecond)
+}