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

Signed-off-by: Erik Sipsma <erik@sipsma.dev>
---
 cache/manager.go        | 10 +++++++---
 cache/manager_dagger.go | 12 ++++++++++--
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/cache/manager.go b/cache/manager.go
index d2823d2bc0402..04cd9b7d95c28 100644
--- a/cache/manager.go
+++ b/cache/manager.go
@@ -1440,14 +1440,18 @@ func (cm *cacheManager) DiskUsage(ctx context.Context, opt client.DiskUsageInfo)
 			func(d *client.UsageInfo) {
 				eg.Go(func() error {
 					cm.mu.Lock()
-					ref, err := cm.get(ctx, d.ID, nil, NoUpdateLastUsed)
+					// TODO: note about this change, do upstream too? May be real bug fix
+					// previous code was getting equalImmutable of mutable refs and setting size to 0
+					// ref, err := cm.get(ctx, d.ID, nil, NoUpdateLastUsed)
+					rec, err := cm.getRecord(ctx, d.ID, NoUpdateLastUsed)
 					cm.mu.Unlock()
 					if err != nil {
 						d.Size = 0
 						return nil
 					}
-					defer ref.Release(context.TODO())
-					s, err := ref.size(ctx)
+					// defer ref.Release(context.TODO())
+					// s, err := ref.size(ctx)
+					s, err := rec.size(ctx)
 					if err != nil {
 						return err
 					}
diff --git a/cache/manager_dagger.go b/cache/manager_dagger.go
index 48236c160a065..fb99c8f4a42a7 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 {