Support scheduling by labels (#244)

internal/command/create/vm.go

@@ -22,6 +22,7 @@ var headless bool
 var username string
 var password string
 var resources map[string]string
+var labels map[string]string
 var restartPolicy string
 var startupScript string
 var hostDirsRaw []string
@@ -47,6 +48,8 @@ func newCreateVMCommand() *cobra.Command {
 		"SSH password to use when executing a startup script on the VM")
 	command.PersistentFlags().StringToStringVar(&resources, "resources", map[string]string{},
 		"resources to request for this VM")
+	command.PersistentFlags().StringToStringVar(&labels, "labels", map[string]string{},
+		"labels required by this VM")
 	command.PersistentFlags().StringVar(&restartPolicy, "restart-policy", string(v1.RestartPolicyNever),
 		fmt.Sprintf("restart policy for this VM: specify %q to never restart or %q "+
 			"to only restart when the VM fails", v1.RestartPolicyNever, v1.RestartPolicyOnFailure))
@@ -96,6 +99,7 @@ func runCreateVM(cmd *cobra.Command, args []string) error {
 		Headless:   headless,
 		Username:   username,
 		Password:   password,
+		Labels:     labels,
 		HostDirs:   hostDirs,
 	}
 

internal/command/worker/run.go

@@ -28,6 +28,7 @@ var bootstrapTokenRaw string
 var bootstrapTokenStdin bool
 var logFilePath string
 var stringToStringResources map[string]string
+var labels map[string]string
 var noPKI bool
 var defaultCPU uint64
 var defaultMemory uint64
@@ -51,6 +52,8 @@ func newRunCommand() *cobra.Command {
 		"optional path to a file where logs (up to 100 Mb) will be written.")
 	cmd.PersistentFlags().StringToStringVar(&stringToStringResources, "resources", map[string]string{},
 		"resources that this worker provides")
+	cmd.PersistentFlags().StringToStringVar(&labels, "labels", map[string]string{},
+		"labels that this worker supports")
 	cmd.PersistentFlags().BoolVar(&noPKI, "no-pki", false,
 		"do not use the host's root CA set and instead validate the Controller's presented "+
 			"certificate using a bootstrap token (or manually via fingerprint, "+
@@ -122,6 +125,7 @@ func runWorker(cmd *cobra.Command, args []string) (err error) {
 		controllerClient,
 		worker.WithName(name),
 		worker.WithResources(resources),
+		worker.WithLabels(labels),
 		worker.WithDefaultCPUAndMemory(defaultCPU, defaultMemory),
 		worker.WithLogger(logger),
 	)

internal/controller/api_workers.go

@@ -98,6 +98,7 @@ func (controller *Controller) createWorker(ctx *gin.Context) responder.Responder
 
 		dbWorker.LastSeen = worker.LastSeen
 		dbWorker.Resources = worker.Resources
+		dbWorker.Labels = worker.Labels
 		dbWorker.DefaultCPU = worker.DefaultCPU
 		dbWorker.DefaultMemory = worker.DefaultMemory
 

internal/controller/scheduler/scheduler.go

@@ -235,7 +235,8 @@ NextVM:
 
 			if worker.Offline(scheduler.workerOfflineTimeout) ||
 				worker.SchedulingPaused ||
-				!resourcesRemaining.CanFit(unscheduledVM.Resources) {
+				!resourcesRemaining.CanFit(unscheduledVM.Resources) ||
+				!worker.Labels.Contains(unscheduledVM.Labels) {
 				continue NextWorker
 			}
 

internal/tests/labels_test.go

@@ -0,0 +1,79 @@
+package tests_test
+
+import (
+	"context"
+	"github.com/cirruslabs/orchard/internal/tests/devcontroller"
+	"github.com/cirruslabs/orchard/internal/tests/wait"
+	v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
+	"github.com/stretchr/testify/require"
+	"testing"
+	"time"
+)
+
+func TestLabels(t *testing.T) {
+	ctx := context.Background()
+
+	// Create a development environment
+	devClient, _, _ := devcontroller.StartIntegrationTestEnvironmentWithAdditionalOpts(t,
+		false, nil,
+		true, nil,
+	)
+
+	// Create a worker that doesn't have any labels
+	_, err := devClient.Workers().Create(ctx, v1.Worker{
+		Meta: v1.Meta{
+			Name: "worker-without-labels",
+		},
+		Resources: map[string]uint64{
+			v1.ResourceTartVMs: 2,
+		},
+	})
+	require.NoError(t, err)
+
+	// Create a VM that requests a "role=test" label
+	vmName := "test-vm"
+
+	require.NoError(t, devClient.VMs().Create(ctx, &v1.VM{
+		Meta: v1.Meta{
+			Name: vmName,
+		},
+		Image:  "example.com/doesnt/matter:latest",
+		Labels: map[string]string{"role": "test"},
+		Status: v1.VMStatusPending,
+	}))
+
+	// Ensure that this VM doesn't get assigned in 30 seconds
+	require.False(t, wait.Wait(30*time.Second, func() bool {
+		vm, err := devClient.VMs().Get(context.Background(), vmName)
+		require.NoError(t, err)
+
+		t.Logf("Waiting for the VM %s to be assigned", vmName)
+
+		return vm.Worker != ""
+	}), "VM %s was not expected to be assigned to any worker, but was assigned to some worker", vmName)
+
+	// Now create one more worker that has the required labels
+	const workerWithLabelsName = "worker-with-labels"
+
+	_, err = devClient.Workers().Create(ctx, v1.Worker{
+		Meta: v1.Meta{
+			Name: workerWithLabelsName,
+		},
+		Resources: map[string]uint64{
+			v1.ResourceTartVMs: 2,
+		},
+		Labels: map[string]string{"role": "test"},
+	})
+	require.NoError(t, err)
+
+	// Wait for the VM to be assigned to the new worker
+	require.True(t, wait.Wait(30*time.Second, func() bool {
+		vm, err := devClient.VMs().Get(context.Background(), vmName)
+		require.NoError(t, err)
+
+		t.Logf("Waiting for the VM %s to be assigned to a worker", vmName)
+
+		return vm.Worker == workerWithLabelsName
+	}), "VM was %s expected to be assigned to the worker %q, but was assigned "+
+		"to another worker (or no worker)", vmName, workerWithLabelsName)
+}

internal/worker/option.go

@@ -19,6 +19,12 @@ func WithResources(resources v1.Resources) Option {
 	}
 }
 
+func WithLabels(labels v1.Labels) Option {
+	return func(worker *Worker) {
+		worker.labels = labels
+	}
+}
+
 func WithDefaultCPUAndMemory(defaultCPU uint64, defaultMemory uint64) Option {
 	return func(worker *Worker) {
 		worker.defaultCPU = defaultCPU

internal/worker/worker.go

@@ -32,6 +32,7 @@ type Worker struct {
 	client        *client.Client
 	pollTicker    *time.Ticker
 	resources     v1.Resources
+	labels        v1.Labels
 
 	defaultCPU    uint64
 	defaultMemory uint64
@@ -195,6 +196,7 @@ func (worker *Worker) registerWorker(ctx context.Context) error {
 			Name: worker.name,
 		},
 		Resources:     worker.resources,
+		Labels:        worker.labels,
 		LastSeen:      time.Now(),
 		MachineID:     platformUUID,
 		DefaultCPU:    worker.defaultCPU,

pkg/resource/v1/labels.go

@@ -0,0 +1,13 @@
+package v1
+
+type Labels map[string]string
+
+func (labels Labels) Contains(other Labels) bool {
+	for label, value := range other {
+		if labels[label] != value {
+			return false
+		}
+	}
+
+	return true
+}

pkg/resource/v1/labels_test.go

@@ -0,0 +1,45 @@
+package v1_test
+
+import (
+	v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
+	"github.com/stretchr/testify/require"
+	"testing"
+)
+
+func TestLabelsMatch(t *testing.T) {
+	// Two nil labels
+	a := v1.Labels(nil)
+	b := v1.Labels(nil)
+	require.True(t, a.Contains(b))
+	require.True(t, b.Contains(a))
+
+	// Two empty labels
+	a = map[string]string{}
+	b = map[string]string{}
+	require.True(t, a.Contains(b))
+	require.True(t, b.Contains(a))
+
+	// Two identical labels
+	a = map[string]string{"foo": "bar"}
+	b = map[string]string{"foo": "bar"}
+	require.True(t, a.Contains(b))
+	require.True(t, b.Contains(a))
+
+	// Supersets against nil labels
+	a = v1.Labels(nil)
+	b = map[string]string{"baz": "qux", "foo": "bar"}
+	require.False(t, a.Contains(b))
+	require.True(t, b.Contains(a))
+
+	// Superset against empty labels
+	a = map[string]string{}
+	b = map[string]string{"baz": "qux", "foo": "bar"}
+	require.False(t, a.Contains(b))
+	require.True(t, b.Contains(a))
+
+	// Superset against subset labels
+	a = map[string]string{"foo": "bar"}
+	b = map[string]string{"baz": "qux", "foo": "bar"}
+	require.False(t, a.Contains(b))
+	require.True(t, b.Contains(a))
+}

pkg/resource/v1/v1.go

@@ -63,6 +63,9 @@ type VM struct {
 	// Resources required by this VM.
 	Resources Resources `json:"resources,omitempty"`
 
+	// Labels required by this VM.
+	Labels Labels `json:"labels,omitempty"`
+
 	// HostDir is a list of host directories to be mounted to the VM.
 	HostDirs []HostDir `json:"hostDirs,omitempty"`
 

pkg/resource/v1/worker.go

@@ -14,6 +14,9 @@ type Worker struct {
 	// Resources available on this Worker.
 	Resources Resources `json:"resources,omitempty"`
 
+	// Labels that this Worker supports.
+	Labels Labels `json:"labels,omitempty"`
+
 	// DefaultCPU is the amount of CPUs to assign to a VM
 	// when it doesn't explicitly request a specific amount.
 	DefaultCPU uint64 `json:"defaultCPU,omitempty"`