Fix #4757: Race condition in exemplarSchemaMap (#4774)

pkg/connection/connection_lifecycle_test.go

@@ -0,0 +1,109 @@
+package connection
+
+import (
+	"sync"
+	"testing"
+)
+
+// TestExemplarSchemaMapConcurrentAccess tests concurrent access to exemplarSchemaMap
+// This test demonstrates issue #4757 - race condition when writing to exemplarSchemaMap
+// without proper mutex protection.
+func TestExemplarSchemaMapConcurrentAccess(t *testing.T) {
+	// Create a refreshConnectionState with initialized exemplarSchemaMap
+	state := &refreshConnectionState{
+		exemplarSchemaMap:    make(map[string]string),
+		exemplarSchemaMapMut: sync.Mutex{},
+	}
+
+	// Number of concurrent goroutines
+	numGoroutines := 10
+	numIterations := 100
+
+	var wg sync.WaitGroup
+	wg.Add(numGoroutines)
+
+	// Launch multiple goroutines that will concurrently read and write to exemplarSchemaMap
+	for i := 0; i < numGoroutines; i++ {
+		go func(id int) {
+			defer wg.Done()
+
+			for j := 0; j < numIterations; j++ {
+				pluginName := "aws"
+				connectionName := "connection"
+
+				// Simulate the FIXED pattern in executeUpdateForConnections
+				// Read with mutex (line 581-591)
+				state.exemplarSchemaMapMut.Lock()
+				_, haveExemplarSchema := state.exemplarSchemaMap[pluginName]
+				state.exemplarSchemaMapMut.Unlock()
+
+				// FIXED: Write with mutex protection (line 602-604)
+				if !haveExemplarSchema {
+					// Now properly protected with mutex
+					state.exemplarSchemaMapMut.Lock()
+					state.exemplarSchemaMap[pluginName] = connectionName
+					state.exemplarSchemaMapMut.Unlock()
+				}
+			}
+		}(i)
+	}
+
+	// Wait for all goroutines to complete
+	wg.Wait()
+
+	// Verify the map has an entry (basic sanity check)
+	state.exemplarSchemaMapMut.Lock()
+	if len(state.exemplarSchemaMap) == 0 {
+		t.Error("Expected exemplarSchemaMap to have at least one entry")
+	}
+	state.exemplarSchemaMapMut.Unlock()
+}
+
+// TestExemplarSchemaMapRaceCondition specifically tests the race condition pattern
+// found in refresh_connections_state.go:601 - now FIXED
+func TestExemplarSchemaMapRaceCondition(t *testing.T) {
+	// This test now PASSES with -race flag after the bug fix
+	state := &refreshConnectionState{
+		exemplarSchemaMap:    make(map[string]string),
+		exemplarSchemaMapMut: sync.Mutex{},
+	}
+
+	plugins := []string{"aws", "azure", "gcp", "github", "slack"}
+
+	var wg sync.WaitGroup
+
+	// Simulate multiple connections being processed concurrently
+	for _, plugin := range plugins {
+		for i := 0; i < 5; i++ {
+			wg.Add(1)
+			go func(p string, connNum int) {
+				defer wg.Done()
+
+				// This simulates the FIXED code pattern in executeUpdateForConnections
+				state.exemplarSchemaMapMut.Lock()
+				_, haveExemplar := state.exemplarSchemaMap[p]
+				state.exemplarSchemaMapMut.Unlock()
+
+				// FIXED: This write is now protected by the mutex
+				if !haveExemplar {
+					// No more race condition!
+					state.exemplarSchemaMapMut.Lock()
+					state.exemplarSchemaMap[p] = p + "_connection"
+					state.exemplarSchemaMapMut.Unlock()
+				}
+			}(plugin, i)
+		}
+	}
+
+	wg.Wait()
+
+	// Verify all plugins are in the map
+	state.exemplarSchemaMapMut.Lock()
+	defer state.exemplarSchemaMapMut.Unlock()
+
+	for _, plugin := range plugins {
+		if _, ok := state.exemplarSchemaMap[plugin]; !ok {
+			t.Errorf("Expected plugin %s to be in exemplarSchemaMap", plugin)
+		}
+	}
+}

pkg/connection/refresh_connections_state.go

@@ -598,7 +598,10 @@ func (s *refreshConnectionState) executeUpdateForConnections(ctx context.Context
 			// we can clone this plugin, add to exemplarSchemaMap
 			// (AFTER executing the update query)
 			if !haveExemplarSchema && connectionState.CanCloneSchema() {
+				// Fix #4757: Protect map write with mutex to prevent race condition
+				s.exemplarSchemaMapMut.Lock()
 				s.exemplarSchemaMap[connectionState.Plugin] = connectionName
+				s.exemplarSchemaMapMut.Unlock()
 			}
 		}
 	}