Fix race while room close || user add

pkg/com/com.go

@@ -21,7 +21,6 @@ func NewNetMap[K stringer, T NetClient[K]]() NetMap[K, T] {
 func (m *NetMap[K, T]) Add(client T) bool         { return m.Put(client.Id(), client) }
 func (m *NetMap[K, T]) Empty() bool               { return m.Map.Len() == 0 }
 func (m *NetMap[K, T]) Remove(client T)           { m.Map.Remove(client.Id()) }
-func (m *NetMap[K, T]) RemoveL(client T) int      { return m.Map.RemoveL(client.Id()) }
 func (m *NetMap[K, T]) Reset()                    { m.Map = Map[K, T]{m: make(map[K]T, 10)} }
 func (m *NetMap[K, T]) RemoveDisconnect(client T) { client.Disconnect(); m.Remove(client) }
 func (m *NetMap[K, T]) Find(id string) T {

pkg/com/map.go

@@ -92,14 +92,6 @@ func (m *Map[K, V]) Pop(key K) V {
 	return val
 }
 
-// RemoveL removes the key and returns the new length of the map.
-func (m *Map[K, _]) RemoveL(key K) int {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-	delete(m.m, key)
-	return len(m.m)
-}
-
 // Clear empties the map.
 func (m *Map[K, V]) Clear() {
 	m.mu.Lock()

pkg/com/map_test.go

@@ -53,6 +53,36 @@ func TestMap_Base(t *testing.T) {
 	}
 }
 
+func TestMap_Remove(t *testing.T) {
+	m := Map[int, int]{m: make(map[int]int)}
+
+	m.Put(1, 10)
+	m.Put(2, 20)
+
+	if m.Len() == 0 {
+		t.Error("should not be empty after puts")
+	}
+
+	m.Remove(1)
+	if m.Len() == 0 {
+		t.Error("should not be empty, key 2 remains")
+	}
+	if m.Len() != 1 {
+		t.Errorf("expected len 1, got %v", m.Len())
+	}
+
+	m.Remove(2)
+	if m.Len() != 0 {
+		t.Error("should be empty after removing last key")
+	}
+
+	// Remove of non-existent key is a no-op.
+	m.Remove(99)
+	if m.Len() != 0 {
+		t.Error("should still be empty after no-op remove")
+	}
+}
+
 func TestMap_Concurrency(t *testing.T) {
 	m := Map[int, int]{m: make(map[int]int)}
 	for i := range 100 {

pkg/worker/room/room.go

@@ -26,7 +26,7 @@ type SessionManager[T Session] interface {
 	Add(T) bool
 	Empty() bool
 	Find(string) T
-	RemoveL(T) int
+	Remove(T)
 	// Reset used for proper cleanup of the resources if needed.
 	Reset()
 	Values() iter.Seq[T]
@@ -127,32 +127,36 @@ func (r *Router[T]) FindRoom(id string) *Room[T] {
 	return nil
 }
 
-func (r *Router[T]) Remove(user T) {
-	if left := r.users.RemoveL(user); left == 0 {
-		r.Close()
-		r.SetRoom(nil) // !to remove
-	}
-}
-
-func (r *Router[T]) AddUser(user T)           { r.users.Add(user) }
-func (r *Router[T]) Close()                   { r.mu.Lock(); r.room.Close(); r.room = nil; r.mu.Unlock() }
+func (r *Router[T]) AddUser(user T)           { r.mu.Lock(); r.users.Add(user); r.mu.Unlock() }
 func (r *Router[T]) FindUser(uid string) T    { return r.users.Find(uid) }
+func (r *Router[T]) Users() SessionManager[T] { return r.users }
 func (r *Router[T]) Room() *Room[T]           { r.mu.Lock(); defer r.mu.Unlock(); return r.room }
 func (r *Router[T]) SetRoom(room *Room[T])    { r.mu.Lock(); r.room = room; r.mu.Unlock() }
 func (r *Router[T]) HasRoom() bool            { r.mu.Lock(); defer r.mu.Unlock(); return r.room != nil }
-func (r *Router[T]) Users() SessionManager[T] { return r.users }
+func (r *Router[T]) Close()                   { r.mu.Lock(); r.close(); r.mu.Unlock() }
 func (r *Router[T]) Reset() {
 	r.mu.Lock()
-	if r.room != nil {
-		r.room.Close()
-		r.room = nil
-	}
+	r.close()
 	for u := range r.users.Values() {
 		u.Disconnect()
 	}
 	r.users.Reset()
 	r.mu.Unlock()
 }
+func (r *Router[T]) Remove(user T) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.users.Remove(user)
+	if r.users.Empty() {
+		r.close()
+	}
+}
+func (r *Router[T]) close() {
+	if r.room != nil {
+		r.room.Close()
+		r.room = nil
+	}
+}
 
 type AppSession struct {
 	Session

pkg/worker/room/router_test.go

@@ -1,6 +1,7 @@
 package room
 
 import (
+	"sync"
 	"testing"
 	"time"
 
@@ -77,6 +78,65 @@ func TestRouterReset(t *testing.T) {
 	}
 }
 
+func TestRouterRemove(t *testing.T) {
+	router := newTestRouter()
+	router.room = &Room[*tSession]{}
+	u1 := &tSession{id: "u1"}
+	u2 := &tSession{id: "u2"}
+	router.AddUser(u1)
+	router.AddUser(u2)
+
+	closed := false
+	router.room.HandleClose = func() { closed = true }
+
+	router.Remove(u1)
+	if router.Users().Empty() {
+		t.Error("should have user u2")
+	}
+	if closed {
+		t.Error("room should not close while wiht users")
+	}
+
+	router.Remove(u2)
+	if !router.Users().Empty() {
+		t.Error("should be empty after last user removed")
+	}
+	if !closed {
+		t.Error("room should close when last user leaves")
+	}
+}
+
+func TestRouterRemove_Race(t *testing.T) {
+	router := newTestRouter()
+	router.room = &Room[*tSession]{}
+
+	var emptyAtClose bool
+	router.room.HandleClose = func() {
+		emptyAtClose = router.Users().Empty()
+	}
+
+	var wg sync.WaitGroup
+	n := 5
+	wg.Add(n)
+
+	for range n {
+		go func() {
+			defer wg.Done()
+			for range 100 {
+				u := &tSession{id: "u"}
+				router.AddUser(u)
+				router.Remove(u)
+			}
+		}()
+	}
+
+	wg.Wait()
+
+	if closed := router.Room() == nil; closed && !emptyAtClose {
+		t.Error("room closed while users still present")
+	}
+}
+
 func newTestRouter() *Router[*tSession] {
 	u := com.NewNetMap[sKey, *tSession]()
 	return &Router[*tSession]{users: &u}