diff --git a/spinifex/daemon/daemon.go b/spinifex/daemon/daemon.go index 7801f2d22..341ab8502 100644 --- a/spinifex/daemon/daemon.go +++ b/spinifex/daemon/daemon.go @@ -1869,7 +1869,13 @@ func (d *Daemon) respondWithVolumeAttachment(msg *nats.Msg, volumeID, instanceID func (rm *ResourceManager) canAllocate(instanceType *ec2.InstanceTypeInfo, count int) int { rm.mu.RLock() defer rm.mu.RUnlock() + return rm.canAllocateLocked(instanceType, count) +} +// canAllocateLocked is the lock-free body of canAllocate. Caller must already +// hold rm.mu for read or write. Extracted so allocate can re-check capacity +// while holding the write lock without dropping it. +func (rm *ResourceManager) canAllocateLocked(instanceType *ec2.InstanceTypeInfo, count int) int { // GPU capacity is managed exclusively by gpuManager.Claim; don't double-gate // on host CPU/memory which is always abundant on GPU-class hardware. if instancetypes.IsGPUType(instanceType) { @@ -1886,17 +1892,21 @@ func (rm *ResourceManager) canAllocate(instanceType *ec2.InstanceTypeInfo, count ) } -// allocate reserves resources for an instance and updates NATS subscriptions +// allocate reserves resources for one instance and updates NATS subscriptions. +// Check and commit run under a single write-lock acquisition; without this, +// two concurrent callers could both observe free capacity through the read +// lock and then both commit, overcommitting the host. Multi-instance launch +// paths loop on allocate per VM, relying on this per-call atomicity. func (rm *ResourceManager) allocate(instanceType *ec2.InstanceTypeInfo) error { - if rm.canAllocate(instanceType, 1) < 1 { + rm.mu.Lock() + if rm.canAllocateLocked(instanceType, 1) < 1 { + rm.mu.Unlock() instanceTypeName := "" if instanceType.InstanceType != nil { instanceTypeName = *instanceType.InstanceType } return fmt.Errorf("insufficient resources for instance type %s", instanceTypeName) } - - rm.mu.Lock() vCPUs := instanceTypeVCPUs(instanceType) memoryGB := float64(instanceTypeMemoryMiB(instanceType)) / 1024.0 rm.allocatedVCPU += int(vCPUs) diff --git a/spinifex/daemon/daemon_test.go b/spinifex/daemon/daemon_test.go index 44fe29731..2cc3385e9 100644 --- a/spinifex/daemon/daemon_test.go +++ b/spinifex/daemon/daemon_test.go @@ -21,6 +21,7 @@ import ( "runtime" "strings" "sync" + "sync/atomic" "testing" "time" @@ -937,6 +938,84 @@ func TestCanAllocate_CountEdgeCases(t *testing.T) { }) } +// TestAllocate_NoOvercommitUnderContention exercises the TOCTOU window that +// the pre-fix code had between canAllocate's RUnlock and the subsequent +// mu.Lock(): N concurrent allocate() callers on a pool with capacity for +// exactly N-1 must produce exactly N-1 successes and 1 insufficient-capacity +// failure. Run under -race to also catch any regression that re-introduces +// a check-then-commit gap. Pre-fix this test would intermittently overcommit; +// post-fix the per-call write-lock acquisition makes it deterministic. +func TestAllocate_NoOvercommitUnderContention(t *testing.T) { + rm, err := NewResourceManager(nil, nil) + require.NoError(t, err) + + var microType *ec2.InstanceTypeInfo + for key, it := range rm.instanceTypes { + if strings.HasSuffix(key, ".micro") { + microType = it + break + } + } + require.NotNil(t, microType, "test requires a .micro instance type") + + vCPUs := int(instanceTypeVCPUs(microType)) + memGB := float64(instanceTypeMemoryMiB(microType)) / 1024.0 + require.Greater(t, vCPUs, 0, "micro type must report vCPU count") + require.Greater(t, memGB, float64(0), "micro type must report memory") + + const capacityFor = 8 // slots + const goroutines = capacityFor + 1 + rm.mu.Lock() + rm.hostVCPU = vCPUs * capacityFor + rm.hostMemGB = memGB * float64(capacityFor) + rm.reservedVCPU = 0 + rm.reservedMem = 0 + rm.allocatedVCPU = 0 + rm.allocatedMem = 0 + rm.mu.Unlock() + + require.Equal(t, capacityFor, rm.canAllocate(microType, goroutines), + "setup sanity: pool should have exactly capacityFor slots") + + var ( + wg sync.WaitGroup + startGate = make(chan struct{}) + successes int64 + insufficient int64 + otherErrs int64 + errsMu sync.Mutex + seenErrs []string + ) + for range goroutines { + wg.Go(func() { + <-startGate + if err := rm.allocate(microType); err == nil { + atomic.AddInt64(&successes, 1) + } else if strings.Contains(err.Error(), "insufficient resources") { + atomic.AddInt64(&insufficient, 1) + } else { + atomic.AddInt64(&otherErrs, 1) + errsMu.Lock() + seenErrs = append(seenErrs, err.Error()) + errsMu.Unlock() + } + }) + } + close(startGate) + wg.Wait() + + assert.Equal(t, int64(0), otherErrs, "unexpected error classes: %v", seenErrs) + assert.Equal(t, int64(capacityFor), successes, "must commit exactly capacityFor allocations") + assert.Equal(t, int64(1), insufficient, "exactly one caller must observe insufficient capacity") + + rm.mu.RLock() + finalVCPU := rm.allocatedVCPU + finalMem := rm.allocatedMem + rm.mu.RUnlock() + assert.Equal(t, vCPUs*capacityFor, finalVCPU, "allocated vCPU must not exceed schedulable pool") + assert.InDelta(t, memGB*float64(capacityFor), finalMem, 0.001, "allocated memory must not exceed schedulable pool") +} + // TestDescribeInstances_ReservationGrouping tests that instances are grouped by reservation ID func TestDescribeInstances_ReservationGrouping(t *testing.T) { natsURL := sharedNATSURL diff --git a/spinifex/handlers/ec2/instance/service_impl.go b/spinifex/handlers/ec2/instance/service_impl.go index 55ec26892..29e822035 100644 --- a/spinifex/handlers/ec2/instance/service_impl.go +++ b/spinifex/handlers/ec2/instance/service_impl.go @@ -565,6 +565,9 @@ func (s *InstanceServiceImpl) PrepareRunInstances(input *ec2.RunInstancesInput, launchCount := allocatableCount slog.Info("PrepareRunInstances: count determined", "min", minCount, "max", maxCount, "launching", launchCount) + // Allocate is atomic per call (check + commit under a single write lock), + // so this loop can never overcommit. Under contention it may allocate + // fewer than allocatableCount; the < minCount branch below rolls back. var allocatedCount int for i := 0; i < launchCount; i++ { if err := s.resourceMgr.Allocate(instanceType); err != nil {