From 03a1b4b8ab49ea9aaa98a1b871ad57ccda76b997 Mon Sep 17 00:00:00 2001
From: Taahir Ahmed <taahm@google.com>
Date: Fri, 20 Oct 2023 18:05:33 -0700
Subject: [PATCH] LRUExpireCache: Allow removing multiple keys under lock

Kubernetes-commit: e83baddbb1a7b8f88db357fcefc8c7f3dde4856c
---
 pkg/util/cache/lruexpirecache.go      | 13 +++++++++++++
 pkg/util/cache/lruexpirecache_test.go | 14 ++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/pkg/util/cache/lruexpirecache.go b/pkg/util/cache/lruexpirecache.go
index 1328dd612..ad486d580 100644
--- a/pkg/util/cache/lruexpirecache.go
+++ b/pkg/util/cache/lruexpirecache.go
@@ -136,6 +136,19 @@ func (c *LRUExpireCache) Remove(key interface{}) {
 	delete(c.entries, key)
 }
 
+// RemoveAll removes all keys that match predicate.
+func (c *LRUExpireCache) RemoveAll(predicate func(key any) bool) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	for key, element := range c.entries {
+		if predicate(key) {
+			c.evictionList.Remove(element)
+			delete(c.entries, key)
+		}
+	}
+}
+
 // Keys returns all unexpired keys in the cache.
 //
 // Keep in mind that subsequent calls to Get() for any of the returned keys
diff --git a/pkg/util/cache/lruexpirecache_test.go b/pkg/util/cache/lruexpirecache_test.go
index f1de853ef..d6bd7464f 100644
--- a/pkg/util/cache/lruexpirecache_test.go
+++ b/pkg/util/cache/lruexpirecache_test.go
@@ -67,6 +67,20 @@ func TestSimpleRemove(t *testing.T) {
 	expectNotEntry(t, c, "long-lived")
 }
 
+func TestSimpleRemoveAll(t *testing.T) {
+	c := NewLRUExpireCache(10)
+	c.Add("long-lived", "12345", 10*time.Hour)
+	c.Add("other-long-lived", "12345", 10*time.Hour)
+	c.RemoveAll(func(k any) bool {
+		return k.(string) == "long-lived"
+	})
+
+	assertKeys(t, c.Keys(), []any{"other-long-lived"})
+
+	expectNotEntry(t, c, "long-lived")
+	expectEntry(t, c, "other-long-lived", "12345")
+}
+
 func TestExpiredGet(t *testing.T) {
 	fakeClock := testingclock.NewFakeClock(time.Now())
 	c := NewLRUExpireCacheWithClock(10, fakeClock)