diff --git a/internal/command/create/vm.go b/internal/command/create/vm.go index 57affb7..7292b96 100644 --- a/internal/command/create/vm.go +++ b/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, } diff --git a/internal/command/worker/run.go b/internal/command/worker/run.go index bc1b59e..d5f30bb 100644 --- a/internal/command/worker/run.go +++ b/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), ) diff --git a/internal/controller/api_workers.go b/internal/controller/api_workers.go index f5d43c4..c30e746 100644 --- a/internal/controller/api_workers.go +++ b/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 diff --git a/internal/controller/scheduler/scheduler.go b/internal/controller/scheduler/scheduler.go index 5412a5b..93ba316 100644 --- a/internal/controller/scheduler/scheduler.go +++ b/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 } diff --git a/internal/tests/labels_test.go b/internal/tests/labels_test.go new file mode 100644 index 0000000..6762cda --- /dev/null +++ b/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) +} diff --git a/internal/worker/option.go b/internal/worker/option.go index 7e87056..aa0a170 100644 --- a/internal/worker/option.go +++ b/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 diff --git a/internal/worker/worker.go b/internal/worker/worker.go index 54db6e2..8dfde99 100644 --- a/internal/worker/worker.go +++ b/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, diff --git a/pkg/resource/v1/labels.go b/pkg/resource/v1/labels.go new file mode 100644 index 0000000..ce0effc --- /dev/null +++ b/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 +} diff --git a/pkg/resource/v1/labels_test.go b/pkg/resource/v1/labels_test.go new file mode 100644 index 0000000..d43bf3b --- /dev/null +++ b/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)) +} diff --git a/pkg/resource/v1/v1.go b/pkg/resource/v1/v1.go index bcdcede..fc31d6b 100644 --- a/pkg/resource/v1/v1.go +++ b/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"` diff --git a/pkg/resource/v1/worker.go b/pkg/resource/v1/worker.go index 02d4bff..b0a3b67 100644 --- a/pkg/resource/v1/worker.go +++ b/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"`