fixup! fixup! fixup! Weekend fun with further optimisations.

autocleanup.go

@@ -253,7 +253,7 @@ type TaskCircuitBreaker struct {
 // NewTaskCircuitBreaker creates a new circuit breaker for background tasks
 // with concurrency limiting capability
 func NewTaskCircuitBreaker(failureThreshold int32, timeout time.Duration, logger *Logger) *TaskCircuitBreaker {
-	maxConcurrent := int32(50) // Default reasonable limit
+	maxConcurrent := int32(20) // CRITICAL FIX: Reduced from 50 to 20 for memory safety
 	return &TaskCircuitBreaker{
 		state:            int32(CircuitBreakerClosed),
 		failureThreshold: failureThreshold,
@@ -831,3 +831,45 @@ func (mm *TaskMemoryMonitor) ForceGC() (before, after TaskMemoryStats, err error
 
 	return before, after, nil
 }
+
+// ShutdownAllTasks gracefully shuts down all background tasks
+// CRITICAL FIX: Ensures proper termination of all goroutines in production
+func ShutdownAllTasks() {
+	registry := GetGlobalTaskRegistry()
+
+	registry.mu.Lock()
+	tasks := make([]*BackgroundTask, 0, len(registry.tasks))
+	for _, task := range registry.tasks {
+		tasks = append(tasks, task)
+	}
+	registry.mu.Unlock()
+
+	// Stop all tasks in parallel
+	var wg sync.WaitGroup
+	for _, task := range tasks {
+		wg.Add(1)
+		go func(t *BackgroundTask) {
+			defer wg.Done()
+			if t != nil {
+				t.Stop()
+			}
+		}(task)
+	}
+
+	// Wait with timeout
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		// All tasks stopped successfully
+	case <-time.After(10 * time.Second):
+		// Timeout - tasks may still be running
+		if registry.logger != nil {
+			registry.logger.Errorf("Timeout waiting for all background tasks to stop")
+		}
+	}
+}

error_recovery_test.go

@@ -652,15 +652,15 @@ func TestRetryExecutorBackoffAndJitter(t *testing.T) {
 	t.Run("Retry with timeout context", func(t *testing.T) {
 		config := RetryConfig{
 			MaxAttempts:     5,
-			InitialDelay:    50 * time.Millisecond,
-			MaxDelay:        200 * time.Millisecond,
+			InitialDelay:    30 * time.Millisecond,  // Reduced from 50ms
+			MaxDelay:        100 * time.Millisecond, // Reduced from 200ms
 			BackoffFactor:   2.0,
 			EnableJitter:    false,
 			RetryableErrors: []string{"timeout", "temporary failure"},
 		}
 		re := NewRetryExecutor(config, logger)
 
-		ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond)
+		ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) // Reduced from 150ms
 		defer cancel()
 
 		attempts := 0
@@ -681,9 +681,15 @@ func TestRetryExecutorBackoffAndJitter(t *testing.T) {
 			t.Errorf("Expected fewer than %d attempts due to context timeout, got %d", config.MaxAttempts, attempts)
 		}
 
-		// Should have respected context timeout
-		if duration > 200*time.Millisecond {
-			t.Errorf("Operation took too long: %v", duration)
+		// Should have respected context timeout with some buffer for processing
+		maxAllowedDuration := 150 * time.Millisecond // Increased buffer from 200ms to 150ms
+		if duration > maxAllowedDuration {
+			t.Errorf("Operation took too long: %v (max allowed: %v)", duration, maxAllowedDuration)
+		}
+
+		// But should be close to the context timeout
+		if duration < 80*time.Millisecond {
+			t.Errorf("Operation completed too quickly: %v (should be close to context timeout of 100ms)", duration)
 		}
 	})
 

http_client_pool.go

@@ -13,6 +13,8 @@ type SharedTransportPool struct {
 	mu         sync.RWMutex
 	transports map[string]*sharedTransport
 	maxConns   int
+	ctx        context.Context
+	cancel     context.CancelFunc
 }
 
 type sharedTransport struct {
@@ -29,12 +31,15 @@ var (
 // GetGlobalTransportPool returns the singleton transport pool instance
 func GetGlobalTransportPool() *SharedTransportPool {
 	globalTransportPoolOnce.Do(func() {
+		ctx, cancel := context.WithCancel(context.Background())
 		globalTransportPool = &SharedTransportPool{
 			transports: make(map[string]*sharedTransport),
 			maxConns:   100, // Total connection limit across all transports
+			ctx:        ctx,
+			cancel:     cancel,
 		}
-		// Start cleanup goroutine
-		go globalTransportPool.cleanupIdleTransports()
+		// Start cleanup goroutine with context cancellation
+		go globalTransportPool.cleanupIdleTransports(ctx)
 	})
 	return globalTransportPool
 }
@@ -103,21 +108,26 @@ func (p *SharedTransportPool) ReleaseTransport(transport *http.Transport) {
 }
 
 // cleanupIdleTransports periodically cleans up unused transports
-func (p *SharedTransportPool) cleanupIdleTransports() {
+func (p *SharedTransportPool) cleanupIdleTransports(ctx context.Context) {
 	ticker := time.NewTicker(1 * time.Minute)
 	defer ticker.Stop()
 
-	for range ticker.C {
-		p.mu.Lock()
-		now := time.Now()
-		for transportKey, shared := range p.transports {
-			// Clean up transports not used for 2 minutes with no references
-			if shared.refCount <= 0 && now.Sub(shared.lastUsed) > 2*time.Minute {
-				shared.transport.CloseIdleConnections()
-				delete(p.transports, transportKey)
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			p.mu.Lock()
+			now := time.Now()
+			for transportKey, shared := range p.transports {
+				// Clean up transports not used for 2 minutes with no references
+				if shared.refCount <= 0 && now.Sub(shared.lastUsed) > 2*time.Minute {
+					shared.transport.CloseIdleConnections()
+					delete(p.transports, transportKey)
+				}
 			}
+			p.mu.Unlock()
 		}
-		p.mu.Unlock()
 	}
 }
 
@@ -132,6 +142,11 @@ func (p *SharedTransportPool) Cleanup() {
 	p.mu.Lock()
 	defer p.mu.Unlock()
 
+	// Stop the cleanup goroutine
+	if p.cancel != nil {
+		p.cancel()
+	}
+
 	for _, shared := range p.transports {
 		shared.transport.CloseIdleConnections()
 	}

memory_monitor.go

@@ -449,3 +449,25 @@ func GetGlobalMemoryMonitor() *MemoryMonitor {
 	})
 	return globalMemoryMonitor
 }
+
+// ResetGlobalMemoryMonitor resets the global memory monitor for testing
+// This should only be used in tests to prevent state pollution between tests
+func ResetGlobalMemoryMonitor() {
+	globalMonitoringMutex.Lock()
+	defer globalMonitoringMutex.Unlock()
+
+	if globalMemoryMonitor != nil {
+		// Stop monitoring if it's active
+		if globalMonitoringStarted {
+			registry := GetGlobalTaskRegistry()
+			if task, exists := registry.GetTask("memory-monitor"); exists {
+				task.Stop()
+			}
+		}
+		globalMemoryMonitor = nil
+	}
+
+	// Reset the singleton state
+	globalMemoryMonitorOnce = sync.Once{}
+	globalMonitoringStarted = false
+}

memory_monitor_singleton_test.go

@@ -9,15 +9,23 @@ import (
 )
 
 func TestMemoryMonitorSingleton(t *testing.T) {
-	// Reset global state
-	globalMonitoringMutex.Lock()
-	globalMonitoringStarted = false
-	globalMonitoringMutex.Unlock()
+	// Use proper cleanup helper
+	TestCleanupHelper(t)
+
+	// Reset global state before test
+	ResetGlobalTaskRegistry()
+	ResetGlobalMemoryMonitor()
 
 	// Clear task registry for clean test
 	registry := GetGlobalTaskRegistry()
 	registry.StopAllTasks()
 
+	// Add cleanup at the end of test
+	defer func() {
+		ResetGlobalMemoryMonitor()
+		ResetGlobalTaskRegistry()
+	}()
+
 	// Test that multiple StartMonitoring calls don't create multiple monitors
 	monitor := GetGlobalMemoryMonitor()
 	ctx, cancel := context.WithCancel(context.Background())
@@ -81,14 +89,22 @@ func TestMemoryMonitorSingleton(t *testing.T) {
 }
 
 func TestMemoryMonitorSingletonTaskCreation(t *testing.T) {
-	// Reset global state
-	globalMonitoringMutex.Lock()
-	globalMonitoringStarted = false
-	globalMonitoringMutex.Unlock()
+	// Use proper cleanup helper
+	TestCleanupHelper(t)
+
+	// Reset global state before test
+	ResetGlobalTaskRegistry()
+	ResetGlobalMemoryMonitor()
 
 	registry := GetGlobalTaskRegistry()
 	registry.StopAllTasks()
 
+	// Add cleanup at the end of test
+	defer func() {
+		ResetGlobalMemoryMonitor()
+		ResetGlobalTaskRegistry()
+	}()
+
 	monitor1 := GetGlobalMemoryMonitor()
 	monitor2 := GetGlobalMemoryMonitor()
 
@@ -122,14 +138,22 @@ func TestMemoryMonitorSingletonTaskCreation(t *testing.T) {
 }
 
 func TestMemoryMonitorTaskRegistryEnforcement(t *testing.T) {
-	// Reset global state
-	globalMonitoringMutex.Lock()
-	globalMonitoringStarted = false
-	globalMonitoringMutex.Unlock()
+	// Use proper cleanup helper
+	TestCleanupHelper(t)
+
+	// Reset global state before test
+	ResetGlobalTaskRegistry()
+	ResetGlobalMemoryMonitor()
 
 	registry := GetGlobalTaskRegistry()
 	registry.StopAllTasks()
 
+	// Add cleanup at the end of test
+	defer func() {
+		ResetGlobalMemoryMonitor()
+		ResetGlobalTaskRegistry()
+	}()
+
 	// Try to create memory-monitor task directly through registry
 	task1, err1 := registry.CreateSingletonTask("memory-monitor", 50*time.Millisecond,
 		func() {}, GetSingletonNoOpLogger(), nil)

session/chunking/chunk_manager.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/gorilla/sessions"
@@ -27,6 +28,12 @@ type TokenConfig struct {
 	RequireJWTFormat  bool
 }
 
+// Global session tracking to prevent memory leaks across all instances
+var (
+	globalSessionCount int64 = 0
+	globalMaxSessions  int64 = 5000 // CRITICAL FIX: Global limit of 5000 total sessions
+)
+
 // Predefined configurations for each token type
 var (
 	AccessTokenConfig = TokenConfig{
@@ -108,8 +115,8 @@ func NewChunkManager(logger Logger) *ChunkManager {
 		logger:      logger,
 		mutex:       &sync.RWMutex{},
 		sessionMap:  make(map[string]*SessionEntry),
-		maxSessions: 1000, // Reasonable limit to prevent memory leaks
-		sessionTTL:  24 * time.Hour,
+		maxSessions: 200,              // CRITICAL FIX: Reduced from 1000 to 200 per instance
+		sessionTTL:  15 * time.Minute, // CRITICAL FIX: Reduced from 24h to 15 minutes
 		lastCleanup: time.Now(),
 	}
 }
@@ -327,21 +334,60 @@ func (cm *ChunkManager) StoreSession(key string, session *sessions.Session) {
 	cm.mutex.Lock()
 	defer cm.mutex.Unlock()
 
-	// Enforce maximum session limit
-	if len(cm.sessionMap) >= cm.maxSessions {
-		// Remove oldest entry
-		var oldestKey string
-		var oldestTime time.Time = time.Now()
+	// CRITICAL FIX: Aggressive session limit enforcement
+	currentLocal := len(cm.sessionMap)
+	currentGlobal := atomic.LoadInt64(&globalSessionCount)
 
+	shouldEvict := false
+	targetCapacity := cm.maxSessions
+
+	// Check global limit first (more critical)
+	if currentGlobal >= globalMaxSessions {
+		shouldEvict = true
+		targetCapacity = cm.maxSessions / 4 // Aggressive reduction to 25%
+	} else if currentGlobal >= globalMaxSessions*8/10 { // 80% of global
+		shouldEvict = true
+		targetCapacity = cm.maxSessions / 2 // Reduce to 50%
+	} else if currentLocal >= cm.maxSessions {
+		shouldEvict = true
+		targetCapacity = cm.maxSessions * 3 / 4 // Reduce to 75%
+	}
+
+	if shouldEvict {
+		// Find oldest sessions to remove
+		type sessionAge struct {
+			key      string
+			lastUsed time.Time
+		}
+
+		sessions := make([]sessionAge, 0, currentLocal)
 		for k, entry := range cm.sessionMap {
-			if entry.LastUsed.Before(oldestTime) {
-				oldestTime = entry.LastUsed
-				oldestKey = k
+			sessions = append(sessions, sessionAge{key: k, lastUsed: entry.LastUsed})
+		}
+
+		// Sort by last used time (oldest first)
+		for i := 0; i < len(sessions)-1; i++ {
+			for j := i + 1; j < len(sessions); j++ {
+				if sessions[i].lastUsed.After(sessions[j].lastUsed) {
+					sessions[i], sessions[j] = sessions[j], sessions[i]
+				}
 			}
 		}
 
-		if oldestKey != "" {
-			delete(cm.sessionMap, oldestKey)
+		// Remove excess sessions
+		excessCount := currentLocal - targetCapacity
+		if excessCount < 0 {
+			excessCount = 0
+		}
+
+		removedCount := int64(0)
+		for i := 0; i < excessCount && i < len(sessions); i++ {
+			delete(cm.sessionMap, sessions[i].key)
+			removedCount++
+		}
+
+		if removedCount > 0 {
+			atomic.AddInt64(&globalSessionCount, -removedCount)
 		}
 	}
 
@@ -350,6 +396,7 @@ func (cm *ChunkManager) StoreSession(key string, session *sessions.Session) {
 		ExpiresAt: time.Now().Add(cm.sessionTTL),
 		LastUsed:  time.Now(),
 	}
+	atomic.AddInt64(&globalSessionCount, 1) // CRITICAL FIX: Track addition
 }
 
 // GetSession retrieves a session from the session map

session_chunk_manager.go

@@ -27,6 +27,12 @@ type TokenConfig struct {
 	RequireJWTFormat  bool
 }
 
+// Global session tracking to prevent memory leaks across all instances
+var (
+	globalSessionCount int64 = 0
+	globalMaxSessions  int64 = 5000 // CRITICAL FIX: Global limit of 5000 total sessions
+)
+
 // Predefined configurations for each token type
 var (
 	AccessTokenConfig = TokenConfig{
@@ -118,8 +124,8 @@ func NewChunkManager(logger *Logger) *ChunkManager {
 		ctx:         ctx,
 		cancel:      cancel,
 		sessionMap:  make(map[string]*SessionEntry),
-		maxSessions: 1000,           // Prevent unbounded growth
-		sessionTTL:  24 * time.Hour, // Auto-expire old sessions
+		maxSessions: 200,              // CRITICAL FIX: Reduced from 1000 to 200 per instance
+		sessionTTL:  15 * time.Minute, // CRITICAL FIX: Reduced from 24h to 15 minutes
 		lastCleanup: time.Now(),
 	}
 
@@ -1132,7 +1138,26 @@ func (cm *ChunkManager) CleanupExpiredSessions(force ...bool) {
 
 // enforceSessionLimit removes oldest sessions when limit is exceeded
 func (cm *ChunkManager) enforceSessionLimit() {
-	if len(cm.sessionMap) <= cm.maxSessions {
+	currentLocal := len(cm.sessionMap)
+	currentGlobal := atomic.LoadInt64(&globalSessionCount)
+
+	// CRITICAL FIX: Aggressive eviction when approaching limits
+	shouldEvict := false
+	targetCapacity := cm.maxSessions
+
+	// Check global limit first (more critical)
+	if currentGlobal >= globalMaxSessions {
+		shouldEvict = true
+		targetCapacity = cm.maxSessions / 4 // Aggressive reduction to 25%
+	} else if currentGlobal >= globalMaxSessions*8/10 { // 80% of global
+		shouldEvict = true
+		targetCapacity = cm.maxSessions / 2 // Reduce to 50%
+	} else if currentLocal >= cm.maxSessions {
+		shouldEvict = true
+		targetCapacity = cm.maxSessions * 3 / 4 // Reduce to 75%
+	}
+
+	if !shouldEvict {
 		return
 	}
 
@@ -1156,18 +1181,30 @@ func (cm *ChunkManager) enforceSessionLimit() {
 		}
 	}
 
-	// Remove excess sessions and track memory
-	excessCount := len(cm.sessionMap) - cm.maxSessions
+	// Remove excess sessions and track memory - CRITICAL FIX: More aggressive
+	excessCount := currentLocal - targetCapacity
+	if excessCount < 0 {
+		excessCount = 0
+	}
+
 	totalBytesFreed := int64(0)
+	removedCount := int64(0)
+
 	for i := 0; i < excessCount && i < len(sessions); i++ {
 		key := sessions[i].key
 		if entry, exists := cm.sessionMap[key]; exists {
 			totalBytesFreed += entry.SizeEstimate
 			atomic.AddInt64(&cm.bytesAllocated, -entry.SizeEstimate)
+			removedCount++
 		}
 		delete(cm.sessionMap, key)
 	}
 
+	// Update global count
+	if removedCount > 0 {
+		atomic.AddInt64(&globalSessionCount, -removedCount)
+	}
+
 	cm.logger.Infof("Enforced session limit: removed %d excess sessions, freed %d bytes",
 		excessCount, totalBytesFreed)
 }

simple_memory_test.go

@@ -144,8 +144,8 @@ func TestChunkManagerInitialization(t *testing.T) {
 			Description: "Verify chunk manager can be created with proper defaults",
 			Input:       nil, // No specific input needed
 			Expected: map[string]interface{}{
-				"maxSessions": 1000,
-				"sessionTTL":  24 * time.Hour,
+				"maxSessions": 200,
+				"sessionTTL":  15 * time.Minute,
 				"notNil":      true,
 			},
 			Timeout: 5 * time.Second,
@@ -164,8 +164,8 @@ func TestChunkManagerInitialization(t *testing.T) {
 			Input:       "singleton_logger",
 			Expected: map[string]interface{}{
 				"hasLogger":   true,
-				"maxSessions": 1000,
-				"sessionTTL":  24 * time.Hour,
+				"maxSessions": 200,
+				"sessionTTL":  15 * time.Minute,
 			},
 			Timeout: 5 * time.Second,
 		},

test_infrastructure.go

@@ -110,6 +110,10 @@ func (g *GlobalTestCleanup) CleanupAll() {
 		runtime.GC() // Force GC to help clean up
 	}
 
+	// Reset all global singletons to prevent state pollution between tests
+	ResetGlobalMemoryMonitor()
+	ResetGlobalTaskRegistry()
+
 	// Give background tasks time to finish cleanup
 	time.Sleep(100 * time.Millisecond)
 	runtime.GC()

universal_cache_singleton.go

@@ -32,12 +32,13 @@ func GetUniversalCacheManager(logger *Logger) *UniversalCacheManager {
 			logger: logger,
 		}
 
-		// Initialize token cache
+		// Initialize token cache - CRITICAL FIX: Reduced from 5000 to 1000
 		universalCacheManager.tokenCache = NewUniversalCache(UniversalCacheConfig{
-			Type:       CacheTypeToken,
-			MaxSize:    5000,
-			DefaultTTL: 1 * time.Hour,
-			Logger:     logger,
+			Type:           CacheTypeToken,
+			MaxSize:        1000,            // CRITICAL FIX: Reduced from 5000 to 1000 items
+			MaxMemoryBytes: 5 * 1024 * 1024, // CRITICAL FIX: Added 5MB memory limit
+			DefaultTTL:     1 * time.Hour,
+			Logger:         logger,
 		})
 
 		// Initialize blacklist cache
@@ -76,12 +77,13 @@ func GetUniversalCacheManager(logger *Logger) *UniversalCacheManager {
 			Logger:     logger,
 		})
 
-		// Initialize session cache
+		// Initialize session cache - CRITICAL FIX: Reduced from 10000 to 2000
 		universalCacheManager.sessionCache = NewUniversalCache(UniversalCacheConfig{
-			Type:       CacheTypeSession,
-			MaxSize:    10000,
-			DefaultTTL: 30 * time.Minute,
-			Logger:     logger,
+			Type:           CacheTypeSession,
+			MaxSize:        2000,            // CRITICAL FIX: Reduced from 10000 to 2000 items
+			MaxMemoryBytes: 5 * 1024 * 1024, // CRITICAL FIX: Added 5MB memory limit
+			DefaultTTL:     30 * time.Minute,
+			Logger:         logger,
 		})
 	})