package websocket import ( "context" "fmt" "net" "net/http" "sync" "git.loafle.net/commons/logging-go" "git.loafle.net/commons/server-go" "github.com/valyala/fasthttp" ) type Server struct { ServerHandler ServerHandler ctx server.ServerCtx stopChan chan struct{} stopWg sync.WaitGroup srw server.ServerReadWriter hs *fasthttp.Server upgrader *Upgrader } 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.(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.srw.ReadwriteHandler = s.ServerHandler s.srw.ServerStopChan = s.stopChan s.srw.ServerStopWg = &s.stopWg s.stopWg.Add(1) return s.handleServer(listener) } func (s *Server) Shutdown(ctx context.Context) error { if s.stopChan == nil { return fmt.Errorf("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) 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.srw.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.(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 } s.upgrader.Upgrade(ctx, responseHeader, func(conn *server.Conn, err error) { if err != nil { s.onError(ctx, fasthttp.StatusInternalServerError, err) return } s.stopWg.Add(1) go s.srw.HandleConnection(servlet, servletCtx, conn) }) } func (s *Server) onError(ctx *fasthttp.RequestCtx, status int, reason error) { s.ServerHandler.(ServerHandler).OnError(s.ctx, ctx, status, reason) }