[#80] Fixed uint32 capacity overflow

maypok86 · Apr 14, 2024 · 9bcb26a · 9bcb26a
1 parent 9d61c67
commit 9bcb26a
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 20 deletions.
diff --git a/internal/core/cache.go b/internal/core/cache.go
@@ -139,7 +139,7 @@ func NewCache[K comparable, V any](c Config[K, V]) *Cache[K, V] {
 	cache := &Cache[K, V]{
 		nodeManager:      nodeManager,
 		hashmap:          hashmap,
-		policy:           s3fifo.NewPolicy[K, V](uint32(c.Capacity)),
+		policy:           s3fifo.NewPolicy[K, V](c.Capacity),
 		expiryPolicy:     expPolicy,
 		readBuffers:      readBuffers,
 		writeBuffer:      queue.NewGrowable[task[K, V]](minWriteBufferCapacity, maxWriteBufferCapacity),
@@ -277,7 +277,7 @@ func (c *Cache[K, V]) SetIfAbsentWithTTL(key K, value V, ttl time.Duration) bool
 
 func (c *Cache[K, V]) set(key K, value V, expiration uint32, onlyIfAbsent bool) bool {
 	cost := c.costFunc(key, value)
-	if cost > c.policy.MaxAvailableCost() {
+	if int(cost) > c.policy.MaxAvailableCost() {
 		c.stats.IncRejectedSets()
 		return false
 	}

diff --git a/internal/core/cache_test.go b/internal/core/cache_test.go
@@ -16,7 +16,7 @@ func TestCache_SetWithCost(t *testing.T) {
 		},
 	})
 
-	goodCost := int(c.policy.MaxAvailableCost())
+	goodCost := c.policy.MaxAvailableCost()
 	badCost := goodCost + 1
 
 	added := c.Set(goodCost, 1)

diff --git a/internal/s3fifo/main.go b/internal/s3fifo/main.go
@@ -22,11 +22,11 @@ const maxReinsertions = 20
 
 type main[K comparable, V any] struct {
 	q       *queue[K, V]
-	cost    uint32
-	maxCost uint32
+	cost    int
+	maxCost int
 }
 
-func newMain[K comparable, V any](maxCost uint32) *main[K, V] {
+func newMain[K comparable, V any](maxCost int) *main[K, V] {
 	return &main[K, V]{
 		q:       newQueue[K, V](),
 		maxCost: maxCost,
@@ -36,7 +36,7 @@ func newMain[K comparable, V any](maxCost uint32) *main[K, V] {
 func (m *main[K, V]) insert(n node.Node[K, V]) {
 	m.q.push(n)
 	n.MarkMain()
-	m.cost += n.Cost()
+	m.cost += int(n.Cost())
 }
 
 func (m *main[K, V]) evict(deleted []node.Node[K, V]) []node.Node[K, V] {
@@ -46,15 +46,15 @@ func (m *main[K, V]) evict(deleted []node.Node[K, V]) []node.Node[K, V] {
 
 		if !n.IsAlive() || n.HasExpired() || n.Frequency() == 0 {
 			n.Unmark()
-			m.cost -= n.Cost()
+			m.cost -= int(n.Cost())
 			return append(deleted, n)
 		}
 
 		// to avoid the worst case O(n), we remove the 20th reinserted consecutive element.
 		reinsertions++
 		if reinsertions >= maxReinsertions {
 			n.Unmark()
-			m.cost -= n.Cost()
+			m.cost -= int(n.Cost())
 			return append(deleted, n)
 		}
 
@@ -65,7 +65,7 @@ func (m *main[K, V]) evict(deleted []node.Node[K, V]) []node.Node[K, V] {
 }
 
 func (m *main[K, V]) remove(n node.Node[K, V]) {
-	m.cost -= n.Cost()
+	m.cost -= int(n.Cost())
 	n.Unmark()
 	m.q.remove(n)
 }

diff --git a/internal/s3fifo/policy.go b/internal/s3fifo/policy.go
@@ -24,12 +24,12 @@ type Policy[K comparable, V any] struct {
 	small                *small[K, V]
 	main                 *main[K, V]
 	ghost                *ghost[K, V]
-	maxCost              uint32
-	maxAvailableNodeCost uint32
+	maxCost              int
+	maxAvailableNodeCost int
 }
 
 // NewPolicy creates a new Policy.
-func NewPolicy[K comparable, V any](maxCost uint32) *Policy[K, V] {
+func NewPolicy[K comparable, V any](maxCost int) *Policy[K, V] {
 	smallMaxCost := maxCost / 10
 	mainMaxCost := maxCost - smallMaxCost
 
@@ -95,7 +95,7 @@ func (p *Policy[K, V]) Delete(n node.Node[K, V]) {
 }
 
 // MaxAvailableCost returns the maximum available cost of the node.
-func (p *Policy[K, V]) MaxAvailableCost() uint32 {
+func (p *Policy[K, V]) MaxAvailableCost() int {
 	return p.maxAvailableNodeCost
 }
 

diff --git a/internal/s3fifo/small.go b/internal/s3fifo/small.go
@@ -22,12 +22,12 @@ type small[K comparable, V any] struct {
 	q       *queue[K, V]
 	main    *main[K, V]
 	ghost   *ghost[K, V]
-	cost    uint32
-	maxCost uint32
+	cost    int
+	maxCost int
 }
 
 func newSmall[K comparable, V any](
-	maxCost uint32,
+	maxCost int,
 	main *main[K, V],
 	ghost *ghost[K, V],
 ) *small[K, V] {
@@ -42,7 +42,7 @@ func newSmall[K comparable, V any](
 func (s *small[K, V]) insert(n node.Node[K, V]) {
 	s.q.push(n)
 	n.MarkSmall()
-	s.cost += n.Cost()
+	s.cost += int(n.Cost())
 }
 
 func (s *small[K, V]) evict(deleted []node.Node[K, V]) []node.Node[K, V] {
@@ -51,7 +51,7 @@ func (s *small[K, V]) evict(deleted []node.Node[K, V]) []node.Node[K, V] {
 	}
 
 	n := s.q.pop()
-	s.cost -= n.Cost()
+	s.cost -= int(n.Cost())
 	n.Unmark()
 	if !n.IsAlive() || n.HasExpired() {
 		return append(deleted, n)
@@ -70,7 +70,7 @@ func (s *small[K, V]) evict(deleted []node.Node[K, V]) []node.Node[K, V] {
 }
 
 func (s *small[K, V]) remove(n node.Node[K, V]) {
-	s.cost -= n.Cost()
+	s.cost -= int(n.Cost())
 	n.Unmark()
 	s.q.remove(n)
 }