websocket_fasthttp/socket.go

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)
}