acceptor.go 1.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. package connect
  2. import (
  3. "context"
  4. "fmt"
  5. "go.uber.org/zap"
  6. "net/http"
  7. "net/url"
  8. "time"
  9. "git.bvbej.com/bvbej/base-golang/pkg/websocket/peer"
  10. "github.com/gorilla/websocket"
  11. )
  12. var _ WsAcceptor = (*wsAcceptor)(nil)
  13. var upgrader = websocket.Upgrader{
  14. CheckOrigin: func(r *http.Request) bool {
  15. return true
  16. },
  17. }
  18. type WsAcceptor interface {
  19. Start(addr string, sessMgr *peer.SessionManager) error
  20. Stop(sessMgr *peer.SessionManager)
  21. }
  22. type wsAcceptor struct {
  23. server *http.Server
  24. logger *zap.Logger
  25. addr string
  26. }
  27. func NewWsAcceptor(loggers *zap.Logger) WsAcceptor {
  28. return &wsAcceptor{logger: loggers}
  29. }
  30. func (ws *wsAcceptor) Start(addr string, sessMgr *peer.SessionManager) error {
  31. urlObj, err := url.Parse(addr)
  32. if err != nil {
  33. return fmt.Errorf("websocket urlparse failed. url(%s) %v", addr, err)
  34. }
  35. if urlObj.Path == "" {
  36. return fmt.Errorf("websocket start failed. expect path in url to listen addr:%s", addr)
  37. }
  38. http.HandleFunc(urlObj.Path, func(w http.ResponseWriter, r *http.Request) {
  39. c, upgradeErr := upgrader.Upgrade(w, r, nil)
  40. if upgradeErr != nil {
  41. ws.logger.Sugar().Errorf("upgrade http failed: %s", upgradeErr)
  42. return
  43. }
  44. sessMgr.Register <- peer.NewSession(newConnection(c, sessMgr))
  45. })
  46. ws.server = &http.Server{
  47. Addr: urlObj.Host,
  48. Handler: nil,
  49. }
  50. ws.addr = addr
  51. err = ws.server.ListenAndServe()
  52. if err != nil && err != http.ErrServerClosed {
  53. return fmt.Errorf("websocket ListenAndServe addr:%s failed:%v", addr, err)
  54. }
  55. return nil
  56. }
  57. func (ws *wsAcceptor) Stop(sessMgr *peer.SessionManager) {
  58. sessMgr.CloseAllSession()
  59. ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
  60. defer cancel()
  61. if err := ws.server.Shutdown(ctx); err != nil {
  62. ws.logger.Sugar().Errorf("server shutdown err:[%s]", err)
  63. }
  64. }