Clear race conditions

autocleanup.go

@@ -13,6 +13,7 @@ import (
 // It supports both internal and external WaitGroup coordination for complex cleanup scenarios.
 type BackgroundTask struct {
 	stopChan   chan struct{}
+	doneChan   chan struct{} // Signals when the task goroutine has completed
 	taskFunc   func()
 	logger     *Logger
 	externalWG *sync.WaitGroup
@@ -21,9 +22,10 @@ type BackgroundTask struct {
 	interval   time.Duration
 	stopOnce   sync.Once
 	startOnce  sync.Once
-	mu         sync.Mutex
-	stopped    bool
-	started    bool
+	// Use atomic fields to avoid race conditions
+	stopped    int32 // 1 = stopped, 0 = not stopped
+	started    int32 // 1 = started, 0 = not started
+	doneClosed int32 // 1 = doneChan closed, 0 = not closed
 }
 
 // NewBackgroundTask creates a new background task with the specified configuration.
@@ -46,6 +48,7 @@ func NewBackgroundTask(name string, interval time.Duration, taskFunc func(), log
 		name:       name,
 		interval:   interval,
 		stopChan:   make(chan struct{}),
+		doneChan:   make(chan struct{}),
 		taskFunc:   taskFunc,
 		logger:     logger,
 		externalWG: externalWG,
@@ -58,13 +61,15 @@ func NewBackgroundTask(name string, interval time.Duration, taskFunc func(), log
 // This method is safe to call multiple times - only the first call will start the task.
 func (bt *BackgroundTask) Start() {
 	bt.startOnce.Do(func() {
-		bt.mu.Lock()
-		defer bt.mu.Unlock()
-
-		if bt.stopped {
+		// Check if already stopped using atomic operation
+		if atomic.LoadInt32(&bt.stopped) == 1 {
 			if bt.logger != nil {
 				bt.logger.Infof("Attempted to start already stopped task: %s", bt.name)
 			}
+			// Close doneChan since the task won't run
+			if atomic.CompareAndSwapInt32(&bt.doneClosed, 0, 1) {
+				close(bt.doneChan)
+			}
 			return
 		}
 
@@ -74,13 +79,17 @@ func (bt *BackgroundTask) Start() {
 			if bt.logger != nil {
 				bt.logger.Debugf("Cannot start task %s: %v (circuit breaker protection working as expected)", bt.name, err)
 			}
+			// Close doneChan since the task won't run
+			if atomic.CompareAndSwapInt32(&bt.doneClosed, 0, 1) {
+				close(bt.doneChan)
+			}
 			return
 		}
 
 		// Reserve the task slot immediately when starting
 		registry.cb.OnTaskStart(bt.name)
 
-		bt.started = true
+		atomic.StoreInt32(&bt.started, 1)
 		bt.internalWG.Add(1)
 		if bt.externalWG != nil {
 			bt.externalWG.Add(1)
@@ -94,9 +103,17 @@ func (bt *BackgroundTask) Start() {
 // This method is safe to call multiple times.
 func (bt *BackgroundTask) Stop() {
 	bt.stopOnce.Do(func() {
-		bt.mu.Lock()
-		bt.stopped = true
-		bt.mu.Unlock()
+		// Set stopped flag atomically
+		atomic.StoreInt32(&bt.stopped, 1)
+
+		// Check if the task was actually started
+		if atomic.LoadInt32(&bt.started) == 0 {
+			// Task was never started, close doneChan to unblock any waiters
+			if atomic.CompareAndSwapInt32(&bt.doneClosed, 0, 1) {
+				close(bt.doneChan)
+			}
+			return
+		}
 
 		// Safe close with panic recovery
 		func() {
@@ -111,21 +128,19 @@ func (bt *BackgroundTask) Stop() {
 			close(bt.stopChan)
 		}()
 
-		// Wait for the goroutine to finish with timeout
-		done := make(chan struct{})
-		go func() {
-			bt.internalWG.Wait()
-			close(done)
-		}()
-
+		// Wait for the task goroutine to complete using doneChan
+		// This avoids the race condition with WaitGroup
 		select {
-		case <-done:
+		case <-bt.doneChan:
 			// Normal completion
 		case <-time.After(5 * time.Second):
 			if bt.logger != nil {
 				bt.logger.Errorf("Timeout waiting for background task %s to stop", bt.name)
 			}
 		}
+
+		// Wait for the internal WaitGroup synchronously after doneChan signals
+		bt.internalWG.Wait()
 	})
 }
 
@@ -140,6 +155,11 @@ func (bt *BackgroundTask) run() {
 		// Register task completion with circuit breaker
 		registry.cb.OnTaskComplete(bt.name)
 
+		// Close doneChan to signal that the task has completed
+		if atomic.CompareAndSwapInt32(&bt.doneClosed, 0, 1) {
+			close(bt.doneChan)
+		}
+
 		bt.internalWG.Done()
 		if bt.externalWG != nil {
 			bt.externalWG.Done()

autocleanup_circuit_breaker_test.go

@@ -143,7 +143,8 @@ func TestBackgroundTaskCircuitBreaker_SingletonPattern(t *testing.T) {
 
 	// Simulate multiple attempts to start cleanup tasks (like from different middleware instances)
 	var wg sync.WaitGroup
-	tasks := make([]*BackgroundTask, 0, numAttempts)
+	tasks := make([]*BackgroundTask, numAttempts)
+	var tasksMu sync.Mutex
 
 	for i := 0; i < numAttempts; i++ {
 		wg.Add(1)
@@ -175,7 +176,9 @@ func TestBackgroundTaskCircuitBreaker_SingletonPattern(t *testing.T) {
 			// Here we test that even if multiple start, resource usage is controlled
 			task.Start()
 
-			tasks = append(tasks, task)
+			tasksMu.Lock()
+			tasks[id] = task
+			tasksMu.Unlock()
 		}(i)
 	}
 

error_recovery.go

@@ -874,13 +874,18 @@ func (gd *GracefulDegradation) executeFallback(serviceName string) (interface{},
 
 // startHealthCheckRoutine starts the background health check routine
 func (gd *GracefulDegradation) startHealthCheckRoutine() {
-	gd.healthCheckTask = NewBackgroundTask(
+	task := NewBackgroundTask(
 		"graceful-degradation-health-check",
 		gd.config.HealthCheckInterval,
 		gd.performHealthChecks,
 		gd.BaseRecoveryMechanism.logger,
 	)
-	gd.healthCheckTask.Start()
+
+	gd.mutex.Lock()
+	gd.healthCheckTask = task
+	gd.mutex.Unlock()
+
+	task.Start()
 }
 
 // performHealthChecks runs health checks for all registered services
@@ -940,9 +945,13 @@ func (gd *GracefulDegradation) Close() {
 		}
 
 		// Stop health check task
-		if gd.healthCheckTask != nil {
-			gd.healthCheckTask.Stop()
-			gd.healthCheckTask = nil
+		gd.mutex.Lock()
+		task := gd.healthCheckTask
+		gd.mutex.Unlock()
+
+		if task != nil {
+			task.Stop()
+			// Don't set to nil to avoid race conditions
 		}
 
 		gd.logger.Info("GracefulDegradation shut down successfully")

memory_monitor.go

@@ -134,16 +134,19 @@ func (mm *MemoryMonitor) GetCurrentStats() *MemoryStats {
 
 	// Calculate GC frequency
 	gcFrequency := 0.0
-	if mm.lastStats != nil {
-		timeDiff := now.Sub(mm.lastStats.Timestamp).Minutes()
+	mm.mu.RLock()
+	lastStats := mm.lastStats
+	lastGCCount := mm.lastGCCount
+	mm.mu.RUnlock()
+
+	if lastStats != nil {
+		timeDiff := now.Sub(lastStats.Timestamp).Minutes()
 		if timeDiff > 0 {
-			gcDiff := float64(memStats.NumGC - mm.lastGCCount)
+			gcDiff := float64(memStats.NumGC - lastGCCount)
 			gcFrequency = gcDiff / timeDiff
 		}
 	}
 
-	mm.lastGCCount = memStats.NumGC
-
 	stats := &MemoryStats{
 		HeapAllocBytes:    memStats.HeapAlloc,
 		HeapSysBytes:      memStats.HeapSys,
@@ -174,6 +177,7 @@ func (mm *MemoryMonitor) GetCurrentStats() *MemoryStats {
 
 	mm.mu.Lock()
 	mm.lastStats = stats
+	mm.lastGCCount = memStats.NumGC
 	mm.mu.Unlock()
 
 	return stats
@@ -234,32 +238,54 @@ func (mm *MemoryMonitor) updateGoroutineTracking(stats *MemoryStats) {
 
 	// Check for potential goroutine leak
 	if stats.NumGoroutines > mm.baselineGoroutines+int(mm.alertThresholds.GoroutineCount) {
-		if !mm.goroutineLeakAlert {
+		mm.mu.Lock()
+		wasAlert := mm.goroutineLeakAlert
+		if !wasAlert {
 			mm.goroutineLeakAlert = true
+		}
+		mm.mu.Unlock()
+		if !wasAlert {
 			mm.logger.Error("Potential goroutine leak detected: %d goroutines (baseline: %d)",
 				stats.NumGoroutines, mm.baselineGoroutines)
 		}
 	} else {
+		mm.mu.Lock()
 		mm.goroutineLeakAlert = false
+		mm.mu.Unlock()
 	}
 }
 
 // updateHeapGrowthTracking monitors heap growth rate
 func (mm *MemoryMonitor) updateHeapGrowthTracking(stats *MemoryStats) {
-	if mm.lastStats != nil {
-		timeDiff := stats.Timestamp.Sub(mm.lastStats.Timestamp).Seconds()
-		if timeDiff > 0 {
-			heapDiff := float64(stats.HeapAllocBytes) - float64(mm.lastStats.HeapAllocBytes)
-			mm.heapGrowthRate = heapDiff / timeDiff // bytes per second
+	mm.mu.RLock()
+	lastStats := mm.lastStats
+	mm.mu.RUnlock()
 
-			growthRateMB := mm.heapGrowthRate / (1024 * 1024)
+	if lastStats != nil {
+		timeDiff := stats.Timestamp.Sub(lastStats.Timestamp).Seconds()
+		if timeDiff > 0 {
+			heapDiff := float64(stats.HeapAllocBytes) - float64(lastStats.HeapAllocBytes)
+			heapGrowthRate := heapDiff / timeDiff // bytes per second
+
+			mm.mu.Lock()
+			mm.heapGrowthRate = heapGrowthRate
+			mm.mu.Unlock()
+
+			growthRateMB := heapGrowthRate / (1024 * 1024)
 			if growthRateMB > mm.alertThresholds.HeapGrowthRateMB {
-				if !mm.suspiciousGrowth {
+				mm.mu.Lock()
+				wasSuspicious := mm.suspiciousGrowth
+				if !wasSuspicious {
 					mm.suspiciousGrowth = true
+				}
+				mm.mu.Unlock()
+				if !wasSuspicious {
 					mm.logger.Error("Suspicious heap growth rate: %.2f MB/sec", growthRateMB)
 				}
 			} else {
+				mm.mu.Lock()
 				mm.suspiciousGrowth = false
+				mm.mu.Unlock()
 			}
 		}
 	}