github.com/ranjib/nomad@v0.1.1-0.20160225204057-97751b02f70b/client/driver/env/env.go

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"

	// Prefix for passing both dynamic and static port allocations to
	// tasks.
	// E.g. $NOMAD_IP_1=127.0.0.1:1 or $NOMAD_IP_http=127.0.0.1:80
	AddrPrefix = "NOMAD_ADDR_"

	// Prefix for passing the host port when a portmap is specified.
	HostPortPrefix = "NOMAD_HOST_PORT_"

	// Prefix for passing task meta data.
	MetaPrefix = "NOMAD_META_"
)

// The node values that can be interpreted.
const (
	nodeIdKey    = "node.unique.id"
	nodeDcKey    = "node.datacenter"
	nodeNameKey  = "node.unique.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
	AllocDir string
	TaskDir  string
	CpuLimit int
	MemLimit int
	Node     *structs.Node
	Networks []*structs.NetworkResource
	PortMap  map[string]int

	// 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 _, network := range t.Networks {
		for label, value := range network.MapLabelToValues(t.PortMap) {
			IPPort := fmt.Sprintf("%s:%d", network.IP, value)
			t.TaskEnv[fmt.Sprintf("%s%s", AddrPrefix, label)] = IPPort

			// Pass an explicit port mapping to the environment
			if port, ok := t.PortMap[label]; ok {
				t.TaskEnv[fmt.Sprintf("%s%s", HostPortPrefix, 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 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) SetNetworks(networks []*structs.NetworkResource) *TaskEnvironment {
	t.Networks = networks
	return t
}

func (t *TaskEnvironment) clearNetworks() *TaskEnvironment {
	t.Networks = nil
	return t
}

func (t *TaskEnvironment) SetPortMap(portMap map[string]int) *TaskEnvironment {
	t.PortMap = portMap
	return t
}

func (t *TaskEnvironment) clearPortMap() *TaskEnvironment {
	t.PortMap = 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
}

// Appends the given environment variables.
func (t *TaskEnvironment) AppendEnvvars(m map[string]string) *TaskEnvironment {
	if t.Env == nil {
		t.Env = make(map[string]string, len(m))
	}

	for k, v := range m {
		t.Env[k] = v
	}
	return t
}

func (t *TaskEnvironment) ClearEnvvars() *TaskEnvironment {
	t.Env = nil
	return t
}