Expose snapshot reference counts

Author: sjmiller609

cmd/api/api/snapshots.go

@@ -336,6 +336,7 @@ func snapshotToOAPI(snapshot instances.Snapshot) oapi.Snapshot {
 		SourceHypervisor:   sourceHypervisor,
 		CreatedAt:          snapshot.CreatedAt,
 		SizeBytes:          snapshot.SizeBytes,
+		RefCount:           snapshot.RefCount,
 		Name:               lo.ToPtr(snapshot.Name),
 	}
 	if snapshot.CompressionState != "" {

lib/instances/snapshot.go

@@ -32,6 +32,13 @@ func (m *manager) listSnapshots(ctx context.Context, filter *ListSnapshotsFilter
 	if err != nil {
 		return nil, fmt.Errorf("list snapshots: %w", err)
 	}
+	counts, err := m.snapshotForkCounts()
+	if err != nil {
+		return nil, err
+	}
+	for i := range snapshots {
+		snapshots[i].RefCount = counts[snapshots[i].Id]
+	}
 	return snapshots, nil
 }
 
@@ -44,6 +51,11 @@ func (m *manager) getSnapshot(ctx context.Context, snapshotID string) (*Snapshot
 		}
 		return nil, err
 	}
+	refCount, err := m.countSnapshotForks(snapshotID)
+	if err != nil {
+		return nil, err
+	}
+	snapshot.RefCount = refCount
 	return snapshot, nil
 }
 
@@ -679,26 +691,34 @@ func (m *manager) listSnapshotRecords() ([]snapshotRecord, error) {
 }
 
 func (m *manager) countSnapshotForks(snapshotID string) (int, error) {
-	metaFiles, err := m.listMetadataFiles()
+	counts, err := m.snapshotForkCounts()
 	if err != nil {
 		return 0, err
 	}
+	return counts[snapshotID], nil
+}
+
+func (m *manager) snapshotForkCounts() (map[string]int, error) {
+	metaFiles, err := m.listMetadataFiles()
+	if err != nil {
+		return nil, err
+	}
 
-	count := 0
+	counts := make(map[string]int)
 	for _, metaPath := range metaFiles {
 		content, err := os.ReadFile(metaPath)
 		if err != nil {
-			return 0, fmt.Errorf("read instance metadata %s: %w", metaPath, err)
+			return nil, fmt.Errorf("read instance metadata %s: %w", metaPath, err)
 		}
 		var meta metadata
 		if err := json.Unmarshal(content, &meta); err != nil {
-			return 0, fmt.Errorf("unmarshal instance metadata %s: %w", metaPath, err)
+			return nil, fmt.Errorf("unmarshal instance metadata %s: %w", metaPath, err)
 		}
-		if meta.ForkOfSnapshot == snapshotID {
-			count++
+		if meta.ForkOfSnapshot != "" {
+			counts[meta.ForkOfSnapshot]++
 		}
 	}
-	return count, nil
+	return counts, nil
 }
 
 func snapshotMemHardlinkSource(srcDir string) (absPath, relSlash string, ok bool) {

lib/instances/snapshot_test.go

@@ -43,6 +43,7 @@ func TestStoppedSnapshotLifecycleAndForkAfterSourceDeletion(t *testing.T) {
 	got, err := mgr.GetSnapshot(ctx, snap.Id)
 	require.NoError(t, err)
 	require.Equal(t, snap.Id, got.Id)
+	require.Equal(t, 0, got.RefCount)
 
 	forked, err := mgr.ForkSnapshot(ctx, snap.Id, ForkSnapshotRequest{
 		Name:             "snapshot-stopped-fork",
@@ -57,6 +58,15 @@ func TestStoppedSnapshotLifecycleAndForkAfterSourceDeletion(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, snap.Id, forkMeta.ForkOfSnapshot)
 
+	got, err = mgr.GetSnapshot(ctx, snap.Id)
+	require.NoError(t, err)
+	require.Equal(t, 1, got.RefCount)
+
+	snapshots, err := mgr.ListSnapshots(ctx, &ListSnapshotsFilter{SourceInstanceID: &sourceID})
+	require.NoError(t, err)
+	require.Len(t, snapshots, 1)
+	require.Equal(t, 1, snapshots[0].RefCount)
+
 	err = mgr.DeleteSnapshot(ctx, snap.Id)
 	require.Error(t, err)
 	assert.ErrorIs(t, err, ErrInvalidState)