fix: resolve 6 critical/high-priority concurrency and performance issues

- Fix WebSocket goroutine leaks by adding per-connection context cancellation
- Fix race conditions in connection state management (set conn to nil after close)
- Fix subscription panic risks by removing panic recovery and adding grace period
- Add context cancellation support to stream functions for proper cleanup
- Fix heartbeat goroutine management to prevent accumulation
- Optimize rate limiter with timestamp-based implementation (eliminate background goroutine)

Test coverage: Added comprehensive tests for goroutine leaks, race conditions, and subscription panics.

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

Author: donowu

go.mod

@@ -35,6 +35,7 @@ require (
 	github.com/supranational/blst v0.3.16-0.20250831170142-f48500c1fdbe // indirect
 	github.com/tklauser/go-sysconf v0.3.12 // indirect
 	github.com/tklauser/numcpus v0.6.1 // indirect
+	go.uber.org/goleak v1.3.0 // indirect
 	golang.org/x/crypto v0.45.0 // indirect
 	golang.org/x/sync v0.12.0 // indirect
 	golang.org/x/sys v0.38.0 // indirect

go.sum

@@ -195,6 +195,8 @@ github.com/urfave/cli/v2 v2.27.5 h1:WoHEJLdsXr6dDWoJgMq/CboDmyY/8HMMH1fTECbih+w=
 github.com/urfave/cli/v2 v2.27.5/go.mod h1:3Sevf16NykTbInEnD0yKkjDAeZDS0A6bzhBH5hrMvTQ=
 github.com/xrash/smetrics v0.0.0-20240521201337-686a1a2994c1 h1:gEOO8jv9F4OT7lGCjxCBTO/36wtF6j2nSip77qHd4x4=
 github.com/xrash/smetrics v0.0.0-20240521201337-686a1a2994c1/go.mod h1:Ohn+xnUBiLI6FVj/9LpzZWtj1/D6lUovWYBkxHVV3aM=
+go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
+go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
 golang.org/x/crypto v0.45.0 h1:jMBrvKuj23MTlT0bQEOBcAE0mjg8mK9RXFhRH6nyF3Q=
 golang.org/x/crypto v0.45.0/go.mod h1:XTGrrkGJve7CYK7J8PEww4aY7gM3qMCElcJQ8n8JdX4=
 golang.org/x/exp v0.0.0-20230626212559-97b1e661b5df h1:UA2aFVmmsIlefxMk29Dp2juaUSth8Pyn3Tq5Y5mJGME=

pkg/clob/impl.go

@@ -376,9 +376,18 @@ func (c *clientImpl) startHeartbeats() {
 	if c.httpClient == nil || c.signer == nil || c.apiKey == nil || c.heartbeat == nil {
 		return
 	}
+
+	// Stop old heartbeat goroutine if it exists
 	if c.heartbeatStop != nil {
-		close(c.heartbeatStop)
+		oldStop := c.heartbeatStop
+		c.heartbeatStop = nil
+		close(oldStop)
+		// Add a small delay to allow old goroutine to exit
+		c.heartbeatMu.Unlock()
+		time.Sleep(50 * time.Millisecond)
+		c.heartbeatMu.Lock()
 	}
+
 	stop := make(chan struct{})
 	c.heartbeatStop = stop
 	interval := c.heartbeatInterval

pkg/clob/stream.go

@@ -32,22 +32,35 @@ func StreamDataWithCursor[T any](ctx context.Context, cursor string, fetch Strea
 		}
 
 		for cursor != clobtypes.EndCursor {
+			// Check context before each fetch operation
 			if err := ctx.Err(); err != nil {
 				out <- StreamResult[T]{Err: err}
 				return
 			}
+
+			// Make fetch operation cancellable by passing context
 			items, next, err := fetch(ctx, cursor)
 			if err != nil {
 				out <- StreamResult[T]{Err: err}
 				return
 			}
+
 			for _, item := range items {
+				// Check context before sending each item
 				if err := ctx.Err(); err != nil {
 					out <- StreamResult[T]{Err: err}
 					return
 				}
-				out <- StreamResult[T]{Item: item}
+
+				// Use select to make send cancellable
+				select {
+				case out <- StreamResult[T]{Item: item}:
+				case <-ctx.Done():
+					out <- StreamResult[T]{Err: ctx.Err()}
+					return
+				}
 			}
+
 			if next == "" || next == cursor {
 				return
 			}

pkg/clob/ws/goroutine_leak_test.go

@@ -0,0 +1,243 @@
+package ws
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/gorilla/websocket"
+	"go.uber.org/goleak"
+)
+
+// TestWebSocketGoroutineLeaks tests that goroutines are properly cleaned up
+// during reconnection scenarios to prevent goroutine leaks.
+func TestWebSocketGoroutineLeaks_Reconnection(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+
+	// Create a mock WebSocket server that closes connections immediately
+	upgrader := websocket.Upgrader{}
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		// Close immediately to trigger reconnection
+		conn.Close()
+	}))
+	defer server.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(server.URL, "http")
+
+	client, err := NewClient(wsURL, nil, nil)
+	if err != nil {
+		t.Fatalf("failed to create client: %v", err)
+	}
+
+	// Set short timeouts to speed up test
+	impl := client.(*clientImpl)
+	impl.reconnectDelay = 10 * time.Millisecond
+	impl.reconnectMaxDelay = 50 * time.Millisecond
+	impl.reconnectMax = 2
+	impl.setReadTimeout(100 * time.Millisecond)
+
+	// Wait for reconnection attempts
+	time.Sleep(200 * time.Millisecond)
+
+	// Close the client
+	if err := client.Close(); err != nil {
+		t.Fatalf("failed to close client: %v", err)
+	}
+
+	// Give goroutines time to exit
+	time.Sleep(100 * time.Millisecond)
+
+	// goleak.VerifyNone will check for leaked goroutines at test end
+}
+
+// TestWebSocketGoroutineLeaks_MultipleReconnections tests that multiple
+// reconnection cycles don't accumulate goroutines.
+func TestWebSocketGoroutineLeaks_MultipleReconnections(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+
+	connectionCount := 0
+	upgrader := websocket.Upgrader{}
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		connectionCount++
+		// Close after a short delay
+		time.AfterFunc(20*time.Millisecond, func() {
+			conn.Close()
+		})
+	}))
+	defer server.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(server.URL, "http")
+
+	client, err := NewClient(wsURL, nil, nil)
+	if err != nil {
+		t.Fatalf("failed to create client: %v", err)
+	}
+
+	impl := client.(*clientImpl)
+	impl.reconnectDelay = 10 * time.Millisecond
+	impl.reconnectMaxDelay = 50 * time.Millisecond
+	impl.reconnectMax = 5
+	impl.setReadTimeout(50 * time.Millisecond)
+
+	// Wait for multiple reconnection cycles
+	time.Sleep(500 * time.Millisecond)
+
+	if err := client.Close(); err != nil {
+		t.Fatalf("failed to close client: %v", err)
+	}
+
+	time.Sleep(100 * time.Millisecond)
+
+	if connectionCount < 2 {
+		t.Logf("Warning: only %d connections made, expected multiple reconnections", connectionCount)
+	}
+}
+
+// TestWebSocketGoroutineLeaks_CloseWhileReading tests that closing the client
+// while a read is in progress doesn't leak goroutines.
+func TestWebSocketGoroutineLeaks_CloseWhileReading(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+
+	upgrader := websocket.Upgrader{}
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		// Keep connection open but don't send data
+		// This simulates a hanging read
+		time.Sleep(5 * time.Second)
+		conn.Close()
+	}))
+	defer server.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(server.URL, "http")
+
+	client, err := NewClient(wsURL, nil, nil)
+	if err != nil {
+		t.Fatalf("failed to create client: %v", err)
+	}
+
+	impl := client.(*clientImpl)
+	impl.setReadTimeout(100 * time.Millisecond)
+
+	// Give readLoop time to start
+	time.Sleep(50 * time.Millisecond)
+
+	// Close while read is potentially in progress
+	if err := client.Close(); err != nil {
+		t.Fatalf("failed to close client: %v", err)
+	}
+
+	time.Sleep(100 * time.Millisecond)
+}
+
+// TestWebSocketGoroutineLeaks_PingLoopCleanup tests that pingLoop goroutines
+// are properly cleaned up when connections are closed.
+func TestWebSocketGoroutineLeaks_PingLoopCleanup(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+
+	upgrader := websocket.Upgrader{}
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		// Keep connection alive and respond to pings
+		for {
+			_, msg, err := conn.ReadMessage()
+			if err != nil {
+				return
+			}
+			if string(msg) == "PING" {
+				conn.WriteMessage(websocket.TextMessage, []byte("PONG"))
+			}
+		}
+	}))
+	defer server.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(server.URL, "http")
+
+	client, err := NewClient(wsURL, nil, nil)
+	if err != nil {
+		t.Fatalf("failed to create client: %v", err)
+	}
+
+	impl := client.(*clientImpl)
+	impl.heartbeatInterval = 20 * time.Millisecond
+	impl.disablePing = false
+
+	// Wait for ping loop to start and send some pings
+	time.Sleep(100 * time.Millisecond)
+
+	if err := client.Close(); err != nil {
+		t.Fatalf("failed to close client: %v", err)
+	}
+
+	time.Sleep(100 * time.Millisecond)
+}
+
+// TestWebSocketGoroutineLeaks_SubscriptionCleanup tests that subscription
+// goroutines are properly cleaned up.
+func TestWebSocketGoroutineLeaks_SubscriptionCleanup(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+
+	upgrader := websocket.Upgrader{}
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		defer conn.Close()
+
+		// Read subscription requests and keep connection alive
+		for {
+			_, _, err := conn.ReadMessage()
+			if err != nil {
+				return
+			}
+		}
+	}))
+	defer server.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(server.URL, "http")
+
+	client, err := NewClient(wsURL, nil, nil)
+	if err != nil {
+		t.Fatalf("failed to create client: %v", err)
+	}
+
+	// Create multiple subscriptions
+	ctx := context.Background()
+	stream1, err := client.SubscribeOrderbookStream(ctx, []string{"asset1"})
+	if err != nil {
+		t.Fatalf("failed to subscribe: %v", err)
+	}
+
+	stream2, err := client.SubscribePricesStream(ctx, []string{"asset2"})
+	if err != nil {
+		t.Fatalf("failed to subscribe: %v", err)
+	}
+
+	// Close streams
+	stream1.Close()
+	stream2.Close()
+
+	// Close client
+	if err := client.Close(); err != nil {
+		t.Fatalf("failed to close client: %v", err)
+	}
+
+	time.Sleep(100 * time.Millisecond)
+}