diff --git a/libp2p/protocols/pubsub/timedcache.nim b/libp2p/protocols/pubsub/timedcache.nim index 4a55a33d4e..9187069cd4 100644 --- a/libp2p/protocols/pubsub/timedcache.nim +++ b/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) diff --git a/tests/pubsub/testtimedcache.nim b/tests/pubsub/testtimedcache.nim index 0dcf07f314..66dec1b6f1 100644 --- a/tests/pubsub/testtimedcache.nim +++ b/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