From f1f79a573263678a961ca62a61ce8eb1e833e479 Mon Sep 17 00:00:00 2001
From: "Han Verstraete (OpenFaaS Ltd)" <han@openfaas.com>
Date: Thu, 30 May 2024 16:15:56 +0200
Subject: [PATCH] Support garbage collection of expired tokens from cache

Signed-off-by: Han Verstraete (OpenFaaS Ltd) <han@openfaas.com>
---
 tokencache.go | 41 ++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 40 insertions(+), 1 deletion(-)

diff --git a/tokencache.go b/tokencache.go
index 7d1f121..9935245 100644
--- a/tokencache.go
+++ b/tokencache.go
@@ -1,24 +1,31 @@
 package sdk
 
-import "sync"
+import (
+	"context"
+	"sync"
+	"time"
+)
 
 type TokenCache interface {
 	Get(key string) (*Token, bool)
 	Set(key string, token *Token)
 }
 
+// MemoryTokenCache is a basic in-memory token cache implementation.
 type MemoryTokenCache struct {
 	tokens map[string]*Token
 
 	lock sync.RWMutex
 }
 
+// NewMemoryTokenCache creates a new in memory token cache instance.
 func NewMemoryTokenCache() *MemoryTokenCache {
 	return &MemoryTokenCache{
 		tokens: map[string]*Token{},
 	}
 }
 
+// Set adds or updates a token with the given key in the cache.
 func (c *MemoryTokenCache) Set(key string, token *Token) {
 	c.lock.Lock()
 	defer c.lock.Unlock()
@@ -26,6 +33,8 @@ func (c *MemoryTokenCache) Set(key string, token *Token) {
 	c.tokens[key] = token
 }
 
+// Get retrieves the token associated with the given key from the cache. The bool
+// return value will be false if no matching key is found, and true otherwise.
 func (c *MemoryTokenCache) Get(key string) (*Token, bool) {
 	c.lock.RLock()
 	token, ok := c.tokens[key]
@@ -41,3 +50,33 @@ func (c *MemoryTokenCache) Get(key string) (*Token, bool) {
 
 	return token, ok
 }
+
+// StartGC starts garbage collection of expired tokens.
+func (c *MemoryTokenCache) StartGC(ctx context.Context, gcInterval time.Duration) {
+	if gcInterval <= 0 {
+		return
+	}
+
+	ticker := time.NewTicker(gcInterval)
+
+	for {
+		select {
+		case <-ticker.C:
+			c.clearExpired()
+		case <-ctx.Done():
+			ticker.Stop()
+			return
+		}
+	}
+}
+
+// clearExpired removes all expired tokens from the cache.
+func (c *MemoryTokenCache) clearExpired() {
+	for key, token := range c.tokens {
+		if token.Expired() {
+			c.lock.Lock()
+			delete(c.tokens, key)
+			c.lock.Unlock()
+		}
+	}
+}