Spread only on same machine size. (#625)

Author: Gerrit91

cmd/metal-api/internal/datastore/machine.go

@@ -422,7 +422,7 @@ func (rs *RethinkStore) UpdateMachine(oldMachine *metal.Machine, newMachine *met
 // FindWaitingMachine returns an available, not allocated, waiting and alive machine of given size within the given partition.
 // TODO: the algorithm can be optimized / shortened by using a rethinkdb join command and then using .Sample(1)
 // but current implementation should have a slightly better readability.
-func (rs *RethinkStore) FindWaitingMachine(ctx context.Context, projectid, partitionid string, size metal.Size, placementTags []string) (*metal.Machine, error) {
+func (rs *RethinkStore) FindWaitingMachine(ctx context.Context, projectid, partitionid string, size metal.Size, placementTags []string, role metal.Role) (*metal.Machine, error) {
 	q := *rs.machineTable()
 	q = q.Filter(map[string]interface{}{
 		"allocation":  nil,
@@ -473,6 +473,7 @@ func (rs *RethinkStore) FindWaitingMachine(ctx context.Context, projectid, parti
 	var partitionMachines metal.Machines
 	err = rs.SearchMachines(&MachineSearchQuery{
 		PartitionID: &partitionid,
+		SizeID:      &size.ID,
 	}, &partitionMachines)
 	if err != nil {
 		return nil, err
@@ -487,12 +488,12 @@ func (rs *RethinkStore) FindWaitingMachine(ctx context.Context, projectid, parti
 		return nil, err
 	}
 
-	ok := checkSizeReservations(available, projectid, partitionMachines.WithSize(size.ID).ByProjectID(), reservations)
+	ok := checkSizeReservations(available, projectid, partitionMachines.ByProjectID(), reservations)
 	if !ok {
 		return nil, errors.New("no machine available")
 	}
 
-	projectMachines := partitionMachines.ByProjectID()[projectid]
+	projectMachines := partitionMachines.WithRole(role).ByProjectID()[projectid]
 
 	spreadCandidates := spreadAcrossRacks(available, projectMachines, placementTags)
 	if len(spreadCandidates) == 0 {
@@ -523,8 +524,6 @@ func checkSizeReservations(available metal.Machines, projectid string, machinesB
 	)
 
 	for _, r := range reservations {
-		r := r
-
 		// sum up the amount of reservations
 		amount += r.Amount
 

cmd/metal-api/internal/datastore/machine_integration_test.go

@@ -5,6 +5,7 @@ package datastore
 
 import (
 	"context"
+	"fmt"
 	"log/slog"
 	"os"
 	"strconv"
@@ -993,8 +994,6 @@ func Test_FindWaitingMachine_NoConcurrentModificationErrors(t *testing.T) {
 			},
 		},
 	} {
-		initEntity := initEntity
-
 		err := sharedDS.createEntity(initEntity.table, initEntity.entity)
 		require.NoError(t, err)
 
@@ -1004,8 +1003,7 @@ func Test_FindWaitingMachine_NoConcurrentModificationErrors(t *testing.T) {
 		}()
 	}
 
-	for i := 0; i < 100; i++ {
-		i := i
+	for i := range 100 {
 		wg.Add(1)
 
 		log := root.With("worker", i)
@@ -1014,7 +1012,7 @@ func Test_FindWaitingMachine_NoConcurrentModificationErrors(t *testing.T) {
 			defer wg.Done()
 
 			for {
-				machine, err := sharedDS.FindWaitingMachine(context.Background(), "project", "partition", size, nil)
+				machine, err := sharedDS.FindWaitingMachine(context.Background(), "project", "partition", size, nil, metal.RoleMachine)
 				if err != nil {
 					if metal.IsConflict(err) {
 						t.Errorf("concurrent modification occurred, shared mutex is not working")
@@ -1061,3 +1059,121 @@ func Test_FindWaitingMachine_NoConcurrentModificationErrors(t *testing.T) {
 
 	assert.Equal(t, 100, count)
 }
+
+func Test_FindWaitingMachine_RackSpreadingDistribution(t *testing.T) {
+	var (
+		partitionID = "partition"
+		projectID   = "project"
+		size1       = metal.Size{Base: metal.Base{ID: "1"}}
+		fiveRacks   = func(i int) string {
+			return "rack-" + strconv.FormatInt(int64((i%5)+1), 10)
+		}
+	)
+
+	defer func() {
+		_, err := sharedDS.machineTable().Delete().RunWrite(sharedDS.session)
+		require.NoError(t, err)
+		_, err = sharedDS.eventTable().Delete().RunWrite(sharedDS.session)
+		require.NoError(t, err)
+	}()
+
+	for i := range 200 {
+		err := sharedDS.createEntity(sharedDS.machineTable(), &metal.Machine{
+			Base: metal.Base{
+				ID: strconv.Itoa(i),
+			},
+			PartitionID: partitionID,
+			SizeID:      size1.ID,
+			State: metal.MachineState{
+				Value: metal.AvailableState,
+			},
+			Waiting:      true,
+			PreAllocated: false,
+			RackID:       fiveRacks(i),
+		})
+		require.NoError(t, err)
+
+		err = sharedDS.createEntity(sharedDS.eventTable(), &metal.ProvisioningEventContainer{
+			Base: metal.Base{
+				ID: strconv.Itoa(i),
+			},
+			Liveliness: metal.MachineLivelinessAlive,
+		})
+		require.NoError(t, err)
+	}
+
+	// just allocate some machines with different specs that should not influence later allocs
+	for i, spec := range []struct {
+		role metal.Role
+		size string
+	}{
+		{role: metal.RoleFirewall, size: "firewall"},
+		{role: metal.RoleFirewall, size: "firewall"},
+		{role: metal.RoleMachine, size: "machine"},
+		{role: metal.RoleMachine, size: "machine"},
+		// just to prove that it affects the algorithm:
+		// {role: metal.RoleMachine, size: size1.ID},
+	} {
+		err := sharedDS.createEntity(sharedDS.machineTable(), &metal.Machine{
+			Base: metal.Base{
+				ID: "allocated-" + strconv.Itoa(i),
+			},
+			PartitionID: partitionID,
+			SizeID:      spec.size,
+			State: metal.MachineState{
+				Value: metal.AvailableState,
+			},
+			Allocation: &metal.MachineAllocation{
+				Project: projectID,
+				Role:    spec.role,
+			},
+			RackID: fiveRacks(i),
+		})
+		require.NoError(t, err)
+
+		err = sharedDS.createEntity(sharedDS.eventTable(), &metal.ProvisioningEventContainer{
+			Base: metal.Base{
+				ID: "allocated-" + strconv.Itoa(i),
+			},
+			Liveliness: metal.MachineLivelinessAlive,
+		})
+		require.NoError(t, err)
+	}
+
+	for range 100 {
+		machine, err := sharedDS.FindWaitingMachine(context.Background(), projectID, partitionID, size1, nil, metal.RoleMachine)
+		require.NoError(t, err)
+
+		newMachine := *machine
+		newMachine.PreAllocated = false
+		newMachine.Allocation = &metal.MachineAllocation{
+			Project: projectID,
+		}
+		newMachine.Allocation.Role = metal.RoleMachine
+		newMachine.SizeID = size1.ID
+
+		err = sharedDS.updateEntity(sharedDS.machineTable(), &newMachine, machine)
+		if err != nil {
+			t.Errorf("unable to update machine: %s", err)
+		}
+
+		t.Logf("machine %s allocated in %s", newMachine.ID, newMachine.RackID)
+	}
+
+	var ms metal.Machines
+	err := sharedDS.SearchMachines(&MachineSearchQuery{AllocationProject: &projectID, SizeID: &size1.ID, PartitionID: &partitionID}, &ms)
+	require.NoError(t, err)
+
+	require.Len(t, ms, 100)
+
+	machinesByRack := map[string]int{}
+	for _, m := range ms {
+		machinesByRack[m.RackID]++
+	}
+
+	for id, count := range machinesByRack {
+		assert.Equal(t, 100/5, count, "uneven machine distribution in %s", id)
+	}
+
+	fmt.Println(machinesByRack)
+}

cmd/metal-api/internal/metal/machine.go

@@ -334,8 +334,6 @@ func (ms Machines) WithSize(id string) Machines {
 	var res Machines
 
 	for _, m := range ms {
-		m := m
-
 		if m.SizeID != id {
 			continue
 		}
@@ -350,8 +348,6 @@ func (ms Machines) WithPartition(id string) Machines {
 	var res Machines
 
 	for _, m := range ms {
-		m := m
-
 		if m.PartitionID != id {
 			continue
 		}
@@ -362,6 +358,24 @@ func (ms Machines) WithPartition(id string) Machines {
 	return res
 }
 
+func (ms Machines) WithRole(role Role) Machines {
+	var res Machines
+
+	for _, m := range ms {
+		if m.Allocation == nil {
+			continue
+		}
+
+		if m.Allocation.Role != role {
+			continue
+		}
+
+		res = append(res, m)
+	}
+
+	return res
+}
+
 // MachineNetwork stores the Network details of the machine
 type MachineNetwork struct {
 	NetworkID           string   `rethinkdb:"networkid" json:"networkid"`

cmd/metal-api/internal/service/machine-service.go

@@ -1434,7 +1434,7 @@ func findWaitingMachine(ctx context.Context, ds *datastore.RethinkStore, allocat
 		return nil, fmt.Errorf("partition cannot be found: %w", err)
 	}
 
-	machine, err := ds.FindWaitingMachine(ctx, allocationSpec.ProjectID, partition.ID, *size, allocationSpec.PlacementTags)
+	machine, err := ds.FindWaitingMachine(ctx, allocationSpec.ProjectID, partition.ID, *size, allocationSpec.PlacementTags, allocationSpec.Role)
 	if err != nil {
 		return nil, err
 	}