exec pipeline (#462)

Signed-off-by: David Chung <david.chung@docker.com>

Author: David Chung

pkg/leader/etcd/v3/etcd_test.go

@@ -33,7 +33,7 @@ func TestWithRealEtcd(t *testing.T) {
 		}
 	}
 
-	defer etcd.StopContainer.Run(containerName)
+	defer etcd.StopContainer.Start(containerName)
 
 	t.Run("AmILeader", testAmILeader)
 }

pkg/store/etcd/v3/etcd_test.go

@@ -34,7 +34,7 @@ func TestWithRealEtcd(t *testing.T) {
 		}
 	}
 
-	defer etcd.StopContainer.Run(containerName)
+	defer etcd.StopContainer.Start(containerName)
 
 	t.Run("SaveLoad", testSaveLoad)
 }

pkg/util/etcd/v3/etcd.go

@@ -62,7 +62,7 @@ func LocalIP() string {
 var (
 
 	// RunContainer is a command that shells out to Docker to run the etcd server in a container
-	RunContainer exec.Command = `
+	RunContainer = exec.Command(`
 docker run --rm -d \
        -v /usr/share/ca-certificates/:/etc/ssl/certs \
        -p 4001:4001 \
@@ -78,18 +78,21 @@ docker run --rm -d \
        -listen-peer-urls http://0.0.0.0:2380 \
        -initial-cluster-token etcd-cluster-1 \
        -initial-cluster etcd0=http://{{ arg 1 }}:2380 \
-       -initial-cluster-state new`
+       -initial-cluster-state new
+`)
 
 	// StopContainer stops the etcd container
-	StopContainer exec.Command = `docker stop {{ arg 1 }}`
+	StopContainer = exec.Command(`docker stop {{ arg 1 }}`)
 
 	// LsMembers lists the members in the cluster
-	LsMembers exec.Command = `
+	LsMembers = exec.Command(`
 docker run --rm -e ETCDCTL_API=3 \
-       quay.io/coreos/etcd etcdctl --endpoints={{ arg 1 }}:2379 member list`
+       quay.io/coreos/etcd etcdctl --endpoints={{ arg 1 }}:2379 member list
+`)
 
 	// Get fetches a value via etcdctl
-	Get exec.Command = `
+	Get = exec.Command(`
 docker run --rm -e ETCDCTL_API=3 \
-       quay.io/coreos/etcd etcdctl --endpoints={{ arg 1 }}:2379 get --print-value-only {{ arg 2 }}`
+       quay.io/coreos/etcd etcdctl --endpoints={{ arg 1 }}:2379 get --print-value-only {{ arg 2 }}
+`)
 )

pkg/util/etcd/v3/etcd_test.go

@@ -90,6 +90,6 @@ func TestRunInContainer(t *testing.T) {
 	}
 
 	log.Infoln("Stopping etcd")
-	err = StopContainer.Run(containerName)
+	err = StopContainer.Start(containerName)
 	require.NoError(t, err)
 }

pkg/util/exec/exec.go

@@ -1,6 +1,7 @@
 package exec
 
 import (
+	"fmt"
 	"io"
 	"os"
 	"os/exec"
@@ -12,75 +13,41 @@ import (
 
 var log = logutil.New("module", "util/exec")
 
-// Command is the template which is rendered before it's executed
-type Command string
-
-// Output runs the command until completion and returns the results
-func (c Command) Output(args ...string) ([]byte, error) {
-	return c.builder().Output(args...)
-}
-
-// Start runs the command without blocking
-func (c Command) Start(args ...string) error {
-	return c.builder().Start(args...)
-}
-
-// Run does a Cmd.Run on the command
-func (c Command) Run(args ...string) error {
-	return c.builder().Run(args...)
-}
-
-// String returns the interpolated version of the command
-func (c Command) String(args ...string) (string, error) {
-	p, err := c.builder().generate(args...)
-	if err == nil {
-		return strings.Join(p, " "), nil
+// Command returns a fluent builder for running a command where the command string
+// can have template functions and arguments
+func Command(s string) *Builder {
+	return &Builder{
+		command: s,
+		funcs:   map[string]interface{}{},
+		args:    map[interface{}]interface{}{},
 	}
-	return string(c), err
-}
-
-// WithOptions adds the template options
-func (c Command) WithOptions(options template.Options) *Builder {
-	return c.builder().WithOptions(options)
-}
-
-// WithFunc adds a function that can be used in the template
-func (c Command) WithFunc(name string, function interface{}) *Builder {
-	return c.builder().WithFunc(name, function)
-}
-
-// WithContext sets the context for the template
-func (c Command) WithContext(context interface{}) *Builder {
-	return c.builder().WithContext(context)
-}
-
-// InheritEnvs determines whether the process should inherit the envs of the parent
-func (c Command) InheritEnvs(v bool) *Builder {
-	return c.builder().InheritEnvs(v)
-}
-
-// NewCommand creates an instance of the command builder to allow detailed configuration
-func NewCommand(s string) *Builder {
-	return Command(s).builder()
 }
 
 // Builder collects options until it's run
 type Builder struct {
-	command     Command
+	command     string
 	options     template.Options
 	inheritEnvs bool
 	envs        []string
 	funcs       map[string]interface{}
+	args        map[interface{}]interface{}
 	context     interface{}
 	rendered    string // rendered command string
 	cmd         *exec.Cmd
 }
 
-func (c Command) builder() *Builder {
-	return &Builder{
-		command: c,
-		funcs:   map[string]interface{}{},
+// WithArg sets the arg key, value pair that can be accessed via the 'arg' function
+func (b *Builder) WithArg(key string, value interface{}) *Builder {
+	b.args[key] = value
+	return b
+}
+
+// WithArgs adds the command line args array
+func (b *Builder) WithArgs(args ...interface{}) *Builder {
+	for i, arg := range args {
+		b.args[i+1] = arg
 	}
+	return b
 }
 
 // InheritEnvs determines whether the process should inherit the envs of the parent
@@ -113,24 +80,6 @@ func (b *Builder) WithContext(context interface{}) *Builder {
 	return b
 }
 
-// Output runs the command until completion and returns the results
-func (b *Builder) Output(args ...string) ([]byte, error) {
-	run, err := b.exec(args...)
-	if err != nil {
-		return nil, err
-	}
-	return run.Output()
-}
-
-// Start runs the command without blocking
-func (b *Builder) Start(args ...string) error {
-	run, err := b.exec(args...)
-	if err != nil {
-		return err
-	}
-	return run.Start()
-}
-
 // Step is something you do with the processes streams
 type Step func(stdin io.WriteCloser, stdout io.ReadCloser, stderr io.ReadCloser) error
 
@@ -154,10 +103,10 @@ func (t *Thenable) Then(then Step) *Thenable {
 
 // Done returns the final function
 func (t *Thenable) Done() Step {
-	steps := t.steps
+	all := t.steps
 	return func(stdin io.WriteCloser, stdout, stderr io.ReadCloser) error {
-		for _, step := range steps {
-			if err := step(stdin, stdout, stderr); err != nil {
+		for _, next := range all {
+			if err := next(stdin, stdout, stderr); err != nil {
 				return err
 			}
 		}
@@ -200,57 +149,89 @@ func MergeOutput(out io.Writer) Step {
 
 // StartWithStreams starts the the process and then calls the function which allows
 // the streams to be wired.  Calling the provided function blocks.
-func (b *Builder) StartWithStreams(f Step,
-	args ...string) error {
+func (b *Builder) StartWithStreams(f Step, args ...interface{}) error {
 
-	_, err := b.exec(args...)
-	if err != nil {
+	if err := b.Prepare(args...); err != nil {
 		return err
 	}
 
-	pOut, err := b.cmd.StdoutPipe()
-	if err != nil {
-		return err
-	}
-	pErr, err := b.cmd.StderrPipe()
-	if err != nil {
-		return err
+	run := func() error { return nil }
+	if f != nil {
+		pIn, err := b.cmd.StdinPipe()
+		if err != nil {
+			return err
+		}
+
+		pOut, err := b.cmd.StdoutPipe()
+		if err != nil {
+			return err
+		}
+
+		pErr, err := b.cmd.StderrPipe()
+		if err != nil {
+			return err
+		}
+
+		run = func() error {
+			return f(pIn, pOut, pErr)
+		}
 	}
-	pIn, err := b.cmd.StdinPipe()
-	if err != nil {
+
+	if err := b.cmd.Start(); err != nil {
 		return err
 	}
 
-	err = b.cmd.Start()
-	if err != nil {
+	return run()
+}
+
+// Start does a Cmd.Start on the command
+func (b *Builder) Start(args ...interface{}) error {
+	if err := b.Prepare(args...); err != nil {
 		return err
 	}
+	return b.StartWithStreams(nil, args...)
+}
 
-	return f(pIn, pOut, pErr)
+// Wait waits for the command to complete
+func (b *Builder) Wait() error {
+	return b.cmd.Wait()
 }
 
-// Run does a Cmd.Run on the command
-func (b *Builder) Run(args ...string) error {
-	run, err := b.exec(args...)
-	if err != nil {
-		return err
+// Output runs the command until completion and returns the results
+func (b *Builder) Output(args ...interface{}) ([]byte, error) {
+	if err := b.Prepare(args...); err != nil {
+		return nil, err
 	}
-	return run.Run()
+	return b.cmd.Output()
 }
 
-func (b *Builder) generate(args ...string) ([]string, error) {
+func (b *Builder) generate(args ...interface{}) ([]string, error) {
+	// also index the args by index
+	for i, v := range args {
+		b.args[i+1] = v
+	}
+
 	ct, err := template.NewTemplate("str://"+string(b.command), template.Options{})
 	if err != nil {
 		return nil, err
 	}
 	for k, v := range b.funcs {
 		ct.AddFunc(k, v)
 	}
-	ct.AddFunc("arg", func(i int) interface{} {
-		return args[i-1] // starts at 1
+	ct.AddFunc("arg", func(i interface{}) interface{} {
+		if i, is := i.(int); is {
+			if len(args) > i {
+				return args[i-1] // starts at 1
+			}
+		}
+		return b.args[i]
 	})
 	ct.AddFunc("argv", func() interface{} {
-		return args
+		argv := []string{}
+		for _, arg := range args {
+			argv = append(argv, fmt.Sprintf("%v", arg))
+		}
+		return argv
 	})
 	cmd, err := ct.Render(b.context)
 	if err != nil {
@@ -267,16 +248,39 @@ func (b *Builder) generate(args ...string) ([]string, error) {
 	}
 	return command, nil
 }
-func (b *Builder) exec(args ...string) (*exec.Cmd, error) {
+
+// Prepare generates the command based on the input args. This is the step before actual Start or Run
+func (b *Builder) Prepare(args ...interface{}) error {
 	command, err := b.generate(args...)
 	if err != nil {
-		return nil, err
+		return err
 	}
 	log.Debug("exec", "command", command)
 	b.cmd = exec.Command(command[0], command[1:]...)
 	if b.inheritEnvs {
 		b.cmd.Env = append(os.Environ(), b.envs...)
 	}
+	return nil
+}
+
+// Stdin takes the input from the writer
+func (b *Builder) Stdin(f func(w io.Writer) error) error {
+	input, err := b.cmd.StdinPipe()
+	if err != nil {
+		return err
+	}
+	defer input.Close()
+	return f(input)
+}
+
+// StdoutTo connects the stdout of this to the next stage
+func (b *Builder) StdoutTo(next *Builder) {
+	r, w := io.Pipe()
+	b.cmd.Stdout = w
+	next.cmd.Stdin = r
+}
 
-	return b.cmd, nil
+// Stdout sets the stdout
+func (b *Builder) Stdout(w io.Writer) {
+	b.cmd.Stdout = w
 }