Replace LRU cache with a simpler map that expires unused items.

Author: sbruens

cmd/outline-ss-server/expiring_map.go

@@ -0,0 +1,109 @@
+// Copyright 2024 The Outline Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// ExpiringMap is a thread-safe, generic map that automatically removes
+// key-value pairs after a specified duration of inactivity.
+// It employs reference counting to ensure safe concurrent access and prevent
+// premature deletion during cleanup.
+type ExpiringMap[K comparable, V any] struct {
+	data       map[K]*item[V]
+	mu         sync.RWMutex
+	expiryTime time.Duration
+	done       chan struct{}
+}
+
+type item[V any] struct {
+	value      V
+	lastAccess time.Time
+	refCount   int32
+}
+
+func NewExpiringMap[K comparable, V any](expiryTime time.Duration) *ExpiringMap[K, V] {
+	em := &ExpiringMap[K, V]{
+		data:       make(map[K]*item[V]),
+		expiryTime: expiryTime,
+		done:       make(chan struct{}),
+	}
+	go em.cleanupLoop()
+	return em
+}
+
+func (em *ExpiringMap[K, V]) Set(key K, value V) {
+	em.mu.Lock()
+	defer em.mu.Unlock()
+
+	em.data[key] = &item[V]{
+		value:      value,
+		lastAccess: time.Now(),
+		refCount:   0,
+	}
+}
+
+func (em *ExpiringMap[K, V]) Get(key K) (V, bool) {
+	em.mu.RLock()
+	item, ok := em.data[key]
+	if !ok {
+		em.mu.RUnlock()
+		var zeroValue V
+		return zeroValue, false
+	}
+
+	atomic.AddInt32(&item.refCount, 1)
+	em.mu.RUnlock()
+
+	em.mu.Lock()
+	defer em.mu.Unlock()
+
+	atomic.AddInt32(&item.refCount, -1)
+
+	item.lastAccess = time.Now()
+	return item.value, true
+}
+
+func (em *ExpiringMap[K, V]) cleanup() {
+	em.mu.Lock()
+	defer em.mu.Unlock()
+
+	for key, item := range em.data {
+		if time.Since(item.lastAccess) > em.expiryTime && atomic.LoadInt32(&item.refCount) == 0 {
+			delete(em.data, key)
+		}
+	}
+}
+
+func (em *ExpiringMap[K, V]) cleanupLoop() {
+	ticker := time.NewTicker(em.expiryTime / 2)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			em.cleanup()
+		case <-em.done:
+			return
+		}
+	}
+}
+
+func (em *ExpiringMap[K, V]) StopCleanup() {
+	close(em.done)
+}

cmd/outline-ss-server/expiring_map_test.go

@@ -0,0 +1,89 @@
+// Copyright 2024 The Outline Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestExpiringMap(t *testing.T) {
+	t.Run("BasicSetGet", func(t *testing.T) {
+		em := NewExpiringMap[string, int](2 * time.Second)
+		em.Set("key1", 10)
+		val, ok := em.Get("key1")
+		require.True(t, ok)
+		require.Equal(t, 10, val)
+
+		time.Sleep(3 * time.Second)
+
+		_, ok = em.Get("key1")
+		require.False(t, ok)
+	})
+
+	t.Run("Expiration", func(t *testing.T) {
+		em := NewExpiringMap[string, int](2 * time.Second)
+		em.Set("a", 1)
+
+		time.Sleep(3 * time.Second)
+
+		_, ok := em.Get("a")
+		require.False(t, ok, "Expected `a` to have been evicted")
+		em.StopCleanup()
+	})
+
+	t.Run("Concurrency", func(t *testing.T) {
+		em := NewExpiringMap[int, string](2 * time.Second)
+		var wg sync.WaitGroup
+		for i := 0; i < 100; i++ {
+			wg.Add(1)
+			go func(i int) {
+				defer wg.Done()
+				em.Set(i, fmt.Sprintf("value%d", i))
+
+				time.Sleep(time.Duration(i) * time.Millisecond)
+
+				val, ok := em.Get(i)
+				require.True(t, ok)
+				require.Equal(t, fmt.Sprintf("value%d", i), val)
+			}(i)
+		}
+		wg.Wait()
+	})
+}
+
+func BenchmarkExpiringMap(b *testing.B) {
+	b.Run("Set", func(b *testing.B) {
+		em := NewExpiringMap[int64, int64](10 * time.Second)
+		for i := 0; i < b.N; i++ {
+			em.Set(int64(i), int64(i))
+		}
+	})
+
+	b.Run("Get", func(b *testing.B) {
+		em := NewExpiringMap[int64, int64](10 * time.Second)
+		for i := 0; i < b.N; i++ {
+			em.Set(int64(i), int64(i))
+		}
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			em.Get(int64(i))
+		}
+	})
+}