diff options
Diffstat (limited to 'decaymap')
| -rw-r--r-- | decaymap/decaymap.go | 87 | ||||
| -rw-r--r-- | decaymap/decaymap_test.go | 31 |
2 files changed, 118 insertions, 0 deletions
diff --git a/decaymap/decaymap.go b/decaymap/decaymap.go new file mode 100644 index 0000000..edcbd1a --- /dev/null +++ b/decaymap/decaymap.go @@ -0,0 +1,87 @@ +package decaymap + +import ( + "sync" + "time" +) + +func Zilch[T any]() T { + var zero T + return zero +} + +// Impl is a lazy key->value map. It's a wrapper around a map and a mutex. If values exceed their time-to-live, they are pruned at Get time. +type Impl[K comparable, V any] struct { + data map[K]decayMapEntry[V] + lock sync.RWMutex +} + +type decayMapEntry[V any] struct { + Value V + expiry time.Time +} + +// New creates a new DecayMap of key type K and value type V. +// +// Key types must be comparable to work with maps. +func New[K comparable, V any]() *Impl[K, V] { + return &Impl[K, V]{ + data: make(map[K]decayMapEntry[V]), + } +} + +// expire forcibly expires a key by setting its time-to-live one second in the past. +func (m *Impl[K, V]) expire(key K) bool { + m.lock.RLock() + val, ok := m.data[key] + m.lock.RUnlock() + + if !ok { + return false + } + + m.lock.Lock() + val.expiry = time.Now().Add(-1 * time.Second) + m.data[key] = val + m.lock.Unlock() + + return true +} + +// Get gets a value from the DecayMap by key. +// +// If a value has expired, forcibly delete it if it was not updated. +func (m *Impl[K, V]) Get(key K) (V, bool) { + m.lock.RLock() + value, ok := m.data[key] + m.lock.RUnlock() + + if !ok { + return Zilch[V](), false + } + + if time.Now().After(value.expiry) { + m.lock.Lock() + // Since previously reading m.data[key], the value may have been updated. + // Delete the entry only if the expiry time is still the same. + if m.data[key].expiry == value.expiry { + delete(m.data, key) + } + m.lock.Unlock() + + return Zilch[V](), false + } + + return value.Value, true +} + +// Set sets a key value pair in the map. +func (m *Impl[K, V]) Set(key K, value V, ttl time.Duration) { + m.lock.Lock() + defer m.lock.Unlock() + + m.data[key] = decayMapEntry[V]{ + Value: value, + expiry: time.Now().Add(ttl), + } +} diff --git a/decaymap/decaymap_test.go b/decaymap/decaymap_test.go new file mode 100644 index 0000000..c930e08 --- /dev/null +++ b/decaymap/decaymap_test.go @@ -0,0 +1,31 @@ +package decaymap + +import ( + "testing" + "time" +) + +func TestImpl(t *testing.T) { + dm := New[string, string]() + + dm.Set("test", "hi", 5*time.Minute) + + val, ok := dm.Get("test") + if !ok { + t.Error("somehow the test key was not set") + } + + if val != "hi" { + t.Errorf("wanted value %q, got: %q", "hi", val) + } + + ok = dm.expire("test") + if !ok { + t.Error("somehow could not force-expire the test key") + } + + _, ok = dm.Get("test") + if ok { + t.Error("got value even though it was supposed to be expired") + } +} |
