github.com/huiliang/nomad@v0.2.1-0.20151124023127-7a8b664699ff/client/driver/environment/vars.go (about)

     1  package environment
     2  
     3  import (
     4  	"fmt"
     5  	"strconv"
     6  	"strings"
     7  )
     8  
     9  // A set of environment variables that are exported by each driver.
    10  const (
    11  	// The path to the alloc directory that is shared across tasks within a task
    12  	// group.
    13  	AllocDir = "NOMAD_ALLOC_DIR"
    14  
    15  	// The path to the tasks local directory where it can store data that is
    16  	// persisted to the alloc is removed.
    17  	TaskLocalDir = "NOMAD_TASK_DIR"
    18  
    19  	// The tasks memory limit in MBs.
    20  	MemLimit = "NOMAD_MEMORY_LIMIT"
    21  
    22  	// The tasks limit in MHz.
    23  	CpuLimit = "NOMAD_CPU_LIMIT"
    24  
    25  	// The IP address for the task.
    26  	TaskIP = "NOMAD_IP"
    27  
    28  	// Prefix for passing both dynamic and static port allocations to
    29  	// tasks.
    30  	// E.g. $NOMAD_PORT_1 or $NOMAD_PORT_http
    31  	PortPrefix = "NOMAD_PORT_"
    32  
    33  	// Prefix for passing task meta data.
    34  	MetaPrefix = "NOMAD_META_"
    35  )
    36  
    37  var (
    38  	nomadVars = []string{AllocDir, TaskLocalDir, MemLimit, CpuLimit, TaskIP, PortPrefix, MetaPrefix}
    39  )
    40  
    41  type TaskEnvironment map[string]string
    42  
    43  func NewTaskEnivornment() TaskEnvironment {
    44  	return make(map[string]string)
    45  }
    46  
    47  // ParseFromList parses a list of strings with NAME=value pairs and returns a
    48  // TaskEnvironment.
    49  func ParseFromList(envVars []string) (TaskEnvironment, error) {
    50  	t := NewTaskEnivornment()
    51  
    52  	for _, pair := range envVars {
    53  		// Start the search from the second byte to skip a possible leading
    54  		// "=". Cmd.exe on Windows creates some special environment variables
    55  		// that start with an "=" and they can be properly retrieved by OS
    56  		// functions so we should handle them properly here.
    57  		idx := strings.Index(pair[1:], "=")
    58  		if idx == -1 {
    59  			return nil, fmt.Errorf("Couldn't parse environment variable: %v", pair)
    60  		}
    61  		idx++ // adjust for slice offset above
    62  		t[pair[:idx]] = pair[idx+1:]
    63  	}
    64  
    65  	return t, nil
    66  }
    67  
    68  // Returns a list of strings with NAME=value pairs.
    69  func (t TaskEnvironment) List() []string {
    70  	env := []string{}
    71  	for k, v := range t {
    72  		env = append(env, fmt.Sprintf("%s=%s", k, v))
    73  	}
    74  
    75  	return env
    76  }
    77  
    78  func (t TaskEnvironment) Map() map[string]string {
    79  	return t
    80  }
    81  
    82  func (t TaskEnvironment) SetAllocDir(dir string) {
    83  	t[AllocDir] = dir
    84  }
    85  
    86  func (t TaskEnvironment) ClearAllocDir() {
    87  	delete(t, AllocDir)
    88  }
    89  
    90  func (t TaskEnvironment) SetTaskLocalDir(dir string) {
    91  	t[TaskLocalDir] = dir
    92  }
    93  
    94  func (t TaskEnvironment) ClearTaskLocalDir() {
    95  	delete(t, TaskLocalDir)
    96  }
    97  
    98  func (t TaskEnvironment) SetMemLimit(limit int) {
    99  	t[MemLimit] = strconv.Itoa(limit)
   100  }
   101  
   102  func (t TaskEnvironment) ClearMemLimit() {
   103  	delete(t, MemLimit)
   104  }
   105  
   106  func (t TaskEnvironment) SetCpuLimit(limit int) {
   107  	t[CpuLimit] = strconv.Itoa(limit)
   108  }
   109  
   110  func (t TaskEnvironment) ClearCpuLimit() {
   111  	delete(t, CpuLimit)
   112  }
   113  
   114  func (t TaskEnvironment) SetTaskIp(ip string) {
   115  	t[TaskIP] = ip
   116  }
   117  
   118  func (t TaskEnvironment) ClearTaskIp() {
   119  	delete(t, TaskIP)
   120  }
   121  
   122  // Takes a map of port labels to their port value.
   123  func (t TaskEnvironment) SetPorts(ports map[string]int) {
   124  	for label, port := range ports {
   125  		t[fmt.Sprintf("%s%s", PortPrefix, label)] = strconv.Itoa(port)
   126  	}
   127  }
   128  
   129  func (t TaskEnvironment) ClearPorts() {
   130  	for k, _ := range t {
   131  		if strings.HasPrefix(k, PortPrefix) {
   132  			delete(t, k)
   133  		}
   134  	}
   135  }
   136  
   137  // Takes a map of meta values to be passed to the task. The keys are capatilized
   138  // when the environent variable is set.
   139  func (t TaskEnvironment) SetMeta(m map[string]string) {
   140  	for k, v := range m {
   141  		t[fmt.Sprintf("%s%s", MetaPrefix, strings.ToUpper(k))] = v
   142  	}
   143  }
   144  
   145  func (t TaskEnvironment) ClearMeta() {
   146  	for k, _ := range t {
   147  		if strings.HasPrefix(k, MetaPrefix) {
   148  			delete(t, k)
   149  		}
   150  	}
   151  }
   152  
   153  func (t TaskEnvironment) SetEnvvars(m map[string]string) {
   154  	for k, v := range m {
   155  		t[k] = v
   156  	}
   157  }
   158  
   159  func (t TaskEnvironment) ClearEnvvars() {
   160  OUTER:
   161  	for k, _ := range t {
   162  		for _, nomadPrefix := range nomadVars {
   163  			if strings.HasPrefix(k, nomadPrefix) {
   164  				continue OUTER
   165  			}
   166  		}
   167  		delete(t, k)
   168  	}
   169  }