lib/cache: add SetCanExpire to keep entries which are still busy

An entry is expired when it hasn't been used for a while, but the user
of the cache may know it is still doing something in the background.
SetCanExpire sets a function which is asked before expiring an entry
and can keep it in the cache until it is ready to go.

The function is called without the cache locked as it may take a
while, so an entry which is used in the meantime isn't expired.

(cherry picked from commit 49d8566ec5)

lib/cache/cache.go

@@ -17,6 +17,7 @@ type Cache struct {
 	expireDuration time.Duration // expire the cache entry when it is older than this
 	expireInterval time.Duration // interval to run the cache expire
 	finalize       func(value any)
+	canExpire      func(value any) bool // if set, asked before expiring a value
 }
 
 // New creates a new cache with the default expire duration and interval
@@ -222,8 +223,29 @@ func (c *Cache) cacheExpire() {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	now := time.Now()
+	expired := func(entry *cacheEntry) bool {
+		return entry.pinCount <= 0 && now.Sub(entry.lastUsed) > c.expireDuration
+	}
+	candidates := map[string]*cacheEntry{}
 	for key, entry := range c.cache {
-		if entry.pinCount <= 0 && now.Sub(entry.lastUsed) > c.expireDuration {
+		if expired(entry) {
+			candidates[key] = entry
+		}
+	}
+	if c.canExpire != nil && len(candidates) != 0 {
+		// canExpire may block so call it without the lock
+		canExpire := c.canExpire
+		c.mu.Unlock()
+		for key, entry := range candidates {
+			if !canExpire(entry.value) {
+				delete(candidates, key)
+			}
+		}
+		c.mu.Lock()
+	}
+	for key, entry := range candidates {
+		// The entry may have been used or replaced while the lock was released
+		if c.cache[key] == entry && expired(entry) {
 			c.finalize(entry.value)
 			delete(c.cache, key)
 		}
@@ -261,6 +283,19 @@ func (c *Cache) SetFinalizer(finalize func(any)) {
 	c.mu.Unlock()
 }
 
+// SetCanExpire sets a function which is asked whether a value which
+// hasn't been used recently can be expired from the cache. If it
+// returns false the value is kept and asked about again the next time
+// the cache expiry runs.
+//
+// It is called without the cache locked so may use the cache. It
+// isn't called for values removed in other ways, e.g. Delete or Clear.
+func (c *Cache) SetCanExpire(canExpire func(value any) bool) {
+	c.mu.Lock()
+	c.canExpire = canExpire
+	c.mu.Unlock()
+}
+
 // EntriesWithPinCount returns the number of pinned and unpinned entries in the cache
 //
 // Each entry is counted only once, regardless of entry.pinCount

lib/cache/cache_test.go

@@ -383,3 +383,64 @@ func TestCacheFinalize(t *testing.T) {
 	c.cacheExpire() // "ok" and "new" fall out of cache
 	assert.Equal(t, 6, numCalled)
 }
+
+func TestCacheCanExpire(t *testing.T) {
+	c := New()
+	finalized := 0
+	c.SetFinalizer(func(v any) {
+		finalized++
+	})
+	busy := true
+	c.SetCanExpire(func(v any) bool {
+		// Must be callable without the lock
+		assert.Equal(t, 2, c.Entries())
+		return v != "busy" || !busy
+	})
+	create := func(path string) (any, bool, error) {
+		return path, true, nil
+	}
+	c.expireDuration = 1 * time.Millisecond
+	_, _ = c.Get("busy", create)
+	_, _ = c.Get("pinned", create)
+	c.Pin("pinned")
+	time.Sleep(2 * time.Millisecond)
+
+	// Neither the busy nor the pinned entry expire
+	c.cacheExpire()
+	assert.Equal(t, 2, c.Entries())
+	assert.Equal(t, 0, finalized)
+	assert.Equal(t, true, c.expireRunning)
+
+	// The busy entry expires as soon as it is no longer busy
+	busy = false
+	c.cacheExpire()
+	assert.Equal(t, 1, c.Entries())
+	assert.Equal(t, 1, finalized)
+	_, found := c.GetMaybe("pinned")
+	assert.True(t, found)
+}
+
+func TestCacheCanExpireUsed(t *testing.T) {
+	c := New()
+	finalized := 0
+	c.SetFinalizer(func(v any) {
+		finalized++
+	})
+	create := func(path string) (any, bool, error) {
+		return path, true, nil
+	}
+	// Use the entry while it is being asked about
+	c.SetCanExpire(func(v any) bool {
+		_, _ = c.Get("/", create)
+		return true
+	})
+	c.expireDuration = 1 * time.Second
+	_, _ = c.Get("/", create)
+	c.mu.Lock()
+	c.cache["/"].lastUsed = time.Now().Add(-time.Minute)
+	c.mu.Unlock()
+
+	c.cacheExpire()
+	assert.Equal(t, 1, c.Entries())
+	assert.Equal(t, 0, finalized)
+}