fix: add decay map

Author: jmattheis

decaymap/decaymap.go

@@ -0,0 +1,63 @@
+package decaymap
+
+import (
+	"sync"
+	"time"
+)
+
+// DecayMap is a coarse-grained paired map that bulk-frees outdated entries.
+type DecayMap[K comparable, V any] struct {
+	mu         sync.RWMutex
+	maps       [2]map[K]V
+	epoch      time.Time
+	lastInsert time.Time
+	generation uint64
+	period     time.Duration
+}
+
+func NewDecayMap[K comparable, V any](epoch time.Time, period time.Duration) *DecayMap[K, V] {
+	return &DecayMap[K, V]{
+		maps:       [2]map[K]V{make(map[K]V), make(map[K]V)},
+		epoch:      epoch,
+		generation: 0,
+		period:     period,
+	}
+}
+
+// Attempts to retrieve a value from the store, does not guarantee the value is unexpired.
+func (m *DecayMap[K, V]) Get(key K) (res V, ok bool) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	res, ok = m.maps[0][key]
+	if ok {
+		return
+	}
+	res, ok = m.maps[1][key]
+	return
+}
+
+// Sets a value in the store, overwriting the existing value if exists.
+//
+// Inserting values in non monotonic order is a no op.
+func (m *DecayMap[K, V]) Set(now time.Time, key K, value V) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if now.Before(m.lastInsert) {
+		return
+	}
+	m.lastInsert = now
+	curGen := uint64(now.Sub(m.epoch) / m.period)
+	genDelta := curGen - m.generation
+	if genDelta >= 2 {
+		// both are outdated, clear both
+		m.maps[0] = make(map[K]V)
+		m.maps[1] = make(map[K]V)
+		m.generation = curGen
+	} else if genDelta == 1 {
+		// one is outdated, clear the outdated one
+		m.maps[curGen&1] = make(map[K]V)
+		m.generation = curGen
+	}
+	delete(m.maps[(curGen&1)^1], key)
+	m.maps[curGen&1][key] = value
+}

decaymap/decaymap_test.go

@@ -0,0 +1,39 @@
+
+package decaymap
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDecayMap2(t *testing.T) {
+	epoch := time.Now()
+	dm := NewDecayMap[string, string](epoch, 10*time.Millisecond)
+
+	now := epoch
+	for ts := 0; ts < 100; ts++ {
+		dm.Set(now, fmt.Sprintf("key%d", ts), fmt.Sprintf("value%d", ts))
+		for backts := 0; backts < ts; backts++ {
+			res, ok := dm.Get(fmt.Sprintf("key%d", backts))
+			if ts-backts <= 10 {
+				assert.True(t, ok)
+				assert.Equal(t, fmt.Sprintf("value%d", backts), res)
+			} else if ts-backts >= 20 {
+				assert.False(t, ok)
+				assert.Equal(t, "", res)
+			}
+		}
+		now = now.Add(1 * time.Millisecond)
+	}
+
+	now = now.Add(20 * time.Millisecond)
+	dm.Set(now, "dummy", "dummy") // rachet internal state
+	for ts := 0; ts < 100; ts++ {
+		res, ok := dm.Get(fmt.Sprintf("key%d", ts))
+		assert.False(t, ok)
+		assert.Equal(t, "", res)
+	}
+}