mirror of
https://github.com/kemko/nomad.git
synced 2026-01-06 18:35:44 +03:00
Add folder
This commit is contained in:
Alex Dadgar
265
client/driver/env/env.go
vendored
Normal file
265
client/driver/env/env.go
vendored
Normal file
@@ -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
|
||||
}
|
||||
174
client/driver/env/env_test.go
vendored
Normal file
174
client/driver/env/env_test.go
vendored
Normal file
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user