ing

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()