Merge pull request #14 from arthhhhh23/fix/ttl-map-orphaned-elements

fix: map ttl reinsert head or tail

Author: C-Pro

locker_test.go

@@ -93,6 +93,10 @@ func TestLockerRPanics(t *testing.T) {
 		t.Errorf("expected panic (Delete on RLocked)")
 	}
 
+	if !panics(func() { tx.SetIfPresent(1, 1) }) {
+		t.Errorf("expected panic (SetIfPresent on RLocked)")
+	}
+
 	tx.Unlock()
 	if !panics(func() { tx.Unlock() }) {
 		t.Errorf("expected panic (Unlock on already unlocked)")
@@ -117,6 +121,10 @@ func TestLockerRPanics(t *testing.T) {
 	if !panics(func() { tx.Snapshot() }) {
 		t.Errorf("expected panic (Snapshot on already unlocked)")
 	}
+
+	if !panics(func() { tx.SetIfPresent(1, 1) }) {
+		t.Errorf("expected panic (SetIfPresent on already unlocked)")
+	}
 }
 
 func TestLockerRWPanics(t *testing.T) {

map_ttl.go

@@ -190,29 +190,46 @@ func (c *MapTTLCache[K, V]) Len() int {
 }
 
 func (c *MapTTLCache[K, V]) set(key K, value V) {
+	ts := c.now()
 	val := ttlRec[K, V]{
 		value:     value,
 		prev:      c.tail,
-		timestamp: c.now(),
+		timestamp: ts,
 	}
 
 	if c.head == c.zero {
 		c.head = key
 		c.tail = key
-		val.prev = c.zero
 		c.data[key] = val
 		return
 	}
 
+	// If it's already the tail, we only need to update the value and timestamp
+	if c.tail == key {
+		rec := c.data[c.tail]
+		rec.timestamp = ts
+		rec.value = value
+		c.data[c.tail] = rec
+		return
+	}
+
 	// If the record for this key already exists
-	// and is somewhere in the middle of the list
+	// and is not already the tail of the list,
 	// removing it before adding to the tail.
-	if rec, ok := c.data[key]; ok && key != c.tail {
-		prev := c.data[rec.prev]
+	if rec, ok := c.data[key]; ok {
 		next := c.data[rec.next]
-		prev.next = rec.next
-		next.prev = rec.prev
-		c.data[rec.prev] = prev
+
+		// edge case: the current head becomes the new tail
+		if key == c.head {
+			c.head = rec.next
+			next.prev = c.zero
+		} else {
+			prev := c.data[rec.prev]
+			prev.next = rec.next
+			c.data[rec.prev] = prev
+			next.prev = rec.prev
+		}
+
 		c.data[rec.next] = next
 	}
 

ttl_test.go

@@ -2,7 +2,9 @@ package geche
 
 import (
 	"context"
+	"math/rand"
 	"strconv"
+	"sync"
 	"testing"
 	"time"
 )
@@ -174,6 +176,132 @@ func TestTTLScenario(t *testing.T) {
 	}
 }
 
+func TestHeadTailLogicConcurrent(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	m := NewMapTTLCache[string, string](ctx, time.Millisecond, time.Hour)
+
+	pool := make([]string, 50)
+	for i := range pool {
+		pool[i] = randomString(10)
+	}
+
+	wg := sync.WaitGroup{}
+	for i := 0; i < 1000; i++ {
+		idx := rand.Intn(len(pool))
+		wg.Add(1)
+
+		go func() {
+			defer wg.Done()
+
+			_, err := m.Get(pool[idx])
+			if err != nil {
+				m.Set(pool[idx], pool[idx])
+			}
+		}()
+	}
+
+	wg.Wait()
+	time.Sleep(time.Millisecond * 3)
+
+	var expectedHeadKey, expectedTailKey string
+	var expectedHeadValue, expectedTailValue ttlRec[string, string]
+
+	keys := map[string]struct{}{}
+
+	for k, v := range m.data {
+		keys[k] = struct{}{}
+		if expectedTailKey == "" {
+			expectedTailKey = k
+			expectedHeadKey = k
+			expectedHeadValue = v
+			expectedTailValue = v
+		} else if expectedTailValue.timestamp.Before(v.timestamp) {
+			expectedTailKey = k
+			expectedTailValue = v
+		} else if expectedHeadValue.timestamp.After(v.timestamp) {
+			expectedHeadKey = k
+			expectedHeadValue = v
+		}
+	}
+
+	for k, v := range m.data {
+		if _, ok := keys[v.next]; k != m.tail && !ok {
+			t.Errorf("expected key %q not found in data", v.next)
+		}
+
+		if _, ok := keys[v.prev]; k != m.head && !ok {
+			t.Errorf("expected key %q not found in data", v.prev)
+		}
+	}
+
+	if m.head != expectedHeadKey {
+		t.Errorf("expected head key %q, but got %v", expectedHeadKey, m.head)
+	}
+
+	if m.tail != expectedTailKey {
+		t.Errorf("expected tail key %q, but got %v", expectedTailKey, m.tail)
+	}
+
+	if err := m.cleanup(); err != nil {
+		t.Errorf("unexpected error in cleanup: %v", err)
+	}
+
+	if m.Len() != 0 {
+		t.Errorf("expected clean to have %d elements, but got %d", 0, m.Len())
+	}
+
+	if m.tail != m.zero {
+		t.Errorf("expected tail to be zero, but got %v", m.tail)
+	}
+
+	if m.head != m.zero {
+		t.Errorf("expected head to be zero, but got %v", m.head)
+	}
+}
+
+func TestReinsertHead(t *testing.T) {
+	c := NewMapTTLCache[string, string](context.Background(), time.Millisecond, time.Second)
+	c.Set("k1", "v1")
+	c.Set("k2", "v2")
+	c.Set("k3", "v3")
+	c.Set("k1", "v2")
+	time.Sleep(2 * time.Millisecond)
+	if err := c.cleanup(); err != nil {
+		t.Errorf("unexpected cleanup error: %v", err)
+	}
+
+	if c.Len() != 0 {
+		t.Errorf("expected cache data len to be 0 but got %d", c.Len())
+	}
+}
+
+func TestReinsertTail(t *testing.T) {
+	c := NewMapTTLCache[string, string](context.Background(), time.Millisecond, time.Second)
+	c.Set("k1", "v1")
+	c.Set("k2", "v2")
+	c.Set("k3", "v3")
+	c.Set("k3", "v4")
+	time.Sleep(2 * time.Millisecond)
+
+	if c.data["k3"].next != "" {
+		t.Errorf("expected tail next to be zero")
+	}
+
+	if c.data["k3"].prev != "k2" {
+		t.Errorf("expected tail prev to be k2, but got %s", c.data["k3"].prev)
+	}
+
+	if err := c.cleanup(); err != nil {
+		t.Errorf("unexpected cleanup error: %v", err)
+	}
+
+	if c.Len() != 0 {
+		t.Errorf("expected cache data len to be 0 but got %d", c.Len())
+	}
+}
+
 func TestSetIfPresentResetsTTL(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()

updater_test.go

@@ -103,6 +103,47 @@ func TestUpdaterScenario(t *testing.T) {
 	}
 }
 
+func TestUpdaterSetIfPresent(t *testing.T) {
+	u := NewCacheUpdater[string, string](
+		NewMapCache[string, string](),
+		updateFn,
+		2,
+	)
+
+	if u.Len() != 0 {
+		t.Errorf("expected length to be 0, but got %d", u.Len())
+	}
+
+	s, inserted := u.SetIfPresent("test", "test")
+	if inserted {
+		t.Error("expected not to insert the value")
+	}
+
+	if s != "" {
+		t.Errorf("expected to get empty string, but got %q", s)
+	}
+
+	u.Set("test", "test")
+
+	s, inserted = u.SetIfPresent("test", "test2")
+	if !inserted {
+		t.Error("expected to insert the value")
+	}
+
+	if s != "test" {
+		t.Errorf("expected to get %q, but got %q", "test", s)
+	}
+
+	v, err := u.Get("test")
+	if err != nil {
+		t.Errorf("unexpected error in Get: %v", err)
+	}
+
+	if v != "test2" {
+		t.Errorf("expected to get %q, but got %q", "test2", v)
+	}
+}
+
 func TestUpdaterErr(t *testing.T) {
 	u := NewCacheUpdater[string, string](
 		NewMapCache[string, string](),