Add envconsul-like support and refactor environment handling

api/tasks.go

@@ -336,6 +336,7 @@ type Template struct {
 	Perms        *string        `mapstructure:"perms"`
 	LeftDelim    *string        `mapstructure:"left_delimiter"`
 	RightDelim   *string        `mapstructure:"right_delimiter"`
+	Envvars      *bool          `mapstructure:"env"`
 }
 
 func (tmpl *Template) Canonicalize() {
@@ -373,6 +374,9 @@ func (tmpl *Template) Canonicalize() {
 	if tmpl.RightDelim == nil {
 		tmpl.RightDelim = helper.StringToPtr("}}")
 	}
+	if tmpl.Envvars == nil {
+		tmpl.Envvars = helper.BoolToPtr(false)
+	}
 }
 
 type Vault struct {

client/alloc_runner_test.go

@@ -135,13 +135,12 @@ func TestAllocRunner_RetryArtifact(t *testing.T) {
 }
 
 func TestAllocRunner_TerminalUpdate_Destroy(t *testing.T) {
-	ctestutil.ExecCompatible(t)
 	upd, ar := testAllocRunner(false)
 
 	// Ensure task takes some time
 	task := ar.alloc.Job.TaskGroups[0].Tasks[0]
-	task.Config["command"] = "/bin/sleep"
-	task.Config["args"] = []string{"10"}
+	task.Driver = "mock_driver"
+	task.Config["run_for"] = "10s"
 	go ar.Run()
 
 	testutil.WaitForResult(func() (bool, error) {
@@ -234,13 +233,12 @@ func TestAllocRunner_TerminalUpdate_Destroy(t *testing.T) {
 }
 
 func TestAllocRunner_Destroy(t *testing.T) {
-	ctestutil.ExecCompatible(t)
 	upd, ar := testAllocRunner(false)
 
 	// Ensure task takes some time
 	task := ar.alloc.Job.TaskGroups[0].Tasks[0]
-	task.Config["command"] = "/bin/sleep"
-	task.Config["args"] = []string{"10"}
+	task.Driver = "mock_driver"
+	task.Config["run_for"] = "10s"
 	go ar.Run()
 	start := time.Now()
 
@@ -269,7 +267,7 @@ func TestAllocRunner_Destroy(t *testing.T) {
 
 			return nil
 		}); err != nil {
-			return false, fmt.Errorf("state not destroyed")
+			return false, fmt.Errorf("state not destroyed: %v", err)
 		}
 
 		// Check the alloc directory was cleaned
@@ -290,13 +288,12 @@ func TestAllocRunner_Destroy(t *testing.T) {
 }
 
 func TestAllocRunner_Update(t *testing.T) {
-	ctestutil.ExecCompatible(t)
 	_, ar := testAllocRunner(false)
 
 	// Ensure task takes some time
 	task := ar.alloc.Job.TaskGroups[0].Tasks[0]
-	task.Config["command"] = "/bin/sleep"
-	task.Config["args"] = []string{"10"}
+	task.Driver = "mock_driver"
+	task.Config["run_for"] = "10s"
 	go ar.Run()
 	defer ar.Destroy()
 
@@ -612,6 +609,7 @@ func TestAllocRunner_RestoreOldState(t *testing.T) {
 
 	logger := testLogger()
 	conf := config.DefaultConfig()
+	conf.Node = mock.Node()
 	conf.StateDir = os.TempDir()
 	conf.AllocDir = os.TempDir()
 	tmp, err := ioutil.TempFile("", "state-db")

client/client.go

@@ -949,7 +949,7 @@ func (c *Client) setupDrivers() error {
 
 	var avail []string
 	var skipped []string
-	driverCtx := driver.NewDriverContext("", "", c.config, c.config.Node, c.logger, nil, nil)
+	driverCtx := driver.NewDriverContext("", "", c.config, c.config.Node, c.logger, nil)
 	for name := range driver.BuiltinDrivers {
 		// Skip fingerprinting drivers that are not in the whitelist if it is
 		// enabled.

client/consul_template.go

@@ -1,7 +1,10 @@
 package client
 
 import (
+	"bufio"
+	"bytes"
 	"fmt"
+	"io"
 	"math/rand"
 	"os"
 	"path/filepath"
@@ -76,7 +79,7 @@ type TaskTemplateManager struct {
 
 func NewTaskTemplateManager(hook TaskHooks, tmpls []*structs.Template,
 	config *config.Config, vaultToken, taskDir string,
-	taskEnv *env.TaskEnvironment) (*TaskTemplateManager, error) {
+	envBuilder *env.Builder) (*TaskTemplateManager, error) {
 
 	// Check pre-conditions
 	if hook == nil {
@@ -85,7 +88,7 @@ func NewTaskTemplateManager(hook TaskHooks, tmpls []*structs.Template,
 		return nil, fmt.Errorf("Invalid config given")
 	} else if taskDir == "" {
 		return nil, fmt.Errorf("Invalid task directory given")
-	} else if taskEnv == nil {
+	} else if envBuilder == nil {
 		return nil, fmt.Errorf("Invalid task environment given")
 	}
 
@@ -114,14 +117,14 @@ func NewTaskTemplateManager(hook TaskHooks, tmpls []*structs.Template,
 	}
 
 	// Build the consul-template runner
-	runner, lookup, err := templateRunner(tmpls, config, vaultToken, taskDir, taskEnv)
+	runner, lookup, err := templateRunner(tmpls, config, vaultToken, taskDir, envBuilder.Build())
 	if err != nil {
 		return nil, err
 	}
 	tm.runner = runner
 	tm.lookup = lookup
 
-	go tm.run()
+	go tm.run(envBuilder, taskDir)
 	return tm, nil
 }
 
@@ -144,7 +147,7 @@ func (tm *TaskTemplateManager) Stop() {
 }
 
 // run is the long lived loop that handles errors and templates being rendered
-func (tm *TaskTemplateManager) run() {
+func (tm *TaskTemplateManager) run(envBuilder *env.Builder, taskDir string) {
 	// Runner is nil if there is no templates
 	if tm.runner == nil {
 		// Unblock the start if there is nothing to do
@@ -192,6 +195,12 @@ WAIT:
 		}
 	}
 
+	for _, t := range tm.templates {
+		if err := loadTemplateEnv(envBuilder, taskDir, t); err != nil {
+			tm.hook.Kill("consul-template", err.Error(), true)
+			return
+		}
+	}
 	allRenderedTime = time.Now()
 	tm.hook.UnblockStart("consul-template")
 
@@ -243,6 +252,11 @@ WAIT:
 				}
 
 				for _, tmpl := range tmpls {
+					if err := loadTemplateEnv(envBuilder, taskDir, tmpl); err != nil {
+
+						tm.hook.Kill("consul-template", err.Error(), true)
+						return
+					}
 					switch tmpl.ChangeMode {
 					case structs.TemplateChangeModeSignal:
 						signals[tmpl.ChangeSignal] = struct{}{}
@@ -317,7 +331,7 @@ func (tm *TaskTemplateManager) allTemplatesNoop() bool {
 // lookup by destination to the template. If no templates are given, a nil
 // template runner and lookup is returned.
 func templateRunner(tmpls []*structs.Template, config *config.Config,
-	vaultToken, taskDir string, taskEnv *env.TaskEnvironment) (
+	vaultToken, taskDir string, taskEnv *env.TaskEnv) (
 	*manager.Runner, map[string][]*structs.Template, error) {
 
 	if len(tmpls) == 0 {
@@ -350,7 +364,7 @@ func templateRunner(tmpls []*structs.Template, config *config.Config,
 	}
 
 	// Set Nomad's environment variables
-	runner.Env = taskEnv.Build().EnvMapAll()
+	runner.Env = taskEnv.All()
 
 	// Build the lookup
 	idMap := runner.TemplateConfigMapping()
@@ -368,9 +382,7 @@ func templateRunner(tmpls []*structs.Template, config *config.Config,
 
 // parseTemplateConfigs converts the tasks templates into consul-templates
 func parseTemplateConfigs(tmpls []*structs.Template, taskDir string,
-	taskEnv *env.TaskEnvironment, allowAbs bool) (map[ctconf.TemplateConfig]*structs.Template, error) {
-	// Build the task environment
-	taskEnv.Build()
+	taskEnv *env.TaskEnv, allowAbs bool) (map[ctconf.TemplateConfig]*structs.Template, error) {
 
 	ctmpls := make(map[ctconf.TemplateConfig]*structs.Template, len(tmpls))
 	for _, tmpl := range tmpls {
@@ -492,3 +504,60 @@ func runnerConfig(config *config.Config, vaultToken string) (*ctconf.Config, err
 	conf.Finalize()
 	return conf, nil
 }
+
+// loadTemplateEnv loads task environment variables from templates.
+func loadTemplateEnv(builder *env.Builder, taskDir string, t *structs.Template) error {
+	if !t.Envvars {
+		return nil
+	}
+	f, err := os.Open(filepath.Join(taskDir, t.DestPath))
+	if err != nil {
+		return fmt.Errorf("error opening env template: %v", err)
+	}
+	defer f.Close()
+
+	// Parse environment fil
+	vars, err := parseEnvFile(f)
+	if err != nil {
+		return fmt.Errorf("error parsing env template %q: %v", t.DestPath, err)
+	}
+
+	// Set the environment variables
+	builder.SetTemplateEnv(vars)
+	return nil
+}
+
+// parseEnvFile and return a map of the environment variables suitable for
+// TaskEnvironment.AppendEnvvars or an error.
+//
+// See nomad/structs#Template.Envvars comment for format.
+func parseEnvFile(r io.Reader) (map[string]string, error) {
+	vars := make(map[string]string, 50)
+	lines := 0
+	scanner := bufio.NewScanner(r)
+	for scanner.Scan() {
+		lines++
+		buf := scanner.Bytes()
+		if len(buf) == 0 {
+			// Skip empty lines
+			continue
+		}
+		if buf[0] == '#' {
+			// Skip lines starting with a #
+			continue
+		}
+		n := bytes.IndexByte(buf, '=')
+		if n == -1 {
+			return nil, fmt.Errorf("line %d: no '=' sign: %q", lines, string(buf))
+		}
+		if len(buf) > n {
+			vars[string(buf[0:n])] = string(buf[n+1 : len(buf)])
+		} else {
+			vars[string(buf[0:n])] = ""
+		}
+	}
+	if err := scanner.Err(); err != nil {
+		return nil, err
+	}
+	return vars, nil
+}