diff --git a/pkg/com/com.go b/pkg/com/com.go index 8b475622..31207876 100644 --- a/pkg/com/com.go +++ b/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 { diff --git a/pkg/com/map.go b/pkg/com/map.go index ce2c5cd5..36c7c808 100644 --- a/pkg/com/map.go +++ b/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() diff --git a/pkg/com/map_test.go b/pkg/com/map_test.go index 15af76c4..252340c4 100644 --- a/pkg/com/map_test.go +++ b/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 { diff --git a/pkg/worker/room/room.go b/pkg/worker/room/room.go index 9ff9576e..ccc9f15b 100644 --- a/pkg/worker/room/room.go +++ b/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 diff --git a/pkg/worker/room/router_test.go b/pkg/worker/room/router_test.go index 1c9858c3..c9f73971 100644 --- a/pkg/worker/room/router_test.go +++ b/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}