package websocket_fasthttp import ( "io" "net" "sync" "time" "git.loafle.net/commons_go/websocket_fasthttp/websocket" ) func newSocket(id string, conn *websocket.Conn, sh SocketHandler) *Socket { s := retainSocket() s.Conn = conn s.sh = sh s.id = id s.SetReadLimit(sh.GetMaxMessageSize()) if 0 < sh.GetReadTimeout() { s.SetReadDeadline(time.Now().Add(sh.GetReadTimeout() * time.Second)) } return s } type Socket struct { *websocket.Conn sh SocketHandler id string attributes map[interface{}]interface{} sc *SocketConn } func (s *Socket) ID() string { return s.id } func (s *Socket) GetAttribute(key interface{}) interface{} { if nil == s.attributes { return nil } return s.attributes[key] } func (s *Socket) SetAttribute(key interface{}, value interface{}) { if nil == s.attributes { s.attributes = make(map[interface{}]interface{}) } s.attributes[key] = value } func (s *Socket) WaitRequest() (*SocketConn, error) { if nil != s.sc { releaseSocketConn(s.sc) s.sc = nil } var mt int var err error var r io.Reader if mt, r, err = s.NextReader(); nil != err { return nil, err } s.sc = retainSocketConn() s.sc.s = s s.sc.MessageType = mt s.sc.r = r return s.sc, nil } func (s *Socket) NextWriter(messageType int) (io.WriteCloser, error) { if 0 < s.sh.GetWriteTimeout() { s.SetWriteDeadline(time.Now().Add(s.sh.GetWriteTimeout() * time.Second)) } return s.Conn.NextWriter(messageType) } func (s *Socket) WriteMessage(messageType int, data []byte) error { if 0 < s.sh.GetWriteTimeout() { s.SetWriteDeadline(time.Now().Add(s.sh.GetWriteTimeout() * time.Second)) } return s.Conn.WriteMessage(messageType, data) } func (s *Socket) Close() error { err := s.Conn.Close() releaseSocket(s) return err } type SocketConn struct { net.Conn s *Socket MessageType int r io.Reader wc io.WriteCloser } func (sc *SocketConn) Read(b []byte) (n int, err error) { return sc.r.Read(b) } func (sc *SocketConn) Write(b []byte) (n int, err error) { if nil == sc.wc { var err error if sc.wc, err = sc.s.NextWriter(sc.MessageType); nil != err { return 0, err } } return sc.wc.Write(b) } func (sc *SocketConn) Close() error { var err error if sc.wc != nil { err = sc.wc.Close() } sc.s.sc = nil releaseSocketConn(sc) return err } func (sc *SocketConn) LocalAddr() net.Addr { return sc.s.LocalAddr() } func (sc *SocketConn) RemoteAddr() net.Addr { return sc.s.RemoteAddr() } func (sc *SocketConn) SetDeadline(t time.Time) error { if err := sc.s.SetReadDeadline(t); nil != err { return err } if err := sc.s.SetWriteDeadline(t); nil != err { return err } return nil } func (sc *SocketConn) SetReadDeadline(t time.Time) error { return sc.s.SetReadDeadline(t) } func (sc *SocketConn) SetWriteDeadline(t time.Time) error { return sc.s.SetWriteDeadline(t) } var socketPool sync.Pool func retainSocket() *Socket { v := socketPool.Get() if v == nil { return &Socket{} } return v.(*Socket) } func releaseSocket(s *Socket) { s.sh = nil s.sc = nil s.id = "" socketPool.Put(s) } var socketConnPool sync.Pool func retainSocketConn() *SocketConn { v := socketConnPool.Get() if v == nil { return &SocketConn{} } return v.(*SocketConn) } func releaseSocketConn(sc *SocketConn) { sc.s = nil sc.r = nil sc.wc = nil socketConnPool.Put(sc) }