diff --git a/internal/list.go b/internal/list.go index 875e179..949a18a 100644 --- a/internal/list.go +++ b/internal/list.go @@ -18,7 +18,8 @@ const ( // The zero value for List is an empty list ready to use. type List[K comparable, V any] struct { root Entry[K, V] // sentinel list element, only &root, root.prev, and root.next are used - len int // current list length excluding (this) sentinel element + len int // current list length(sum of costs) excluding (this) sentinel element + count int // count of entries in list capacity uint bounded bool listType uint8 // 1 tinylfu list, 2 timerwheel list @@ -97,6 +98,7 @@ func (l *List[K, V]) insert(e, at *Entry[K, V]) *Entry[K, V] { e.next(l.listType).setPrev(e, l.listType) if l.bounded { l.len += int(e.cost.Load()) + l.count += 1 } return evicted } @@ -122,6 +124,7 @@ func (l *List[K, V]) remove(e *Entry[K, V]) { } if l.bounded { l.len -= int(e.cost.Load()) + l.count -= 1 } } diff --git a/internal/list_test.go b/internal/list_test.go index 3eb462e..0912f91 100644 --- a/internal/list_test.go +++ b/internal/list_test.go @@ -55,6 +55,24 @@ func TestList(t *testing.T) { } +func TestListCountCost(t *testing.T) { + l := NewList[string, string](100, LIST_PROBATION) + require.Equal(t, uint(100), l.capacity) + require.Equal(t, LIST_PROBATION, l.listType) + for i := 0; i < 5; i++ { + evicted := l.PushFront(NewEntry(fmt.Sprintf("%d", i), "", 20, 0)) + require.Nil(t, evicted) + } + require.Equal(t, 100, l.len) + require.Equal(t, 5, l.count) + for i := 0; i < 3; i++ { + entry := l.PopTail() + require.NotNil(t, entry) + } + require.Equal(t, 40, l.len) + require.Equal(t, 2, l.count) +} + func TestWheelList(t *testing.T) { l := NewList[string, string](5, WHEEL_LIST) require.Equal(t, uint(5), l.capacity) diff --git a/internal/tlfu.go b/internal/tlfu.go index 8ebc0c5..1592299 100644 --- a/internal/tlfu.go +++ b/internal/tlfu.go @@ -96,7 +96,7 @@ func (t *TinyLfu[K, V]) Set(entry *Entry[K, V]) *Entry[K, V] { return entry } } else { - count := t.slru.probation.len + t.slru.protected.len + count := t.slru.probation.count + t.slru.protected.count t.sketch.ensureCapacity(uint(count + count/100)) } evicted := t.slru.insert(entry)