From 7fc428cdfc931263fe4926e30444413b1962d60e Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Wed, 20 Nov 2024 22:47:41 -0800
Subject: [PATCH] solver: release unreferenced cache keys after gc

Previously this routine only ran after user ran
prune command or on reboot of the daemon.

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 control/control.go         | 24 ++++++++++++++++--------
 solver/cachemanager.go     |  5 +++++
 solver/combinedcache.go    | 12 ++++++++++++
 solver/llbsolver/bridge.go |  7 +++++++
 solver/types.go            |  2 ++
 5 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/control/control.go b/control/control.go
index 01408839aa86..9f1c4ce6708f 100644
--- a/control/control.go
+++ b/control/control.go
@@ -76,14 +76,15 @@ type Opt struct {
 
 type Controller struct { // TODO: ControlService
 	// buildCount needs to be 64bit aligned
-	buildCount       int64
-	opt              Opt
-	solver           *llbsolver.Solver
-	history          *llbsolver.HistoryQueue
-	cache            solver.CacheManager
-	gatewayForwarder *controlgateway.GatewayForwarder
-	throttledGC      func()
-	gcmu             sync.Mutex
+	buildCount                   int64
+	opt                          Opt
+	solver                       *llbsolver.Solver
+	history                      *llbsolver.HistoryQueue
+	cache                        solver.CacheManager
+	gatewayForwarder             *controlgateway.GatewayForwarder
+	throttledGC                  func()
+	throttledReleaseUnreferenced func()
+	gcmu                         sync.Mutex
 	tracev1.UnimplementedTraceServiceServer
 }
 
@@ -124,6 +125,8 @@ func NewController(opt Opt) (*Controller, error) {
 		gatewayForwarder: gatewayForwarder,
 	}
 	c.throttledGC = throttle.After(time.Minute, c.gc)
+	// use longer interval for releaseUnreferencedCache deleting links quickly is less important
+	c.throttledReleaseUnreferenced = throttle.After(5*time.Minute, func() { c.releaseUnreferencedCache(context.TODO()) })
 
 	defer func() {
 		time.AfterFunc(time.Second, c.throttledGC)
@@ -194,6 +197,10 @@ func (c *Controller) DiskUsage(ctx context.Context, r *controlapi.DiskUsageReque
 	return resp, nil
 }
 
+func (c *Controller) releaseUnreferencedCache(ctx context.Context) error {
+	return c.cache.ReleaseUnreferenced(ctx)
+}
+
 func (c *Controller) Prune(req *controlapi.PruneRequest, stream controlapi.Control_PruneServer) error {
 	if atomic.LoadInt64(&c.buildCount) == 0 {
 		imageutil.CancelCacheLeases()
@@ -634,6 +641,7 @@ func (c *Controller) gc() {
 	<-done
 	if size > 0 {
 		bklog.G(ctx).Debugf("gc cleaned up %d bytes", size)
+		go c.throttledReleaseUnreferenced()
 	}
 }
 
diff --git a/solver/cachemanager.go b/solver/cachemanager.go
index bca2c65e72ed..f9d95a507f8f 100644
--- a/solver/cachemanager.go
+++ b/solver/cachemanager.go
@@ -42,8 +42,13 @@ type cacheManager struct {
 }
 
 func (c *cacheManager) ReleaseUnreferenced(ctx context.Context) error {
+	visited := map[string]struct{}{}
 	return c.backend.Walk(func(id string) error {
 		return c.backend.WalkResults(id, func(cr CacheResult) error {
+			if _, ok := visited[cr.ID]; ok {
+				return nil
+			}
+			visited[cr.ID] = struct{}{}
 			if !c.results.Exists(ctx, cr.ID) {
 				c.backend.Release(cr.ID)
 			}
diff --git a/solver/combinedcache.go b/solver/combinedcache.go
index 211d30f0e9df..db99aaf0da43 100644
--- a/solver/combinedcache.go
+++ b/solver/combinedcache.go
@@ -33,6 +33,18 @@ func (cm *combinedCacheManager) ID() string {
 	return cm.id
 }
 
+func (cm *combinedCacheManager) ReleaseUnreferenced(ctx context.Context) error {
+	eg, ctx := errgroup.WithContext(ctx)
+	for _, c := range cm.cms {
+		func(c CacheManager) {
+			eg.Go(func() error {
+				return c.ReleaseUnreferenced(ctx)
+			})
+		}(c)
+	}
+	return eg.Wait()
+}
+
 func (cm *combinedCacheManager) Query(inp []CacheKeyWithSelector, inputIndex Index, dgst digest.Digest, outputIndex Index) ([]*CacheKey, error) {
 	eg, _ := errgroup.WithContext(context.TODO())
 	keys := make(map[string]*CacheKey, len(cm.cms))
diff --git a/solver/llbsolver/bridge.go b/solver/llbsolver/bridge.go
index 2f8fce8bf137..bf682ec6ce02 100644
--- a/solver/llbsolver/bridge.go
+++ b/solver/llbsolver/bridge.go
@@ -425,6 +425,13 @@ func (lcm *lazyCacheManager) Save(key *solver.CacheKey, s solver.Result, created
 	return lcm.main.Save(key, s, createdAt)
 }
 
+func (lcm *lazyCacheManager) ReleaseUnreferenced(ctx context.Context) error {
+	if err := lcm.wait(); err != nil {
+		return err
+	}
+	return lcm.main.ReleaseUnreferenced(ctx)
+}
+
 func (lcm *lazyCacheManager) wait() error {
 	<-lcm.waitCh
 	return lcm.err
diff --git a/solver/types.go b/solver/types.go
index 2ea7e1188896..56a53d0fd912 100644
--- a/solver/types.go
+++ b/solver/types.go
@@ -267,4 +267,6 @@ type CacheManager interface {
 
 	// Save saves a result based on a cache key
 	Save(key *CacheKey, s Result, createdAt time.Time) (*ExportableCacheKey, error)
+
+	ReleaseUnreferenced(context.Context) error
 }