diff --git a/client/driver/env/env.go b/client/driver/env/env.go new file mode 100644 index 000000000..d5ab3e9fc --- /dev/null +++ b/client/driver/env/env.go @@ -0,0 +1,265 @@ +package env + +import ( + "fmt" + "strconv" + "strings" + + hargs "github.com/hashicorp/nomad/helper/args" + "github.com/hashicorp/nomad/nomad/structs" +) + +// A set of environment variables that are exported by each driver. +const ( + // The path to the alloc directory that is shared across tasks within a task + // group. + AllocDir = "NOMAD_ALLOC_DIR" + + // The path to the tasks local directory where it can store data that is + // persisted to the alloc is removed. + TaskLocalDir = "NOMAD_TASK_DIR" + + // The tasks memory limit in MBs. + MemLimit = "NOMAD_MEMORY_LIMIT" + + // The tasks limit in MHz. + CpuLimit = "NOMAD_CPU_LIMIT" + + // The IP address for the task. + TaskIP = "NOMAD_IP" + + // Prefix for passing both dynamic and static port allocations to + // tasks. + // E.g. $NOMAD_PORT_1 or $NOMAD_PORT_http + PortPrefix = "NOMAD_PORT_" + + // Prefix for passing task meta data. + MetaPrefix = "NOMAD_META_" +) + +// The node values that can be interpreted. +const ( + nodeIdKey = "node.id" + nodeDcKey = "node.datacenter" + nodeNameKey = "node.name" + nodeClassKey = "node.class" + + // Prefixes used for lookups. + nodeAttributePrefix = "attr." + nodeMetaPrefix = "meta." +) + +// 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 + ports map[string]int + allocDir string + taskDir string + cpuLimit int + memLimit int + ip string + node *structs.Node + + // taskEnv is the variables that will be set in the tasks environment + taskEnv map[string]string + + // nodeValues is the values that are allowed for interprolation from the + // node. + nodeValues map[string]string +} + +func NewTaskEnvironment(node *structs.Node) *TaskEnvironment { + return &TaskEnvironment{node: node} +} + +// ParseAndReplace takes the user supplied args replaces any instance of an +// environment variable or nomad variable in the args with the actual value. +func (t *TaskEnvironment) ParseAndReplace(args []string) []string { + replaced := make([]string, len(args)) + for i, arg := range args { + replaced[i] = hargs.ReplaceEnv(arg, t.taskEnv, t.nodeValues) + } + + return replaced +} + +// ReplaceEnv takes an arg and replaces all occurences of environment variables +// and nomad variables. If the variable is found in the passed map it is +// replaced, otherwise the original string is returned. +func (t *TaskEnvironment) ReplaceEnv(arg string) string { + return hargs.ReplaceEnv(arg, t.taskEnv, t.nodeValues) +} + +// Build must be called after all the tasks environment values have been set. +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 ports + for label, port := range t.ports { + t.taskEnv[fmt.Sprintf("%s%s", PortPrefix, label)] = strconv.Itoa(port) + } + + // Build the directories + if t.allocDir != "" { + t.taskEnv[AllocDir] = t.allocDir + } + if t.taskDir != "" { + t.taskEnv[TaskLocalDir] = t.taskDir + } + + // Build the resource limits + if t.memLimit != 0 { + t.taskEnv[MemLimit] = strconv.Itoa(t.memLimit) + } + if t.cpuLimit != 0 { + t.taskEnv[CpuLimit] = strconv.Itoa(t.cpuLimit) + } + + // Build the IP + if t.ip != "" { + t.taskEnv[TaskIP] = t.ip + } + + // Build the node + if t.node != nil { + // Set up the node values. + t.nodeValues[nodeIdKey] = t.node.ID + t.nodeValues[nodeDcKey] = t.node.Datacenter + t.nodeValues[nodeNameKey] = t.node.Name + t.nodeValues[nodeClassKey] = t.node.NodeClass + + // Set up the attributes. + for k, v := range t.node.Attributes { + t.nodeValues[fmt.Sprintf("%s%s", nodeAttributePrefix, k)] = v + } + + // Set up the meta. + for k, v := range t.node.Meta { + t.nodeValues[fmt.Sprintf("%s%s", nodeMetaPrefix, k)] = v + } + } + + // Interpret the environment variables + interpreted := make(map[string]string, len(t.env)) + for k, v := range t.env { + interpreted[k] = hargs.ReplaceEnv(v, t.nodeValues, t.taskEnv) + } + + for k, v := range interpreted { + t.taskEnv[k] = v + } + + return t +} + +// EnvList returns a list of strings with NAME=value pairs. +func (t *TaskEnvironment) EnvList() []string { + env := []string{} + for k, v := range t.taskEnv { + env = append(env, fmt.Sprintf("%s=%s", k, v)) + } + + return env +} + +// EnvMap returns a copy of the tasks environment variables. +func (t *TaskEnvironment) EnvMap() map[string]string { + m := make(map[string]string, len(t.taskEnv)) + for k, v := range t.taskEnv { + m[k] = v + } + + return m +} + +// Builder methods to build the TaskEnvironment +func (t *TaskEnvironment) SetAllocDir(dir string) *TaskEnvironment { + t.allocDir = dir + return t +} + +func (t *TaskEnvironment) ClearAllocDir() *TaskEnvironment { + t.allocDir = "" + return t +} + +func (t *TaskEnvironment) SetTaskLocalDir(dir string) *TaskEnvironment { + t.taskDir = dir + return t +} + +func (t *TaskEnvironment) ClearTaskLocalDir() *TaskEnvironment { + t.taskDir = "" + return t +} + +func (t *TaskEnvironment) SetMemLimit(limit int) *TaskEnvironment { + t.memLimit = limit + return t +} + +func (t *TaskEnvironment) ClearMemLimit() *TaskEnvironment { + t.memLimit = 0 + return t +} + +func (t *TaskEnvironment) SetCpuLimit(limit int) *TaskEnvironment { + t.cpuLimit = limit + return t +} + +func (t *TaskEnvironment) ClearCpuLimit() *TaskEnvironment { + t.cpuLimit = 0 + return t +} + +func (t *TaskEnvironment) SetTaskIp(ip string) *TaskEnvironment { + t.ip = ip + return t +} + +func (t *TaskEnvironment) ClearTaskIp() *TaskEnvironment { + t.ip = "" + return t +} + +// Takes a map of port labels to their port value. +func (t *TaskEnvironment) SetPorts(ports map[string]int) *TaskEnvironment { + t.ports = ports + return t +} + +func (t *TaskEnvironment) ClearPorts() *TaskEnvironment { + t.ports = nil + return t +} + +// 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 + return t +} + +func (t *TaskEnvironment) ClearMeta() *TaskEnvironment { + t.meta = nil + return t +} + +func (t *TaskEnvironment) SetEnvvars(m map[string]string) *TaskEnvironment { + t.env = m + return t +} + +func (t *TaskEnvironment) ClearEnvvars() *TaskEnvironment { + t.env = nil + return t +} diff --git a/client/driver/env/env_test.go b/client/driver/env/env_test.go new file mode 100644 index 000000000..43a8eb65d --- /dev/null +++ b/client/driver/env/env_test.go @@ -0,0 +1,174 @@ +package env + +import ( + "fmt" + "reflect" + "sort" + "testing" + + "github.com/hashicorp/nomad/nomad/mock" +) + +const ( + // Node values that tests can rely on + metaKey = "instance" + metaVal = "t2-micro" + attrKey = "arch" + attrVal = "amd64" + nodeName = "test node" + nodeClass = "test class" + + // Environment variable values that tests can rely on + envOneKey = "NOMAD_IP" + envOneVal = "127.0.0.1" + envTwoKey = "NOMAD_PORT_WEB" + envTwoVal = ":80" +) + +func testTaskEnvironment() *TaskEnvironment { + n := mock.Node() + n.Attributes = map[string]string{ + attrKey: attrVal, + } + n.Meta = map[string]string{ + metaKey: metaVal, + } + n.Name = nodeName + n.NodeClass = nodeClass + + envVars := map[string]string{ + envOneKey: envOneVal, + envTwoKey: envTwoVal, + } + return NewTaskEnvironment(n).SetEnvvars(envVars).Build() +} + +func TestEnvironment_ParseAndReplace_Env(t *testing.T) { + env := testTaskEnvironment() + + input := []string{fmt.Sprintf(`"$%v"!`, envOneKey), fmt.Sprintf("$%s$%s", envOneKey, envTwoKey)} + act := env.ParseAndReplace(input) + exp := []string{fmt.Sprintf(`"%s"!`, envOneVal), fmt.Sprintf("%s%s", envOneVal, envTwoVal)} + + if !reflect.DeepEqual(act, exp) { + t.Fatalf("ParseAndReplace(%v) returned %#v; want %#v", input, act, exp) + } +} + +func TestEnvironment_ParseAndReplace_Meta(t *testing.T) { + input := []string{fmt.Sprintf("$%v%v", nodeMetaPrefix, metaKey)} + exp := []string{metaVal} + env := testTaskEnvironment() + act := env.ParseAndReplace(input) + + if !reflect.DeepEqual(act, exp) { + t.Fatalf("ParseAndReplace(%v) returned %#v; want %#v", input, act, exp) + } +} + +func TestEnvironment_ParseAndReplace_Attr(t *testing.T) { + input := []string{fmt.Sprintf("$%v%v", nodeAttributePrefix, attrKey)} + exp := []string{attrVal} + env := testTaskEnvironment() + act := env.ParseAndReplace(input) + + if !reflect.DeepEqual(act, exp) { + t.Fatalf("ParseAndReplace(%v) returned %#v; want %#v", input, act, exp) + } +} + +func TestEnvironment_ParseAndReplace_Node(t *testing.T) { + input := []string{fmt.Sprintf("$%v", nodeNameKey), fmt.Sprintf("$%v", nodeClassKey)} + exp := []string{nodeName, nodeClass} + env := testTaskEnvironment() + act := env.ParseAndReplace(input) + + if !reflect.DeepEqual(act, exp) { + t.Fatalf("ParseAndReplace(%v) returned %#v; want %#v", input, act, exp) + } +} + +func TestEnvironment_ParseAndReplace_Mixed(t *testing.T) { + input := []string{ + fmt.Sprintf("$%v$%v%v", nodeNameKey, nodeAttributePrefix, attrKey), + fmt.Sprintf("$%v$%v%v", nodeClassKey, nodeMetaPrefix, metaKey), + fmt.Sprintf("$%v$%v", envTwoKey, nodeClassKey), + } + exp := []string{ + fmt.Sprintf("%v%v", nodeName, attrVal), + fmt.Sprintf("%v%v", nodeClass, metaVal), + fmt.Sprintf("%v%v", envTwoVal, nodeClass), + } + env := testTaskEnvironment() + act := env.ParseAndReplace(input) + + if !reflect.DeepEqual(act, exp) { + t.Fatalf("ParseAndReplace(%v) returned %#v; want %#v", input, act, exp) + } +} + +func TestEnvironment_ReplaceEnv_Mixed(t *testing.T) { + input := fmt.Sprintf("$%v$%v%v", nodeNameKey, nodeAttributePrefix, attrKey) + exp := fmt.Sprintf("%v%v", nodeName, attrVal) + env := testTaskEnvironment() + act := env.ReplaceEnv(input) + + if act != exp { + t.Fatalf("ParseAndReplace(%v) returned %#v; want %#v", input, act, exp) + } +} + +func TestEnvironment_AsList(t *testing.T) { + n := mock.Node() + env := NewTaskEnvironment(n). + SetTaskIp("127.0.0.1").SetPorts(map[string]int{"http": 80}). + SetMeta(map[string]string{"foo": "baz"}).Build() + + act := env.EnvList() + exp := []string{"NOMAD_IP=127.0.0.1", "NOMAD_PORT_http=80", "NOMAD_META_FOO=baz"} + sort.Strings(act) + sort.Strings(exp) + if !reflect.DeepEqual(act, exp) { + t.Fatalf("env.List() returned %v; want %v", act, exp) + } +} + +func TestEnvironment_ClearEnvvars(t *testing.T) { + n := mock.Node() + env := NewTaskEnvironment(n). + SetTaskIp("127.0.0.1"). + SetEnvvars(map[string]string{"foo": "baz", "bar": "bang"}).Build() + + act := env.EnvList() + exp := []string{"NOMAD_IP=127.0.0.1", "bar=bang", "foo=baz"} + sort.Strings(act) + sort.Strings(exp) + if !reflect.DeepEqual(act, exp) { + t.Fatalf("env.List() returned %v; want %v", act, exp) + } + + // Clear the environent variables. + env.ClearEnvvars().Build() + + act = env.EnvList() + exp = []string{"NOMAD_IP=127.0.0.1"} + sort.Strings(act) + sort.Strings(exp) + if !reflect.DeepEqual(act, exp) { + t.Fatalf("env.List() returned %v; want %v", act, exp) + } +} + +func TestEnvironment_Interprolate(t *testing.T) { + env := testTaskEnvironment(). + SetEnvvars(map[string]string{"test": "$node.class", "test2": "$attr.arch"}). + Build() + + act := env.EnvList() + exp := []string{fmt.Sprintf("test=%s", nodeClass), fmt.Sprintf("test2=%s", attrVal)} + sort.Strings(act) + sort.Strings(exp) + if !reflect.DeepEqual(act, exp) { + t.Fatalf("env.List() returned %v; want %v", act, exp) + } +}