From b63c11cb3b060379f210840f7ea8a51a3cebd533 Mon Sep 17 00:00:00 2001
From: Oleksii Kursov <okursov@roblox.com>
Date: Thu, 9 Jan 2020 13:43:13 -0800
Subject: [PATCH] Add default labels to all Nomad Drone Controller jobs on
 creation

---
 cmd/drone-server/config/config.go    | 19 ++++++++++---------
 cmd/drone-server/inject_scheduler.go |  1 +
 scheduler/nomad/config.go            |  1 +
 scheduler/nomad/nomad.go             |  8 ++++++++
 4 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/cmd/drone-server/config/config.go b/cmd/drone-server/config/config.go
index 1349edc6..9f42d0b7 100644
--- a/cmd/drone-server/config/config.go
+++ b/cmd/drone-server/config/config.go
@@ -143,15 +143,16 @@ type (
 
 	// Nomad configuration.
 	Nomad struct {
-		Enabled     bool     `envconfig:"DRONE_NOMAD_ENABLED"`
-		Datacenters []string `envconfig:"DRONE_NOMAD_DATACENTER" default:"dc1"`
-		Namespace   string   `envconfig:"DRONE_NOMAD_NAMESPACE"`
-		Region      string   `envconfig:"DRONE_NOMAD_REGION"`
-		Prefix      string   `envconfig:"DRONE_NOMAD_JOB_PREFIX" default:"drone-job-"`
-		Image       string   `envconfig:"DRONE_NOMAD_IMAGE"`
-		ImagePull   bool     `envconfig:"DRONE_NOMAD_IMAGE_PULL"`
-		Memory      int      `envconfig:"DRONE_NOMAD_DEFAULT_RAM" default:"1024"`
-		CPU         int      `envconfig:"DRONE_NOMAD_DEFAULT_CPU" default:"500"`
+		Enabled     bool                `envconfig:"DRONE_NOMAD_ENABLED"`
+		Datacenters []string            `envconfig:"DRONE_NOMAD_DATACENTER" default:"dc1"`
+		Namespace   string              `envconfig:"DRONE_NOMAD_NAMESPACE"`
+		Region      string              `envconfig:"DRONE_NOMAD_REGION"`
+		Prefix      string              `envconfig:"DRONE_NOMAD_JOB_PREFIX" default:"drone-job-"`
+		Image       string              `envconfig:"DRONE_NOMAD_IMAGE"`
+		ImagePull   bool                `envconfig:"DRONE_NOMAD_IMAGE_PULL"`
+		Memory      int                 `envconfig:"DRONE_NOMAD_DEFAULT_RAM" default:"1024"`
+		Labels      map[string]string   `envconfig:"DRONE_NOMAD_LABELS"`
+		CPU         int                 `envconfig:"DRONE_NOMAD_DEFAULT_CPU" default:"500"`
 	}
 
 	// License provides license configuration
diff --git a/cmd/drone-server/inject_scheduler.go b/cmd/drone-server/inject_scheduler.go
index b9c4a7fd..26ca6369 100644
--- a/cmd/drone-server/inject_scheduler.go
+++ b/cmd/drone-server/inject_scheduler.go
@@ -87,6 +87,7 @@ func provideNomadScheduler(config config.Config) core.Scheduler {
 	logrus.Info("main: nomad scheduler enabled")
 	sched, err := nomad.FromConfig(nomad.Config{
 		Datacenter:      config.Nomad.Datacenters,
+		Labels:          config.Nomad.Labels,
 		Namespace:       config.Nomad.Namespace,
 		Region:          config.Nomad.Region,
 		DockerImage:     config.Nomad.Image,
diff --git a/scheduler/nomad/config.go b/scheduler/nomad/config.go
index 15473683..4491c3b2 100644
--- a/scheduler/nomad/config.go
+++ b/scheduler/nomad/config.go
@@ -17,6 +17,7 @@ package nomad
 // Config is the configuration for the Nomad scheduler.
 type Config struct {
 	Datacenter       []string
+	Labels           map[string]string
 	Namespace        string
 	Region           string
 	DockerImage      string
diff --git a/scheduler/nomad/nomad.go b/scheduler/nomad/nomad.go
index 07732399..d2fcfe57 100644
--- a/scheduler/nomad/nomad.go
+++ b/scheduler/nomad/nomad.go
@@ -157,6 +157,14 @@ func (s *nomadScheduler) Schedule(ctx context.Context, stage *core.Stage) error
 		})
 	}
 
+	for k, v := range s.config.Labels {
+        job.Constraints = append(job.Constraints, &api.Constraint{
+            LTarget: fmt.Sprintf("${meta.%s}", k),
+            RTarget: v,
+            Operand: "=",
+        })
+    }
+
 	log := logrus.WithFields(logrus.Fields{
 		"stage-id":     stage.ID,
 		"stage-number": stage.Number,