websocket_fasthttp/socket.go

188 lines
3.3 KiB
Go
Raw Normal View History

2017-11-08 10:15:09 +00:00
package websocket_fasthttp
import (
"io"
"net"
"sync"
"time"
"git.loafle.net/commons_go/websocket_fasthttp/websocket"
)
2017-11-10 08:33:40 +00:00
func newSocket(id string, conn *websocket.Conn, sh SocketHandler) *Socket {
2017-11-08 10:15:09 +00:00
s := retainSocket()
s.Conn = conn
s.sh = sh
2017-11-09 09:10:19 +00:00
s.id = id
2017-11-08 10:15:09 +00:00
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
2017-11-10 08:33:40 +00:00
id string
2017-11-09 09:10:19 +00:00
attributes map[interface{}]interface{}
2017-11-08 10:15:09 +00:00
sc *SocketConn
}
2017-11-10 08:33:40 +00:00
func (s *Socket) ID() string {
2017-11-09 09:10:19 +00:00
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
}
2017-11-08 10:15:09 +00:00
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
}
2017-11-08 10:24:05 +00:00
s.sc = retainSocketConn()
s.sc.s = s
s.sc.MessageType = mt
s.sc.r = r
2017-11-08 10:15:09 +00:00
2017-11-08 10:24:05 +00:00
return s.sc, nil
2017-11-08 10:15:09 +00:00
}
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
2017-11-08 10:26:34 +00:00
releaseSocketConn(sc)
2017-11-08 10:15:09 +00:00
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
2017-11-10 08:33:40 +00:00
s.id = ""
2017-11-08 10:15:09 +00:00
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)
}