Merge remote-tracking branch 'origin/master' into ws-extension-updates

extensions/websocket/internal/wsutil/manager.go

@@ -58,6 +58,7 @@ type SocketManager struct {
 	sockets                  *xsync.MapOf[string, *xsync.MapOf[string, SocketConnection]]
 	idToRoom                 *xsync.MapOf[string, string]
 	listeners                []chan SocketEvent
+	listenersQueuedForDelete *xsync.MapOf[chan SocketEvent, bool]
 	goroutinesRunning        atomic.Int32
 	opts                     *opts.ExtensionOpts
 	lock                     sync.Mutex
@@ -129,6 +130,7 @@ func NewSocketManager(opts *opts.ExtensionOpts) *SocketManager {
 	return &SocketManager{
 		sockets:                  xsync.NewMapOf[string, *xsync.MapOf[string, SocketConnection]](),
 		idToRoom:                 xsync.NewMapOf[string, string](),
+		listenersQueuedForDelete: xsync.NewMapOf[chan SocketEvent, bool](),
 		opts:                     opts,
 		goroutinesRunning:        atomic.Int32{},
 	}
@@ -185,6 +187,10 @@ func (manager *SocketManager) Listen(listener chan SocketEvent) {
 	}
 }
 
+func (manager *SocketManager) RemoveListener(listener chan SocketEvent) {
+	manager.listenersQueuedForDelete.Store(listener, true)
+}
+
 func (manager *SocketManager) dispatch(event SocketEvent) {
 	done := make(chan struct{}, 1)
 	go func() {
@@ -197,9 +203,25 @@ func (manager *SocketManager) dispatch(event SocketEvent) {
 			}
 		}
 	}()
+
+	newListener := make([]chan SocketEvent, 0)
 	for _, listener := range manager.listeners {
-		listener <- event
+		if _, ok := manager.listenersQueuedForDelete.Load(listener); !ok {
+			newListener = append(newListener, listener)
 		}
+	}
+	manager.listeners = newListener
+
+	wg := sync.WaitGroup{}
+	for _, listener := range manager.listeners {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			listener <- event
+		}()
+	}
+
+	wg.Wait()
 	done <- struct{}{}
 }