Merge pull request #985 from hashicorp/b-interpret-job-tg-meta

client: Add job and task group meta interpretation
hashicorp · Mar 25, 2016 · b253416 · b253416
2 parents aea4a2f + be6da2a
commit b253416
Show file tree

Hide file tree

Showing 4 changed files with 90 additions and 37 deletions.
diff --git a/client/driver/driver.go b/client/driver/driver.go
@@ -121,8 +121,11 @@ func NewExecContext(alloc *allocdir.AllocDir, allocID string) *ExecContext {
 func GetTaskEnv(allocDir *allocdir.AllocDir, node *structs.Node,
 	task *structs.Task, alloc *structs.Allocation) (*env.TaskEnvironment, error) {
 
+	tg := alloc.Job.LookupTaskGroup(alloc.TaskGroup)
 	env := env.NewTaskEnvironment(node).
-		SetMeta(task.Meta).
+		SetTaskMeta(task.Meta).
+		SetTaskGroupMeta(tg.Meta).
+		SetJobMeta(alloc.Job.Meta).
 		SetEnvvars(task.Env).
 		SetTaskName(task.Name)
 

diff --git a/client/driver/driver_test.go b/client/driver/driver_test.go
@@ -124,21 +124,25 @@ func TestDriver_GetTaskEnv(t *testing.T) {
 		t.Fatalf("GetTaskEnv() failed: %v", err)
 	}
 	exp := map[string]string{
-		"NOMAD_CPU_LIMIT":       "1000",
-		"NOMAD_MEMORY_LIMIT":    "500",
-		"NOMAD_ADDR_one":        "1.2.3.4:80",
-		"NOMAD_ADDR_two":        "1.2.3.4:443",
-		"NOMAD_ADDR_three":      "1.2.3.4:8080",
-		"NOMAD_ADDR_four":       "1.2.3.4:12345",
-		"NOMAD_ADDR_admin":      "1.2.3.4:8081",
-		"NOMAD_ADDR_web":        "1.2.3.4:8086",
-		"NOMAD_META_CHOCOLATE":  "cake",
-		"NOMAD_META_STRAWBERRY": "icecream",
-		"HELLO":                 "world",
-		"lorem":                 "ipsum",
-		"NOMAD_ALLOC_ID":        alloc.ID,
-		"NOMAD_ALLOC_NAME":      alloc.Name,
-		"NOMAD_TASK_NAME":       task.Name,
+		"NOMAD_CPU_LIMIT":               "1000",
+		"NOMAD_MEMORY_LIMIT":            "500",
+		"NOMAD_ADDR_one":                "1.2.3.4:80",
+		"NOMAD_ADDR_two":                "1.2.3.4:443",
+		"NOMAD_ADDR_three":              "1.2.3.4:8080",
+		"NOMAD_ADDR_four":               "1.2.3.4:12345",
+		"NOMAD_ADDR_admin":              "1.2.3.4:8081",
+		"NOMAD_ADDR_web":                "1.2.3.4:8086",
+		"NOMAD_META_CHOCOLATE":          "cake",
+		"NOMAD_META_STRAWBERRY":         "icecream",
+		"NOMAD_META_ELB_CHECK_INTERVAL": "30s",
+		"NOMAD_META_ELB_CHECK_TYPE":     "http",
+		"NOMAD_META_ELB_CHECK_MIN":      "3",
+		"NOMAD_META_OWNER":              "armon",
+		"HELLO":                         "world",
+		"lorem":                         "ipsum",
+		"NOMAD_ALLOC_ID":                alloc.ID,
+		"NOMAD_ALLOC_NAME":              alloc.Name,
+		"NOMAD_TASK_NAME":               task.Name,
 	}
 
 	act := env.EnvMap()

diff --git a/client/driver/env/env.go b/client/driver/env/env.go
@@ -67,19 +67,21 @@ const (
 // TaskEnvironment is used to expose information to a task via environment
 // variables and provide interpolation of Nomad variables.
 type TaskEnvironment struct {
-	Env        map[string]string
-	Meta       map[string]string
-	AllocDir   string
-	TaskDir    string
-	CpuLimit   int
-	MemLimit   int
-	TaskName   string
-	AllocIndex int
-	AllocId    string
-	AllocName  string
-	Node       *structs.Node
-	Networks   []*structs.NetworkResource
-	PortMap    map[string]int
+	Env           map[string]string
+	TaskMeta      map[string]string
+	TaskGroupMeta map[string]string
+	JobMeta       map[string]string
+	AllocDir      string
+	TaskDir       string
+	CpuLimit      int
+	MemLimit      int
+	TaskName      string
+	AllocIndex    int
+	AllocId       string
+	AllocName     string
+	Node          *structs.Node
+	Networks      []*structs.NetworkResource
+	PortMap       map[string]int
 
 	// taskEnv is the variables that will be set in the tasks environment
 	TaskEnv map[string]string
@@ -116,9 +118,11 @@ func (t *TaskEnvironment) Build() *TaskEnvironment {
 	t.NodeValues = make(map[string]string)
 	t.TaskEnv = make(map[string]string)
 
-	// Build the task metadata
-	for k, v := range t.Meta {
-		t.TaskEnv[fmt.Sprintf("%s%s", MetaPrefix, strings.ToUpper(k))] = v
+	// Build the meta with the following precedence: task, task group, job.
+	for _, meta := range []map[string]string{t.JobMeta, t.TaskGroupMeta, t.TaskMeta} {
+		for k, v := range meta {
+			t.TaskEnv[fmt.Sprintf("%s%s", MetaPrefix, strings.ToUpper(k))] = v
+		}
 	}
 
 	// Build the ports
@@ -279,13 +283,33 @@ func (t *TaskEnvironment) clearPortMap() *TaskEnvironment {
 
 // Takes a map of meta values to be passed to the task. The keys are capatilized
 // when the environent variable is set.
-func (t *TaskEnvironment) SetMeta(m map[string]string) *TaskEnvironment {
-	t.Meta = m
+func (t *TaskEnvironment) SetTaskMeta(m map[string]string) *TaskEnvironment {
+	t.TaskMeta = m
+	return t
+}
+
+func (t *TaskEnvironment) ClearTaskMeta() *TaskEnvironment {
+	t.TaskMeta = nil
+	return t
+}
+
+func (t *TaskEnvironment) SetTaskGroupMeta(m map[string]string) *TaskEnvironment {
+	t.TaskGroupMeta = m
+	return t
+}
+
+func (t *TaskEnvironment) ClearTaskGroupMeta() *TaskEnvironment {
+	t.TaskGroupMeta = nil
+	return t
+}
+
+func (t *TaskEnvironment) SetJobMeta(m map[string]string) *TaskEnvironment {
+	t.JobMeta = m
 	return t
 }
 
-func (t *TaskEnvironment) ClearMeta() *TaskEnvironment {
-	t.Meta = nil
+func (t *TaskEnvironment) ClearJobMeta() *TaskEnvironment {
+	t.JobMeta = nil
 	return t
 }
 

diff --git a/client/driver/env/env_test.go b/client/driver/env/env_test.go
@@ -140,14 +140,16 @@ func TestEnvironment_AsList(t *testing.T) {
 	env := NewTaskEnvironment(n).
 		SetNetworks(networks).
 		SetPortMap(portMap).
-		SetMeta(map[string]string{"foo": "baz"}).Build()
+		SetTaskGroupMeta(map[string]string{"foo": "bar", "baz": "bam"}).
+		SetTaskMeta(map[string]string{"foo": "baz"}).Build()
 
 	act := env.EnvList()
 	exp := []string{
 		"NOMAD_ADDR_http=127.0.0.1:80",
 		"NOMAD_ADDR_https=127.0.0.1:443",
 		"NOMAD_HOST_PORT_https=443",
 		"NOMAD_META_FOO=baz",
+		"NOMAD_META_BAZ=bam",
 	}
 	sort.Strings(act)
 	sort.Strings(exp)
@@ -225,3 +227,23 @@ func TestEnvironment_AppendHostEnvVars(t *testing.T) {
 		t.Fatalf("Didn't filter environment variable %q", skip)
 	}
 }
+
+func TestEnvironment_MetaPrecedence(t *testing.T) {
+	n := mock.Node()
+	env := NewTaskEnvironment(n).
+		SetJobMeta(map[string]string{"foo": "job", "bar": "job", "baz": "job"}).
+		SetTaskGroupMeta(map[string]string{"foo": "tg", "bar": "tg"}).
+		SetTaskMeta(map[string]string{"foo": "task"}).Build()
+
+	act := env.EnvList()
+	exp := []string{
+		"NOMAD_META_FOO=task",
+		"NOMAD_META_BAR=tg",
+		"NOMAD_META_BAZ=job",
+	}
+	sort.Strings(act)
+	sort.Strings(exp)
+	if !reflect.DeepEqual(act, exp) {
+		t.Fatalf("env.List() returned %v; want %v", act, exp)
+	}
+}