Merge branch 'main' into schema-update/v0.64.2-7ee8e196a121397a

Author: lpcox

internal/launcher/connection_pool.go

@@ -206,8 +206,12 @@ func (p *SessionConnectionPool) Stop() {
 
 // Get retrieves a connection from the pool
 func (p *SessionConnectionPool) Get(backendID, sessionID string) (*mcp.Connection, bool) {
-	p.mu.RLock()
-	defer p.mu.RUnlock()
+	// Use a write lock for the entire operation to atomically read and update
+	// metadata. The previous pattern of acquiring a read lock, manually unlocking,
+	// and then acquiring a write lock introduced a race window in which another
+	// goroutine could delete the connection between the two lock acquisitions.
+	p.mu.Lock()
+	defer p.mu.Unlock()
 
 	key := ConnectionKey{BackendID: backendID, SessionID: sessionID}
 	metadata, exists := p.connections[key]
@@ -225,14 +229,9 @@ func (p *SessionConnectionPool) Get(backendID, sessionID string) (*mcp.Connectio
 	logPool.Printf("Reusing connection: backend=%s, session=%s, requests=%d",
 		backendID, sessionID, metadata.RequestCount)
 
-	// Update last used time and state (need write lock for this)
-	p.mu.RUnlock()
-	p.mu.Lock()
 	metadata.LastUsedAt = time.Now()
 	metadata.RequestCount++
 	metadata.State = ConnectionStateActive
-	p.mu.Unlock()
-	p.mu.RLock()
 
 	return metadata.Connection, true
 }

internal/launcher/connection_pool_test.go

@@ -242,6 +242,48 @@ func TestConnectionPoolConcurrency(t *testing.T) {
 	assert.Equal(t, 1000, metadata.RequestCount)
 }
 
+// TestConnectionPoolConcurrencyWithDeletes verifies that concurrent Get and Delete
+// operations do not race. Previously, Get used a manual RUnlock/Lock/RLock upgrade
+// that created a window in which another goroutine could delete the connection,
+// causing Get to update and return a deleted connection.
+func TestConnectionPoolConcurrencyWithDeletes(t *testing.T) {
+	ctx := context.Background()
+	pool := NewSessionConnectionPool(ctx)
+
+	mockConn := &mcp.Connection{}
+	pool.Set("backend1", "session1", mockConn)
+
+	done := make(chan bool)
+
+	// Goroutines that continuously Get the connection
+	for i := 0; i < 5; i++ {
+		go func() {
+			for j := 0; j < 200; j++ {
+				pool.Get("backend1", "session1")
+			}
+			done <- true
+		}()
+	}
+
+	// Goroutines that interleave Set and Delete operations
+	for i := 0; i < 3; i++ {
+		go func() {
+			for j := 0; j < 50; j++ {
+				pool.Set("backend1", "session1", mockConn)
+				pool.Delete("backend1", "session1")
+			}
+			done <- true
+		}()
+	}
+
+	for i := 0; i < 8; i++ {
+		<-done
+	}
+	// No assertion needed — the goal is to detect races via -race flag.
+	// If the race condition is present, this test will fail non-deterministically
+	// under the Go race detector (go test -race).
+}
+
 func TestConnectionPoolDeleteNonExistent(t *testing.T) {
 	ctx := context.Background()
 	pool := NewSessionConnectionPool(ctx)