diff --git a/socket.go b/socket.go index 715db84..624723d 100644 --- a/socket.go +++ b/socket.go @@ -3,6 +3,7 @@ package websocket_fasthttp import ( "io" "net" + "sync" "time" "git.loafle.net/commons_go/websocket_fasthttp/websocket" @@ -147,7 +148,7 @@ type Socket interface { } func newSocket(socketHandler SocketHandler, socketCTX SocketContext, conn *websocket.Conn, id string) Socket { - s := &fasthttpWebSocket{} + s := retainSocket() s.Conn = conn s.sh = socketHandler s.id = id @@ -219,6 +220,7 @@ func (s *fasthttpWebSocket) WriteMessage(messageType int, data []byte) error { func (s *fasthttpWebSocket) Close() error { err := s.Conn.Close() + releaseSocket(s) return err } @@ -278,23 +280,23 @@ func (s *fasthttpWebSocket) Close() error { // return sc.s.SetWriteDeadline(t) // } -// var socketPool sync.Pool +var socketPool sync.Pool -// func retainSocket() *fasthttpWebSocket { -// v := socketPool.Get() -// if v == nil { -// return &fasthttpWebSocket{} -// } -// return v.(*fasthttpWebSocket) -// } +func retainSocket() *fasthttpWebSocket { + v := socketPool.Get() + if v == nil { + return &fasthttpWebSocket{} + } + return v.(*fasthttpWebSocket) +} -// func releaseSocket(s *fasthttpWebSocket) { -// s.sh = nil -// s.ctx = nil -// s.id = "" +func releaseSocket(s *fasthttpWebSocket) { + s.sh = nil + s.ctx = nil + s.id = "" -// socketPool.Put(s) -// } + socketPool.Put(s) +} // var socketConnPool sync.Pool