165 lines
2.9 KiB
Go
165 lines
2.9 KiB
Go
package websocket_fasthttp
|
|
|
|
import (
|
|
"io"
|
|
"net"
|
|
"sync"
|
|
"time"
|
|
|
|
"git.loafle.net/commons_go/websocket_fasthttp/websocket"
|
|
)
|
|
|
|
func newSocket(conn *websocket.Conn, sh SocketHandler) *Socket {
|
|
s := retainSocket()
|
|
s.Conn = conn
|
|
s.sh = sh
|
|
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
|
|
|
|
sc *SocketConn
|
|
}
|
|
|
|
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
|
|
|
|
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)
|
|
}
|