2018-04-04 04:01:26 +00:00
|
|
|
package websocket
|
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
|
|
|
"fmt"
|
|
|
|
"net"
|
|
|
|
"net/http"
|
|
|
|
"sync"
|
|
|
|
"time"
|
|
|
|
|
|
|
|
"git.loafle.net/commons/logging-go"
|
|
|
|
"git.loafle.net/commons/server-go"
|
|
|
|
"github.com/valyala/fasthttp"
|
|
|
|
)
|
|
|
|
|
|
|
|
type Server struct {
|
|
|
|
ServerHandler ServerHandler
|
|
|
|
|
|
|
|
ctx server.ServerCtx
|
|
|
|
|
|
|
|
hs *fasthttp.Server
|
|
|
|
upgrader *Upgrader
|
|
|
|
|
|
|
|
connections sync.Map
|
|
|
|
stopChan chan struct{}
|
|
|
|
stopWg sync.WaitGroup
|
|
|
|
}
|
|
|
|
|
|
|
|
func (s *Server) ListenAndServe() error {
|
|
|
|
var (
|
|
|
|
err error
|
|
|
|
listener net.Listener
|
|
|
|
)
|
|
|
|
if nil == s.ServerHandler {
|
|
|
|
return fmt.Errorf("Server: server handler must be specified")
|
|
|
|
}
|
|
|
|
s.ServerHandler.Validate()
|
|
|
|
|
|
|
|
if s.stopChan != nil {
|
|
|
|
return fmt.Errorf(s.serverMessage("already running. Stop it before starting it again"))
|
|
|
|
}
|
|
|
|
|
|
|
|
s.ctx = s.ServerHandler.ServerCtx()
|
|
|
|
if nil == s.ctx {
|
|
|
|
return fmt.Errorf(s.serverMessage("ServerCtx is nil"))
|
|
|
|
}
|
|
|
|
|
|
|
|
s.hs = &fasthttp.Server{
|
|
|
|
Handler: s.httpHandler,
|
|
|
|
Name: s.ServerHandler.GetName(),
|
|
|
|
Concurrency: s.ServerHandler.GetConcurrency(),
|
|
|
|
ReadBufferSize: s.ServerHandler.GetReadBufferSize(),
|
|
|
|
WriteBufferSize: s.ServerHandler.GetWriteBufferSize(),
|
|
|
|
ReadTimeout: s.ServerHandler.GetReadTimeout(),
|
|
|
|
WriteTimeout: s.ServerHandler.GetWriteTimeout(),
|
|
|
|
}
|
|
|
|
|
|
|
|
s.upgrader = &Upgrader{
|
|
|
|
HandshakeTimeout: s.ServerHandler.GetHandshakeTimeout(),
|
|
|
|
ReadBufferSize: s.ServerHandler.GetReadBufferSize(),
|
|
|
|
WriteBufferSize: s.ServerHandler.GetWriteBufferSize(),
|
|
|
|
CheckOrigin: s.ServerHandler.CheckOrigin,
|
|
|
|
Error: s.onError,
|
|
|
|
EnableCompression: s.ServerHandler.IsEnableCompression(),
|
|
|
|
}
|
|
|
|
|
|
|
|
if err = s.ServerHandler.Init(s.ctx); nil != err {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
|
|
|
if listener, err = s.ServerHandler.Listener(s.ctx); nil != err {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
|
|
|
s.stopChan = make(chan struct{})
|
|
|
|
s.stopWg.Add(1)
|
|
|
|
return s.handleServer(listener)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (s *Server) Shutdown(ctx context.Context) error {
|
|
|
|
if s.stopChan == nil {
|
|
|
|
return fmt.Errorf(s.serverMessage("server must be started before stopping it"))
|
|
|
|
}
|
|
|
|
close(s.stopChan)
|
|
|
|
s.stopWg.Wait()
|
|
|
|
|
|
|
|
s.ServerHandler.Destroy(s.ctx)
|
|
|
|
|
|
|
|
s.stopChan = nil
|
|
|
|
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func (s *Server) ConnectionSize() int {
|
|
|
|
var sz int
|
|
|
|
s.connections.Range(func(k, v interface{}) bool {
|
|
|
|
sz++
|
|
|
|
return true
|
|
|
|
})
|
|
|
|
return sz
|
|
|
|
}
|
|
|
|
|
|
|
|
func (s *Server) serverMessage(msg string) string {
|
|
|
|
return fmt.Sprintf("Server[%s]: %s", s.ServerHandler.GetName(), msg)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (s *Server) handleServer(listener net.Listener) error {
|
|
|
|
var (
|
|
|
|
err error
|
|
|
|
)
|
|
|
|
|
|
|
|
errChan := make(chan error)
|
|
|
|
|
|
|
|
defer func() {
|
|
|
|
if nil != listener {
|
|
|
|
listener.Close()
|
|
|
|
}
|
|
|
|
s.stopWg.Done()
|
|
|
|
}()
|
|
|
|
|
|
|
|
go func() {
|
|
|
|
if err := s.hs.Serve(listener); nil != err {
|
|
|
|
errChan <- err
|
|
|
|
return
|
|
|
|
}
|
|
|
|
close(errChan)
|
|
|
|
}()
|
|
|
|
|
|
|
|
select {
|
|
|
|
case err, _ := <-errChan:
|
|
|
|
if nil != err {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
defer func() {
|
|
|
|
s.ServerHandler.OnStop(s.ctx)
|
|
|
|
|
|
|
|
logging.Logger().Infof(s.serverMessage("Stopped"))
|
|
|
|
}()
|
|
|
|
|
|
|
|
if err = s.ServerHandler.OnStart(s.ctx); nil != err {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
|
|
|
logging.Logger().Infof(s.serverMessage("Started"))
|
|
|
|
|
|
|
|
select {
|
|
|
|
case <-s.stopChan:
|
|
|
|
listener.Close()
|
|
|
|
listener = nil
|
|
|
|
}
|
|
|
|
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func (s *Server) httpHandler(ctx *fasthttp.RequestCtx) {
|
|
|
|
path := string(ctx.Path())
|
|
|
|
var (
|
|
|
|
servlet Servlet
|
|
|
|
err error
|
|
|
|
)
|
|
|
|
|
|
|
|
if 0 < s.ServerHandler.GetConcurrency() {
|
|
|
|
sz := s.ConnectionSize()
|
|
|
|
if sz >= s.ServerHandler.GetConcurrency() {
|
|
|
|
logging.Logger().Warnf(s.serverMessage(fmt.Sprintf("max connections size %d, refuse", sz)))
|
|
|
|
s.onError(ctx, fasthttp.StatusServiceUnavailable, err)
|
|
|
|
return
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
if servlet = s.ServerHandler.Servlet(path); nil == servlet {
|
|
|
|
s.onError(ctx, fasthttp.StatusInternalServerError, err)
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
var responseHeader *fasthttp.ResponseHeader
|
|
|
|
servletCtx := servlet.ServletCtx(s.ctx)
|
|
|
|
|
|
|
|
if responseHeader, err = servlet.Handshake(servletCtx, ctx); nil != err {
|
|
|
|
s.onError(ctx, http.StatusNotAcceptable, fmt.Errorf("Handshake err: %v", err))
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
2018-04-04 05:47:10 +00:00
|
|
|
s.upgrader.Upgrade(ctx, responseHeader, func(conn *server.Conn, err error) {
|
2018-04-04 04:01:26 +00:00
|
|
|
if err != nil {
|
|
|
|
s.onError(ctx, fasthttp.StatusInternalServerError, err)
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
s.stopWg.Add(1)
|
|
|
|
go s.handleConnection(servlet, servletCtx, conn)
|
|
|
|
})
|
|
|
|
}
|
|
|
|
|
2018-04-04 05:47:10 +00:00
|
|
|
func (s *Server) handleConnection(servlet Servlet, servletCtx server.ServletCtx, conn *server.Conn) {
|
2018-04-04 04:01:26 +00:00
|
|
|
addr := conn.RemoteAddr()
|
|
|
|
|
|
|
|
defer func() {
|
2018-04-04 05:31:10 +00:00
|
|
|
if nil != conn {
|
|
|
|
conn.Close()
|
|
|
|
}
|
2018-04-04 04:01:26 +00:00
|
|
|
servlet.OnDisconnect(servletCtx)
|
|
|
|
logging.Logger().Infof(s.serverMessage(fmt.Sprintf("Client[%s] has been disconnected", addr)))
|
|
|
|
s.stopWg.Done()
|
|
|
|
}()
|
|
|
|
|
|
|
|
logging.Logger().Infof(s.serverMessage(fmt.Sprintf("Client[%s] has been connected", addr)))
|
2018-04-04 05:31:10 +00:00
|
|
|
|
2018-04-04 04:01:26 +00:00
|
|
|
s.connections.Store(conn, true)
|
2018-04-04 05:31:10 +00:00
|
|
|
defer s.connections.Delete(conn)
|
|
|
|
|
2018-04-04 04:01:26 +00:00
|
|
|
servlet.OnConnect(servletCtx, conn)
|
2018-04-04 08:36:23 +00:00
|
|
|
conn.SetCloseHandler(func(code int, text string) error {
|
|
|
|
logging.Logger().Debugf("close")
|
|
|
|
return nil
|
|
|
|
})
|
2018-04-04 04:01:26 +00:00
|
|
|
|
2018-04-04 05:31:10 +00:00
|
|
|
stopChan := make(chan struct{})
|
|
|
|
servletDoneChan := make(chan struct{})
|
2018-04-04 04:01:26 +00:00
|
|
|
|
|
|
|
readChan := make(chan []byte)
|
|
|
|
writeChan := make(chan []byte)
|
|
|
|
|
2018-04-04 05:31:10 +00:00
|
|
|
readerDoneChan := make(chan struct{})
|
|
|
|
writerDoneChan := make(chan struct{})
|
|
|
|
|
|
|
|
go servlet.Handle(servletCtx, stopChan, servletDoneChan, readChan, writeChan)
|
|
|
|
go handleRead(s, conn, stopChan, readerDoneChan, readChan)
|
|
|
|
go handleWrite(s, conn, stopChan, writerDoneChan, writeChan)
|
2018-04-04 04:01:26 +00:00
|
|
|
|
|
|
|
select {
|
2018-04-04 05:31:10 +00:00
|
|
|
case <-readerDoneChan:
|
|
|
|
close(stopChan)
|
|
|
|
<-writerDoneChan
|
|
|
|
<-servletDoneChan
|
|
|
|
case <-writerDoneChan:
|
|
|
|
close(stopChan)
|
|
|
|
<-readerDoneChan
|
|
|
|
<-servletDoneChan
|
|
|
|
case <-servletDoneChan:
|
|
|
|
close(stopChan)
|
|
|
|
<-readerDoneChan
|
|
|
|
<-writerDoneChan
|
2018-04-04 04:01:26 +00:00
|
|
|
case <-s.stopChan:
|
2018-04-04 05:31:10 +00:00
|
|
|
close(stopChan)
|
|
|
|
<-readerDoneChan
|
|
|
|
<-writerDoneChan
|
|
|
|
<-servletDoneChan
|
2018-04-04 04:01:26 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func (s *Server) onError(ctx *fasthttp.RequestCtx, status int, reason error) {
|
|
|
|
s.ServerHandler.OnError(s.ctx, ctx, status, reason)
|
|
|
|
}
|
|
|
|
|
2018-04-04 08:36:23 +00:00
|
|
|
func handleRead(s *Server, conn *server.Conn, stopChan <-chan struct{}, doneChan chan<- struct{}, readChan chan []byte) {
|
2018-04-04 05:31:10 +00:00
|
|
|
defer func() {
|
2018-04-04 08:36:23 +00:00
|
|
|
doneChan <- struct{}{}
|
2018-04-04 05:31:10 +00:00
|
|
|
}()
|
|
|
|
|
2018-04-04 08:36:23 +00:00
|
|
|
if 0 < s.ServerHandler.GetMaxMessageSize() {
|
|
|
|
conn.SetReadLimit(s.ServerHandler.GetMaxMessageSize())
|
|
|
|
}
|
|
|
|
if 0 < s.ServerHandler.GetReadTimeout() {
|
|
|
|
conn.SetReadDeadline(time.Now().Add(s.ServerHandler.GetReadTimeout()))
|
|
|
|
}
|
2018-04-04 04:01:26 +00:00
|
|
|
conn.SetPongHandler(func(string) error {
|
|
|
|
conn.SetReadDeadline(time.Now().Add(s.ServerHandler.GetPongTimeout()))
|
|
|
|
return nil
|
|
|
|
})
|
|
|
|
|
2018-04-04 05:31:10 +00:00
|
|
|
var (
|
|
|
|
message []byte
|
|
|
|
err error
|
|
|
|
)
|
|
|
|
|
2018-04-04 04:01:26 +00:00
|
|
|
for {
|
2018-04-04 05:31:10 +00:00
|
|
|
readMessageChan := make(chan struct{})
|
|
|
|
|
|
|
|
go func() {
|
|
|
|
_, message, err = conn.ReadMessage()
|
|
|
|
close(readMessageChan)
|
|
|
|
}()
|
|
|
|
|
|
|
|
select {
|
|
|
|
case <-s.stopChan:
|
|
|
|
<-readMessageChan
|
2018-04-04 08:36:23 +00:00
|
|
|
return
|
2018-04-04 05:31:10 +00:00
|
|
|
case <-readMessageChan:
|
2018-04-04 04:01:26 +00:00
|
|
|
}
|
2018-04-04 05:31:10 +00:00
|
|
|
|
|
|
|
if nil != err {
|
2018-04-04 08:36:23 +00:00
|
|
|
if server.IsUnexpectedCloseError(err, server.CloseGoingAway, server.CloseAbnormalClosure) {
|
|
|
|
logging.Logger().Debugf(s.serverMessage(fmt.Sprintf("Read error %v", err)))
|
2018-04-04 05:31:10 +00:00
|
|
|
}
|
2018-04-04 08:36:23 +00:00
|
|
|
return
|
2018-04-04 05:31:10 +00:00
|
|
|
}
|
|
|
|
|
2018-04-04 04:01:26 +00:00
|
|
|
readChan <- message
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2018-04-04 08:36:23 +00:00
|
|
|
func handleWrite(s *Server, conn *server.Conn, stopChan <-chan struct{}, doneChan chan<- struct{}, writeChan chan []byte) {
|
2018-04-04 05:31:10 +00:00
|
|
|
defer func() {
|
2018-04-04 08:36:23 +00:00
|
|
|
doneChan <- struct{}{}
|
2018-04-04 05:31:10 +00:00
|
|
|
}()
|
|
|
|
|
2018-04-04 04:01:26 +00:00
|
|
|
ticker := time.NewTicker(s.ServerHandler.GetPingPeriod())
|
|
|
|
defer func() {
|
|
|
|
ticker.Stop()
|
|
|
|
}()
|
|
|
|
for {
|
|
|
|
select {
|
|
|
|
case message, ok := <-writeChan:
|
2018-04-04 08:36:23 +00:00
|
|
|
if 0 < s.ServerHandler.GetWriteTimeout() {
|
|
|
|
conn.SetWriteDeadline(time.Now().Add(s.ServerHandler.GetWriteTimeout()))
|
|
|
|
}
|
2018-04-04 04:01:26 +00:00
|
|
|
if !ok {
|
2018-04-04 05:47:10 +00:00
|
|
|
conn.WriteMessage(server.CloseMessage, []byte{})
|
2018-04-04 04:01:26 +00:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
2018-04-04 05:47:10 +00:00
|
|
|
w, err := conn.NextWriter(server.TextMessage)
|
2018-04-04 04:01:26 +00:00
|
|
|
if err != nil {
|
|
|
|
return
|
|
|
|
}
|
|
|
|
w.Write(message)
|
|
|
|
|
|
|
|
if err := w.Close(); nil != err {
|
|
|
|
return
|
|
|
|
}
|
|
|
|
case <-ticker.C:
|
|
|
|
conn.SetWriteDeadline(time.Now().Add(s.ServerHandler.GetPingTimeout()))
|
2018-04-04 05:47:10 +00:00
|
|
|
if err := conn.WriteMessage(server.PingMessage, nil); nil != err {
|
2018-04-04 04:01:26 +00:00
|
|
|
return
|
|
|
|
}
|
2018-04-04 05:31:10 +00:00
|
|
|
case <-s.stopChan:
|
2018-04-04 08:36:23 +00:00
|
|
|
return
|
2018-04-04 04:01:26 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|