diff --git a/server.go b/server.go index c3d6a9b..121b9be 100644 --- a/server.go +++ b/server.go @@ -1,6 +1,7 @@ package overflow_gateway_websocket import ( + "context" "net/http" "sync" @@ -16,6 +17,7 @@ type Server interface { } type server struct { + _ctx context.Context _option *ServerOptions _upgrader *websocket.Upgrader _handlers map[string]*SocketOptions @@ -25,8 +27,9 @@ type server struct { _removeSocketCh chan Socket } -func NewServer(o *ServerOptions) Server { +func NewServer(ctx context.Context, o *ServerOptions) Server { s := &server{ + _ctx: ctx, _option: o.Validate(), _handlers: make(map[string]*SocketOptions, 1), _sockets: make(map[string]Socket, 100), @@ -91,6 +94,8 @@ func (s *server) onConnection(ctx *fasthttp.RequestCtx) { func (s *server) listenHandler() { for { select { + case <-s._ctx.Done(): + return // Add new a socket case soc := <-s._addSocketCh: s._socketsMtx.Lock()