Harden burst-proof accounting paths

Author: sjmiller609

lib/images/disk_usage.go

@@ -9,29 +9,52 @@ import (
 
 // totalReadyImageBytesFromMetadata sums ready image sizes directly from metadata.json files.
 // This is conservative for admission control and disk accounting: if metadata says an
-// image is ready, we count its recorded size without re-validating the rootfs path.
+// image is ready, we count its recorded size without re-validating the rootfs path. If
+// the metadata file is unreadable or malformed, we fall back to counting any rootfs disk
+// files found in the digest directory so we do not undercount host disk usage.
 func totalReadyImageBytesFromMetadata(imagesDir string) (int64, error) {
 	var total int64
 
 	err := filepath.Walk(imagesDir, func(path string, info os.FileInfo, err error) error {
 		if err != nil {
-			return nil
+			if os.IsNotExist(err) {
+				return nil
+			}
+			return err
 		}
 		if info.IsDir() || info.Name() != "metadata.json" {
 			return nil
 		}
 
 		data, err := os.ReadFile(path)
 		if err != nil {
-			return nil
+			rootfsBytes, fallbackErr := totalRootfsBytesInDigestDir(filepath.Dir(path))
+			if fallbackErr == nil {
+				total += rootfsBytes
+				return nil
+			}
+			return fmt.Errorf("read image metadata %s: %w", path, err)
 		}
 
 		var meta imageMetadata
 		if err := json.Unmarshal(data, &meta); err != nil {
-			return nil
+			rootfsBytes, fallbackErr := totalRootfsBytesInDigestDir(filepath.Dir(path))
+			if fallbackErr == nil {
+				total += rootfsBytes
+				return nil
+			}
+			return fmt.Errorf("unmarshal image metadata %s: %w", path, err)
 		}
 		if meta.Status == StatusReady && meta.SizeBytes > 0 {
 			total += meta.SizeBytes
+			return nil
+		}
+		if meta.Status == StatusReady {
+			rootfsBytes, err := totalRootfsBytesInDigestDir(filepath.Dir(path))
+			if err != nil {
+				return fmt.Errorf("stat ready image rootfs for %s: %w", path, err)
+			}
+			total += rootfsBytes
 		}
 		return nil
 	})
@@ -50,7 +73,10 @@ func totalOCICacheBlobBytesFromFilesystem(blobDir string) (int64, error) {
 
 	err := filepath.Walk(blobDir, func(path string, info os.FileInfo, err error) error {
 		if err != nil {
-			return nil
+			if os.IsNotExist(err) {
+				return nil
+			}
+			return err
 		}
 		if info.IsDir() {
 			return nil
@@ -117,3 +143,32 @@ func (m *manager) computeDiskUsageTotals() (int64, int64, error) {
 	}
 	return readyImageBytes, ociCacheBytes, nil
 }
+
+func totalRootfsBytesInDigestDir(digestDir string) (int64, error) {
+	rootfsPaths, err := filepath.Glob(filepath.Join(digestDir, "rootfs.*"))
+	if err != nil {
+		return 0, err
+	}
+	if len(rootfsPaths) == 0 {
+		return 0, os.ErrNotExist
+	}
+
+	var total int64
+	for _, rootfsPath := range rootfsPaths {
+		info, err := os.Stat(rootfsPath)
+		if err != nil {
+			if os.IsNotExist(err) {
+				continue
+			}
+			return 0, err
+		}
+		if info.IsDir() {
+			continue
+		}
+		total += info.Size()
+	}
+	if total == 0 {
+		return 0, os.ErrNotExist
+	}
+	return total, nil
+}

lib/images/disk_usage_test.go

@@ -0,0 +1,37 @@
+package images
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestTotalReadyImageBytesFromMetadata_UsesRootfsFallbackForMalformedMetadata(t *testing.T) {
+	t.Parallel()
+
+	imagesDir := t.TempDir()
+	digestDir := filepath.Join(imagesDir, "docker.io", "library", "alpine", "sha256deadbeef")
+	require.NoError(t, os.MkdirAll(digestDir, 0o755))
+	require.NoError(t, os.WriteFile(filepath.Join(digestDir, "metadata.json"), []byte("{not-json"), 0o644))
+	require.NoError(t, os.WriteFile(filepath.Join(digestDir, "rootfs.erofs"), []byte("rootfs-data"), 0o644))
+
+	total, err := totalReadyImageBytesFromMetadata(imagesDir)
+	require.NoError(t, err)
+	require.Equal(t, int64(len("rootfs-data")), total)
+}
+
+func TestTotalReadyImageBytesFromMetadata_UsesRootfsFallbackForReadyImageWithoutSize(t *testing.T) {
+	t.Parallel()
+
+	imagesDir := t.TempDir()
+	digestDir := filepath.Join(imagesDir, "docker.io", "library", "alpine", "sha256deadbeef")
+	require.NoError(t, os.MkdirAll(digestDir, 0o755))
+	require.NoError(t, os.WriteFile(filepath.Join(digestDir, "metadata.json"), []byte(`{"status":"ready","size_bytes":0}`), 0o644))
+	require.NoError(t, os.WriteFile(filepath.Join(digestDir, "rootfs.erofs"), []byte("another-rootfs"), 0o644))
+
+	total, err := totalReadyImageBytesFromMetadata(imagesDir)
+	require.NoError(t, err)
+	require.Equal(t, int64(len("another-rootfs")), total)
+}

lib/instances/admission_allocations.go

@@ -47,6 +47,9 @@ func (m *manager) ensureAdmissionAllocations() error {
 		if err != nil {
 			continue
 		}
+		// Cold-start cache hydration uses socket existence as the ground-truth
+		// signal for "currently active" because persisted metadata can outlive a
+		// crashed or already-stopped hypervisor process.
 		allocations[id] = m.allocationFromStoredMetadata(&meta.StoredMetadata, admissionSocketActive(&meta.StoredMetadata))
 	}
 
@@ -71,6 +74,9 @@ func (m *manager) syncAdmissionAllocation(meta *metadata) {
 		m.admissionAllocations = make(map[string]resources.InstanceAllocation)
 	}
 
+	// Incremental cache updates trust metadata because lifecycle paths update
+	// visibility explicitly once boot/restore has succeeded, avoiding repeated
+	// filesystem/socket probes on the hot path.
 	m.admissionAllocations[meta.Id] = m.allocationFromStoredMetadata(&meta.StoredMetadata, admissionMetadataActive(&meta.StoredMetadata))
 }
 

lib/instances/create.go

@@ -511,6 +511,10 @@ func (m *manager) createInstance(
 		return nil, err
 	}
 	startVMSpanEnd(nil)
+	// Mark the instance visible before releasing its pending reservation so we
+	// never create an undercount window. The tiny overlap is intentionally
+	// over-conservative: concurrent admissions may briefly see both visible and
+	// pending usage for this instance, but they will not oversubscribe the host.
 	m.setAdmissionAllocationActive(stored, true)
 	if reservedResources {
 		m.resourceValidator.FinishAllocation(id)

lib/instances/restore.go

@@ -261,6 +261,10 @@ func (m *manager) restoreInstance(
 		return nil, fmt.Errorf("resume vm failed: %w", err)
 	}
 	resumeSpanEnd(nil)
+	// Mark the instance visible before releasing its pending reservation so we
+	// never create an undercount window. The tiny overlap is intentionally
+	// over-conservative: concurrent admissions may briefly see both visible and
+	// pending usage for this instance, but they will not oversubscribe the host.
 	m.setAdmissionAllocationActive(stored, true)
 	if reservedResources {
 		m.resourceValidator.FinishAllocation(id)

lib/instances/start.go

@@ -196,6 +196,10 @@ func (m *manager) startInstance(
 		return nil, err
 	}
 	startVMSpanEnd(nil)
+	// Mark the instance visible before releasing its pending reservation so we
+	// never create an undercount window. The tiny overlap is intentionally
+	// over-conservative: concurrent admissions may briefly see both visible and
+	// pending usage for this instance, but they will not oversubscribe the host.
 	m.setAdmissionAllocationActive(stored, true)
 	if reservedResources {
 		m.resourceValidator.FinishAllocation(id)

lib/resources/resource.go

@@ -443,8 +443,8 @@ func (m *Manager) GetFullStatus(ctx context.Context) (*FullResourceStatus, error
 
 // CanAllocate checks if the requested amount can be allocated for a resource type.
 func (m *Manager) CanAllocate(ctx context.Context, rt ResourceType, amount int64) (bool, error) {
-	m.mu.Lock()
-	defer m.mu.Unlock()
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 
 	usage, err := m.collectAdmissionUsageLocked(ctx)
 	if err != nil {
@@ -713,8 +713,8 @@ func (m *Manager) validateAllocationLocked(ctx context.Context, excludeID string
 // Returns nil if allocation is allowed, or a detailed error describing
 // which resource is insufficient and the current capacity/usage.
 func (m *Manager) ValidateAllocation(ctx context.Context, vcpus int, memoryBytes int64, networkDownloadBps int64, networkUploadBps int64, diskIOBps int64, diskBytes int64, needsGPU bool) error {
-	m.mu.Lock()
-	defer m.mu.Unlock()
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 
 	req := newPendingAllocation(vcpus, memoryBytes, networkDownloadBps, networkUploadBps, diskIOBps, diskBytes, needsGPU)
 	return m.validateAllocationLocked(ctx, "", req)