github.com/ryanslade/nomad@v0.2.4-0.20160128061903-fc95782f2089/client/driver/env/env.go (about) 1 package env 2 3 import ( 4 "fmt" 5 "strconv" 6 "strings" 7 8 hargs "github.com/hashicorp/nomad/helper/args" 9 "github.com/hashicorp/nomad/nomad/structs" 10 ) 11 12 // A set of environment variables that are exported by each driver. 13 const ( 14 // The path to the alloc directory that is shared across tasks within a task 15 // group. 16 AllocDir = "NOMAD_ALLOC_DIR" 17 18 // The path to the tasks local directory where it can store data that is 19 // persisted to the alloc is removed. 20 TaskLocalDir = "NOMAD_TASK_DIR" 21 22 // The tasks memory limit in MBs. 23 MemLimit = "NOMAD_MEMORY_LIMIT" 24 25 // The tasks limit in MHz. 26 CpuLimit = "NOMAD_CPU_LIMIT" 27 28 // The IP address for the task. 29 TaskIP = "NOMAD_IP" 30 31 // Prefix for passing both dynamic and static port allocations to 32 // tasks. 33 // E.g. $NOMAD_PORT_1 or $NOMAD_PORT_http 34 PortPrefix = "NOMAD_PORT_" 35 36 // Prefix for passing task meta data. 37 MetaPrefix = "NOMAD_META_" 38 ) 39 40 // The node values that can be interpreted. 41 const ( 42 nodeIdKey = "node.unique.id" 43 nodeDcKey = "node.datacenter" 44 nodeNameKey = "node.unique.name" 45 nodeClassKey = "node.class" 46 47 // Prefixes used for lookups. 48 nodeAttributePrefix = "attr." 49 nodeMetaPrefix = "meta." 50 ) 51 52 // TaskEnvironment is used to expose information to a task via environment 53 // variables and provide interpolation of Nomad variables. 54 type TaskEnvironment struct { 55 env map[string]string 56 meta map[string]string 57 ports map[string]int 58 allocDir string 59 taskDir string 60 cpuLimit int 61 memLimit int 62 ip string 63 node *structs.Node 64 65 // taskEnv is the variables that will be set in the tasks environment 66 taskEnv map[string]string 67 68 // nodeValues is the values that are allowed for interprolation from the 69 // node. 70 nodeValues map[string]string 71 } 72 73 func NewTaskEnvironment(node *structs.Node) *TaskEnvironment { 74 return &TaskEnvironment{node: node} 75 } 76 77 // ParseAndReplace takes the user supplied args replaces any instance of an 78 // environment variable or nomad variable in the args with the actual value. 79 func (t *TaskEnvironment) ParseAndReplace(args []string) []string { 80 replaced := make([]string, len(args)) 81 for i, arg := range args { 82 replaced[i] = hargs.ReplaceEnv(arg, t.taskEnv, t.nodeValues) 83 } 84 85 return replaced 86 } 87 88 // ReplaceEnv takes an arg and replaces all occurences of environment variables 89 // and nomad variables. If the variable is found in the passed map it is 90 // replaced, otherwise the original string is returned. 91 func (t *TaskEnvironment) ReplaceEnv(arg string) string { 92 return hargs.ReplaceEnv(arg, t.taskEnv, t.nodeValues) 93 } 94 95 // Build must be called after all the tasks environment values have been set. 96 func (t *TaskEnvironment) Build() *TaskEnvironment { 97 t.nodeValues = make(map[string]string) 98 t.taskEnv = make(map[string]string) 99 100 // Build the task metadata 101 for k, v := range t.meta { 102 t.taskEnv[fmt.Sprintf("%s%s", MetaPrefix, strings.ToUpper(k))] = v 103 } 104 105 // Build the ports 106 for label, port := range t.ports { 107 t.taskEnv[fmt.Sprintf("%s%s", PortPrefix, label)] = strconv.Itoa(port) 108 } 109 110 // Build the directories 111 if t.allocDir != "" { 112 t.taskEnv[AllocDir] = t.allocDir 113 } 114 if t.taskDir != "" { 115 t.taskEnv[TaskLocalDir] = t.taskDir 116 } 117 118 // Build the resource limits 119 if t.memLimit != 0 { 120 t.taskEnv[MemLimit] = strconv.Itoa(t.memLimit) 121 } 122 if t.cpuLimit != 0 { 123 t.taskEnv[CpuLimit] = strconv.Itoa(t.cpuLimit) 124 } 125 126 // Build the IP 127 if t.ip != "" { 128 t.taskEnv[TaskIP] = t.ip 129 } 130 131 // Build the node 132 if t.node != nil { 133 // Set up the node values. 134 t.nodeValues[nodeIdKey] = t.node.ID 135 t.nodeValues[nodeDcKey] = t.node.Datacenter 136 t.nodeValues[nodeNameKey] = t.node.Name 137 t.nodeValues[nodeClassKey] = t.node.NodeClass 138 139 // Set up the attributes. 140 for k, v := range t.node.Attributes { 141 t.nodeValues[fmt.Sprintf("%s%s", nodeAttributePrefix, k)] = v 142 } 143 144 // Set up the meta. 145 for k, v := range t.node.Meta { 146 t.nodeValues[fmt.Sprintf("%s%s", nodeMetaPrefix, k)] = v 147 } 148 } 149 150 // Interpret the environment variables 151 interpreted := make(map[string]string, len(t.env)) 152 for k, v := range t.env { 153 interpreted[k] = hargs.ReplaceEnv(v, t.nodeValues, t.taskEnv) 154 } 155 156 for k, v := range interpreted { 157 t.taskEnv[k] = v 158 } 159 160 return t 161 } 162 163 // EnvList returns a list of strings with NAME=value pairs. 164 func (t *TaskEnvironment) EnvList() []string { 165 env := []string{} 166 for k, v := range t.taskEnv { 167 env = append(env, fmt.Sprintf("%s=%s", k, v)) 168 } 169 170 return env 171 } 172 173 // EnvMap returns a copy of the tasks environment variables. 174 func (t *TaskEnvironment) EnvMap() map[string]string { 175 m := make(map[string]string, len(t.taskEnv)) 176 for k, v := range t.taskEnv { 177 m[k] = v 178 } 179 180 return m 181 } 182 183 // Builder methods to build the TaskEnvironment 184 func (t *TaskEnvironment) SetAllocDir(dir string) *TaskEnvironment { 185 t.allocDir = dir 186 return t 187 } 188 189 func (t *TaskEnvironment) ClearAllocDir() *TaskEnvironment { 190 t.allocDir = "" 191 return t 192 } 193 194 func (t *TaskEnvironment) SetTaskLocalDir(dir string) *TaskEnvironment { 195 t.taskDir = dir 196 return t 197 } 198 199 func (t *TaskEnvironment) ClearTaskLocalDir() *TaskEnvironment { 200 t.taskDir = "" 201 return t 202 } 203 204 func (t *TaskEnvironment) SetMemLimit(limit int) *TaskEnvironment { 205 t.memLimit = limit 206 return t 207 } 208 209 func (t *TaskEnvironment) ClearMemLimit() *TaskEnvironment { 210 t.memLimit = 0 211 return t 212 } 213 214 func (t *TaskEnvironment) SetCpuLimit(limit int) *TaskEnvironment { 215 t.cpuLimit = limit 216 return t 217 } 218 219 func (t *TaskEnvironment) ClearCpuLimit() *TaskEnvironment { 220 t.cpuLimit = 0 221 return t 222 } 223 224 func (t *TaskEnvironment) SetTaskIp(ip string) *TaskEnvironment { 225 t.ip = ip 226 return t 227 } 228 229 func (t *TaskEnvironment) ClearTaskIp() *TaskEnvironment { 230 t.ip = "" 231 return t 232 } 233 234 // Takes a map of port labels to their port value. 235 func (t *TaskEnvironment) SetPorts(ports map[string]int) *TaskEnvironment { 236 t.ports = ports 237 return t 238 } 239 240 func (t *TaskEnvironment) ClearPorts() *TaskEnvironment { 241 t.ports = nil 242 return t 243 } 244 245 // Takes a map of meta values to be passed to the task. The keys are capatilized 246 // when the environent variable is set. 247 func (t *TaskEnvironment) SetMeta(m map[string]string) *TaskEnvironment { 248 t.meta = m 249 return t 250 } 251 252 func (t *TaskEnvironment) ClearMeta() *TaskEnvironment { 253 t.meta = nil 254 return t 255 } 256 257 func (t *TaskEnvironment) SetEnvvars(m map[string]string) *TaskEnvironment { 258 t.env = m 259 return t 260 } 261 262 // Appends the given environment variables. 263 func (t *TaskEnvironment) AppendEnvvars(m map[string]string) *TaskEnvironment { 264 if t.env == nil { 265 t.env = make(map[string]string, len(m)) 266 } 267 268 for k, v := range m { 269 t.env[k] = v 270 } 271 return t 272 } 273 274 func (t *TaskEnvironment) ClearEnvvars() *TaskEnvironment { 275 t.env = nil 276 return t 277 }