Deal with the memory growth issue.

* TokenBlacklist limit is set to 1000
* Increased token cleanup frequency
This commit is contained in:
2025-02-06 23:34:05 +00:00
parent 09daa1025c
commit 241cb1c209
2 changed files with 63 additions and 5 deletions
+28
View File
@@ -292,12 +292,16 @@ type TokenBlacklist struct {
// mutex protects concurrent access to the blacklist
mutex sync.RWMutex
// maxSize is the maximum number of tokens in the blacklist
maxSize int
}
// NewTokenBlacklist creates a new TokenBlacklist instance.
func NewTokenBlacklist() *TokenBlacklist {
return &TokenBlacklist{
blacklist: make(map[string]time.Time),
maxSize: 1000, // Limit the size to prevent unbounded growth
}
}
@@ -305,6 +309,30 @@ func NewTokenBlacklist() *TokenBlacklist {
func (tb *TokenBlacklist) Add(tokenID string, expiration time.Time) {
tb.mutex.Lock()
defer tb.mutex.Unlock()
// Clean up expired tokens if we're at capacity
if len(tb.blacklist) >= tb.maxSize {
now := time.Now()
for token, exp := range tb.blacklist {
if now.After(exp) {
delete(tb.blacklist, token)
}
}
// If still at capacity after cleanup, remove oldest token
if len(tb.blacklist) >= tb.maxSize {
var oldestToken string
var oldestTime time.Time
first := true
for token, exp := range tb.blacklist {
if first || exp.Before(oldestTime) {
oldestToken = token
oldestTime = exp
first = false
}
}
delete(tb.blacklist, oldestToken)
}
}
tb.blacklist[tokenID] = expiration
}
+35 -5
View File
@@ -664,12 +664,42 @@ func (t *TraefikOidc) buildAuthURL(redirectURL, state, nonce string) string {
// startTokenCleanup starts the token cleanup goroutine
func (t *TraefikOidc) startTokenCleanup() {
ticker := time.NewTicker(1 * time.Minute)
ctx, cancel := context.WithCancel(context.Background())
ticker := time.NewTicker(30 * time.Second) // Increased frequency to prevent memory buildup
go func() {
for range ticker.C {
t.logger.Debug("Cleaning up token cache")
t.tokenCache.Cleanup()
t.tokenBlacklist.Cleanup()
defer ticker.Stop()
defer cancel()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
t.logger.Debug("Starting token cleanup cycle")
// Run cleanup in a separate goroutine with timeout
cleanupCtx, cleanupCancel := context.WithTimeout(ctx, 10*time.Second)
done := make(chan struct{})
go func() {
defer close(done)
t.tokenCache.Cleanup()
t.tokenBlacklist.Cleanup()
}()
// Wait for cleanup to complete or timeout
select {
case <-cleanupCtx.Done():
if cleanupCtx.Err() == context.DeadlineExceeded {
t.logger.Error("Token cleanup cycle timed out")
}
case <-done:
t.logger.Debug("Token cleanup cycle completed successfully")
}
cleanupCancel()
}
}
}()
}