Add burst-proof resource reservations

Author: sjmiller609

lib/instances/admission.go

@@ -0,0 +1,22 @@
+package instances
+
+func volumeOverlayReservationBytes(volumes []VolumeAttachment) int64 {
+	var total int64
+	for _, vol := range volumes {
+		if vol.Overlay {
+			total += vol.OverlaySize
+		}
+	}
+	return total
+}
+
+func requestedDiskReservationBytes(overlaySize int64, volumes []VolumeAttachment) int64 {
+	return overlaySize + volumeOverlayReservationBytes(volumes)
+}
+
+func storedDiskReservationBytes(stored *StoredMetadata) int64 {
+	if stored == nil {
+		return 0
+	}
+	return requestedDiskReservationBytes(stored.OverlaySize, stored.Volumes)
+}

lib/instances/admission_test.go

@@ -0,0 +1,33 @@
+package instances
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRequestedDiskReservationBytes(t *testing.T) {
+	t.Parallel()
+
+	diskBytes := requestedDiskReservationBytes(10, []VolumeAttachment{
+		{VolumeID: "base-only", Overlay: false, OverlaySize: 100},
+		{VolumeID: "overlay-a", Overlay: true, OverlaySize: 20},
+		{VolumeID: "overlay-b", Overlay: true, OverlaySize: 30},
+	})
+
+	assert.Equal(t, int64(60), diskBytes)
+}
+
+func TestStoredDiskReservationBytes(t *testing.T) {
+	t.Parallel()
+
+	diskBytes := storedDiskReservationBytes(&StoredMetadata{
+		OverlaySize: 15,
+		Volumes: []VolumeAttachment{
+			{VolumeID: "base-only", Overlay: false, OverlaySize: 100},
+			{VolumeID: "overlay", Overlay: true, OverlaySize: 25},
+		},
+	})
+
+	assert.Equal(t, int64(40), diskBytes)
+}

lib/instances/create.go

@@ -182,13 +182,22 @@ func (m *manager) createInstance(
 		return nil, fmt.Errorf("total memory %d (size + hotplug_size) exceeds maximum allowed %d per instance", totalMemory, m.limits.MaxMemoryPerInstance)
 	}
 
-	// Validate aggregate resource limits via ResourceValidator (if configured)
+	diskBytes := requestedDiskReservationBytes(overlaySize, req.Volumes)
+	reservedResources := false
+
+	// Reserve aggregate resources for this create while it is in flight.
 	if m.resourceValidator != nil {
 		needsGPU := req.GPU != nil && req.GPU.Profile != ""
-		if err := m.resourceValidator.ValidateAllocation(ctx, vcpus, totalMemory, req.NetworkBandwidthDownload, req.NetworkBandwidthUpload, req.DiskIOBps, needsGPU); err != nil {
-			log.ErrorContext(ctx, "resource validation failed", "error", err)
+		if err := m.resourceValidator.ReserveAllocation(ctx, id, vcpus, totalMemory, req.NetworkBandwidthDownload, req.NetworkBandwidthUpload, req.DiskIOBps, diskBytes, needsGPU); err != nil {
+			log.ErrorContext(ctx, "resource reservation failed", "error", err)
 			return nil, fmt.Errorf("%w: %v", ErrInsufficientResources, err)
 		}
+		reservedResources = true
+		defer func() {
+			if reservedResources {
+				m.resourceValidator.FinishAllocation(id)
+			}
+		}()
 	}
 
 	if req.Env == nil {
@@ -492,6 +501,10 @@ func (m *manager) createInstance(
 		return nil, err
 	}
 	startVMSpanEnd(nil)
+	if reservedResources {
+		m.resourceValidator.FinishAllocation(id)
+		reservedResources = false
+	}
 
 	// 20. Persist runtime metadata updates after VM boot.
 	meta = &metadata{StoredMetadata: *stored}

lib/instances/manager.go

@@ -84,7 +84,12 @@ type ResourceValidator interface {
 	// ValidateAllocation checks if the requested resources are available.
 	// Returns nil if allocation is allowed, or a detailed error describing
 	// which resource is insufficient and the current capacity/usage.
-	ValidateAllocation(ctx context.Context, vcpus int, memoryBytes int64, networkDownloadBps int64, networkUploadBps int64, diskIOBps int64, needsGPU bool) error
+	ValidateAllocation(ctx context.Context, vcpus int, memoryBytes int64, networkDownloadBps int64, networkUploadBps int64, diskIOBps int64, diskBytes int64, needsGPU bool) error
+	// ReserveAllocation tentatively reserves resources for an in-flight operation.
+	// Call FinishAllocation once the operation fails or becomes visible to resource accounting.
+	ReserveAllocation(ctx context.Context, instanceID string, vcpus int, memoryBytes int64, networkDownloadBps int64, networkUploadBps int64, diskIOBps int64, diskBytes int64, needsGPU bool) error
+	// FinishAllocation removes any pending reservation for the given instance ID.
+	FinishAllocation(instanceID string)
 }
 
 type manager struct {

lib/instances/restore.go

@@ -59,13 +59,21 @@ func (m *manager) restoreInstance(
 	}
 
 	// 2b. Validate aggregate resource limits before allocating resources (if configured)
+	reservedResources := false
 	if m.resourceValidator != nil {
 		needsGPU := stored.GPUProfile != ""
 		totalMemory := stored.Size + stored.HotplugSize
-		if err := m.resourceValidator.ValidateAllocation(ctx, stored.Vcpus, totalMemory, stored.NetworkBandwidthDownload, stored.NetworkBandwidthUpload, stored.DiskIOBps, needsGPU); err != nil {
-			log.ErrorContext(ctx, "resource validation failed for restore", "instance_id", id, "error", err)
+		diskBytes := storedDiskReservationBytes(stored)
+		if err := m.resourceValidator.ReserveAllocation(ctx, id, stored.Vcpus, totalMemory, stored.NetworkBandwidthDownload, stored.NetworkBandwidthUpload, stored.DiskIOBps, diskBytes, needsGPU); err != nil {
+			log.ErrorContext(ctx, "resource reservation failed for restore", "instance_id", id, "error", err)
 			return nil, fmt.Errorf("%w: %v", ErrInsufficientResources, err)
 		}
+		reservedResources = true
+		defer func() {
+			if reservedResources {
+				m.resourceValidator.FinishAllocation(id)
+			}
+		}()
 	}
 
 	// 3. Get snapshot directory
@@ -253,6 +261,10 @@ func (m *manager) restoreInstance(
 		return nil, fmt.Errorf("resume vm failed: %w", err)
 	}
 	resumeSpanEnd(nil)
+	if reservedResources {
+		m.resourceValidator.FinishAllocation(id)
+		reservedResources = false
+	}
 
 	// Forked standby restores may allocate a fresh identity while the guest memory snapshot
 	// still has the source VM's old IP configuration. Reconfigure guest networking after

lib/instances/start.go

@@ -61,13 +61,21 @@ func (m *manager) startInstance(
 	}
 
 	// 2b. Validate aggregate resource limits before allocating resources (if configured)
+	reservedResources := false
 	if m.resourceValidator != nil {
 		needsGPU := stored.GPUProfile != ""
 		totalMemory := stored.Size + stored.HotplugSize
-		if err := m.resourceValidator.ValidateAllocation(ctx, stored.Vcpus, totalMemory, stored.NetworkBandwidthDownload, stored.NetworkBandwidthUpload, stored.DiskIOBps, needsGPU); err != nil {
-			log.ErrorContext(ctx, "resource validation failed for start", "instance_id", id, "error", err)
+		diskBytes := storedDiskReservationBytes(stored)
+		if err := m.resourceValidator.ReserveAllocation(ctx, id, stored.Vcpus, totalMemory, stored.NetworkBandwidthDownload, stored.NetworkBandwidthUpload, stored.DiskIOBps, diskBytes, needsGPU); err != nil {
+			log.ErrorContext(ctx, "resource reservation failed for start", "instance_id", id, "error", err)
 			return nil, fmt.Errorf("%w: %v", ErrInsufficientResources, err)
 		}
+		reservedResources = true
+		defer func() {
+			if reservedResources {
+				m.resourceValidator.FinishAllocation(id)
+			}
+		}()
 	}
 
 	// 3. Get image info (needed for buildHypervisorConfig)
@@ -188,6 +196,10 @@ func (m *manager) startInstance(
 		return nil, err
 	}
 	startVMSpanEnd(nil)
+	if reservedResources {
+		m.resourceValidator.FinishAllocation(id)
+		reservedResources = false
+	}
 
 	// Success - release cleanup stack (prevent cleanup)
 	cu.Release()