diff --git a/autocleanup.go b/autocleanup.go index dd7bff8..1b1b4a2 100644 --- a/autocleanup.go +++ b/autocleanup.go @@ -521,53 +521,22 @@ func (tr *TaskRegistry) GetTaskCount() int { func (tr *TaskRegistry) CreateSingletonTask(name string, interval time.Duration, taskFunc func(), logger *Logger, wg *sync.WaitGroup) (*BackgroundTask, error) { - tr.mu.Lock() - defer tr.mu.Unlock() - - // Strict singleton enforcement: check if ANY task with similar name pattern exists - for taskName := range tr.tasks { - if strings.Contains(taskName, "cleanup") && strings.Contains(name, "cleanup") { - if tr.logger != nil { - tr.logger.Debug("Singleton enforcement: cleanup task %s already exists, rejecting %s", taskName, name) - } - return nil, fmt.Errorf("singleton cleanup task already exists: %s (requested: %s)", taskName, name) - } - if strings.Contains(taskName, "singleton") && strings.Contains(name, "singleton") { - if tr.logger != nil { - tr.logger.Debug("Singleton enforcement: singleton task %s already exists, rejecting %s", taskName, name) - } - return nil, fmt.Errorf("singleton task already exists: %s (requested: %s)", taskName, name) - } - if strings.Contains(taskName, "memory-monitor") && strings.Contains(name, "memory-monitor") { - if tr.logger != nil { - tr.logger.Debug("Singleton enforcement: memory-monitor task %s already exists, rejecting %s", taskName, name) - } - return nil, fmt.Errorf("singleton memory-monitor task already exists: %s (requested: %s)", taskName, name) - } + // Delegate to the singleton resource manager instead + rm := GetResourceManager() + err := rm.RegisterBackgroundTask(name, interval, taskFunc) + if err != nil { + return nil, err } - // Check if exact task already exists - if existing, exists := tr.tasks[name]; exists { - if tr.logger != nil { - tr.logger.Debug("Singleton task %s already exists, returning existing task", name) - } - return existing, nil + // Start the task if not already running + if !rm.IsTaskRunning(name) { + rm.StartBackgroundTask(name) } - // Check circuit breaker - if err := tr.cb.CanCreateTask(name); err != nil { - tr.cb.OnTaskFailure(name, err) - return nil, fmt.Errorf("circuit breaker prevented singleton task creation: %w", err) - } - - // Create new task (execution tracking handled in run() method) - task := NewBackgroundTask(name, interval, taskFunc, logger, wg) - tr.tasks[name] = task - tr.cb.OnTaskSuccess(name) - - if tr.logger != nil { - tr.logger.Info("Created singleton background task: %s", name) - } + // Get the task from resource manager's internal registry + rm.tasksMu.RLock() + task := rm.tasks[name] + rm.tasksMu.RUnlock() return task, nil } diff --git a/main.go b/main.go index 48b32e4..1d523af 100644 --- a/main.go +++ b/main.go @@ -350,6 +350,12 @@ func mergeScopes(defaultScopes, userScopes []string) []string { // - The configured TraefikOidc handler ready to process requests. // - An error if essential configuration is missing or invalid (e.g., short encryption key). func New(ctx context.Context, next http.Handler, config *Config, name string) (http.Handler, error) { + return NewWithContext(ctx, config, next, name) +} + +// NewWithContext creates a new TraefikOidc middleware instance with proper context handling. +// This is the preferred constructor that ensures proper goroutine lifecycle management. +func NewWithContext(ctx context.Context, config *Config, next http.Handler, name string) (*TraefikOidc, error) { if config == nil { config = CreateConfig() } @@ -377,7 +383,14 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h goroutineWG := &sync.WaitGroup{} cacheManager := GetGlobalCacheManager(goroutineWG) - pluginCtx, cancelFunc := context.WithCancel(context.Background()) + // Use provided context instead of creating new one + var pluginCtx context.Context + var cancelFunc context.CancelFunc + if ctx != nil { + pluginCtx, cancelFunc = context.WithCancel(ctx) + } else { + pluginCtx, cancelFunc = context.WithCancel(context.Background()) + } t := &TraefikOidc{ next: next, @@ -503,7 +516,14 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h logger.Debugf("TraefikOidc.New: Final t.scopes initialized to: %v", t.scopes) t.providerURL = config.ProviderURL - // Initialize metadata in background but track the goroutine + + // Use singleton resource manager for metadata initialization + rm := GetResourceManager() + + // Add reference for this instance + rm.AddReference(name) + + // Initialize metadata in a goroutine with proper tracking if t.goroutineWG != nil { t.goroutineWG.Add(1) } @@ -520,6 +540,14 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h t.initializeMetadata(config.ProviderURL) }() + // Setup cleanup hook for when context is cancelled + if pluginCtx != nil { + go func() { + <-pluginCtx.Done() + t.Close() + }() + } + return t, nil } @@ -575,91 +603,42 @@ func (t *TraefikOidc) updateMetadataEndpoints(metadata *ProviderMetadata) { // Parameters: // - providerURL: The base URL of the OIDC provider, used for subsequent refresh attempts. func (t *TraefikOidc) startMetadataRefresh(providerURL string) { - ticker := time.NewTicker(2 * time.Hour) + // Use singleton resource manager for metadata refresh + rm := GetResourceManager() + taskName := "singleton-metadata-refresh" - // Track this goroutine for proper cleanup - if t.goroutineWG != nil { - t.goroutineWG.Add(1) + // Create refresh function + refreshFunc := func() { + if t.metadataCache == nil || t.httpClient == nil { + return + } + + metadata, err := t.metadataCache.GetMetadata(providerURL, t.httpClient, t.logger) + if err != nil { + t.safeLogErrorf("Failed to refresh provider metadata: %v", err) + return + } + + if metadata != nil { + t.updateMetadataEndpoints(metadata) + t.safeLogDebug("Successfully refreshed provider metadata") + } } - go func() { - defer func() { - // Ensure ticker is always stopped - ticker.Stop() + // Register as singleton task - will return existing if already registered + err := rm.RegisterBackgroundTask(taskName, 2*time.Hour, refreshFunc) + if err != nil { + t.logger.Errorf("Failed to register metadata refresh task: %v", err) + return + } - // Always call Done() even if panic occurs - if t.goroutineWG != nil { - t.goroutineWG.Done() - } - - // Recover from panics to prevent goroutine leaks - if r := recover(); r != nil { - t.logger.Errorf("Metadata refresh goroutine panic recovered: %v", r) - } - }() - - consecutiveFailures := 0 - for { - select { - case <-ticker.C: - // Add timeout for metadata refresh operations - refreshCtx, cancel := context.WithTimeout(t.ctx, 30*time.Second) - - func() { - defer cancel() // Ensure context is always canceled - - if consecutiveFailures >= 3 { - t.safeLogDebug("Skipping metadata refresh due to consecutive failures") - consecutiveFailures++ - if consecutiveFailures > 10 { - consecutiveFailures = 3 - } - return - } - - t.safeLogDebug("Refreshing OIDC metadata") - var metadata *ProviderMetadata - var err error - if t.errorRecoveryManager != nil { - metadata, err = t.metadataCache.GetMetadataWithRecovery(providerURL, t.httpClient, t.logger, t.errorRecoveryManager) - } else { - metadata, err = t.metadataCache.GetMetadata(providerURL, t.httpClient, t.logger) - } - - // Check for context cancellation - select { - case <-refreshCtx.Done(): - t.safeLogDebug("Metadata refresh timeout or cancellation") - consecutiveFailures++ - return - default: - } - - if err != nil { - consecutiveFailures++ - t.safeLogErrorf("Failed to refresh metadata (attempt %d): %v", consecutiveFailures, err) - return - } - - if metadata != nil { - t.updateMetadataEndpoints(metadata) - t.logger.Debug("Successfully refreshed metadata") - consecutiveFailures = 0 - } else { - consecutiveFailures++ - t.logger.Error("Received nil metadata during refresh") - } - }() - - case <-t.metadataRefreshStopChan: - t.logger.Debug("Metadata refresh goroutine stopped.") - return - case <-t.ctx.Done(): - t.logger.Debug("Metadata refresh goroutine stopped due to context cancellation.") - return - } - } - }() + // Start the task if not already running + if !rm.IsTaskRunning(taskName) { + rm.StartBackgroundTask(taskName) + t.logger.Debug("Started singleton metadata refresh task") + } else { + t.logger.Debug("Metadata refresh task already running, skipping duplicate") + } } // ServeHTTP implements the main middleware logic for processing HTTP requests. @@ -684,21 +663,8 @@ func (t *TraefikOidc) ServeHTTP(rw http.ResponseWriter, req *http.Request) { if !t.metadataRefreshStarted && t.providerURL != "" { t.metadataRefreshStarted = true - if t.goroutineWG != nil { - t.goroutineWG.Add(1) - } - go func() { - defer func() { - if t.goroutineWG != nil { - t.goroutineWG.Done() - } - // Recover from panics to prevent goroutine leaks - if r := recover(); r != nil { - t.logger.Errorf("Start metadata refresh goroutine panic recovered: %v", r) - } - }() - t.startMetadataRefresh(t.providerURL) - }() + // Metadata refresh is handled by singleton resource manager + t.startMetadataRefresh(t.providerURL) } } t.firstRequestMutex.Unlock() @@ -1645,160 +1611,54 @@ func (t *TraefikOidc) startTokenCleanup() { return } - // Capture values to avoid race conditions + // Use singleton resource manager for token cleanup + rm := GetResourceManager() + taskName := "singleton-token-cleanup" + + // Capture values for the cleanup function tokenCache := t.tokenCache jwkCache := t.jwkCache sessionManager := t.sessionManager logger := t.logger - goroutineWG := t.goroutineWG - stopChan := t.tokenCleanupStopChan - ctx := t.ctx cleanupInterval := 1 * time.Minute if isTestMode() { cleanupInterval = 50 * time.Millisecond // Fast interval for tests } - ticker := time.NewTicker(cleanupInterval) - if goroutineWG != nil { - goroutineWG.Add(1) + // Create cleanup function + cleanupFunc := func() { + if logger != nil && !isTestMode() { + logger.Debug("Starting token cleanup cycle") + } + if tokenCache != nil { + tokenCache.Cleanup() + } + if jwkCache != nil { + jwkCache.Cleanup() + } + if sessionManager != nil { + sessionManager.PeriodicChunkCleanup() + if logger != nil && !isTestMode() { + logger.Debug("Running session health monitoring") + } + } } - // Ensure cleanup channels are available - if stopChan == nil && ctx == nil { - if logger != nil && !isTestMode() { - logger.Debug("No cleanup mechanism available for token cleanup goroutine, skipping") - } + // Register as singleton task - will return existing if already registered + err := rm.RegisterBackgroundTask(taskName, cleanupInterval, cleanupFunc) + if err != nil { + logger.Errorf("Failed to register token cleanup task: %v", err) return } - go func() { - defer func() { - if goroutineWG != nil { - goroutineWG.Done() - } - ticker.Stop() - - if r := recover(); r != nil { - // Only log if context is valid and not cancelled to prevent logging after test completion - if logger != nil && ctx != nil && ctx.Err() == nil && !isTestMode() { - logger.Errorf("Token cleanup goroutine panic recovered: %v", r) - } - } - }() - - // Build dynamic select cases - for { - if stopChan != nil && ctx != nil { - // Both channels available - select { - case <-ticker.C: - if logger != nil && !isTestMode() { - logger.Debug("Starting token cleanup cycle") - } - if tokenCache != nil { - tokenCache.Cleanup() - } - if jwkCache != nil { - jwkCache.Cleanup() - } - if sessionManager != nil { - // Check if context is cancelled before calling session cleanup to prevent logging after test completion - if ctx == nil || ctx.Err() == nil { - sessionManager.PeriodicChunkCleanup() - } - if logger != nil && !isTestMode() && ctx != nil && ctx.Err() == nil { - logger.Debug("Running session health monitoring") - } - } - case <-stopChan: - if logger != nil { - logger.Debug("Token cleanup goroutine stopped.") - } - return - case <-ctx.Done(): - if logger != nil { - logger.Debug("Token cleanup goroutine stopped due to context cancellation.") - } - return - } - } else if stopChan != nil { - // Only stopChan available - select { - case <-ticker.C: - if logger != nil && !isTestMode() { - logger.Debug("Starting token cleanup cycle") - } - if tokenCache != nil { - tokenCache.Cleanup() - } - if jwkCache != nil { - jwkCache.Cleanup() - } - if sessionManager != nil { - // Check if context is cancelled before calling session cleanup to prevent logging after test completion - if ctx == nil || ctx.Err() == nil { - sessionManager.PeriodicChunkCleanup() - } - if logger != nil && !isTestMode() && ctx != nil && ctx.Err() == nil { - logger.Debug("Running session health monitoring") - } - } - case <-stopChan: - if logger != nil { - logger.Debug("Token cleanup goroutine stopped.") - } - return - } - } else if ctx != nil { - // Only ctx available - select { - case <-ticker.C: - if logger != nil && !isTestMode() { - logger.Debug("Starting token cleanup cycle") - } - if tokenCache != nil { - tokenCache.Cleanup() - } - if jwkCache != nil { - jwkCache.Cleanup() - } - if sessionManager != nil { - // Check if context is cancelled before calling session cleanup to prevent logging after test completion - if ctx == nil || ctx.Err() == nil { - sessionManager.PeriodicChunkCleanup() - } - if logger != nil && !isTestMode() && ctx != nil && ctx.Err() == nil { - logger.Debug("Running session health monitoring") - } - } - case <-ctx.Done(): - if logger != nil { - logger.Debug("Token cleanup goroutine stopped due to context cancellation.") - } - return - } - } else { - // Neither available - just use ticker (this should not happen due to the check above) - <-ticker.C - if logger != nil && !isTestMode() { - logger.Debug("Starting token cleanup cycle") - } - if tokenCache != nil { - tokenCache.Cleanup() - } - if jwkCache != nil { - jwkCache.Cleanup() - } - if sessionManager != nil { - sessionManager.PeriodicChunkCleanup() - if logger != nil { - logger.Debug("Running session health monitoring") - } - } - } - } - }() + // Start the task if not already running + if !rm.IsTaskRunning(taskName) { + rm.StartBackgroundTask(taskName) + logger.Debug("Started singleton token cleanup task") + } else { + logger.Debug("Token cleanup task already running, skipping duplicate") + } } // RevokeToken revokes a token locally by adding it to the blacklist cache. @@ -2551,11 +2411,22 @@ func (t *TraefikOidc) Close() error { t.shutdownOnce.Do(func() { t.safeLogDebug("Closing TraefikOidc plugin instance") + // Get resource manager for cleanup + rm := GetResourceManager() + + // Stop singleton tasks related to this instance + rm.StopBackgroundTask("singleton-token-cleanup") + rm.StopBackgroundTask("singleton-metadata-refresh") + + // Remove reference for this instance + rm.RemoveReference(t.name) + if t.cancelFunc != nil { t.cancelFunc() t.safeLogDebug("Context cancellation signaled to all goroutines") } + // Clean up legacy stop channels if they exist if t.tokenCleanupStopChan != nil { close(t.tokenCleanupStopChan) t.safeLogDebug("tokenCleanupStopChan closed") diff --git a/middleware/auth_middleware.go b/middleware/auth_middleware.go index c6d6099..af72f55 100644 --- a/middleware/auth_middleware.go +++ b/middleware/auth_middleware.go @@ -163,21 +163,9 @@ func (m *AuthMiddleware) ServeHTTP(rw http.ResponseWriter, req *http.Request) { if !m.metadataRefreshStarted && m.providerURL != "" { m.metadataRefreshStarted = true - if m.goroutineWG != nil { - m.goroutineWG.Add(1) - } - go func() { - defer func() { - if m.goroutineWG != nil { - m.goroutineWG.Done() - } - // Recover from panics to prevent goroutine leaks - if r := recover(); r != nil { - m.logger.Errorf("Start metadata refresh goroutine panic recovered: %v", r) - } - }() - m.startMetadataRefreshFunc(m.providerURL) - }() + // Metadata refresh is now handled by singleton resource manager + // Just call the function directly - it will use the singleton internally + m.startMetadataRefreshFunc(m.providerURL) } } m.firstRequestMutex.Unlock() diff --git a/singleton_resources.go b/singleton_resources.go new file mode 100644 index 0000000..8c7c888 --- /dev/null +++ b/singleton_resources.go @@ -0,0 +1,543 @@ +package traefikoidc + +import ( + "context" + "fmt" + "net/http" + "sync" + "sync/atomic" + "time" +) + +var ( + globalResourceManager *ResourceManager + resourceManagerOnce sync.Once + resourceManagerMutex sync.Mutex +) + +// ResourceManager manages shared resources across all middleware instances +// to prevent duplication and goroutine leaks when Traefik recreates middleware +type ResourceManager struct { + // HTTP clients shared across instances + httpClients map[string]*http.Client + clientsMu sync.RWMutex + + // Caches shared across instances + caches map[string]interface{} + cachesMu sync.RWMutex + + // Background tasks registry + tasks map[string]*BackgroundTask + tasksMu sync.RWMutex + + // Goroutine pools for controlled concurrency + pools map[string]*GoroutinePool + poolsMu sync.RWMutex + + // Reference counting for cleanup + references map[string]*int32 + referencesMu sync.RWMutex + + // Logger + logger *Logger + + // Shutdown coordination + shutdownOnce sync.Once + shutdownChan chan struct{} + wg sync.WaitGroup +} + +// GetResourceManager returns the global singleton ResourceManager instance +func GetResourceManager() *ResourceManager { + resourceManagerOnce.Do(func() { + globalResourceManager = &ResourceManager{ + httpClients: make(map[string]*http.Client), + caches: make(map[string]interface{}), + tasks: make(map[string]*BackgroundTask), + pools: make(map[string]*GoroutinePool), + references: make(map[string]*int32), + logger: GetSingletonNoOpLogger(), + shutdownChan: make(chan struct{}), + } + }) + return globalResourceManager +} + +// GetHTTPClient returns a shared HTTP client for the given key +func (rm *ResourceManager) GetHTTPClient(key string) *http.Client { + rm.clientsMu.RLock() + client, exists := rm.httpClients[key] + rm.clientsMu.RUnlock() + + if exists { + return client + } + + rm.clientsMu.Lock() + defer rm.clientsMu.Unlock() + + // Double-check after acquiring write lock + if client, exists := rm.httpClients[key]; exists { + return client + } + + // Create new HTTP client with reasonable defaults + client = &http.Client{ + Timeout: 30 * time.Second, + Transport: &http.Transport{ + MaxIdleConns: 100, + MaxIdleConnsPerHost: 10, + IdleConnTimeout: 90 * time.Second, + DisableKeepAlives: false, + }, + } + + rm.httpClients[key] = client + return client +} + +// GetCache returns a shared cache for the given key +func (rm *ResourceManager) GetCache(key string) interface{} { + rm.cachesMu.RLock() + cache, exists := rm.caches[key] + rm.cachesMu.RUnlock() + + if exists { + return cache + } + + rm.cachesMu.Lock() + defer rm.cachesMu.Unlock() + + // Double-check after acquiring write lock + if cache, exists := rm.caches[key]; exists { + return cache + } + + // Create cache based on key type + // Use global cache manager for proper singleton caches + cacheManager := GetGlobalCacheManager(&rm.wg) + switch key { + case "metadata-cache": + cache = cacheManager.GetSharedMetadataCache() + case "token-cache": + cache = cacheManager.GetSharedTokenCache() + case "jwk-cache": + cache = cacheManager.GetSharedJWKCache() + default: + // Generic cache implementation + cache = NewGenericCache(1*time.Hour, rm.logger) + } + + rm.caches[key] = cache + return cache +} + +// RegisterBackgroundTask registers a singleton background task +func (rm *ResourceManager) RegisterBackgroundTask(name string, interval time.Duration, taskFunc func()) error { + rm.tasksMu.Lock() + defer rm.tasksMu.Unlock() + + // Check if task already exists + if _, exists := rm.tasks[name]; exists { + if rm.logger != nil { + rm.logger.Debugf("Background task %s already registered", name) + } + // Return existing task without error for idempotency + return nil + } + + // Create new task with WaitGroup for proper cleanup + task := NewBackgroundTask(name, interval, taskFunc, rm.logger, &rm.wg) + rm.tasks[name] = task + + if rm.logger != nil { + rm.logger.Infof("Registered singleton background task: %s", name) + } + + return nil +} + +// StartBackgroundTask starts a registered background task +func (rm *ResourceManager) StartBackgroundTask(name string) error { + rm.tasksMu.RLock() + task, exists := rm.tasks[name] + rm.tasksMu.RUnlock() + + if !exists { + return fmt.Errorf("task %s not registered", name) + } + + task.Start() + return nil +} + +// StopBackgroundTask stops a running background task +func (rm *ResourceManager) StopBackgroundTask(name string) error { + rm.tasksMu.RLock() + task, exists := rm.tasks[name] + rm.tasksMu.RUnlock() + + if !exists { + return fmt.Errorf("task %s not registered", name) + } + + task.Stop() + return nil +} + +// IsTaskRunning checks if a background task is running +func (rm *ResourceManager) IsTaskRunning(name string) bool { + rm.tasksMu.RLock() + task, exists := rm.tasks[name] + rm.tasksMu.RUnlock() + + if !exists { + return false + } + + // Check if task has been started and not stopped + return atomic.LoadInt32(&task.started) == 1 && atomic.LoadInt32(&task.stopped) == 0 +} + +// GetGoroutinePool returns a shared goroutine pool for controlled concurrency +func (rm *ResourceManager) GetGoroutinePool(key string, maxWorkers int) *GoroutinePool { + rm.poolsMu.RLock() + pool, exists := rm.pools[key] + rm.poolsMu.RUnlock() + + if exists { + return pool + } + + rm.poolsMu.Lock() + defer rm.poolsMu.Unlock() + + // Double-check after acquiring write lock + if pool, exists := rm.pools[key]; exists { + return pool + } + + // Create new pool + pool = NewGoroutinePool(maxWorkers, rm.logger) + rm.pools[key] = pool + + return pool +} + +// AddReference increments the reference count for a given instance +func (rm *ResourceManager) AddReference(instanceID string) { + rm.referencesMu.Lock() + defer rm.referencesMu.Unlock() + + if count, exists := rm.references[instanceID]; exists { + atomic.AddInt32(count, 1) + } else { + initial := int32(1) + rm.references[instanceID] = &initial + } + + if rm.logger != nil { + rm.logger.Debugf("Added reference for instance %s", instanceID) + } +} + +// RemoveReference decrements the reference count and triggers cleanup if needed +func (rm *ResourceManager) RemoveReference(instanceID string) { + rm.referencesMu.Lock() + defer rm.referencesMu.Unlock() + + if count, exists := rm.references[instanceID]; exists { + newCount := atomic.AddInt32(count, -1) + if newCount <= 0 { + delete(rm.references, instanceID) + if rm.logger != nil { + rm.logger.Debugf("Removed last reference for instance %s", instanceID) + } + // Trigger cleanup for this instance if needed + rm.cleanupInstance(instanceID) + } + } +} + +// GetReferenceCount returns the current reference count for an instance +func (rm *ResourceManager) GetReferenceCount(instanceID string) int32 { + rm.referencesMu.RLock() + defer rm.referencesMu.RUnlock() + + if count, exists := rm.references[instanceID]; exists { + return atomic.LoadInt32(count) + } + return 0 +} + +// cleanupInstance performs cleanup for a specific instance when its reference count reaches zero +func (rm *ResourceManager) cleanupInstance(instanceID string) { + // Instance-specific cleanup logic + if rm.logger != nil { + rm.logger.Infof("Cleaning up resources for instance %s", instanceID) + } + + // Clean up any instance-specific resources + // This is a hook for future instance-specific cleanup needs +} + +// Shutdown gracefully shuts down all managed resources +func (rm *ResourceManager) Shutdown(ctx context.Context) error { + var err error + + rm.shutdownOnce.Do(func() { + close(rm.shutdownChan) + + if rm.logger != nil { + rm.logger.Info("Starting ResourceManager shutdown") + } + + // Stop all background tasks + rm.tasksMu.RLock() + tasks := make([]*BackgroundTask, 0, len(rm.tasks)) + for _, task := range rm.tasks { + tasks = append(tasks, task) + } + rm.tasksMu.RUnlock() + + for _, task := range tasks { + task.Stop() + } + + // Shutdown all goroutine pools + rm.poolsMu.RLock() + pools := make([]*GoroutinePool, 0, len(rm.pools)) + for _, pool := range rm.pools { + pools = append(pools, pool) + } + rm.poolsMu.RUnlock() + + for _, pool := range pools { + if shutdownErr := pool.Shutdown(ctx); shutdownErr != nil && err == nil { + err = shutdownErr + } + } + + // Wait for all goroutines with timeout + done := make(chan struct{}) + go func() { + rm.wg.Wait() + close(done) + }() + + select { + case <-done: + if rm.logger != nil { + rm.logger.Info("ResourceManager shutdown completed successfully") + } + case <-ctx.Done(): + err = fmt.Errorf("shutdown timeout: %w", ctx.Err()) + if rm.logger != nil { + rm.logger.Errorf("ResourceManager shutdown timeout: %v", err) + } + } + }) + + return err +} + +// GoroutinePool provides a pool of workers for controlled concurrency +type GoroutinePool struct { + maxWorkers int + taskQueue chan func() + workerWG sync.WaitGroup + shutdownOnce sync.Once + shutdownChan chan struct{} + logger *Logger + started int32 +} + +// NewGoroutinePool creates a new goroutine pool with the specified max workers +func NewGoroutinePool(maxWorkers int, logger *Logger) *GoroutinePool { + pool := &GoroutinePool{ + maxWorkers: maxWorkers, + taskQueue: make(chan func(), maxWorkers*2), // Buffer for queuing + shutdownChan: make(chan struct{}), + logger: logger, + } + + // Start workers + for i := 0; i < maxWorkers; i++ { + pool.workerWG.Add(1) + go pool.worker(i) + } + + atomic.StoreInt32(&pool.started, 1) + + if logger != nil { + logger.Infof("Created goroutine pool with %d workers", maxWorkers) + } + + return pool +} + +// worker is the main loop for a pool worker +func (p *GoroutinePool) worker(id int) { + defer p.workerWG.Done() + + for { + select { + case task := <-p.taskQueue: + if task != nil { + // Execute task with panic recovery + func() { + defer func() { + if r := recover(); r != nil { + if p.logger != nil { + p.logger.Errorf("Worker %d panic recovered: %v", id, r) + } + } + }() + task() + }() + } + case <-p.shutdownChan: + if p.logger != nil { + p.logger.Debugf("Worker %d shutting down", id) + } + return + } + } +} + +// Submit submits a task to the pool +func (p *GoroutinePool) Submit(task func()) error { + if atomic.LoadInt32(&p.started) == 0 { + return fmt.Errorf("pool is shutdown") + } + + select { + case p.taskQueue <- task: + return nil + case <-p.shutdownChan: + return fmt.Errorf("pool is shutting down") + default: + // Queue is full, try with a small timeout + select { + case p.taskQueue <- task: + return nil + case <-time.After(100 * time.Millisecond): + return fmt.Errorf("task queue is full") + case <-p.shutdownChan: + return fmt.Errorf("pool is shutting down") + } + } +} + +// Wait waits for all submitted tasks to complete +func (p *GoroutinePool) Wait() { + // Drain the task queue + for len(p.taskQueue) > 0 { + time.Sleep(10 * time.Millisecond) + } +} + +// Shutdown gracefully shuts down the pool +func (p *GoroutinePool) Shutdown(ctx context.Context) error { + var err error + + p.shutdownOnce.Do(func() { + atomic.StoreInt32(&p.started, 0) + close(p.shutdownChan) + + // Wait for workers to finish with context timeout + done := make(chan struct{}) + go func() { + p.workerWG.Wait() + close(done) + }() + + select { + case <-done: + if p.logger != nil { + p.logger.Debug("Goroutine pool shutdown completed") + } + case <-ctx.Done(): + err = fmt.Errorf("pool shutdown timeout: %w", ctx.Err()) + if p.logger != nil { + p.logger.Errorf("Goroutine pool shutdown timeout: %v", err) + } + } + }) + + return err +} + +// GenericCache provides a simple cache implementation for testing +type GenericCache struct { + data map[string]interface{} + mu sync.RWMutex + ttl time.Duration + logger *Logger + stopChan chan struct{} +} + +// NewGenericCache creates a new generic cache +func NewGenericCache(ttl time.Duration, logger *Logger) *GenericCache { + cache := &GenericCache{ + data: make(map[string]interface{}), + ttl: ttl, + logger: logger, + stopChan: make(chan struct{}), + } + + // Start cleanup routine + go cache.cleanupRoutine() + + return cache +} + +// Get retrieves a value from the cache +func (gc *GenericCache) Get(key string) (interface{}, bool) { + gc.mu.RLock() + defer gc.mu.RUnlock() + + val, exists := gc.data[key] + return val, exists +} + +// Set stores a value in the cache +func (gc *GenericCache) Set(key string, value interface{}) { + gc.mu.Lock() + defer gc.mu.Unlock() + + gc.data[key] = value +} + +// Delete removes a value from the cache +func (gc *GenericCache) Delete(key string) { + gc.mu.Lock() + defer gc.mu.Unlock() + + delete(gc.data, key) +} + +// cleanupRoutine periodically cleans up the cache +func (gc *GenericCache) cleanupRoutine() { + ticker := time.NewTicker(gc.ttl) + defer ticker.Stop() + + for { + select { + case <-ticker.C: + gc.mu.Lock() + // Simple cleanup - clear all data after TTL + // In production, you'd track individual entry TTLs + gc.data = make(map[string]interface{}) + gc.mu.Unlock() + case <-gc.stopChan: + return + } + } +} + +// Stop stops the cleanup routine +func (gc *GenericCache) Stop() { + close(gc.stopChan) +} diff --git a/singleton_resources_test.go b/singleton_resources_test.go new file mode 100644 index 0000000..dd58207 --- /dev/null +++ b/singleton_resources_test.go @@ -0,0 +1,482 @@ +package traefikoidc + +import ( + "context" + "fmt" + "runtime" + "sync" + "sync/atomic" + "testing" + "time" +) + +// TestSingletonResourceManager tests the singleton resource manager implementation +func TestSingletonResourceManager(t *testing.T) { + t.Run("SingletonInstance", func(t *testing.T) { + // Test that GetResourceManager returns the same instance + rm1 := GetResourceManager() + rm2 := GetResourceManager() + + if rm1 != rm2 { + t.Error("GetResourceManager did not return singleton instance") + } + }) + + t.Run("ThreadSafeInitialization", func(t *testing.T) { + // Reset singleton for test + resetResourceManagerForTesting() + + const numGoroutines = 100 + instances := make([]*ResourceManager, numGoroutines) + var wg sync.WaitGroup + + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + instances[idx] = GetResourceManager() + }(i) + } + + wg.Wait() + + // Verify all instances are the same + first := instances[0] + for i := 1; i < numGoroutines; i++ { + if instances[i] != first { + t.Errorf("Instance %d differs from first instance", i) + } + } + }) + + t.Run("SharedHTTPClient", func(t *testing.T) { + rm := GetResourceManager() + + client1 := rm.GetHTTPClient("test-client-1") + client2 := rm.GetHTTPClient("test-client-1") + + if client1 != client2 { + t.Error("GetHTTPClient did not return same client for same key") + } + + client3 := rm.GetHTTPClient("test-client-2") + if client1 == client3 { + t.Error("GetHTTPClient returned same client for different keys") + } + }) + + t.Run("SharedCache", func(t *testing.T) { + rm := GetResourceManager() + + cache1 := rm.GetCache("test-cache-1") + cache2 := rm.GetCache("test-cache-1") + + if cache1 != cache2 { + t.Error("GetCache did not return same cache for same key") + } + }) + + t.Run("SingletonTaskRegistry", func(t *testing.T) { + rm := GetResourceManager() + + err := rm.RegisterBackgroundTask("test-task", 1*time.Second, func() { + // Test task + }) + + if err != nil { + t.Errorf("Failed to register task: %v", err) + } + + // Try to register same task again - should return existing + err = rm.RegisterBackgroundTask("test-task", 1*time.Second, func() { + // Duplicate task + }) + + if err != nil { + t.Errorf("Failed to handle duplicate task registration: %v", err) + } + }) + + t.Run("ReferenceCountingCleanup", func(t *testing.T) { + rm := GetResourceManager() + + // Add reference + rm.AddReference("test-instance-1") + + // Get reference count + if rm.GetReferenceCount("test-instance-1") != 1 { + t.Error("Reference count should be 1") + } + + // Add another reference + rm.AddReference("test-instance-1") + if rm.GetReferenceCount("test-instance-1") != 2 { + t.Error("Reference count should be 2") + } + + // Remove reference + rm.RemoveReference("test-instance-1") + if rm.GetReferenceCount("test-instance-1") != 1 { + t.Error("Reference count should be 1 after removal") + } + + // Remove last reference + rm.RemoveReference("test-instance-1") + if rm.GetReferenceCount("test-instance-1") != 0 { + t.Error("Reference count should be 0 after removing all references") + } + }) + + t.Run("GracefulShutdown", func(t *testing.T) { + rm := GetResourceManager() + + // Register a task + taskExecuted := false + err := rm.RegisterBackgroundTask("shutdown-test-task", 100*time.Millisecond, func() { + taskExecuted = true + }) + + if err != nil { + t.Errorf("Failed to register task: %v", err) + } + + // Start the task + rm.StartBackgroundTask("shutdown-test-task") + + // Wait for task to execute at least once + time.Sleep(150 * time.Millisecond) + + if !taskExecuted { + t.Error("Task was not executed") + } + + // Shutdown + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + err = rm.Shutdown(ctx) + if err != nil { + t.Errorf("Shutdown failed: %v", err) + } + + // Verify task is stopped + if rm.IsTaskRunning("shutdown-test-task") { + t.Error("Task should be stopped after shutdown") + } + }) +} + +// TestContextAwareGoroutineManagement tests context-aware goroutine management +func TestContextAwareGoroutineManagement(t *testing.T) { + t.Run("GoroutineCleanupOnContextCancel", func(t *testing.T) { + initialGoroutines := runtime.NumGoroutine() + + ctx, cancel := context.WithCancel(context.Background()) + + // Create a TraefikOidc instance with context + config := &Config{ + ProviderURL: "https://example.com", + ClientID: "test-client", + ClientSecret: "test-secret", + } + + plugin, err := NewWithContext(ctx, config, nil, "test") + if err != nil { + t.Fatalf("Failed to create plugin: %v", err) + } + + // Wait for goroutines to start + time.Sleep(100 * time.Millisecond) + + midGoroutines := runtime.NumGoroutine() + if midGoroutines <= initialGoroutines { + t.Error("No goroutines were created") + } + + // Cancel context + cancel() + + // Wait for cleanup + time.Sleep(500 * time.Millisecond) + + finalGoroutines := runtime.NumGoroutine() + + // Allow for some system goroutines + tolerance := 5 + if finalGoroutines > initialGoroutines+tolerance { + t.Errorf("Goroutine leak detected: initial=%d, final=%d", initialGoroutines, finalGoroutines) + } + + // Ensure plugin is properly closed + plugin.Close() + }) + + t.Run("NoGoroutineLeakOnMultipleInstances", func(t *testing.T) { + initialGoroutines := runtime.NumGoroutine() + + configs := []Config{ + {ProviderURL: "https://example1.com", ClientID: "client1", ClientSecret: "secret1"}, + {ProviderURL: "https://example2.com", ClientID: "client2", ClientSecret: "secret2"}, + {ProviderURL: "https://example3.com", ClientID: "client3", ClientSecret: "secret3"}, + } + + var plugins []*TraefikOidc + var cancels []context.CancelFunc + + // Create multiple instances + for i, config := range configs { + ctx, cancel := context.WithCancel(context.Background()) + cancels = append(cancels, cancel) + + plugin, err := NewWithContext(ctx, &config, nil, fmt.Sprintf("test-%d", i)) + if err != nil { + t.Fatalf("Failed to create plugin %d: %v", i, err) + } + plugins = append(plugins, plugin) + } + + // Wait for all goroutines to start + time.Sleep(200 * time.Millisecond) + + midGoroutines := runtime.NumGoroutine() + + // Cancel all contexts + for _, cancel := range cancels { + cancel() + } + + // Close all plugins + for _, plugin := range plugins { + plugin.Close() + } + + // Wait for cleanup + time.Sleep(500 * time.Millisecond) + + finalGoroutines := runtime.NumGoroutine() + + // Check for leaks + tolerance := 5 + if finalGoroutines > initialGoroutines+tolerance { + t.Errorf("Goroutine leak with multiple instances: initial=%d, mid=%d, final=%d", + initialGoroutines, midGoroutines, finalGoroutines) + } + }) + + t.Run("SingletonTasksAcrossInstances", func(t *testing.T) { + rm := GetResourceManager() + + // Register singleton cleanup task + var cleanupCount int32 + err := rm.RegisterBackgroundTask("singleton-cleanup", 100*time.Millisecond, func() { + atomic.AddInt32(&cleanupCount, 1) + }) + + if err != nil { + t.Fatalf("Failed to register singleton task: %v", err) + } + + // Start the task + rm.StartBackgroundTask("singleton-cleanup") + + // Create multiple plugin instances + var plugins []*TraefikOidc + for i := 0; i < 3; i++ { + ctx := context.Background() + config := &Config{ + ProviderURL: fmt.Sprintf("https://example%d.com", i), + ClientID: fmt.Sprintf("client%d", i), + ClientSecret: fmt.Sprintf("secret%d", i), + } + + plugin, err := NewWithContext(ctx, config, nil, fmt.Sprintf("test-%d", i)) + if err != nil { + t.Fatalf("Failed to create plugin %d: %v", i, err) + } + plugins = append(plugins, plugin) + } + + // Wait for cleanup to run multiple times + time.Sleep(350 * time.Millisecond) + + // Check that cleanup ran but not excessively (should be singleton) + count := atomic.LoadInt32(&cleanupCount) + if count < 2 || count > 5 { + t.Errorf("Unexpected cleanup count: %d (expected 2-5 for singleton)", count) + } + + // Cleanup + for _, plugin := range plugins { + plugin.Close() + } + + rm.StopBackgroundTask("singleton-cleanup") + }) +} + +// TestResourcePooling tests resource pooling implementation +func TestResourcePooling(t *testing.T) { + t.Run("GoroutinePoolLimiting", func(t *testing.T) { + rm := GetResourceManager() + + // Configure pool with max workers + pool := rm.GetGoroutinePool("test-pool", 5) // Max 5 workers + + if pool == nil { + t.Fatal("Failed to get goroutine pool") + } + + // Submit more tasks than pool size + var taskCount int32 + var runningCount int32 + maxRunning := int32(0) + + for i := 0; i < 20; i++ { + err := pool.Submit(func() { + atomic.AddInt32(&taskCount, 1) + current := atomic.AddInt32(&runningCount, 1) + + // Track max concurrent tasks + for { + oldMax := atomic.LoadInt32(&maxRunning) + if current <= oldMax || atomic.CompareAndSwapInt32(&maxRunning, oldMax, current) { + break + } + } + + time.Sleep(50 * time.Millisecond) + atomic.AddInt32(&runningCount, -1) + }) + + if err != nil { + t.Errorf("Failed to submit task %d: %v", i, err) + } + } + + // Wait for all tasks to complete + pool.Wait() + + // Verify all tasks executed + if atomic.LoadInt32(&taskCount) != 20 { + t.Errorf("Expected 20 tasks to execute, got %d", taskCount) + } + + // Verify concurrency was limited + if atomic.LoadInt32(&maxRunning) > 5 { + t.Errorf("Max concurrent tasks exceeded pool size: %d > 5", maxRunning) + } + }) + + t.Run("PoolShutdown", func(t *testing.T) { + rm := GetResourceManager() + pool := rm.GetGoroutinePool("shutdown-pool", 3) + + // Submit tasks + var completed int32 + for i := 0; i < 10; i++ { + pool.Submit(func() { + time.Sleep(10 * time.Millisecond) + atomic.AddInt32(&completed, 1) + }) + } + + // Shutdown pool + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + err := pool.Shutdown(ctx) + if err != nil { + t.Errorf("Pool shutdown failed: %v", err) + } + + // Try to submit after shutdown - should fail + err = pool.Submit(func() { + t.Error("Task should not execute after shutdown") + }) + + if err == nil { + t.Error("Expected error when submitting to shutdown pool") + } + }) + + t.Run("ResourceReuse", func(t *testing.T) { + rm := GetResourceManager() + + // Get same pool multiple times + pool1 := rm.GetGoroutinePool("reuse-pool", 3) + pool2 := rm.GetGoroutinePool("reuse-pool", 3) + + if pool1 != pool2 { + t.Error("Expected same pool instance for same key") + } + + // Get HTTP client multiple times + client1 := rm.GetHTTPClient("reuse-client") + client2 := rm.GetHTTPClient("reuse-client") + + if client1 != client2 { + t.Error("Expected same HTTP client instance for same key") + } + }) +} + +// TestBackwardCompatibility verifies the changes maintain backward compatibility +func TestBackwardCompatibility(t *testing.T) { + t.Run("LegacyNewFunction", func(t *testing.T) { + // Test that the original New function still works + config := &Config{ + ProviderURL: "https://example.com", + ClientID: "test-client", + ClientSecret: "test-secret", + } + + handler, err := New(context.Background(), nil, config, "test") + if err != nil { + t.Fatalf("Legacy New function failed: %v", err) + } + + if handler == nil { + t.Fatal("Handler should not be nil") + } + + // Cleanup - cast to TraefikOidc if needed + if plugin, ok := handler.(*TraefikOidc); ok { + plugin.Close() + } + }) + + t.Run("ExistingAPICompatibility", func(t *testing.T) { + config := &Config{ + ProviderURL: "https://example.com", + ClientID: "test-client", + ClientSecret: "test-secret", + } + + handler, _ := New(context.Background(), nil, config, "test") + + // Test that the handler works + if handler == nil { + t.Error("Handler should not be nil") + } + + // Cleanup - cast to TraefikOidc if needed + if plugin, ok := handler.(*TraefikOidc); ok { + plugin.Close() + } + }) +} + +// Helper function to reset singleton for testing +func resetResourceManagerForTesting() { + resourceManagerMutex.Lock() + defer resourceManagerMutex.Unlock() + + if globalResourceManager != nil { + ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) + defer cancel() + globalResourceManager.Shutdown(ctx) + } + + resourceManagerOnce = sync.Once{} + globalResourceManager = nil +}