diff --git a/lru/lru_with_expiry_test.go b/lru/lru_with_expiry_test.go index 42812d70..64190859 100644 --- a/lru/lru_with_expiry_test.go +++ b/lru/lru_with_expiry_test.go @@ -45,6 +45,11 @@ func TestTypedGetWithExpiry(t *testing.T) { {name: "string_hit_expiring_big_jump", keyToAdd: "myKey", keyToGet: "myKey", keyExpiry: startTime.Add(time.Hour), setTime: startTime.Add(time.Hour * 24 * 365), expectedOK: false}, {name: "string_hit_expiring_exact_jump", keyToAdd: "myKey", keyToGet: "myKey", keyExpiry: startTime.Add(time.Hour), setTime: startTime.Add(time.Hour), expectedOK: false}, {name: "string_hit_expiring_jump_just_under", keyToAdd: "myKey", keyToGet: "myKey", keyExpiry: startTime.Add(time.Hour), setTime: startTime.Add(time.Hour - time.Nanosecond), expectedOK: true}, + {name: "string_hit_zero_val_no_advance_prior_entry", keyToAdd: "fizzlebat", keyToGet: "fizzlebat", setTime: startTime, expectedOK: true}, + {name: "string_hit_zero_val_big_jump_prior_entry", keyToAdd: "fizzlebat", keyToGet: "fizzlebat", setTime: startTime.Add(time.Hour * 24 * 365), expectedOK: true}, + {name: "string_hit_expiring_big_jump_prior_entry", keyToAdd: "fizzlebat", keyToGet: "fizzlebat", keyExpiry: startTime.Add(time.Hour), setTime: startTime.Add(time.Hour * 24 * 365), expectedOK: false}, + {name: "string_hit_expiring_exact_jump_prior_entry", keyToAdd: "fizzlebat", keyToGet: "fizzlebat", keyExpiry: startTime.Add(time.Hour), setTime: startTime.Add(time.Hour), expectedOK: false}, + {name: "string_hit_expiring_jump_just_under_prior_entry", keyToAdd: "fizzlebat", keyToGet: "fizzlebat", keyExpiry: startTime.Add(time.Hour), setTime: startTime.Add(time.Hour - time.Nanosecond), expectedOK: true}, } for _, tt := range getTests { @@ -52,6 +57,8 @@ func TestTypedGetWithExpiry(t *testing.T) { clk := fake.NewClock(startTime) lru := TypedNew[string, int](0) lru.Clock = clk + // add a pre-existing value that we can update with a different expiry + lru.AddExpiring("fizzlebat", 4321, startTime.Add(time.Minute)) lru.AddExpiring(tt.keyToAdd, 1234, tt.keyExpiry) clk.SetClock(tt.setTime) diff --git a/lru/typed_lru.go b/lru/typed_lru.go index 163c766a..57eb3de3 100644 --- a/lru/typed_lru.go +++ b/lru/typed_lru.go @@ -121,8 +121,19 @@ func (c *TypedCache[K, V]) AddExpiring(key K, value V, expiration time.Time) { if ele, hit := c.cache[key]; hit { c.ll.MoveToFront(ele) ele.value.value = value + // Remove the existing expiration handle from the heap before (possibly) adding a new one. + // this way we don't end up with both expirations in the heap. + if ele.value.hasExpiry { + c.expirations.Remove(ele.value.expHandle) + } if !expiration.IsZero() { ele.value.expHandle = c.expirations.Push(expiration, weak.Make(ele)) + ele.value.hasExpiry = true + ele.value.expiry = expiration.Sub(c.expiryBase) + } else { + ele.value.expHandle = nil + ele.value.hasExpiry = false + ele.value.expiry = 0 } return }