mirror of
https://github.com/lukaszraczylo/traefikoidc.git
synced 2026-10-14 04:07:46 +00:00
fixup! fixup! fixup! fixup! fixup! Weekend fun with further optimisations.
This commit is contained in:
+12
-43
@@ -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
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user