From 3eb806d7b11c9d1debc315378638c2a957b0d4b5 Mon Sep 17 00:00:00 2001
From: Ian Gulliver <github@flamingcow.io>
Date: Mon, 24 Apr 2023 22:30:39 +0000
Subject: [PATCH] Expire old cache entries

---
 justfile        |  2 +-
 potency.go      | 61 ++++++++++++++++++++++++++++++++++++++++++++-----
 potency_test.go | 48 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 104 insertions(+), 7 deletions(-)

diff --git a/justfile b/justfile
index d607f93..c2dd020 100644
--- a/justfile
+++ b/justfile
@@ -11,7 +11,7 @@ tidy:
 test:
 	{{go}} vet ./...
 	golangci-lint run ./...
-	{{go}} test -race -coverprofile=cover.out -timeout=60s -parallel=10 ./...
+	{{go}} test -race -coverprofile=cover.out -timeout=60s ./...
 	{{go}} tool cover -html=cover.out -o=cover.html
 
 todo:
diff --git a/potency.go b/potency.go
index d6b1b53..b2c6130 100644
--- a/potency.go
+++ b/potency.go
@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/gopatchy/jsrest"
 )
@@ -16,15 +17,20 @@ import (
 type Potency struct {
 	handler http.Handler
 
-	// TODO: Expire based on time; probably use age-based linked list and delete at write time
-	cache   map[string]*savedResult
-	cacheMu sync.RWMutex
+	lifetime time.Duration
+
+	cache       map[string]*savedResult
+	cacheOldest *savedResult
+	cacheNewest *savedResult
+	cacheMu     sync.RWMutex
 
 	inProgress   map[string]bool
 	inProgressMu sync.Mutex
 }
 
 type savedResult struct {
+	key string
+
 	method        string
 	url           string
 	requestHeader http.Header
@@ -33,6 +39,9 @@ type savedResult struct {
 	statusCode     int
 	responseHeader http.Header
 	responseBody   []byte
+
+	added time.Time
+	newer *savedResult
 }
 
 var (
@@ -53,6 +62,7 @@ var (
 func NewPotency(handler http.Handler) *Potency {
 	return &Potency{
 		handler:    handler,
+		lifetime:   6 * time.Hour,
 		cache:      map[string]*savedResult{},
 		inProgress: map[string]bool{},
 	}
@@ -71,6 +81,20 @@ func (p *Potency) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+func (p *Potency) SetLifetime(lifetime time.Duration) {
+	p.cacheMu.Lock()
+	defer p.cacheMu.Unlock()
+
+	p.lifetime = lifetime
+}
+
+func (p *Potency) NumCached() int {
+	p.cacheMu.RLock()
+	defer p.cacheMu.RUnlock()
+
+	return len(p.cache)
+}
+
 func (p *Potency) serveHTTP(w http.ResponseWriter, r *http.Request, val string) error {
 	if len(val) < 2 || !strings.HasPrefix(val, `"`) || !strings.HasSuffix(val, `"`) {
 		return jsrest.Errorf(jsrest.ErrBadRequest, "%s (%w)", val, ErrInvalidKey)
@@ -139,6 +163,8 @@ func (p *Potency) serveHTTP(w http.ResponseWriter, r *http.Request, val string)
 	p.handler.ServeHTTP(w, r)
 
 	save := &savedResult{
+		key: key,
+
 		method:        r.Method,
 		url:           r.URL.String(),
 		requestHeader: requestHeader,
@@ -149,7 +175,7 @@ func (p *Potency) serveHTTP(w http.ResponseWriter, r *http.Request, val string)
 		responseBody:   rwi.buf.Bytes(),
 	}
 
-	p.write(key, save)
+	p.write(save)
 
 	return nil
 }
@@ -181,9 +207,32 @@ func (p *Potency) read(key string) *savedResult {
 	return p.cache[key]
 }
 
-func (p *Potency) write(key string, sr *savedResult) {
+func (p *Potency) write(sr *savedResult) {
 	p.cacheMu.Lock()
 	defer p.cacheMu.Unlock()
 
-	p.cache[key] = sr
+	sr.added = time.Now()
+
+	p.cache[sr.key] = sr
+
+	if p.cacheNewest != nil {
+		p.cacheNewest.newer = sr
+	}
+
+	p.cacheNewest = sr
+
+	if p.cacheOldest == nil {
+		p.cacheOldest = sr
+	}
+
+	p.removeExpired()
+}
+
+func (p *Potency) removeExpired() {
+	cutoff := time.Now().Add(-1 * p.lifetime)
+
+	for iter := p.cacheOldest; iter != nil && iter.added.Before(cutoff); iter = iter.newer {
+		delete(p.cache, iter.key)
+		p.cacheOldest = iter
+	}
 }
diff --git a/potency_test.go b/potency_test.go
index 1fdcc0d..e4ce040 100644
--- a/potency_test.go
+++ b/potency_test.go
@@ -123,8 +123,55 @@ func TestPOST(t *testing.T) {
 	require.True(t, resp.IsError())
 }
 
+func TestExpire(t *testing.T) {
+	t.Parallel()
+
+	ts := newTestServer(t)
+	defer ts.shutdown(t)
+
+	ts.pot.SetLifetime(1 * time.Second)
+
+	require.Equal(t, 0, ts.pot.NumCached())
+
+	key1 := uniuri.New()
+
+	resp, err := ts.r().
+		SetHeader("Idempotency-Key", fmt.Sprintf(`"%s"`, key1)).
+		SetBody("test1").
+		Post("")
+	require.NoError(t, err)
+	require.False(t, resp.IsError())
+
+	require.Equal(t, 1, ts.pot.NumCached())
+
+	key2 := uniuri.New()
+
+	resp, err = ts.r().
+		SetHeader("Idempotency-Key", fmt.Sprintf(`"%s"`, key2)).
+		SetBody("test2").
+		Post("")
+	require.NoError(t, err)
+	require.False(t, resp.IsError())
+
+	require.Equal(t, 2, ts.pot.NumCached())
+
+	time.Sleep(1100 * time.Millisecond)
+
+	key3 := uniuri.New()
+
+	resp, err = ts.r().
+		SetHeader("Idempotency-Key", fmt.Sprintf(`"%s"`, key3)).
+		SetBody("test3").
+		Post("")
+	require.NoError(t, err)
+	require.False(t, resp.IsError())
+
+	require.Equal(t, 1, ts.pot.NumCached())
+}
+
 type testServer struct {
 	dir string
+	pot *potency.Potency
 	srv *http.Server
 	rst *resty.Client
 }
@@ -166,6 +213,7 @@ func newTestServer(t *testing.T) *testServer {
 
 	return &testServer{
 		dir: dir,
+		pot: p,
 		srv: srv,
 		rst: rst,
 	}