From 6285ba38bc0e7746c5bdd6932879a32996fa72cb Mon Sep 17 00:00:00 2001
From: Siddhanta Rath <sid.9111@gmail.com>
Date: Fri, 8 Aug 2025 08:26:55 +0530
Subject: [PATCH] fix: prevent raccoon restart by adding RWMutex for connection
 group count map

---
 services/rest/service.go                      |  4 +--
 services/rest/websocket/connection/table.go   |  8 +++--
 .../rest/websocket/connection/table_test.go   | 35 ++++++++++++-------
 3 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/services/rest/service.go b/services/rest/service.go
index 660fbb9e..65db7ea8 100644
--- a/services/rest/service.go
+++ b/services/rest/service.go
@@ -54,9 +54,9 @@ func reportConnectionMetrics(conn connection.Table) {
 	t := time.Tick(config.MetricStatsd.FlushPeriodMs)
 	for {
 		<-t
-		for k, v := range conn.TotalConnectionPerGroup() {
+		conn.RangeConnectionPerGroup(func(k string, v int) {
 			metrics.Gauge("connections_count_current", v, fmt.Sprintf("conn_group=%s", k))
-		}
+		})
 	}
 }
 
diff --git a/services/rest/websocket/connection/table.go b/services/rest/websocket/connection/table.go
index eb340329..4c9cc844 100644
--- a/services/rest/websocket/connection/table.go
+++ b/services/rest/websocket/connection/table.go
@@ -83,6 +83,10 @@ func (t *Table) TotalConnection() int {
 	return len(t.connMap)
 }
 
-func (t *Table) TotalConnectionPerGroup() map[string]int {
-	return t.counter
+func (t *Table) RangeConnectionPerGroup(fn func(string, int)) {
+	t.m.RLock()
+	defer t.m.RUnlock()
+	for k, v := range t.counter {
+		fn(k, v)
+	}
 }
diff --git a/services/rest/websocket/connection/table_test.go b/services/rest/websocket/connection/table_test.go
index 0ef3bb58..a3681201 100644
--- a/services/rest/websocket/connection/table_test.go
+++ b/services/rest/websocket/connection/table_test.go
@@ -8,18 +8,6 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestConnectionPerGroup(t *testing.T) {
-	t.Run("Should return all the group on the table with the count", func(t *testing.T) {
-		table := NewTable(10)
-		table.Store(identification.Identifier{ID: "user1", Group: "group1"})
-		table.Store(identification.Identifier{ID: "user2", Group: "group1"})
-		table.Store(identification.Identifier{ID: "user3", Group: "group1"})
-		table.Store(identification.Identifier{ID: "user1", Group: "group2"})
-		table.Store(identification.Identifier{ID: "user2", Group: "group2"})
-		assert.Equal(t, map[string]int{"group1": 3, "group2": 2}, table.TotalConnectionPerGroup())
-	})
-}
-
 func TestStore(t *testing.T) {
 	t.Run("Should store new connection", func(t *testing.T) {
 		table := NewTable(10)
@@ -146,6 +134,29 @@ func TestStoreBatch(t *testing.T) {
 	})
 }
 
+func TestRangeConnectionPerGroup(t *testing.T) {
+	t.Run("Should iterate over all the group on the table with the count", func(t *testing.T) {
+		//set up the table data
+		table := NewTable(10)
+		table.Store(identification.Identifier{ID: "user1", Group: "group1"})
+		table.Store(identification.Identifier{ID: "user2", Group: "group1"})
+		table.Store(identification.Identifier{ID: "user3", Group: "group1"})
+		table.Store(identification.Identifier{ID: "user1", Group: "group2"})
+		table.Store(identification.Identifier{ID: "user2", Group: "group2"})
+		// Track callback calls
+		called := make(map[string]int)
+		table.RangeConnectionPerGroup(func(k string, v int) {
+			called[k] = v
+		})
+		// Verify all keys are present
+		assert.Equal(t, len(table.counter), len(called))
+		// Verify values match
+		for k, v := range table.counter {
+			assert.Equal(t, v, called[k])
+		}
+	})
+}
+
 func BenchmarkStoreBatch(b *testing.B) {
 	table := NewTable(b.N)
 	for i := 0; i < b.N; i++ {