From 7db2f8d66c352566843c0efa38ec5ac4eb585722 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Tue, 2 Sep 2025 17:03:50 +0100 Subject: [PATCH] fixup! Metadata cache leak fix + profiling --- profiling.go | 84 +++++++++++++++++ profiling_test.go | 228 ++++++++++++++++++++++++++++++++++++++++++++++ session.go | 6 +- 3 files changed, 317 insertions(+), 1 deletion(-) diff --git a/profiling.go b/profiling.go index aa88bf4..0ed439e 100644 --- a/profiling.go +++ b/profiling.go @@ -734,6 +734,90 @@ func (tcp *TokenCompressionProfiler) AnalyzeLeaks(baseline, current *MemorySnaps return analysis } +// MemoryPoolProfiler monitors memory pool usage and detects leaks +type MemoryPoolProfiler struct { + memoryPoolManager *MemoryPoolManager + tokenCompressionPool *TokenCompressionPool + logger *Logger +} + +// NewMemoryPoolProfiler creates a new memory pool profiler +func NewMemoryPoolProfiler(memoryPoolManager *MemoryPoolManager, tokenCompressionPool *TokenCompressionPool, logger *Logger) *MemoryPoolProfiler { + if logger == nil { + logger = newNoOpLogger() + } + return &MemoryPoolProfiler{ + memoryPoolManager: memoryPoolManager, + tokenCompressionPool: tokenCompressionPool, + logger: logger, + } +} + +// TakeSnapshot captures memory pool statistics +func (mpp *MemoryPoolProfiler) TakeSnapshot() (*MemorySnapshot, error) { + snapshot := &MemorySnapshot{ + Timestamp: time.Now(), + CustomMetrics: make(map[string]interface{}), + } + + // Capture runtime stats + runtime.ReadMemStats(&snapshot.RuntimeStats) + + // Add memory pool metrics + if mpp.memoryPoolManager != nil { + snapshot.CustomMetrics["memory_pool_active"] = true + // Note: sync.Pool doesn't expose internal statistics, so we track usage patterns + } + + if mpp.tokenCompressionPool != nil { + snapshot.CustomMetrics["token_compression_pool_active"] = true + } + + return snapshot, nil +} + +// StartProfiling begins profiling (no-op for memory pools) +func (mpp *MemoryPoolProfiler) StartProfiling(config ProfilingConfig) error { + return nil +} + +// StopProfiling ends profiling +func (mpp *MemoryPoolProfiler) StopProfiling() (*MemorySnapshot, error) { + return mpp.TakeSnapshot() +} + +// GetCurrentStats returns current memory statistics +func (mpp *MemoryPoolProfiler) GetCurrentStats() *runtime.MemStats { + stats := &runtime.MemStats{} + runtime.ReadMemStats(stats) + return stats +} + +// AnalyzeLeaks analyzes memory pools for leaks +func (mpp *MemoryPoolProfiler) AnalyzeLeaks(baseline, current *MemorySnapshot) *LeakAnalysis { + analysis := &LeakAnalysis{ + SuspectedLeaks: make([]string, 0), + Recommendations: make([]string, 0), + } + + if baseline == nil || current == nil { + analysis.LeakDescription = "Insufficient memory pool data" + return analysis + } + + // Check for memory leaks in pool operations + memoryIncrease := current.RuntimeStats.Alloc - baseline.RuntimeStats.Alloc + if memoryIncrease > 5*1024*1024 { // 5MB threshold for pool operations + analysis.HasLeak = true + analysis.SuspectedLeaks = append(analysis.SuspectedLeaks, + "Memory pool operations caused significant memory increase") + analysis.Recommendations = append(analysis.Recommendations, + "Check for objects not being returned to memory pools properly") + } + + return analysis +} + // Global profiling manager instance var globalProfilingManager *ProfilingManager var profilingManagerOnce sync.Once diff --git a/profiling_test.go b/profiling_test.go index cab1746..ae1062d 100644 --- a/profiling_test.go +++ b/profiling_test.go @@ -558,3 +558,231 @@ func TestProviderMetadataMemoryLeakDetection(t *testing.T) { float64(intermediate.RuntimeStats.Alloc)/(1024*1024), float64(current.RuntimeStats.Alloc)/(1024*1024)) } + +// TestMemoryPoolLeakDetection tests for memory leaks in memory pool operations +func TestMemoryPoolLeakDetection(t *testing.T) { + logger := NewLogger("debug") + + strictMode := os.Getenv("STRICT_MEMORY_TEST") == "true" + if strictMode { + t.Log("Running in strict memory test mode - will fail on detected leaks") + } else { + t.Log("Running in lenient memory test mode - will log warnings instead of failing") + } + + config := LeakDetectionConfig{ + EnableLeakDetection: true, + LeakThresholdMB: 10, + } + + mto := NewMemoryTestOrchestrator(config, logger) + + // Create memory pool manager and token compression pool + memoryPoolManager := NewMemoryPoolManager() + tokenCompressionPool := NewTokenCompressionPool() + + // Create profiler for memory pools + profiler := NewMemoryPoolProfiler(memoryPoolManager, tokenCompressionPool, logger) + mto.RegisterComponent("memory_pools", profiler) + + // Take initial baseline + baseline, err := profiler.TakeSnapshot() + if err != nil { + t.Fatalf("Failed to take baseline snapshot: %v", err) + } + + initialGoroutines := runtime.NumGoroutine() + + // Phase 1: Simulate various memory pool operations + t.Log("Phase 1: Testing memory pool operations with various patterns...") + + // Test compression buffer pool + for i := 0; i < 100; i++ { + buf := memoryPoolManager.GetCompressionBuffer() + // Simulate some work with the buffer + buf.WriteString(fmt.Sprintf("test data %d", i)) + // Properly return buffer to pool + memoryPoolManager.PutCompressionBuffer(buf) + } + + // Test JWT parsing buffer pool + for i := 0; i < 50; i++ { + jwtBuf := memoryPoolManager.GetJWTParsingBuffer() + // Simulate JWT parsing operations + jwtBuf.HeaderBuf = append(jwtBuf.HeaderBuf, []byte("header")...) + jwtBuf.PayloadBuf = append(jwtBuf.PayloadBuf, []byte("payload")...) + jwtBuf.SignatureBuf = append(jwtBuf.SignatureBuf, []byte("signature")...) + // Properly return buffer to pool + memoryPoolManager.PutJWTParsingBuffer(jwtBuf) + } + + // Test HTTP response buffer pool + for i := 0; i < 75; i++ { + httpBuf := memoryPoolManager.GetHTTPResponseBuffer() + // Simulate HTTP response processing + copy(httpBuf[:min(len(httpBuf), 100)], []byte("http response data")) + // Properly return buffer to pool + memoryPoolManager.PutHTTPResponseBuffer(httpBuf) + } + + // Test string builder pool + for i := 0; i < 60; i++ { + sb := memoryPoolManager.GetStringBuilder() + // Simulate string building operations + sb.WriteString(fmt.Sprintf("built string %d", i)) + _ = sb.String() // Use the result + // Properly return string builder to pool + memoryPoolManager.PutStringBuilder(sb) + } + + // Test token compression pool + for i := 0; i < 40; i++ { + compBuf := tokenCompressionPool.GetCompressionBuffer() + // Simulate compression operations + compBuf.WriteString(fmt.Sprintf("compress data %d", i)) + // Properly return buffer to pool + tokenCompressionPool.PutCompressionBuffer(compBuf) + + decompBuf := tokenCompressionPool.GetDecompressionBuffer() + // Simulate decompression operations + decompBuf.WriteString(fmt.Sprintf("decompress data %d", i)) + // Properly return buffer to pool + tokenCompressionPool.PutDecompressionBuffer(decompBuf) + + sb := tokenCompressionPool.GetStringBuilder() + // Simulate string operations + sb.WriteString(fmt.Sprintf("token string %d", i)) + _ = sb.String() + // Properly return string builder to pool + tokenCompressionPool.PutStringBuilder(sb) + } + + // Take intermediate snapshot + intermediate, err := profiler.TakeSnapshot() + if err != nil { + t.Fatalf("Failed to take intermediate snapshot: %v", err) + } + + // Phase 2: Continue with more intensive operations to test sustained usage + t.Log("Phase 2: Testing sustained memory pool usage...") + + // Simulate mixed operations with varying patterns + for i := 0; i < 200; i++ { + // Mix different pool operations + switch i % 4 { + case 0: + buf := memoryPoolManager.GetCompressionBuffer() + buf.WriteString("mixed operation data") + memoryPoolManager.PutCompressionBuffer(buf) + case 1: + jwtBuf := memoryPoolManager.GetJWTParsingBuffer() + jwtBuf.HeaderBuf = append(jwtBuf.HeaderBuf, []byte("mixed")...) + memoryPoolManager.PutJWTParsingBuffer(jwtBuf) + case 2: + httpBuf := memoryPoolManager.GetHTTPResponseBuffer() + copy(httpBuf[:min(len(httpBuf), 50)], []byte("mixed http")) + memoryPoolManager.PutHTTPResponseBuffer(httpBuf) + case 3: + sb := memoryPoolManager.GetStringBuilder() + sb.WriteString("mixed string building") + _ = sb.String() + memoryPoolManager.PutStringBuilder(sb) + } + } + + // Take final snapshot + current, err := profiler.TakeSnapshot() + if err != nil { + t.Fatalf("Failed to take current snapshot: %v", err) + } + + finalGoroutines := runtime.NumGoroutine() + + // Analyze for leaks + analysis := profiler.AnalyzeLeaks(baseline, current) + + // Assertions for memory leaks + if analysis.HasLeak { + if strictMode { + t.Errorf("Memory leak detected in memory pool operations: %s", analysis.LeakDescription) + for _, leak := range analysis.SuspectedLeaks { + t.Errorf("Suspected leak: %s", leak) + } + } else { + t.Logf("Memory leak warning in memory pool operations: %s", analysis.LeakDescription) + for _, leak := range analysis.SuspectedLeaks { + t.Logf("Suspected leak: %s", leak) + } + } + for _, rec := range analysis.Recommendations { + t.Logf("Recommendation: %s", rec) + } + } + + // Check total memory growth + totalMemoryIncrease := current.RuntimeStats.Alloc - baseline.RuntimeStats.Alloc + if totalMemoryIncrease > 15*1024*1024 { // 15MB threshold for entire test + if strictMode { + t.Errorf("Total memory usage increased by %.2f MB during memory pool operations", float64(totalMemoryIncrease)/(1024*1024)) + } else { + t.Logf("Total memory usage increased by %.2f MB during memory pool operations", float64(totalMemoryIncrease)/(1024*1024)) + } + } + + // Check for gradual memory growth patterns + intermediateMemoryIncrease := intermediate.RuntimeStats.Alloc - baseline.RuntimeStats.Alloc + if intermediateMemoryIncrease > 8*1024*1024 { // 8MB threshold for first phase + if strictMode { + t.Errorf("Memory usage increased by %.2f MB during first phase of memory pool operations", float64(intermediateMemoryIncrease)/(1024*1024)) + } else { + t.Logf("Memory usage increased by %.2f MB during first phase of memory pool operations", float64(intermediateMemoryIncrease)/(1024*1024)) + } + } + + // Check goroutine count stability + goroutineIncrease := finalGoroutines - initialGoroutines + if goroutineIncrease > 3 { // Allow small variance for test environment + if strictMode { + t.Errorf("Goroutine count increased by %d during memory pool operations (initial: %d, final: %d)", + goroutineIncrease, initialGoroutines, finalGoroutines) + } else { + t.Logf("Goroutine count increased by %d during memory pool operations (initial: %d, final: %d)", + goroutineIncrease, initialGoroutines, finalGoroutines) + } + } + + // Phase 3: Test cleanup verification + t.Log("Phase 3: Testing cleanup verification...") + + // Force garbage collection to see if pools are properly managed + runtime.GC() + runtime.GC() // Run twice to ensure cleanup + + time.Sleep(10 * time.Millisecond) // Allow cleanup to complete + + // Take post-cleanup snapshot + postCleanup, err := profiler.TakeSnapshot() + if err != nil { + t.Fatalf("Failed to take post-cleanup snapshot: %v", err) + } + + // Check if memory decreased after cleanup + if postCleanup.RuntimeStats.Alloc < current.RuntimeStats.Alloc { + memoryDecrease := current.RuntimeStats.Alloc - postCleanup.RuntimeStats.Alloc + t.Logf("Memory decreased by %.2f MB after cleanup phase", float64(memoryDecrease)/(1024*1024)) + } else if postCleanup.RuntimeStats.Alloc > current.RuntimeStats.Alloc { + memoryIncrease := postCleanup.RuntimeStats.Alloc - current.RuntimeStats.Alloc + if strictMode { + t.Errorf("Memory increased by %.2f MB after cleanup phase - possible cleanup issues", float64(memoryIncrease)/(1024*1024)) + } else { + t.Logf("Memory increased by %.2f MB after cleanup phase - possible cleanup issues", float64(memoryIncrease)/(1024*1024)) + } + } + + t.Logf("Memory pool leak detection test completed") + t.Logf("Memory usage: baseline=%.2f MB, intermediate=%.2f MB, final=%.2f MB, post-cleanup=%.2f MB", + float64(baseline.RuntimeStats.Alloc)/(1024*1024), + float64(intermediate.RuntimeStats.Alloc)/(1024*1024), + float64(current.RuntimeStats.Alloc)/(1024*1024), + float64(postCleanup.RuntimeStats.Alloc)/(1024*1024)) +} diff --git a/session.go b/session.go index 39c29dd..19660fe 100644 --- a/session.go +++ b/session.go @@ -678,6 +678,8 @@ func (sm *SessionManager) GetSession(r *http.Request) (*SessionData, error) { sessionData.dirty = false var sessionReturned bool + // CRITICAL FIX: Defer logic to return session to pool only on panic, not on successful return + // This prevents memory leaks by ensuring sessions are only pooled when not in use by caller defer func() { if !sessionReturned && sessionData != nil { if r := recover(); r != nil { @@ -743,7 +745,9 @@ func (sm *SessionManager) GetSession(r *http.Request) (*SessionData, error) { sm.getTokenChunkSessions(r, refreshTokenCookie, sessionData.refreshTokenChunks) sm.getTokenChunkSessions(r, idTokenCookie, sessionData.idTokenChunks) - sessionReturned = false + // CRITICAL FIX: Mark session as successfully returned to caller - prevents premature pool return + // This fixes memory leak where sessions were incorrectly returned to pool immediately after success + sessionReturned = true return sessionData, nil }