From f411273fa261fbd36837042dad4abd75e2dea829 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Sun, 14 Sep 2025 18:28:32 +0100 Subject: [PATCH] fixup! fixup! fixup! Weekend fun with further optimisations. --- autocleanup.go | 44 ++++++++++++++++++- error_recovery_test.go | 18 +++++--- http_client_pool.go | 39 +++++++++++------ memory_monitor.go | 22 ++++++++++ memory_monitor_singleton_test.go | 48 +++++++++++++++------ session/chunking/chunk_manager.go | 71 +++++++++++++++++++++++++------ session_chunk_manager.go | 47 +++++++++++++++++--- simple_memory_test.go | 8 ++-- test_infrastructure.go | 4 ++ universal_cache_singleton.go | 22 +++++----- 10 files changed, 261 insertions(+), 62 deletions(-) diff --git a/autocleanup.go b/autocleanup.go index 8666f35..73a462a 100644 --- a/autocleanup.go +++ b/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") + } + } +} diff --git a/error_recovery_test.go b/error_recovery_test.go index 0d6262d..bd7b5ac 100644 --- a/error_recovery_test.go +++ b/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) } }) diff --git a/http_client_pool.go b/http_client_pool.go index e395d50..470535b 100644 --- a/http_client_pool.go +++ b/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() } diff --git a/memory_monitor.go b/memory_monitor.go index 80266e1..794c302 100644 --- a/memory_monitor.go +++ b/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 +} diff --git a/memory_monitor_singleton_test.go b/memory_monitor_singleton_test.go index 7f4d738..4a4364a 100644 --- a/memory_monitor_singleton_test.go +++ b/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) diff --git a/session/chunking/chunk_manager.go b/session/chunking/chunk_manager.go index 7729fcc..30f9c37 100644 --- a/session/chunking/chunk_manager.go +++ b/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 diff --git a/session_chunk_manager.go b/session_chunk_manager.go index fde5f49..ba3c593 100644 --- a/session_chunk_manager.go +++ b/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) } diff --git a/simple_memory_test.go b/simple_memory_test.go index 553149d..f7fd697 100644 --- a/simple_memory_test.go +++ b/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, }, diff --git a/test_infrastructure.go b/test_infrastructure.go index d3860dd..f9bf222 100644 --- a/test_infrastructure.go +++ b/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() diff --git a/universal_cache_singleton.go b/universal_cache_singleton.go index 9d978cd..b1c950b 100644 --- a/universal_cache_singleton.go +++ b/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, }) })