From 31e2bf9c582395b0d3ee4c2906b242824b1c0e77 Mon Sep 17 00:00:00 2001
From: Erik Sipsma <erik@sipsma.dev>
Date: Wed, 31 Jul 2024 22:01:38 -0700
Subject: [PATCH] fix bug that caused mutable refs to have size 0

Signed-off-by: Erik Sipsma <erik@sipsma.dev>
---
 cache/manager_dagger.go | 12 ++++++++++--
 cache/refs.go           |  4 +++-
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/cache/manager_dagger.go b/cache/manager_dagger.go
index 48236c160a06..fb99c8f4a42a 100644
--- a/cache/manager_dagger.go
+++ b/cache/manager_dagger.go
@@ -68,7 +68,7 @@ func (cm *cacheManager) GetOrInitVolume(
 	}
 	parentID := ""
 
-	rec, err := func() (*cacheRecord, error) {
+	rec, err := func() (_ *cacheRecord, rerr error) {
 		cm.mu.Lock()
 		defer cm.mu.Unlock()
 
@@ -90,8 +90,9 @@ func (cm *cacheManager) GetOrInitVolume(
 			if err != nil {
 				return nil, fmt.Errorf("failed to create lease: %w", err)
 			}
+			// TODO: this defer should run outside this function too
 			defer func() {
-				if err != nil {
+				if rerr != nil {
 					ctx := context.WithoutCancel(ctx)
 					if err := cm.LeaseManager.Delete(ctx, leases.Lease{
 						ID: l.ID,
@@ -131,6 +132,10 @@ func (cm *cacheManager) GetOrInitVolume(
 			if err := initializeMetadata(rec.cacheMetadata, rec.parentRefs, opts...); err != nil {
 				return nil, err
 			}
+			// this is needed because for some reason snapshotID is an imageRefOption
+			if err := setImageRefMetadata(rec.cacheMetadata, opts...); err != nil {
+				return nil, fmt.Errorf("failed to append image ref metadata to ref %s: %w", id, err)
+			}
 
 			cm.records[id] = rec
 			return rec, nil
@@ -139,6 +144,9 @@ func (cm *cacheManager) GetOrInitVolume(
 			return nil, fmt.Errorf("failed to get volume cache record: %w", err)
 		}
 	}()
+	if err != nil {
+		return nil, err
+	}
 
 	releaseFunc, err := cm.volumeSnapshotter.Acquire(ctx, id, sharingMode)
 	if err != nil {
diff --git a/cache/refs.go b/cache/refs.go
index a174170c0adf..82235cc88e14 100644
--- a/cache/refs.go
+++ b/cache/refs.go
@@ -335,14 +335,16 @@ func (cr *cacheRecord) size(ctx context.Context) (int64, error) {
 			return s, nil
 		}
 		driverID := cr.getSnapshotID()
+		snapshotter := cr.cm.snapshotterFor(cr.cacheMetadata)
 		if cr.equalMutable != nil {
 			driverID = cr.equalMutable.getSnapshotID()
+			snapshotter = cr.cm.snapshotterFor(cr.equalMutable.cacheMetadata)
 		}
 		cr.mu.Unlock()
 		var usage snapshots.Usage
 		if !cr.getBlobOnly() {
 			var err error
-			usage, err = cr.cm.snapshotterFor(cr.cacheMetadata).Usage(ctx, driverID)
+			usage, err = snapshotter.Usage(ctx, driverID)
 			if err != nil {
 				cr.mu.Lock()
 				isDead := cr.isDead()