fix: prevent raccoon restart by adding RWMutex for connection group count map (#14)

Added RWMutex protection around the connection group counter map to ensure concurrent safety while reporting the metric on the number of connections. It will solve the runtime crash due to concurrent map iteration and map write

Author: siddhanta-rath

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))
-		}
+		})
 	}
 }
 

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)
+	}
 }

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++ {