Use thread safe sync.map

Author: mateeullahmalik

pkg/net/credentials/alts/handshake/handshake_test.go

@@ -434,8 +434,7 @@ func TestHandshakerConcurrentHandshakes(t *testing.T) {
 					"expected %d successful handshakes", maxConcurrentHandshakes*2)
 			}
 
-			// Verify concurrent handshake limits
-			require.LessOrEqual(t, stat.MaxConcurrentCalls, maxConcurrentHandshakes,
+			require.LessOrEqual(t, stat.MaxConcurrentCalls(), maxConcurrentHandshakes,
 				"concurrent handshakes exceeded limit")
 		})
 	}

pkg/net/credentials/alts/testutil/testutil.go

@@ -5,40 +5,46 @@ import (
 	"encoding/binary"
 	"net"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	. "github.com/LumeraProtocol/supernode/pkg/net/credentials/alts/common"
 )
 
 // Stats is used to collect statistics about concurrent handshake calls.
 type Stats struct {
-	mu                 sync.Mutex
-	calls              int
-	MaxConcurrentCalls int
+	stats     sync.Map
+	callCount int64
+	maxCalls  int64
 }
 
 // Update updates the statistics by adding one call.
-func (s *Stats) Update() func() {
-	s.mu.Lock()
-	s.calls++
-	if s.calls > s.MaxConcurrentCalls {
-		s.MaxConcurrentCalls = s.calls
+func (s *Stats) Update() {
+	// Use atomic operations for counters
+	current := atomic.AddInt64(&s.callCount, 1)
+
+	// Update max if needed
+	for {
+		max := atomic.LoadInt64(&s.maxCalls)
+		if current <= max {
+			break
+		}
+		if atomic.CompareAndSwapInt64(&s.maxCalls, max, current) {
+			break
+		}
 	}
-	s.mu.Unlock()
 
-	return func() {
-		s.mu.Lock()
-		s.calls--
-		s.mu.Unlock()
-	}
 }
 
 // Reset resets the statistics.
 func (s *Stats) Reset() {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	s.calls = 0
-	s.MaxConcurrentCalls = 0
+	s.stats = sync.Map{}
+	atomic.StoreInt64(&s.callCount, 0)
+	atomic.StoreInt64(&s.maxCalls, 0)
+}
+
+func (s *Stats) MaxConcurrentCalls() int {
+	return int(atomic.LoadInt64(&s.maxCalls))
 }
 
 // testConn mimics a net.Conn to the peer.
@@ -50,7 +56,7 @@ type testLatencyConn struct {
 // NewTestConnWithReadLatency wraps a net.Conn with artificial read latency
 func NewTestConnWithReadLatency(conn net.Conn, readLatency time.Duration) net.Conn {
 	return &testLatencyConn{
-		Conn:       conn,
+		Conn:        conn,
 		readLatency: readLatency,
 	}
 }
@@ -87,9 +93,9 @@ func NewUnresponsiveTestConn(delay time.Duration) net.Conn {
 
 // Read reads from the in buffer.
 func (c *unresponsiveTestConn) Read([]byte) (n int, err error) {
-    // Wait for delay to simulate network latency
-    time.Sleep(c.delay)
-    // Return empty data (success but zero bytes)
+	// Wait for delay to simulate network latency
+	time.Sleep(c.delay)
+	// Return empty data (success but zero bytes)
 	return 0, nil
 }