feat: add maxSize to TimedCache (#1132)

libp2p/protocols/pubsub/timedcache.nim

@@ -29,6 +29,7 @@ type
     head, tail: TimedEntry[K] # nim linked list doesn't allow inserting at pos
     entries: HashSet[TimedEntry[K]]
     timeout: Duration
+    maxSize: int # Optional max size of the cache, 0 means unlimited
 
 func `==`*[E](a, b: TimedEntry[E]): bool =
   if isNil(a) == isNil(b):
@@ -78,7 +79,18 @@ func put*[K](t: var TimedCache[K], k: K, now = Moment.now()): bool =
   # Puts k in cache, returning true if the item was already present and false
   # otherwise. If the item was already present, its expiry timer will be
   # refreshed.
+  func ensureSizeBound(t: var TimedCache[K]) =
+    if t.maxSize > 0 and t.entries.len() >= t.maxSize and k notin t:
+      if t.head != nil:
+        t.entries.excl(t.head)
+        t.head = t.head.next
+        if t.head != nil:
+          t.head.prev = nil
+        else:
+          t.tail = nil
+
   t.expire(now)
+  t.ensureSizeBound()
 
   let
     previous = t.del(k) # Refresh existing item
@@ -128,5 +140,5 @@ func addedAt*[K](t: var TimedCache[K], k: K): Moment =
 
   default(Moment)
 
-func init*[K](T: type TimedCache[K], timeout: Duration = Timeout): T =
-  T(timeout: timeout)
+func init*[K](T: type TimedCache[K], timeout: Duration = Timeout, maxSize: int = 0): T =
+  T(timeout: timeout, maxSize: maxSize)

tests/pubsub/testtimedcache.nim

@@ -57,3 +57,93 @@ suite "TimedCache":
     for i in 101 .. 100000:
       check:
         i in cache
+
+  test "max size constraint":
+    var cache = TimedCache[int].init(5.seconds, 3) # maxSize = 3
+
+    let now = Moment.now()
+    check:
+      not cache.put(1, now)
+      not cache.put(2, now + 1.seconds)
+      not cache.put(3, now + 2.seconds)
+
+    check:
+      1 in cache
+      2 in cache
+      3 in cache
+
+    check:
+      not cache.put(4, now + 3.seconds) # exceeds maxSize, evicts 1
+
+    check:
+      1 notin cache
+      2 in cache
+      3 in cache
+      4 in cache
+
+    check:
+      not cache.put(5, now + 4.seconds) # exceeds maxSize, evicts 2
+
+    check:
+      1 notin cache
+      2 notin cache
+      3 in cache
+      4 in cache
+      5 in cache
+
+    check:
+      not cache.put(6, now + 5.seconds) # exceeds maxSize, evicts 3
+
+    check:
+      1 notin cache
+      2 notin cache
+      3 notin cache
+      4 in cache
+      5 in cache
+      6 in cache
+
+  test "max size with expiration":
+    var cache = TimedCache[int].init(3.seconds, 2) # maxSize = 2
+
+    let now = Moment.now()
+    check:
+      not cache.put(1, now)
+      not cache.put(2, now + 1.seconds)
+
+    check:
+      1 in cache
+      2 in cache
+
+    check:
+      not cache.put(3, now + 5.seconds) # expires 1 and 2, should only contain 3
+
+    check:
+      1 notin cache
+      2 notin cache
+      3 in cache
+
+  test "max size constraint with refresh":
+    var cache = TimedCache[int].init(5.seconds, 3) # maxSize = 3
+
+    let now = Moment.now()
+    check:
+      not cache.put(1, now)
+      not cache.put(2, now + 1.seconds)
+      not cache.put(3, now + 2.seconds)
+
+    check:
+      1 in cache
+      2 in cache
+      3 in cache
+
+    check:
+      cache.put(1, now + 3.seconds) # refreshes 1, now 2 is the oldest
+
+    check:
+      not cache.put(4, now + 3.seconds) # exceeds maxSize, evicts 2
+
+    check:
+      1 in cache
+      2 notin cache
+      3 in cache
+      4 in cache