From d42abd339d86af19c649a7a36d68ca755af647c0 Mon Sep 17 00:00:00 2001
From: Daniel Ponte <amigan@gmail.com>
Date: Fri, 14 Feb 2025 12:48:45 -0500
Subject: [PATCH] Cache with expiration

---
 internal/cache/cache.go      | 62 ++++++++++++++++++++++++++++++------
 internal/cache/cache_test.go | 31 ++++++++++++++++++
 2 files changed, 83 insertions(+), 10 deletions(-)

diff --git a/internal/cache/cache.go b/internal/cache/cache.go
index a364406..85b7823 100644
--- a/internal/cache/cache.go
+++ b/internal/cache/cache.go
@@ -1,6 +1,9 @@
 package cache
 
-import "sync"
+import (
+	"sync"
+	"time"
+)
 
 type Cache[K comparable, V any] interface {
 	Get(K) (V, bool)
@@ -9,28 +12,67 @@ type Cache[K comparable, V any] interface {
 	Clear()
 }
 
-type inMem[K comparable, V any] struct {
-	sync.RWMutex
-	m map[K]V
+type cacheItem[V any] struct {
+	exp  int64
+	elem V
 }
 
-func New[K comparable, V any]() *inMem[K, V] {
-	return &inMem[K, V]{
-		m: make(map[K]V),
+type inMem[K comparable, V any] struct {
+	sync.RWMutex
+	expiration time.Duration
+	m          map[K]cacheItem[V]
+}
+
+type cacheOpts struct {
+	exp time.Duration
+}
+
+type Option func(*cacheOpts)
+
+func WithExpiration(d time.Duration) Option {
+	return func(o *cacheOpts) {
+		o.exp = d
 	}
 }
 
+func New[K comparable, V any](opts ...Option) *inMem[K, V] {
+	co := cacheOpts{}
+
+	for _, o := range opts {
+		o(&co)
+	}
+
+	c := &inMem[K, V]{
+		m:          make(map[K]cacheItem[V]),
+		expiration: co.exp,
+	}
+
+	return c
+}
+
 func (c *inMem[K, V]) Get(key K) (V, bool) {
 	c.RLock()
 	defer c.RUnlock()
-	v, ok := c.m[key]
-	return v, ok
+
+	v, found := c.m[key]
+	if !found || (v.exp > 0 && time.Now().UnixNano() > v.exp) {
+		return v.elem, false
+	}
+
+	return v.elem, true
 }
 
 func (c *inMem[K, V]) Set(key K, val V) {
 	c.Lock()
 	defer c.Unlock()
-	c.m[key] = val
+
+	var ci cacheItem[V]
+	ci.elem = val
+	if c.expiration > 0 {
+		ci.exp = time.Now().Add(c.expiration).UnixNano()
+	}
+
+	c.m[key] = ci
 }
 
 func (c *inMem[K, V]) Delete(key K) {
diff --git a/internal/cache/cache_test.go b/internal/cache/cache_test.go
index 23079f1..d7102ed 100644
--- a/internal/cache/cache_test.go
+++ b/internal/cache/cache_test.go
@@ -2,6 +2,7 @@ package cache_test
 
 import (
 	"testing"
+	"time"
 
 	"dynatron.me/x/stillbox/internal/cache"
 
@@ -31,3 +32,33 @@ func TestCache(t *testing.T) {
 	assert.False(t, ok)
 	assert.NotEqual(t, "fg", g)
 }
+
+func TestCacheTime(t *testing.T) {
+	c := cache.New[int, string](cache.WithExpiration(100 * time.Millisecond))
+	c.Set(4, "asd")
+	time.Sleep(20 * time.Millisecond)
+	g, ok := c.Get(4)
+	assert.Equal(t, "asd", g)
+	assert.True(t, ok)
+
+	c.Set(2, "gff")
+	time.Sleep(120 * time.Millisecond)
+	g, ok = c.Get(2)
+	assert.False(t, ok)
+
+	_, ok = c.Get(8)
+	assert.False(t, ok)
+
+	c.Set(7, "fg")
+
+	c.Delete(4)
+
+	g, ok = c.Get(4)
+	assert.False(t, ok)
+	assert.NotEqual(t, "asd", g)
+
+	c.Clear()
+	g, ok = c.Get(7)
+	assert.False(t, ok)
+	assert.NotEqual(t, "fg", g)
+}