diff --git a/internal/tlfu.go b/internal/tlfu.go index 84441b3..7348711 100644 --- a/internal/tlfu.go +++ b/internal/tlfu.go @@ -172,8 +172,14 @@ func (t *TinyLfu[K, V]) Set(entry *Entry[K, V]) { if entry.meta.prev == nil { t.window.PushFront(entry) } + t.demoteFromProtected() - t.EvictEntries() + if t.weightedSize > t.capacity { + t.EvictEntries() + } else { + count := t.slru.probation.count + t.slru.protected.count + t.window.count + t.sketch.EnsureCapacity(uint(count)) + } } func (t *TinyLfu[K, V]) Access(item ReadBufItem[K, V]) { diff --git a/internal/tlfu_test.go b/internal/tlfu_test.go index eb2858e..412a655 100644 --- a/internal/tlfu_test.go +++ b/internal/tlfu_test.go @@ -549,3 +549,20 @@ func TestTlfu_AdaptiveAmountRemain(t *testing.T) { require.Equal(t, "998-998>149-110:109-103>101-80:79-0", result) } + +func TestTlfu_SketchResize(t *testing.T) { + hasher := NewHasher[int](nil) + tlfu := NewTinyLfu[int, int](10000, hasher) + + for i := 0; i < 10000; i++ { + tlfu.Set(&Entry[int, int]{key: i, value: i, policyWeight: 1}) + require.True(t, len(tlfu.sketch.Table) >= i, fmt.Sprintf("sketch size %d < %d", len(tlfu.sketch.Table), i)) + } + + size := len(tlfu.sketch.Table) + require.Equal(t, 16384, size) + + for i := 10000; i < 20000; i++ { + require.Equal(t, size, len(tlfu.sketch.Table)) + } +}