Merge pull request #2507 from hashicorp/f-task-env-vars

Add env vars for dc, region
hashicorp · May 4, 2017 · ec3aec5 · ec3aec5
2 parents 0728e14 + 3642434
commit ec3aec5
Show file tree

Hide file tree

Showing 7 changed files with 63 additions and 9 deletions.
diff --git a/client/alloc_runner_test.go b/client/alloc_runner_test.go
@@ -32,6 +32,7 @@ func (m *MockAllocStateUpdater) Update(alloc *structs.Allocation) {
 func testAllocRunnerFromAlloc(alloc *structs.Allocation, restarts bool) (*MockAllocStateUpdater, *AllocRunner) {
 	logger := testLogger()
 	conf := config.DefaultConfig()
+	conf.Node = mock.Node()
 	conf.StateDir = os.TempDir()
 	conf.AllocDir = os.TempDir()
 	upd := &MockAllocStateUpdater{}

diff --git a/client/driver/driver.go b/client/driver/driver.go
@@ -293,6 +293,8 @@ func GetTaskEnv(taskDir *allocdir.TaskDir, node *structs.Node,
 	env := env.NewTaskEnvironment(node).
 		SetTaskMeta(alloc.Job.CombinedTaskMeta(alloc.TaskGroup, task.Name)).
 		SetJobName(alloc.Job.Name).
+		SetDatacenterName(node.Datacenter).
+		SetRegionName(conf.Region).
 		SetEnvvars(task.Env).
 		SetTaskName(task.Name)
 

diff --git a/client/driver/driver_test.go b/client/driver/driver_test.go
@@ -74,6 +74,8 @@ func testConfig() *config.Config {
 	conf.StateDir = os.TempDir()
 	conf.AllocDir = os.TempDir()
 	conf.MaxKillTimeout = 10 * time.Second
+	conf.Region = "global"
+	conf.Node = mock.Node()
 	return conf
 }
 
@@ -88,6 +90,7 @@ type testContext struct {
 // It is up to the caller to call AllocDir.Destroy to cleanup.
 func testDriverContexts(t *testing.T, task *structs.Task) *testContext {
 	cfg := testConfig()
+	cfg.Node = mock.Node()
 	allocDir := allocdir.NewAllocDir(testLogger(), filepath.Join(cfg.AllocDir, structs.GenerateUUID()))
 	if err := allocDir.Build(); err != nil {
 		t.Fatalf("AllocDir.Build() failed: %v", err)
@@ -162,7 +165,7 @@ func setupTaskEnv(t *testing.T, driver string) (*allocdir.TaskDir, map[string]st
 	conf := testConfig()
 	allocDir := allocdir.NewAllocDir(testLogger(), filepath.Join(conf.AllocDir, alloc.ID))
 	taskDir := allocDir.NewTaskDir(task.Name)
-	env, err := GetTaskEnv(taskDir, nil, task, alloc, conf, "")
+	env, err := GetTaskEnv(taskDir, conf.Node, task, alloc, conf, "")
 	if err != nil {
 		t.Fatalf("GetTaskEnv() failed: %v", err)
 	}
@@ -209,6 +212,8 @@ func setupTaskEnv(t *testing.T, driver string) (*allocdir.TaskDir, map[string]st
 		"NOMAD_ALLOC_NAME":              alloc.Name,
 		"NOMAD_TASK_NAME":               task.Name,
 		"NOMAD_JOB_NAME":                alloc.Job.Name,
+		"NOMAD_DC":                      "dc1",
+		"NOMAD_REGION":                  "global",
 	}
 
 	act := env.EnvMap()

diff --git a/client/driver/env/env.go b/client/driver/env/env.go
@@ -48,6 +48,12 @@ const (
 	// AllocIndex is the environment variable for passing the allocation index.
 	AllocIndex = "NOMAD_ALLOC_INDEX"
 
+	// Datacenter is the environment variable for passing the datacenter in which the alloc is running.
+	Datacenter = "NOMAD_DC"
+
+	// Region is the environment variable for passing the region in which the alloc is running.
+	Region = "NOMAD_REGION"
+
 	// AddrPrefix is the prefix for passing both dynamic and static port
 	// allocations to tasks.
 	// E.g $NOMAD_ADDR_http=127.0.0.1:80
@@ -72,10 +78,11 @@ const (
 
 // The node values that can be interpreted.
 const (
-	nodeIdKey    = "node.unique.id"
-	nodeDcKey    = "node.datacenter"
-	nodeNameKey  = "node.unique.name"
-	nodeClassKey = "node.class"
+	nodeIdKey     = "node.unique.id"
+	nodeDcKey     = "node.datacenter"
+	nodeRegionKey = "node.region"
+	nodeNameKey   = "node.unique.name"
+	nodeClassKey  = "node.class"
 
 	// Prefixes used for lookups.
 	nodeAttributePrefix = "attr."
@@ -94,6 +101,8 @@ type TaskEnvironment struct {
 	MemLimit         int
 	TaskName         string
 	AllocIndex       int
+	Datacenter       string
+	Region           string
 	AllocId          string
 	AllocName        string
 	Node             *structs.Node
@@ -195,6 +204,12 @@ func (t *TaskEnvironment) Build() *TaskEnvironment {
 	if t.JobName != "" {
 		t.TaskEnv[JobName] = t.JobName
 	}
+	if t.Datacenter != "" {
+		t.TaskEnv[Datacenter] = t.Datacenter
+	}
+	if t.Region != "" {
+		t.TaskEnv[Region] = t.Region
+	}
 
 	// Build the addr of the other tasks
 	if t.Alloc != nil {
@@ -227,6 +242,7 @@ func (t *TaskEnvironment) Build() *TaskEnvironment {
 		// Set up the node values.
 		t.NodeValues[nodeIdKey] = t.Node.ID
 		t.NodeValues[nodeDcKey] = t.Node.Datacenter
+		t.NodeValues[nodeRegionKey] = t.Region
 		t.NodeValues[nodeNameKey] = t.Node.Name
 		t.NodeValues[nodeClassKey] = t.Node.NodeClass
 
@@ -488,13 +504,13 @@ func (t *TaskEnvironment) SetTaskName(name string) *TaskEnvironment {
 	return t
 }
 
-func (t *TaskEnvironment) SetJobName(name string) *TaskEnvironment {
-	t.JobName = name
+func (t *TaskEnvironment) ClearTaskName() *TaskEnvironment {
+	t.TaskName = ""
 	return t
 }
 
-func (t *TaskEnvironment) ClearTaskName() *TaskEnvironment {
-	t.TaskName = ""
+func (t *TaskEnvironment) SetJobName(name string) *TaskEnvironment {
+	t.JobName = name
 	return t
 }
 
@@ -503,6 +519,26 @@ func (t *TaskEnvironment) ClearJobName() *TaskEnvironment {
 	return t
 }
 
+func (t *TaskEnvironment) SetDatacenterName(name string) *TaskEnvironment {
+	t.Datacenter = name
+	return t
+}
+
+func (t *TaskEnvironment) ClearDatacenterName() *TaskEnvironment {
+	t.Datacenter = ""
+	return t
+}
+
+func (t *TaskEnvironment) SetRegionName(name string) *TaskEnvironment {
+	t.Region = name
+	return t
+}
+
+func (t *TaskEnvironment) ClearRegionName() *TaskEnvironment {
+	t.Region = ""
+	return t
+}
+
 func (t *TaskEnvironment) SetVaultToken(token string, inject bool) *TaskEnvironment {
 	t.VaultToken = token
 	t.InjectVaultToken = inject

diff --git a/client/task_runner_test.go b/client/task_runner_test.go
@@ -74,6 +74,7 @@ func testTaskRunner(t *testing.T, restarts bool) *taskRunnerTestCtx {
 func testTaskRunnerFromAlloc(t *testing.T, restarts bool, alloc *structs.Allocation) *taskRunnerTestCtx {
 	logger := testLogger()
 	conf := config.DefaultConfig()
+	conf.Node = mock.Node()
 	conf.StateDir = os.TempDir()
 	conf.AllocDir = os.TempDir()
 	upd := &MockTaskStateUpdater{}

diff --git a/command/agent/consul/int_test.go b/command/agent/consul/int_test.go
@@ -54,6 +54,7 @@ func TestConsul_Integration(t *testing.T) {
 	defer testconsul.Stop()
 
 	conf := config.DefaultConfig()
+	conf.Node = mock.Node()
 	conf.ConsulConfig.Addr = testconsul.HTTPAddr
 	consulConfig, err := conf.ConsulConfig.ApiConfig()
 	if err != nil {

diff --git a/website/source/docs/runtime/environment.html.md b/website/source/docs/runtime/environment.html.md
@@ -60,6 +60,14 @@ environment variables.
     <td>`NOMAD_JOB_NAME`</td>
     <td>The job's name</td>
   </tr>
+  <tr>
+    <td>`NOMAD_DC`</td>
+    <td>The datacenter in which the allocation is running</td>
+  </tr>
+  <tr>
+    <td>`NOMAD_REGION`</td>
+    <td>The region in which the allocation is running</td>
+  </tr>
   <tr>
     <td>`NOMAD_IP_<label>`</td>
     <td>The IP of the port with the given label</td>