Merge pull request #2223 from hashicorp/f-env-addr

Setting the Addrs of sibling tasks in the env
hashicorp · Jan 23, 2017 · af66ea3 · af66ea3
2 parents 49a8ce1 + bfb01d9
commit af66ea3
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 3 deletions.
diff --git a/client/driver/env/env.go b/client/driver/env/env.go
@@ -101,6 +101,7 @@ type TaskEnvironment struct {
 	VaultToken       string
 	InjectVaultToken bool
 	JobName          string
+	Alloc            *structs.Allocation
 
 	// taskEnv is the variables that will be set in the tasks environment
 	TaskEnv map[string]string
@@ -193,6 +194,28 @@ func (t *TaskEnvironment) Build() *TaskEnvironment {
 		t.TaskEnv[JobName] = t.JobName
 	}
 
+	// Build the addr of the other tasks
+	if t.Alloc != nil {
+		for taskName, resources := range t.Alloc.TaskResources {
+			if taskName == t.TaskName {
+				continue
+			}
+			for _, nw := range resources.Networks {
+				ports := make([]*structs.Port, 0, len(nw.ReservedPorts)+len(nw.DynamicPorts))
+				for _, port := range nw.ReservedPorts {
+					ports = append(ports, &port)
+				}
+				for _, port := range nw.DynamicPorts {
+					ports = append(ports, &port)
+				}
+				for _, p := range ports {
+					key := fmt.Sprintf("%s%s_%s", AddrPrefix, taskName, p.Label)
+					t.TaskEnv[key] = fmt.Sprintf("%s:%d", nw.IP, p.Value)
+				}
+			}
+		}
+	}
+
 	// Build the node
 	if t.Node != nil {
 		// Set up the node values.
@@ -393,6 +416,7 @@ func (t *TaskEnvironment) SetAlloc(alloc *structs.Allocation) *TaskEnvironment {
 	t.AllocId = alloc.ID
 	t.AllocName = alloc.Name
 	t.AllocIndex = alloc.Index()
+	t.Alloc = alloc
 	return t
 }
 

diff --git a/client/driver/env/env_test.go b/client/driver/env/env_test.go
@@ -137,10 +137,13 @@ func TestEnvironment_ReplaceEnv_Mixed(t *testing.T) {
 
 func TestEnvironment_AsList(t *testing.T) {
 	n := mock.Node()
+	a := mock.Alloc()
 	env := NewTaskEnvironment(n).
 		SetNetworks(networks).
 		SetPortMap(portMap).
-		SetTaskMeta(map[string]string{"foo": "baz"}).Build()
+		SetTaskMeta(map[string]string{"foo": "baz"}).
+		SetAlloc(a).
+		SetTaskName("taskA").Build()
 
 	act := env.EnvList()
 	exp := []string{
@@ -153,11 +156,16 @@ func TestEnvironment_AsList(t *testing.T) {
 		"NOMAD_HOST_PORT_http=80",
 		"NOMAD_HOST_PORT_https=8080",
 		"NOMAD_META_FOO=baz",
+		"NOMAD_ADDR_web_main=192.168.0.100:5000",
+		"NOMAD_ADDR_web_http=192.168.0.100:2000",
+		"NOMAD_TASK_NAME=taskA",
 	}
+	allocID := fmt.Sprintf("NOMAD_ALLOC_ID=%s", a.ID)
+	exp = append(exp, allocID)
 	sort.Strings(act)
 	sort.Strings(exp)
 	if !reflect.DeepEqual(act, exp) {
-		t.Fatalf("env.List() returned %v; want %v", act, exp)
+		t.Fatalf("env.List() returned %v;\n want %v", act, exp)
 	}
 }
 

diff --git a/nomad/mock/mock.go b/nomad/mock/mock.go
@@ -281,7 +281,7 @@ func Alloc() *structs.Allocation {
 						IP:            "192.168.0.100",
 						ReservedPorts: []structs.Port{{Label: "main", Value: 5000}},
 						MBits:         50,
-						DynamicPorts:  []structs.Port{{Label: "http"}},
+						DynamicPorts:  []structs.Port{{Label: "http", Value: 2000}},
 					},
 				},
 			},

diff --git a/website/source/docs/runtime/environment.html.md b/website/source/docs/runtime/environment.html.md
@@ -84,6 +84,10 @@ environment variables.
     <td>`VAULT_TOKEN`</td>
     <td>The task's Vault token. See [Vault Integration](/docs/vault-integration/index.html) for more details</td>
   </tr>
+  <tr>
+    <td>`NOMAD_ADDR_<task>_<label`></td>
+    <td>The allocated address, given as IP:Port, for each task, port label</td>
+  </tr>
 </table>
 
 ## Task Identifiers