diff --git a/internal/hub/expirymap/expirymap.go b/internal/hub/expirymap/expirymap.go index db316c45..dcfc962c 100644 --- a/internal/hub/expirymap/expirymap.go +++ b/internal/hub/expirymap/expirymap.go @@ -1,29 +1,33 @@ +// Package expirymap provides a thread-safe map with expiring entries. +// It supports TTL-based expiration with both lazy cleanup on access +// and periodic background cleanup. package expirymap import ( - "reflect" + "sync" "time" "github.com/pocketbase/pocketbase/tools/store" ) -type val[T any] struct { +type val[T comparable] struct { value T expires time.Time } -type ExpiryMap[T any] struct { - store *store.Store[string, *val[T]] - cleanupInterval time.Duration +type ExpiryMap[T comparable] struct { + store *store.Store[string, *val[T]] + stopChan chan struct{} + stopOnce sync.Once } // New creates a new expiry map with custom cleanup interval -func New[T any](cleanupInterval time.Duration) *ExpiryMap[T] { +func New[T comparable](cleanupInterval time.Duration) *ExpiryMap[T] { m := &ExpiryMap[T]{ - store: store.New(map[string]*val[T]{}), - cleanupInterval: cleanupInterval, + store: store.New(map[string]*val[T]{}), + stopChan: make(chan struct{}), } - m.startCleaner() + go m.startCleaner(cleanupInterval) return m } @@ -55,7 +59,7 @@ func (m *ExpiryMap[T]) GetOk(key string) (T, bool) { // GetByValue retrieves a value by value func (m *ExpiryMap[T]) GetByValue(val T) (key string, value T, ok bool) { for key, v := range m.store.GetAll() { - if reflect.DeepEqual(v.value, val) { + if v.value == val { // check if expired if v.expires.Before(time.Now()) { m.store.Remove(key) @@ -75,7 +79,7 @@ func (m *ExpiryMap[T]) Remove(key string) { // RemovebyValue removes a value by value func (m *ExpiryMap[T]) RemovebyValue(value T) (T, bool) { for key, val := range m.store.GetAll() { - if reflect.DeepEqual(val.value, value) { + if val.value == value { m.store.Remove(key) return val.value, true } @@ -84,13 +88,23 @@ func (m *ExpiryMap[T]) RemovebyValue(value T) (T, bool) { } // startCleaner runs the background cleanup process -func (m *ExpiryMap[T]) startCleaner() { - go func() { - tick := time.Tick(m.cleanupInterval) - for range tick { +func (m *ExpiryMap[T]) startCleaner(interval time.Duration) { + tick := time.Tick(interval) + for { + select { + case <-tick: m.cleanup() + case <-m.stopChan: + return } - }() + } +} + +// StopCleaner stops the background cleanup process +func (m *ExpiryMap[T]) StopCleaner() { + m.stopOnce.Do(func() { + close(m.stopChan) + }) } // cleanup removes all expired entries diff --git a/internal/hub/expirymap/expirymap_test.go b/internal/hub/expirymap/expirymap_test.go index 6ac09c43..9fec2012 100644 --- a/internal/hub/expirymap/expirymap_test.go +++ b/internal/hub/expirymap/expirymap_test.go @@ -4,6 +4,7 @@ package expirymap import ( "testing" + "testing/synctest" "time" "github.com/stretchr/testify/assert" @@ -473,3 +474,52 @@ func TestExpiryMap_ValueOperations_Integration(t *testing.T) { assert.Equal(t, "unique", value) assert.Equal(t, "key2", key) } + +func TestExpiryMap_Cleaner(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + em := New[string](time.Second) + defer em.StopCleaner() + + em.Set("test", "value", 500*time.Millisecond) + + // Wait 600ms, value is expired but cleaner hasn't run yet (interval is 1s) + time.Sleep(600 * time.Millisecond) + synctest.Wait() + + // Map should still hold the value in its internal store before lazy access or cleaner + assert.Equal(t, 1, len(em.store.GetAll()), "store should still have 1 item before cleaner runs") + + // Wait another 500ms so cleaner (1s interval) runs + time.Sleep(500 * time.Millisecond) + synctest.Wait() // Wait for background goroutine to process the tick + + assert.Equal(t, 0, len(em.store.GetAll()), "store should be empty after cleaner runs") + }) +} + +func TestExpiryMap_StopCleaner(t *testing.T) { + em := New[string](time.Hour) + + // Initially, stopChan is open, reading would block + select { + case <-em.stopChan: + t.Fatal("stopChan should be open initially") + default: + // success + } + + em.StopCleaner() + + // After StopCleaner, stopChan is closed, reading returns immediately + select { + case <-em.stopChan: + // success + default: + t.Fatal("stopChan was not closed by StopCleaner") + } + + // Calling StopCleaner again should NOT panic thanks to sync.Once + assert.NotPanics(t, func() { + em.StopCleaner() + }) +} diff --git a/internal/hub/systems/systems_test_helpers.go b/internal/hub/systems/systems_test_helpers.go index 8599bca3..53729c01 100644 --- a/internal/hub/systems/systems_test_helpers.go +++ b/internal/hub/systems/systems_test_helpers.go @@ -113,4 +113,5 @@ func (sm *SystemManager) RemoveAllSystems() { for _, system := range sm.systems.GetAll() { sm.RemoveSystem(system.Id) } + sm.smartFetchMap.StopCleaner() }