From 7aabe5f8f47786db68248fdf4b73ff6a23de17a1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joel=20H=C3=B6ner?= <joel@optimyze.cloud>
Date: Tue, 27 Jul 2021 14:16:46 +0200
Subject: [PATCH 1/2] Fix propagation of the expire time

When moving values from the recent queue into the frequent queue, the
previous implementation of the `Get` method would incorrectly reset
the expire time to the default value.
---
 2q.go            | 11 +++++++++--
 simplelru/lru.go | 14 +++++++++++---
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/2q.go b/2q.go
index de7c2a7..cacce7f 100644
--- a/2q.go
+++ b/2q.go
@@ -108,9 +108,16 @@ func (c *TwoQueueCache) Get(key interface{}) (interface{}, bool) {
 
 	// If the value is contained in recent, then we
 	// promote it to frequent
-	if val, ok := c.recent.Peek(key); ok {
+	if val, expire, ok := c.recent.PeekWithExpireTime(key); ok {
 		c.recent.Remove(key)
-		c.frequent.Add(key, val)
+		var expireDuration time.Duration
+		if expire != nil {
+			expireDuration = expire.Sub(time.Now())
+			if expireDuration < 0 {
+				return nil, false
+			}
+		}
+		c.frequent.AddEx(key, val, expireDuration)
 		return val, ok
 	}
 
diff --git a/simplelru/lru.go b/simplelru/lru.go
index d5f0264..ee4fee4 100644
--- a/simplelru/lru.go
+++ b/simplelru/lru.go
@@ -136,13 +136,21 @@ func (c *LRU) Contains(key interface{}) (ok bool) {
 // Returns the key value (or undefined if not found) without updating
 // the "recently used"-ness of the key.
 func (c *LRU) Peek(key interface{}) (value interface{}, ok bool) {
+	v, _, ok := c.PeekWithExpireTime(key)
+	return v, ok
+}
+
+// Returns the key value (or undefined if not found) and its associated expire
+// time without updating the "recently used"-ness of the key.
+func (c *LRU) PeekWithExpireTime(key interface{}) (
+	value interface{}, expire *time.Time, ok bool) {
 	if ent, ok := c.items[key]; ok {
 		if ent.Value.(*entry).IsExpired() {
-			return nil, false
+			return nil, nil, false
 		}
-		return ent.Value.(*entry).value, true
+		return ent.Value.(*entry).value, ent.Value.(*entry).expire, true
 	}
-	return nil, ok
+	return nil, nil, ok
 }
 
 // Remove removes the provided key from the cache, returning if the

From 9fa9c6e47a98a0d16ba0b6d115a021dbe5c73f10 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joel=20H=C3=B6ner?= <joel@optimyze.cloud>
Date: Mon, 4 Oct 2021 15:27:24 +0200
Subject: [PATCH 2/2] Add test case for proper expiry

---
 2q_test.go | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/2q_test.go b/2q_test.go
index 1b0f351..f54e6bd 100644
--- a/2q_test.go
+++ b/2q_test.go
@@ -3,6 +3,7 @@ package lru
 import (
 	"math/rand"
 	"testing"
+	"time"
 )
 
 func Benchmark2Q_Rand(b *testing.B) {
@@ -304,3 +305,27 @@ func Test2Q_Peek(t *testing.T) {
 		t.Errorf("should not have updated recent-ness of 1")
 	}
 }
+
+// Test that values expire as expected
+func Test2Q_Expire(t *testing.T) {
+	l, err := New2Q(100)
+	if err != nil {
+		t.Fatalf("failed to create LRU: %v", err)
+	}
+
+	l.AddEx("hey", "hello", 300*time.Millisecond)
+
+	value, ok := l.Get("hey")
+	if !ok {
+		t.Fatal("failed to read back value")
+	}
+	if value.(string) != "hello" {
+		t.Errorf("expected \"hello\", got %v", value)
+	}
+
+	time.Sleep(500 * time.Millisecond)
+	_, ok = l.Get("hey")
+	if ok {
+		t.Errorf("cached didn't properly expire")
+	}
+}