package websocket_fasthttp import ( "fmt" "net" "net/http" "sync" "git.loafle.net/commons_go/logging" "git.loafle.net/commons_go/websocket_fasthttp/websocket" "github.com/valyala/fasthttp" ) type Server interface { Start() error Stop() } func New(sh ServerHandler) Server { s := &server{ sh: sh, } return s } type server struct { sh ServerHandler httpServer *fasthttp.Server upgrader *websocket.Upgrader listener net.Listener stopChan chan struct{} stopWg sync.WaitGroup } func (s *server) Start() error { if nil == s.sh { panic("Server: server handler must be specified.") } s.sh.Validate() if s.stopChan != nil { panic("Server: server is already running. Stop it before starting it again") } s.httpServer = &fasthttp.Server{ Handler: s.handleRequest, Name: s.sh.GetName(), Concurrency: s.sh.GetConcurrency(), ReadBufferSize: s.sh.GetReadBufferSize(), WriteBufferSize: s.sh.GetWriteBufferSize(), ReadTimeout: s.sh.GetReadTimeout(), WriteTimeout: s.sh.GetWriteTimeout(), } s.upgrader = &websocket.Upgrader{ HandshakeTimeout: s.sh.GetHandshakeTimeout(), ReadBufferSize: s.sh.GetReadBufferSize(), WriteBufferSize: s.sh.GetWriteBufferSize(), CheckOrigin: s.sh.CheckOrigin, Error: s.handleError, EnableCompression: s.sh.IsEnableCompression(), } var err error var listener net.Listener if listener, err = s.sh.Listen(); nil != err { return err } s.listener = newGracefulListener(listener, s.sh.GetMaxStopWaitTime()) s.stopChan = make(chan struct{}) s.sh.OnStart() s.stopWg.Add(1) go handleServer(s) return nil } func (s *server) Stop() { if s.stopChan == nil { panic("Server: server must be started before stopping it") } close(s.stopChan) s.stopWg.Wait() s.stopChan = nil s.sh.OnStop() } func handleServer(s *server) { go func() { defer s.stopWg.Done() if err := s.httpServer.Serve(s.listener); nil != err { logging.Logger().Error(fmt.Sprintf("Server: Server err - %v", err)) } }() select { case <-s.stopChan: s.listener.Close() return } } func (s *server) handleRequest(ctx *fasthttp.RequestCtx) { path := string(ctx.Path()) var socketHandler SocketHandler var err error if socketHandler, err = s.sh.GetSocketHandler(path); nil != err { s.handleError(ctx, fasthttp.StatusNotFound, err) return } var responseHeader *fasthttp.ResponseHeader var socketID interface{} if socketID, responseHeader = socketHandler.Handshake(ctx); nil == socketID { s.handleError(ctx, http.StatusNotAcceptable, fmt.Errorf("Server: Handshake err")) return } s.upgrader.Upgrade(ctx, responseHeader, func(conn *websocket.Conn, err error) { if err != nil { s.handleError(ctx, fasthttp.StatusInternalServerError, err) return } soc := newSocket(socketID, conn, socketHandler) s.stopWg.Add(1) handleConnection(s, soc, socketHandler) }) } func (s *server) handleError(ctx *fasthttp.RequestCtx, status int, reason error) { ctx.Response.Header.Set("Sec-Websocket-Version", "13") ctx.Error(http.StatusText(status), status) s.sh.OnError(ctx, status, reason) } func handleConnection(s *server, soc *Socket, socketHandler SocketHandler) { defer s.stopWg.Done() logging.Logger().Debug(fmt.Sprintf("Server: Client[%s] is connected.", soc.RemoteAddr())) socketHandler.OnConnect(soc) clientStopChan := make(chan struct{}) handleDoneChan := make(chan struct{}) go socketHandler.Handle(soc, clientStopChan, handleDoneChan) select { case <-s.stopChan: close(clientStopChan) <-handleDoneChan case <-handleDoneChan: close(clientStopChan) logging.Logger().Debug(fmt.Sprintf("Server: Client[%s] is disconnected.", soc.RemoteAddr())) socketHandler.OnDisconnect(soc) soc.Close() } }