Repository: puppetlabs/wash
Pull request #257 (2019-05-16 04:15:59)

Refactor Exec's plugin interface
This commit refactors Exec's plugin interface to be more plugin-author
friendly, and simpler to understand. It includes the following changes:

* Rename plugin.ExecResult => plugin.ExecCommand

plugin.ExecCommand better reflects what is being returned.

* Enforce a strict separation between what an Exec caller needs to care
about (plugin.ExecCommand) vs. what a plugin author needs to care about
(plugin.ExecCommandImpl) when implementing Exec

Exec callers should only care about a command's output channel and exit
code. They should not have to worry about context-cancellation cleanup
like ensuring that the command's stopped. That cleanup is the plugin
author's responsibility. On the other hand, it would be annoying for a
plugin author to have to remember to handle context-cancellation
cleanup.

The plugin.ExecCommand interface and the plugin.ExecCommandImpl object
solve both of these problems. Execable#Exec returns a plugin.ExecCommand
object, which only includes OutputCh and ExitCode. These are precisely the
things that Exec callers need to care about. Conversely, the
plugin.ExecCommandImpl implements the plugin.ExecCommand interface, but
also includes additional helpers that are meant to help plugin authors
write their code, including helpers that handle context cancellation
cleanup.

Finally, plugin.ExecCommandImpl replaces the previous ExitCodeCB via
helpers that set the exit-code/exit-code errors. This is more idiomatic
with how one thinks about what exec'ing a command would look like
("Stream the output" => "Get/Set the exit code").


Patch:
diff --git a/api/exec.go b/api/exec.go
index bda28d8..008e709 100644
--- a/api/exec.go
+++ b/api/exec.go
@@ -27,40 +27,6 @@ func sendPacket(ctx context.Context, w *json.Encoder, p *apitypes.ExecPacket) {
 	}
 }
 
-func streamOutput(ctx context.Context, w *json.Encoder, outputCh <-chan plugin.ExecOutputChunk) {
-	if outputCh == nil {
-		return
-	}
-
-	for chunk := range outputCh {
-		packet := apitypes.ExecPacket{TypeField: chunk.StreamID, Timestamp: chunk.Timestamp}
-		if err := chunk.Err; err != nil {
-			packet.Err = newStreamingErrorObj(chunk.StreamID, err.Error())
-		} else {
-			packet.Data = chunk.Data
-		}
-
-		sendPacket(ctx, w, &packet)
-	}
-}
-
-func streamExitCode(ctx context.Context, w *json.Encoder, exitCodeCB func() (int, error)) {
-	if exitCodeCB == nil {
-		return
-	}
-
-	packet := apitypes.ExecPacket{TypeField: apitypes.Exitcode, Timestamp: time.Now()}
-
-	exitCode, err := exitCodeCB()
-	if err != nil {
-		packet.Err = newUnknownErrorObj(fmt.Errorf("could not get the exit code: %v", err))
-	} else {
-		packet.Data = exitCode
-	}
-
-	sendPacket(ctx, w, &packet)
-}
-
 // swagger:route POST /fs/exec exec executeCommand
 //
 // Execute a command on a remote system
@@ -110,24 +76,36 @@ var execHandler handler = func(w http.ResponseWriter, r *http.Request) *errorRes
 	if body.Opts.Input != "" {
 		opts.Stdin = strings.NewReader(body.Opts.Input)
 	}
-	result, err := entry.(plugin.Execable).Exec(ctx, body.Cmd, body.Args, opts)
+	cmd, err := entry.(plugin.Execable).Exec(ctx, body.Cmd, body.Args, opts)
 	if err != nil {
 		return erroredActionResponse(path, plugin.ExecAction(), err.Error())
 	}
 
-	// Setup context cancellation handling
-	if result.CancelFunc != nil {
-		go func() {
-			<-ctx.Done()
-			result.CancelFunc()
-		}()
-	}
-
 	w.WriteHeader(http.StatusOK)
 	fw.Flush()
 
+	// Stream the command's output
 	enc := json.NewEncoder(&streamableResponseWriter{fw})
-	streamOutput(ctx, enc, result.OutputCh)
-	streamExitCode(ctx, enc, result.ExitCodeCB)
+	for chunk := range cmd.OutputCh() {
+		packet := apitypes.ExecPacket{TypeField: chunk.StreamID, Timestamp: chunk.Timestamp}
+		if err := chunk.Err; err != nil {
+			packet.Err = newStreamingErrorObj(chunk.StreamID, err.Error())
+		} else {
+			packet.Data = chunk.Data
+		}
+
+		sendPacket(ctx, enc, &packet)
+	}
+
+	// Now stream its exit code
+	packet := apitypes.ExecPacket{TypeField: apitypes.Exitcode, Timestamp: time.Now()}
+	exitCode, err := cmd.ExitCode()
+	if err != nil {
+		packet.Err = newUnknownErrorObj(fmt.Errorf("could not get the exit code: %v", err))
+	} else {
+		packet.Data = exitCode
+	}
+	sendPacket(ctx, enc, &packet)
+
 	return nil
 }
diff --git a/plugin/aws/ec2Instance.go b/plugin/aws/ec2Instance.go
index 64984ee..c0d78a0 100644
--- a/plugin/aws/ec2Instance.go
+++ b/plugin/aws/ec2Instance.go
@@ -182,10 +182,10 @@ func (inst *ec2Instance) checkLatestConsoleOutput(ctx context.Context) (*ec2Inst
 	return nil, fmt.Errorf("could not access the latest console log: %v", err)
 }
 
-func (inst *ec2Instance) Exec(ctx context.Context, cmd string, args []string, opts plugin.ExecOptions) (plugin.ExecResult, error) {
+func (inst *ec2Instance) Exec(ctx context.Context, cmd string, args []string, opts plugin.ExecOptions) (plugin.ExecCommand, error) {
 	meta, err := inst.Metadata(ctx)
 	if err != nil {
-		return plugin.ExecResult{}, err
+		return nil, err
 	}
 
 	// TODO: scrape default user and authorized keys from console output. Probably only works for Amazon AMIs.
@@ -197,7 +197,7 @@ func (inst *ec2Instance) Exec(ctx context.Context, cmd string, args []string, op
 	} else if ipaddr, ok := meta["PublicIpAddress"]; ok {
 		hostname = ipaddr.(string)
 	} else {
-		return plugin.ExecResult{}, fmt.Errorf("No public interface found for %v", inst)
+		return nil, fmt.Errorf("No public interface found for %v", inst)
 	}
 
 	// Use the default user for Amazon AMIs. See above for ideas on making this more general. Can be
diff --git a/plugin/docker/container.go b/plugin/docker/container.go
index b577c2c..e27ec3b 100644
--- a/plugin/docker/container.go
+++ b/plugin/docker/container.go
@@ -56,9 +56,7 @@ func (c *container) List(ctx context.Context) ([]plugin.Entry, error) {
 	return []plugin.Entry{cm, clf, vol.NewFS("fs", c)}, nil
 }
 
-func (c *container) Exec(ctx context.Context, cmd string, args []string, opts plugin.ExecOptions) (plugin.ExecResult, error) {
-	execResult := plugin.ExecResult{}
-
+func (c *container) Exec(ctx context.Context, cmd string, args []string, opts plugin.ExecOptions) (plugin.ExecCommand, error) {
 	command := append([]string{cmd}, args...)
 	activity.Record(ctx, "Exec %v on %v", command, c.Name())
 
@@ -68,29 +66,18 @@ func (c *container) Exec(ctx context.Context, cmd string, args []string, opts pl
 	}
 	created, err := c.client.ContainerExecCreate(ctx, c.id, cfg)
 	if err != nil {
-		return execResult, err
+		return nil, err
 	}
 
 	resp, err := c.client.ContainerExecAttach(ctx, created.ID, types.ExecStartCheck{})
 	if err != nil {
-		return execResult, err
+		return nil, err
 	}
 
-	// Asynchronously copy container exec output to an exec output channel.
-	outputCh, stdout, stderr := plugin.CreateExecOutputStreams(ctx)
-	go func() {
-		_, err := stdcopy.StdCopy(stdout, stderr, resp.Reader)
-		activity.Record(ctx, "Exec on %v complete: %v", c.Name(), err)
-		stdout.CloseWithError(err)
-		stderr.CloseWithError(err)
-		resp.Close()
-	}()
-
 	// If stdin is supplied, asynchronously copy it to container exec input.
-	var writeErr error
 	if opts.Stdin != nil {
 		go func() {
-			_, writeErr = io.Copy(resp.Conn, opts.Stdin)
+			_, writeErr := io.Copy(resp.Conn, opts.Stdin)
 			// If using a Tty, wait until done reading in case we need to send Ctrl-C; when a Tty is
 			// allocated commands expect user input and will respond to control signals. Otherwise close
 			// input now to ensure commands that depend on EOF execute correctly.
@@ -101,8 +88,8 @@ func (c *container) Exec(ctx context.Context, cmd string, args []string, opts pl
 		}()
 	}
 
-	execResult.OutputCh = outputCh
-	execResult.CancelFunc = func() {
+	execCmd := plugin.NewExecCommand(ctx)
+	execCmd.SetStopFunc(func() {
 		// Close the response on cancellation. Copying will block until there's more to read from the
 		// exec output. For an action with no more output it may never return.
 		if opts.Tty {
@@ -111,24 +98,27 @@ func (c *container) Exec(ctx context.Context, cmd string, args []string, opts pl
 			activity.Record(ctx, "Sent ETX on context termination: %v", err)
 		}
 		resp.Close()
-	}
-	execResult.ExitCodeCB = func() (int, error) {
-		if writeErr != nil {
-			return 0, err
-		}
+	})
+	// Asynchronously copy container exec output, then fetch the exit code once
+	// the copy's finished.
+	go func() {
+		_, err := stdcopy.StdCopy(execCmd.Stdout(), execCmd.Stderr(), resp.Reader)
+		activity.Record(ctx, "Exec on %v complete: %v", c.Name(), err)
+		execCmd.CloseStreamsWithError(err)
+		resp.Close()
 
+		// Command's finished. Now send the exit code.
 		resp, err := c.client.ContainerExecInspect(ctx, created.ID)
 		if err != nil {
-			return 0, err
+			execCmd.SetExitCodeErr(err)
+			return
 		}
-
 		if resp.Running {
-			return 0, fmt.Errorf("the command was marked as 'Running' even though the output streams reached EOF")
+			execCmd.SetExitCodeErr(fmt.Errorf("the command was marked as 'Running' even though the output streams reached EOF"))
+			return
 		}
-
 		activity.Record(ctx, "Exec on %v exited %v", c.Name(), resp.ExitCode)
-		return resp.ExitCode, nil
-	}
-
-	return execResult, nil
+		execCmd.SetExitCode(resp.ExitCode)
+	}()
+	return execCmd, nil
 }
diff --git a/plugin/execCommandImpl.go b/plugin/execCommandImpl.go
new file mode 100644
index 0000000..619decc
--- /dev/null
+++ b/plugin/execCommandImpl.go
@@ -0,0 +1,154 @@
+package plugin
+
+import (
+	"context"
+	"fmt"
+)
+
+/*
+ExecCommandImpl implements the plugin.ExecCommand interface.
+Use plugin.NewExecCommand to create instances of these objects.
+
+ExecCommandImpl provides Stdout/Stderr streams that you can wire-up
+to your plugin's API. These streams ensure that your API's stdout/stderr
+ordering is preserved. For example, if your API sends over Stderr, Stdout,
+and Stderr packets (in that order), then ExecCommand's OutputCh will
+stream them as Stderr, Stdout, and Stderr packets to your clients.
+
+Once the command's finished its execution, use SetExitCode to set the
+command's exit code. If your API fails to retrieve the command's exit
+code, then use SetExitCodeErr to set an appropriate error that ExitCode()
+can return.
+
+ExecCommandImpl also includes helpers that handle context-cancellation
+related cleanup for you. plugin.NewExecCommand ensures that the OutputCh
+is closed upon context-cancellation so that clients streaming your command's
+output are not blocked. SetStopFunc ensures that the exec'ing command is
+stopped upon context-cancellation, which prevents orphaned processes.
+
+See Container#Exec in the Docker plugin and ExecSSH in the transport
+package for examples of how ExecCommandImpl is used.
+*/
+type ExecCommandImpl struct {
+	ctx           context.Context
+	outputCh      chan ExecOutputChunk
+	stdout        *OutputStream
+	stderr        *OutputStream
+	exitCodeCh    chan int
+	exitCodeErrCh chan error
+}
+
+// NewExecCommand creates a new ExecCommandImpl object whose
+// lifetime is tied to the passed-in execution context. This
+// ctx should match what's passed into Execable#Exec. Note that
+// NewExecCommand ensures that the OutputCh is closed when ctx
+// is cancelled.
+func NewExecCommand(ctx context.Context) *ExecCommandImpl {
+	if ctx == nil {
+		panic("plugin.NewExecCommand called with a nil context")
+	}
+	
+	cmd := &ExecCommandImpl{
+		ctx: ctx,
+		exitCodeCh: make(chan int, 1),
+		exitCodeErrCh: make(chan error, 1),
+	}
+
+	// Create the output streams
+	cmd.outputCh = make(chan ExecOutputChunk)
+	closer := &multiCloser{ch: cmd.outputCh, countdown: 2}
+	cmd.stdout = &OutputStream{ctx: cmd.ctx, id: Stdout, ch: cmd.outputCh, closer: closer}
+	cmd.stderr = &OutputStream{ctx: cmd.ctx, id: Stderr, ch: cmd.outputCh, closer: closer}
+
+	// Ensure that the output streams are closed when the context
+	// is cancelled. This guarantees that callers won't be blocked
+	// when they are streaming our command's output.
+	go func() {
+		<-cmd.ctx.Done()
+		cmd.CloseStreamsWithError(ctx.Err())
+	}()
+
+	return cmd
+}
+
+// SetStopFunc sets the function that stops the running command. stopFunc
+// is called when cmd.ctx is cancelled.
+//
+// NOTE: SetStopFunc can be used to avoid orphaned processes.
+func (cmd *ExecCommandImpl) SetStopFunc(stopFunc func()) {
+	// Thankfully, goroutines are cheap. Otherwise, mixing this in with
+	// the goroutine in NewExecCommand heavily complicates things.
+	// For example, we'd have to worry about the possibility that the
+	// NewExecCommand goroutine is invoked before the Execable#Exec
+	// implementation can set a stopFunc, which can happen if the context
+	// is prematurely cancelled. That can result in an orphaned process
+	// in some plugin APIs, which is bad.
+	if stopFunc != nil {
+		go func() {
+			<-cmd.ctx.Done()
+			stopFunc()
+		}()
+	}
+}
+
+// Stdout returns the command's stdout stream. Attach this to your
+// plugin API's stdout stream.
+func (cmd *ExecCommandImpl) Stdout() *OutputStream {
+	return cmd.stdout
+}
+
+// Stderr returns the command's stderr stream. Attach this to your
+// plugin API's stderr stream.
+func (cmd *ExecCommandImpl) Stderr() *OutputStream {
+	return cmd.stderr
+}
+
+// CloseStreamsWithError closes the command's stdout/stderr streams
+// with the given error.
+func (cmd *ExecCommandImpl) CloseStreamsWithError(err error) {
+	cmd.Stdout().CloseWithError(err)
+	cmd.Stderr().CloseWithError(err)
+}
+
+// SetExitCode sets the command's exit code that will be returned by
+// cmd.ExitCode().
+func (cmd *ExecCommandImpl) SetExitCode(exitCode int) {
+	select {
+	case <-cmd.ctx.Done():
+		// Don't send anything if the context is cancelled.
+	default:
+		cmd.exitCodeCh <- exitCode
+	}
+}
+
+// SetExitCodeErr sets the error that will be returned by cmd.ExitCode().
+// Use it if your plugin API fails to get the command's exit code.
+func (cmd *ExecCommandImpl) SetExitCodeErr(err error) {
+	select {
+	case <-cmd.ctx.Done():
+		// Don't send anything if the context is cancelled.
+	default:
+		cmd.exitCodeErrCh <- err
+	}
+}
+
+// OutputCh implements ExecCommand#OutputCh
+func (cmd *ExecCommandImpl) OutputCh() <-chan ExecOutputChunk {
+	return cmd.outputCh
+}
+
+// ExitCode implements ExecCommand#ExitCode
+func (cmd *ExecCommandImpl) ExitCode() (int, error) {
+	select {
+	case <-cmd.ctx.Done():
+		return 0, fmt.Errorf("failed to fetch the command's exit code: %v", cmd.ctx.Err())
+	case exitCode := <-cmd.exitCodeCh:
+		return exitCode, nil
+	case err := <-cmd.exitCodeErrCh:
+		return 0, err
+	}
+}
+
+// sealed implements ExecCommand#sealed
+func (cmd *ExecCommandImpl) sealed() {
+}
\ No newline at end of file
diff --git a/plugin/externalPluginEntry.go b/plugin/externalPluginEntry.go
index 897d104..18555eb 100644
--- a/plugin/externalPluginEntry.go
+++ b/plugin/externalPluginEntry.go
@@ -299,9 +299,7 @@ func (e *externalPluginEntry) Stream(ctx context.Context) (io.ReadCloser, error)
 }
 
 // Exec executes a command on the given entry
-func (e *externalPluginEntry) Exec(ctx context.Context, cmd string, args []string, opts ExecOptions) (ExecResult, error) {
-	execResult := ExecResult{}
-
+func (e *externalPluginEntry) Exec(ctx context.Context, cmd string, args []string, opts ExecOptions) (ExecCommand, error) {
 	// TODO: Figure out how to pass-in opts when we have entries
 	// besides Stdin. Could do something like
 	//   <plugin_script> exec <path> <state> <opts> <cmd> <args...>
@@ -318,42 +316,27 @@ func (e *externalPluginEntry) Exec(ctx context.Context, cmd string, args []strin
 		cmdObj.Stdin = opts.Stdin
 	}
 
+	execCmd := NewExecCommand(ctx)
+
 	// Set-up the output streams
-	outputCh, stdout, stderr := CreateExecOutputStreams(ctx)
-	cmdObj.Stdout = stdout
-	cmdObj.Stderr = stderr
+	cmdObj.Stdout = execCmd.Stdout()
+	cmdObj.Stderr = execCmd.Stderr()
 
 	// Start the command
 	activity.Record(ctx, "Starting command: %v %v", cmdObj.Path, strings.Join(cmdObj.Args, " "))
 	if err := cmdObj.Start(); err != nil {
-		stdout.Close()
-		stderr.Close()
-		return execResult, err
+		return nil, err
 	}
 
 	// Wait for the command to finish
-	var exitCode int
-	var cmdWaitErr error
 	go func() {
 		ec, err := ExitCodeFromErr(cmdObj.Wait())
 		if err != nil {
-			cmdWaitErr = err
-		} else {
-			exitCode = ec
+			execCmd.CloseStreamsWithError(err)
+			return
 		}
-
-		stdout.CloseWithError(err)
-		stderr.CloseWithError(err)
+		execCmd.SetExitCode(ec)
 	}()
-
-	execResult.OutputCh = outputCh
-	execResult.ExitCodeCB = func() (int, error) {
-		if cmdWaitErr != nil {
-			return 0, cmdWaitErr
-		}
-
-		return exitCode, nil
-	}
-
-	return execResult, nil
+	
+	return execCmd, nil
 }
diff --git a/plugin/kubernetes/pod.go b/plugin/kubernetes/pod.go
index 452a02b..cbf7b13 100644
--- a/plugin/kubernetes/pod.go
+++ b/plugin/kubernetes/pod.go
@@ -119,7 +119,7 @@ func (p *pod) Stream(ctx context.Context) (io.ReadCloser, error) {
 	return req.Stream()
 }
 
-func (p *pod) Exec(ctx context.Context, cmd string, args []string, opts plugin.ExecOptions) (plugin.ExecResult, error) {
+func (p *pod) Exec(ctx context.Context, cmd string, args []string, opts plugin.ExecOptions) (plugin.ExecCommand, error) {
 	execRequest := p.client.CoreV1().RESTClient().Post().
 		Resource("pods").
 		Name(p.Name()).
@@ -137,15 +137,12 @@ func (p *pod) Exec(ctx context.Context, cmd string, args []string, opts plugin.E
 		execRequest = execRequest.Param("stdin", "true")
 	}
 
-	execResult := plugin.ExecResult{}
-
 	executor, err := remotecommand.NewSPDYExecutor(p.config, "POST", execRequest.URL())
 	if err != nil {
-		return execResult, errors.Wrap(err, "kubernetes.pod.Exec request")
+		return nil, errors.Wrap(err, "kubernetes.pod.Exec request")
 	}
 
-	outputCh, stdout, stderr := plugin.CreateExecOutputStreams(ctx)
-	exitcode := 0
+	execCmd := plugin.NewExecCommand(ctx)
 
 	// If using a Tty, create an input stream that allows us to send Ctrl-C to end execution;
 	// when a Tty is allocated commands expect user input and will respond to control signals.
@@ -158,33 +155,32 @@ func (p *pod) Exec(ctx context.Context, cmd string, args []string, opts plugin.E
 			stdin = r
 		}
 
-		execResult.CancelFunc = func() {
+		execCmd.SetStopFunc(func() {
 			// Close the response on context cancellation. Copying will block until there's more to
 			// read from the exec output. For an action with no more output it may never return.
 			// Append Ctrl-C to input to signal end of execution.
 			_, err := w.Write([]byte{0x03})
 			activity.Record(ctx, "Sent ETX on context termination: %v", err)
 			w.Close()
-		}
+		})
 	}
 
 	go func() {
-		streamOpts := remotecommand.StreamOptions{Stdout: stdout, Stderr: stderr, Stdin: stdin, Tty: opts.Tty}
+		streamOpts := remotecommand.StreamOptions{
+			Stdout: execCmd.Stdout(),
+			Stderr: execCmd.Stderr(),
+			Stdin: stdin,
+			Tty: opts.Tty,
+		}
 		err = executor.Stream(streamOpts)
 		activity.Record(ctx, "Exec on %v complete: %v", p.Name(), err)
 		if exerr, ok := err.(k8exec.ExitError); ok {
-			exitcode = exerr.ExitStatus()
+			execCmd.SetExitCode(exerr.ExitStatus())
 			err = nil
 		}
 
-		stdout.CloseWithError(err)
-		stderr.CloseWithError(err)
+		execCmd.CloseStreamsWithError(err)
 	}()
 
-	execResult.OutputCh = outputCh
-	execResult.ExitCodeCB = func() (int, error) {
-		return exitcode, nil
-	}
-
-	return execResult, nil
+	return execCmd, nil
 }
diff --git a/plugin/outputStream.go b/plugin/outputStream.go
index b3eb847..364b221 100644
--- a/plugin/outputStream.go
+++ b/plugin/outputStream.go
@@ -13,6 +13,8 @@ type OutputStream struct {
 	id         ExecPacketType
 	ch         chan ExecOutputChunk
 	closer     *multiCloser
+	mux        sync.Mutex
+	closed     bool
 }
 
 func (s *OutputStream) sendData(timestamp time.Time, data string) {
@@ -41,13 +43,19 @@ func (s *OutputStream) Write(data []byte) (int, error) {
 	return len(data), err
 }
 
-// Close ensures the channel is closed when the last OutputStream is closed.
-func (s *OutputStream) Close() {
-	s.closer.Close()
-}
-
-// CloseWithError sends the given error before calling Close()
+// CloseWithError sends the given error before closing the OutputStream.
+// It will noop if the OutputStream's already closed.
 func (s *OutputStream) CloseWithError(err error) {
+	// The lock's necessary because this can be called by multiple threads
+	// (e.g. the goroutine in NewExecCommand + ExecCommandImpl#CloseStreamsWithError)
+	s.mux.Lock()
+	defer s.mux.Unlock()
+	if s.closed {
+		return
+	}
+	defer func() {
+		s.closed = true
+	}()
 	if err != nil {
 		// Avoid re-sending ctx.Err() if it was already sent
 		// by OutputStream#Write
@@ -55,8 +63,7 @@ func (s *OutputStream) CloseWithError(err error) {
 			s.sendError(time.Now(), err)
 		}
 	}
-
-	s.Close()
+	s.closer.Close()
 }
 
 type multiCloser struct {
@@ -72,20 +79,4 @@ func (c *multiCloser) Close() {
 		close(c.ch)
 	}
 	c.mux.Unlock()
-}
-
-// CreateExecOutputStreams creates a pair of writers representing stdout
-// and stderr. They are used to transfer chunks of the Exec'ed cmd's
-// output in the order they're received by the corresponding API. The
-// writers maintain the ordering by writing to a channel.
-//
-// This method returns outputCh, stdout, and stderr, respectively.
-func CreateExecOutputStreams(ctx context.Context) (<-chan ExecOutputChunk, *OutputStream, *OutputStream) {
-	outputCh := make(chan ExecOutputChunk)
-	closer := &multiCloser{ch: outputCh, countdown: 2}
-
-	stdout := &OutputStream{ctx: ctx, id: Stdout, ch: outputCh, closer: closer}
-	stderr := &OutputStream{ctx: ctx, id: Stderr, ch: outputCh, closer: closer}
-
-	return outputCh, stdout, stderr
-}
+}
\ No newline at end of file
diff --git a/plugin/types.go b/plugin/types.go
index 22884e1..ce53808 100644
--- a/plugin/types.go
+++ b/plugin/types.go
@@ -104,25 +104,26 @@ type ExecOutputChunk struct {
 	Err       error
 }
 
-// ExecResult is a struct that contains the result of invoking Execable#exec.
-// Any of these fields can be nil.
+// ExecCommand represents a command that was invoked by a call to Exec.
+// It is a sealed interface, meaning you must use plugin.NewExecCommand
+// to create instances of these objects.
 //
-// OutputCh: contains timestamped chunks of the running command's stdout/stderr. This should be set
-//           to the channel that's returned by plugin.CreateExecOutputStreams(). OutputCh must be
-//           closed to signal that execution is complete.
-// ExitCodeCB: is called after execution completes to determine final status of execution.
-// CancelFunc: cancels a running command. It should noop for a finished command. CancelFunc is
-//             called when the execution context completes to perform necessary termination.
-type ExecResult struct {
-	OutputCh   <-chan ExecOutputChunk
-	ExitCodeCB func() (int, error)
-	CancelFunc func()
+// OutputCh returns a channel containing timestamped chunks of the command's
+// stdout/stderr.
+//
+// ExitCode returns the command's exit code. It will block until the command's
+// exit code is set, or until the execution context is cancelled. ExitCode will
+// return an error if it fails to fetch the command's exit code.
+type ExecCommand interface {
+	OutputCh() <-chan ExecOutputChunk
+	ExitCode() (int, error)
+	sealed()
 }
 
 // Execable is an entry that can have a command run on it.
 type Execable interface {
 	Entry
-	Exec(ctx context.Context, cmd string, args []string, opts ExecOptions) (ExecResult, error)
+	Exec(ctx context.Context, cmd string, args []string, opts ExecOptions) (ExecCommand, error)
 }
 
 // Streamable is an entry that returns a stream of updates.
diff --git a/transport/ssh.go b/transport/ssh.go
index 1593768..ef66405 100644
--- a/transport/ssh.go
+++ b/transport/ssh.go
@@ -97,19 +97,17 @@ type Identity struct {
 // Host *.compute.amazonaws.com
 //   StrictHostKeyChecking no
 // ```
-func ExecSSH(ctx context.Context, id Identity, cmd []string, opts plugin.ExecOptions) (plugin.ExecResult, error) {
-	var result plugin.ExecResult
-
+func ExecSSH(ctx context.Context, id Identity, cmd []string, opts plugin.ExecOptions) (plugin.ExecCommand, error) {
 	// find port, username, etc from .ssh/config
 	port, err := ssh_config.GetStrict(id.Host, "Port")
 	if err != nil {
-		return result, err
+		return nil, err
 	}
 
 	user := id.User
 	if user == "" {
 		if user, err = ssh_config.GetStrict(id.Host, "User"); err != nil {
-			return result, err
+			return nil, err
 		}
 	}
 
@@ -119,30 +117,30 @@ func ExecSSH(ctx context.Context, id Identity, cmd []string, opts plugin.ExecOpt
 
 	strictHostKeyChecking, err := ssh_config.GetStrict(id.Host, "StrictHostKeyChecking")
 	if err != nil {
-		return result, err
+		return nil, err
 	}
 
 	connection, err := sshConnect(id.Host, port, user, strictHostKeyChecking != "no")
 	if err != nil {
-		return result, fmt.Errorf("Failed to connect: %s", err)
+		return nil, fmt.Errorf("Failed to connect: %s", err)
 	}
 
 	// Run command via session
 	session, err := connection.NewSession()
 	if err != nil {
-		return result, fmt.Errorf("Failed to create session: %s", err)
+		return nil, fmt.Errorf("Failed to create session: %s", err)
 	}
 
 	if opts.Tty {
 		// sshd only processes signal codes if a TTY has been allocated. So set one up when requested.
 		modes := ssh.TerminalModes{ssh.ECHO: 0, ssh.TTY_OP_ISPEED: 14400, ssh.TTY_OP_OSPEED: 14400}
 		if err := session.RequestPty("xterm", 40, 80, modes); err != nil {
-			return result, fmt.Errorf("Unable to setup a TTY: %v", err)
+			return nil, fmt.Errorf("Unable to setup a TTY: %v", err)
 		}
 	}
 
-	outputch, stdout, stderr := plugin.CreateExecOutputStreams(ctx)
-	session.Stdin, session.Stdout, session.Stderr = opts.Stdin, stdout, stderr
+	execCmd := plugin.NewExecCommand(ctx)
+	session.Stdin, session.Stdout, session.Stderr = opts.Stdin, execCmd.Stdout(), execCmd.Stderr()
 
 	if opts.Elevate {
 		cmd = append([]string{"sudo"}, cmd...)
@@ -150,22 +148,9 @@ func ExecSSH(ctx context.Context, id Identity, cmd []string, opts plugin.ExecOpt
 
 	cmdStr := shellquote.Join(cmd...)
 	if err := session.Start(cmdStr); err != nil {
-		return result, err
+		return nil, err
 	}
-
-	// Wait for session to complete and stash result. This is done asynchronously because the
-	// exitcode callback won't be called until stdout/stderr are closed, but they're not closed by
-	// session.Start.
-	var exitErr error
-	go func() {
-		exitErr = session.Wait()
-		activity.Record(ctx, "Closing session for %v: %v", id.Host, session.Close())
-		stdout.Close()
-		stderr.Close()
-	}()
-
-	result.OutputCh = outputch
-	result.CancelFunc = func() {
+	execCmd.SetStopFunc(func() {
 		// Close the session on context cancellation. Copying will block until there's more to read
 		// from the exec output. For an action with no more output it may never return.
 		// If a TTY is setup and the session is still open, send Ctrl-C over before closing it.
@@ -173,14 +158,20 @@ func ExecSSH(ctx context.Context, id Identity, cmd []string, opts plugin.ExecOpt
 			activity.Record(ctx, "Sent SIGINT on context termination: %v", session.Signal(ssh.SIGINT))
 		}
 		activity.Record(ctx, "Closing session on context termination for %v: %v", id.Host, session.Close())
-	}
-	result.ExitCodeCB = func() (int, error) {
-		if exitErr == nil {
-			return 0, nil
-		} else if err, ok := exitErr.(*ssh.ExitError); ok {
-			return err.ExitStatus(), nil
+	})
+
+	// Wait for session to complete and stash result.
+	go func() {
+		err := session.Wait()
+		activity.Record(ctx, "Closing session for %v: %v", id.Host, session.Close())
+		execCmd.CloseStreamsWithError(nil)
+		if err == nil {
+			execCmd.SetExitCode(0)
+		} else if exitErr, ok := err.(*ssh.ExitError); ok {
+			execCmd.SetExitCode(exitErr.ExitStatus())
+		} else {
+			execCmd.SetExitCodeErr(err)
 		}
-		return 0, exitErr
-	}
-	return result, nil
+	}()
+	return execCmd, nil
 }
diff --git a/volume/fs.go b/volume/fs.go
index a36cbed..28958c0 100644
--- a/volume/fs.go
+++ b/volume/fs.go
@@ -34,14 +34,14 @@ func (d *FS) List(ctx context.Context) ([]plugin.Entry, error) {
 var errNonZero = fmt.Errorf("Exec exited non-zero")
 
 func exec(ctx context.Context, executor plugin.Execable, cmdline []string) (*bytes.Buffer, error) {
-	result, err := executor.Exec(ctx, cmdline[0], cmdline[1:], plugin.ExecOptions{Elevate: true})
+	cmd, err := executor.Exec(ctx, cmdline[0], cmdline[1:], plugin.ExecOptions{Elevate: true})
 	if err != nil {
 		return nil, err
 	}
 
 	var buf bytes.Buffer
 	var errs []error
-	for chunk := range result.OutputCh {
+	for chunk := range cmd.OutputCh() {
 		if chunk.Err != nil {
 			errs = append(errs, chunk.Err)
 		} else {
@@ -56,7 +56,7 @@ func exec(ctx context.Context, executor plugin.Execable, cmdline []string) (*byt
 		return nil, fmt.Errorf("exec errored: %v", errs)
 	}
 
-	exitcode, err := result.ExitCodeCB()
+	exitcode, err := cmd.ExitCode()
 	if err != nil {
 		return nil, err
 	} else if exitcode != 0 {
@@ -97,25 +97,17 @@ func (d *FS) VolumeOpen(ctx context.Context, path string) (plugin.SizedReader, e
 func (d *FS) VolumeStream(ctx context.Context, path string) (io.ReadCloser, error) {
 	activity.Record(ctx, "Streaming %v on %v", path, plugin.ID(d.executor))
 	execOpts := plugin.ExecOptions{Elevate: true, Tty: true}
-	result, err := d.executor.Exec(ctx, "tail", []string{"-f", path}, execOpts)
+	cmd, err := d.executor.Exec(ctx, "tail", []string{"-f", path}, execOpts)
 	if err != nil {
 		activity.Record(ctx, "Exec error in VolumeRead: %v", err)
 		return nil, err
 	}
 
-	// Setup context cancellation handling
-	if result.CancelFunc != nil {
-		go func() {
-			<-ctx.Done()
-			result.CancelFunc()
-		}()
-	}
-
 	r, w := io.Pipe()
 	go func() {
 		// Exec uses context; if it's canceled, the OutputCh will close. Close the writer.
 		var errs []error
-		for chunk := range result.OutputCh {
+		for chunk := range cmd.OutputCh() {
 			if chunk.Err != nil {
 				activity.Record(ctx, "Error on exec: %v", chunk.Err)
 				errs = append(errs, chunk.Err)


Test patch:
diff --git a/plugin/execCommandImpl_test.go b/plugin/execCommandImpl_test.go
new file mode 100644
index 0000000..f5dc9f3
--- /dev/null
+++ b/plugin/execCommandImpl_test.go
@@ -0,0 +1,182 @@
+package plugin
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/suite"
+)
+
+type ExecCommandImplTestSuite struct {
+	suite.Suite
+	ctx         context.Context
+	cancelFunc  context.CancelFunc
+}
+
+func (suite *ExecCommandImplTestSuite) SetupTest() {
+	suite.ctx, suite.cancelFunc = context.WithCancel(context.Background())
+}
+
+func (suite *ExecCommandImplTestSuite) TearDownTest() {
+	// Ensure that the context is cancelled to avoid dangling
+	// goroutines
+	suite.cancelFunc()
+}
+
+func (suite *ExecCommandImplTestSuite) NewExecCommand() *ExecCommandImpl {
+	return NewExecCommand(suite.ctx)
+}
+
+func (suite *ExecCommandImplTestSuite) TestNewExecCommand_CreatesOutputStreams() {
+	execCmd := suite.NewExecCommand()
+
+	// Our simulated command alternates writing to stdout + stderr
+	expectedChunksCh := make(chan ExecOutputChunk, 1)
+	go func() {
+		// stdout + stderr will be closed before expectedChunksCh,
+		// which means that outputCh should be closed before
+		// expectedChunksCh.
+		defer close(expectedChunksCh)
+		defer execCmd.CloseStreamsWithError(nil)
+
+		writeTo := func(streamName string, stream *OutputStream, data string) {
+			expectedChunksCh <- ExecOutputChunk{StreamID: stream.id, Data: data}
+			_, err := stream.Write([]byte(data))
+			if !suite.NoError(err) {
+				suite.FailNow(
+					fmt.Sprintf("Unexpected error writing to %v: %v", streamName, err),
+				)
+			}
+		}
+		for i := 0; i < 5; i++ {
+			data := strconv.Itoa(i)
+			writeTo("stdout", execCmd.Stdout(), data)
+			writeTo("stderr", execCmd.Stderr(), data)
+		}
+	}()
+
+	for expectedChunk := range expectedChunksCh {
+		timer := time.NewTimer(1 * time.Second)
+		select {
+		case <-timer.C:
+			suite.FailNow(
+				fmt.Sprintf("Timed out while waiting for chunk %v to be sent to the output channel", expectedChunk),
+			)
+		case chunk, ok := <-execCmd.OutputCh():
+			if !ok {
+				suite.FailNow(
+					fmt.Sprintf("Expected chunk %v, but output channel was prematurely closed", expectedChunk),
+				)
+			}
+			EqualChunk(suite.Suite, expectedChunk, chunk)
+		}
+	}
+
+	assertClosedChannel(suite.Suite, execCmd.OutputCh())
+}
+
+func (suite *ExecCommandImplTestSuite) TestNewExecCommand_CancelledContext_ClosesOutputCh() {
+	execCmd := suite.NewExecCommand()
+	suite.cancelFunc()
+	assertSentError(suite.Suite, execCmd.Stdout(), suite.ctx.Err())
+	assertSentError(suite.Suite, execCmd.Stderr(), suite.ctx.Err())
+	assertClosedChannel(suite.Suite, execCmd.OutputCh())
+}
+
+func (suite *ExecCommandImplTestSuite) TestSetStopFunc_CancelledContext_StopsCommand() {
+	execCmd := suite.NewExecCommand()
+	stoppedCh := make(chan bool, 1)
+	time.AfterFunc(1 * time.Second, func() {
+		close(stoppedCh)
+	})
+	execCmd.SetStopFunc(func() {
+		stoppedCh <- true
+	})
+	suite.cancelFunc()
+
+	stopped, ok := <-stoppedCh
+	if !ok {
+		suite.Fail("The command was not stopped")
+	} else {
+		suite.True(stopped)
+	}
+}
+
+func (suite *ExecCommandImplTestSuite) TestSetExitCode_CancelledContext_DoesNotSetExitCode() {
+	execCmd := suite.NewExecCommand()
+	suite.cancelFunc()
+	execCmd.SetExitCode(1)
+	select {
+	case <-execCmd.exitCodeCh:
+		suite.Fail("SetExitCode set the exit code on a cancelled context")
+	default:
+		// Exit code was not sent, so the test passed
+	}
+}
+
+func (suite *ExecCommandImplTestSuite) TestSetExitCode_SetsExitCode() {
+	execCmd := suite.NewExecCommand()
+	execCmd.SetExitCode(1)
+	select {
+	case exitCode := <-execCmd.exitCodeCh:
+		suite.Equal(1, exitCode)
+	default:
+		suite.Fail("SetExitCode did not send the exit code on exitCodeCh")
+	}
+}
+
+func (suite *ExecCommandImplTestSuite) TestSetExitCodeErr_CancelledContext_DoesNotSetExitCodeErr() {
+	execCmd := suite.NewExecCommand()
+	suite.cancelFunc()
+	execCmd.SetExitCodeErr(fmt.Errorf("an error"))
+	select {
+	case <-execCmd.exitCodeErrCh:
+		suite.Fail("SetExitCodeErr set the exit code error on a cancelled context")
+	default:
+		// Exit code error was not sent, so the test passed
+	}
+}
+
+func (suite *ExecCommandImplTestSuite) TestSetExitCodeErr_SetsExitCodeErr() {
+	execCmd := suite.NewExecCommand()
+	expectedErr := fmt.Errorf("error")
+	execCmd.SetExitCodeErr(expectedErr)
+	select {
+	case err := <-execCmd.exitCodeErrCh:
+		suite.Equal(expectedErr, err)
+	default:
+		suite.Fail("SetExitCodeErr did not send the exit code error on exitCodeErrCh")
+	}
+}
+
+func (suite *ExecCommandImplTestSuite) TestExitCode_CancelledContext_ReturnsError() {
+	execCmd := suite.NewExecCommand()
+	suite.cancelFunc()
+	_, err := execCmd.ExitCode()
+	suite.Regexp(suite.ctx.Err(), err)
+}
+
+func (suite *ExecCommandImplTestSuite) TestExitCode_ReturnsExitCodeIfSet() {
+	execCmd := suite.NewExecCommand()
+	execCmd.SetExitCode(1)
+	exitCode, err := execCmd.ExitCode()
+	if suite.NoError(err) {
+		suite.Equal(1, exitCode)
+	}
+}
+
+func (suite *ExecCommandImplTestSuite) TestExitCode_ReturnsExitCodeErrIfSet() {
+	execCmd := suite.NewExecCommand()
+	expectedErr := fmt.Errorf("error")
+	execCmd.SetExitCodeErr(expectedErr)
+	_, err := execCmd.ExitCode()
+	suite.Equal(expectedErr, err)
+}
+
+func TestExecCommandImpl(t *testing.T) {
+	suite.Run(t, new(ExecCommandImplTestSuite))
+}
+
diff --git a/plugin/outputStream_test.go b/plugin/outputStream_test.go
index 877064c..5353818 100644
--- a/plugin/outputStream_test.go
+++ b/plugin/outputStream_test.go
@@ -3,18 +3,34 @@ package plugin
 import (
 	"context"
 	"fmt"
-	"strconv"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/suite"
 )
 
+// NOTE: Some of the assertion helpers in this file take in a suite.Suite
+// object as a parameter because they are also used by the ExecCommandImpl
+// tests.
+
 type OutputStreamTestSuite struct {
 	suite.Suite
 }
 
-func (suite *OutputStreamTestSuite) EqualChunk(expected ExecOutputChunk, actual ExecOutputChunk) bool {
+func newOutputStream(ctx context.Context) *OutputStream {
+	ch := make(chan ExecOutputChunk, 1)
+	return &OutputStream{
+		ctx: ctx,
+		id: Stdout,
+		ch: ch,
+		closer: &multiCloser{
+			ch: ch,
+			countdown: 1,
+		},
+	}
+}
+
+func EqualChunk(suite suite.Suite, expected ExecOutputChunk, actual ExecOutputChunk) bool {
 	expectedStreamName := expected.StreamID
 	eqlStreamName := suite.Equal(
 		expectedStreamName,
@@ -48,22 +64,9 @@ func (suite *OutputStreamTestSuite) EqualChunk(expected ExecOutputChunk, actual
 // for Write pass then so, too, do the tests for WriteWithTimestamp.
 func (suite *OutputStreamTestSuite) TestWrite() {
 	ctx, cancelFunc := context.WithCancel(context.Background())
-	defer func() {
-		// Ensure that the context and its associated resources
-		// are cleaned up
-		select {
-		case <-ctx.Done():
-			// Pass-thru, context was already cancelled
-		default:
-			cancelFunc()
-		}
-	}()
-
-	ch := make(chan ExecOutputChunk, 1)
-	defer close(ch)
-
-	stream := OutputStream{ctx: ctx, id: Stdout, ch: ch}
-
+	defer cancelFunc()
+	stream := newOutputStream(ctx)
+	
 	// Test a successful write
 	data := []byte("data")
 	nw, writeErr := stream.Write(data)
@@ -72,7 +75,8 @@ func (suite *OutputStreamTestSuite) TestWrite() {
 	}
 	select {
 	case chunk := <-stream.ch:
-		suite.EqualChunk(
+		EqualChunk(
+			suite.Suite,
 			ExecOutputChunk{StreamID: Stdout, Data: string(data)},
 			chunk,
 		)
@@ -86,7 +90,8 @@ func (suite *OutputStreamTestSuite) TestWrite() {
 	suite.EqualError(writeErr, ctx.Err().Error(), "Write should have returned the context's error")
 	select {
 	case chunk := <-stream.ch:
-		suite.EqualChunk(
+		EqualChunk(
+			suite.Suite,
 			ExecOutputChunk{StreamID: Stdout, Err: ctx.Err()},
 			chunk,
 		)
@@ -96,7 +101,7 @@ func (suite *OutputStreamTestSuite) TestWrite() {
 	}
 }
 
-func (suite *OutputStreamTestSuite) assertClosedChannel(ch <-chan ExecOutputChunk) {
+func assertClosedChannel(suite suite.Suite, ch <-chan ExecOutputChunk) {
 	timer := time.NewTimer(1 * time.Second)
 	select {
 	case <-timer.C:
@@ -110,119 +115,76 @@ func (suite *OutputStreamTestSuite) assertClosedChannel(ch <-chan ExecOutputChun
 	}
 }
 
-func (suite *OutputStreamTestSuite) TestClose() {
-	ch := make(chan ExecOutputChunk)
-	stream := OutputStream{ch: ch, closer: &multiCloser{ch: ch, countdown: 1}}
-
-	stream.Close()
-	suite.assertClosedChannel(stream.ch)
+func (suite *OutputStreamTestSuite) assertClosed(stream *OutputStream) {
+	assertClosedChannel(suite.Suite, stream.ch)
+	suite.True(stream.closed)
 }
 
-func (suite *OutputStreamTestSuite) TestCloseWithError() {
-	newOutputStream := func(ctx context.Context) OutputStream {
-		ch := make(chan ExecOutputChunk, 1)
-		return OutputStream{ctx: ctx, id: Stdout, ch: ch, closer: &multiCloser{ch: ch, countdown: 1}}
+func assertSentError(suite suite.Suite, stream *OutputStream, err error) {
+	timer := time.NewTimer(1 * time.Second)
+	sentErrorMsg := fmt.Sprintf("Expected the error '%v' to be sent", err)
+	select {
+	case <-timer.C:
+		suite.Fail(sentErrorMsg + ", but timed out while waiting for it")
+	case chunk, ok := <-stream.ch:
+		if !ok {
+			suite.Fail(sentErrorMsg + ", but the channel was closed")
+		} else {
+			EqualChunk(
+				suite,
+				ExecOutputChunk{StreamID: stream.id, Err: err},
+				chunk,
+			)
+		}
 	}
+}
 
-	// Test that if err == nil, then nothing was sent to the channel
+func (suite *OutputStreamTestSuite) TestCloseWithError_NoError() {
 	stream := newOutputStream(context.Background())
 	stream.CloseWithError(nil)
-	suite.assertClosedChannel(stream.ch)
-
-	// Useful assertion for the subsequent tests
-	assertSentError := func(stream OutputStream, err error) {
-		sentErrorMsg := fmt.Sprintf("Expected the error '%v' to be sent", err)
-
-		select {
-		case chunk, ok := <-stream.ch:
-			if !ok {
-				suite.Fail(sentErrorMsg + ", but the channel was closed")
-			} else {
-				suite.EqualChunk(
-					ExecOutputChunk{StreamID: stream.id, Err: err},
-					chunk,
-				)
-			}
-		default:
-			suite.Fail(sentErrorMsg + ", but nothing was sent on the channel")
-		}
-	}
+	// Note that if an error was sent, then assertClosed would fail
+	// because it'd read-in the sent error.
+	suite.assertClosed(stream)
+}
 
-	// Test that if err != nil, then the error is sent to the channel
-	stream = newOutputStream(context.Background())
-	sentErr := fmt.Errorf("an arbitrary error")
-	stream.CloseWithError(sentErr)
-	assertSentError(stream, sentErr)
-	suite.assertClosedChannel(stream.ch)
+func (suite *OutputStreamTestSuite) TestCloseWithError_WithError() {
+	stream := newOutputStream(context.Background())
+	err := fmt.Errorf("an arbitrary error")
+	stream.CloseWithError(err)
+	assertSentError(suite.Suite, stream, err)
+	suite.assertClosed(stream)
+}
 
-	// Test that if err == ctx.Err(), then ctx.Err() is sent if a previous
-	// Write did not send it
+func (suite *OutputStreamTestSuite) TestCloseWithError_ContextError_NotSent() {
 	ctx, cancelFunc := context.WithCancel(context.Background())
-	stream = newOutputStream(ctx)
 	cancelFunc()
+	stream := newOutputStream(ctx)
 	stream.CloseWithError(ctx.Err())
-	assertSentError(stream, ctx.Err())
-	suite.assertClosedChannel(stream.ch)
+	assertSentError(suite.Suite, stream, ctx.Err())
+	suite.assertClosed(stream)
+}
 
-	// Now, test that if err == ctx.Err(), then ctx.Err() is _not_ sent if
-	// a previous Write sent it
-	stream = newOutputStream(ctx)
+func (suite *OutputStreamTestSuite) TestCloseWithError_ContextError_Sent() {
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	cancelFunc()
+	stream := newOutputStream(ctx)
 	stream.sentCtxErr = true
 	stream.CloseWithError(ctx.Err())
-	suite.assertClosedChannel(stream.ch)
-
+	// Note that if CloseWithError sent the passed-in ctx.Err(),
+	// then assertClosed would fail because it'd read-in the sent
+	// ctx.Err()
+	suite.assertClosed(stream)
 }
 
-func (suite *OutputStreamTestSuite) TestCreateExecOutputStreams() {
-	outputCh, stdout, stderr := CreateExecOutputStreams(context.Background())
-
-	// Our simulated command alternates writing to stdout + stderr
-	expectedChunksCh := make(chan ExecOutputChunk, 1)
-	go func() {
-		// stdout + stderr will be closed before expectedChunksCh,
-		// which means that outputCh should be closed before
-		// expectedChunksCh.
-		defer close(expectedChunksCh)
-		defer stdout.Close()
-		defer stderr.Close()
-
-		writeTo := func(streamName string, stream *OutputStream, data string) {
-			expectedChunksCh <- ExecOutputChunk{StreamID: stream.id, Data: data}
-			_, err := stream.Write([]byte(data))
-			if !suite.NoError(err) {
-				suite.FailNow(
-					fmt.Sprintf("Unexpected error writing to %v: %v", streamName, err),
-				)
-			}
-		}
-		for i := 0; i < 5; i++ {
-			data := strconv.Itoa(i)
-			writeTo("stdout", stdout, data)
-			writeTo("stderr", stderr, data)
-		}
-	}()
-
-	for expectedChunk := range expectedChunksCh {
-		timer := time.NewTimer(1 * time.Second)
-		select {
-		case <-timer.C:
-			suite.FailNow(
-				fmt.Sprintf("Timed out while waiting for chunk %v to be sent to the output channel", expectedChunk),
-			)
-		case chunk, ok := <-outputCh:
-			if !ok {
-				suite.FailNow(
-					fmt.Sprintf("Expected chunk %v, but output channel was prematurely closed", expectedChunk),
-				)
-			}
-
-			suite.EqualChunk(expectedChunk, chunk)
-		}
-	}
-
-	// outputCh should be closed, so assert that it is
-	suite.assertClosedChannel(outputCh)
-
+func (suite *OutputStreamTestSuite) TestCloseWithError_ConsecutiveCloses() {
+	// Main thing to test here is to ensure that multiple close calls
+	// won't panic by trying to close stream.ch when it is already
+	// closed
+	stream := newOutputStream(context.Background())
+	stream.CloseWithError(nil)
+	suite.assertClosed(stream)
+	stream.CloseWithError(nil)
+	suite.assertClosed(stream)
 }
 
 func TestOutputStream(t *testing.T) {
diff --git a/volume/fs_test.go b/volume/fs_test.go
index bccd631..a05860e 100644
--- a/volume/fs_test.go
+++ b/volume/fs_test.go
@@ -2,6 +2,7 @@ package volume
 
 import (
 	"context"
+	"fmt"
 	"strings"
 	"testing"
 
@@ -26,33 +27,43 @@ const varLogFixture = `
 
 type fsTestSuite struct {
 	suite.Suite
+	ctx        context.Context
+	cancelFunc context.CancelFunc
 }
 
 func (suite *fsTestSuite) SetupSuite() {
 	plugin.SetTestCache(datastore.NewMemCache())
+	suite.ctx, suite.cancelFunc = context.WithCancel(context.Background())
 }
 
 func (suite *fsTestSuite) TearDownSuite() {
 	plugin.UnsetTestCache()
+	// Cancelling the context ensures that the tests don't leave any
+	// dangling goroutines waiting on a context cancellation. These
+	// goroutines are created by ExecCommandImpl
+	suite.cancelFunc()
 }
 
-func createResult(data string) plugin.ExecResult {
-	outputch := make(chan plugin.ExecOutputChunk, 1)
-	execResult := plugin.ExecResult{
-		OutputCh:   outputch,
-		ExitCodeCB: func() (int, error) { return 0, nil },
-	}
-	outputch <- plugin.ExecOutputChunk{StreamID: plugin.Stdout, Data: data}
-	close(outputch)
-	return execResult
+func (suite *fsTestSuite) createResult(data string) plugin.ExecCommand {
+	cmd := plugin.NewExecCommand(suite.ctx)
+	go func() {
+		_, err := cmd.Stdout().Write([]byte(data))
+		if err != nil {
+			msg := fmt.Sprintf("Unexpected error while setting up mocks: %v", err)
+			panic(msg)
+		}
+		cmd.CloseStreamsWithError(nil)
+		cmd.SetExitCode(0)
+	}()
+	return cmd
 }
 
-func createExec() *mockExecutor {
+func (suite *fsTestSuite) createExec() *mockExecutor {
 	exec := &mockExecutor{EntryBase: plugin.NewEntry("instance")}
 	// Used when recording activity.
 	exec.SetTestID("/instance")
 	cmd := StatCmd("/var/log")
-	exec.On("Exec", mock.Anything, cmd[0], cmd[1:], plugin.ExecOptions{Elevate: true}).Return(createResult(varLogFixture), nil)
+	exec.On("Exec", mock.Anything, cmd[0], cmd[1:], plugin.ExecOptions{Elevate: true}).Return(suite.createResult(varLogFixture), nil)
 	return exec
 }
 
@@ -76,7 +87,7 @@ func (suite *fsTestSuite) find(grp plugin.Group, path string) plugin.Entry {
 }
 
 func (suite *fsTestSuite) TestFSList() {
-	exec := createExec()
+	exec := suite.createExec()
 	fs := NewFS("fs", exec)
 	// ID would normally be set when listing FS within the parent instance.
 	fs.SetTestID("/instance/fs")
@@ -116,7 +127,7 @@ func (suite *fsTestSuite) TestFSList() {
 }
 
 func (suite *fsTestSuite) TestFSRead() {
-	exec := createExec()
+	exec := suite.createExec()
 	fs := NewFS("fs", exec)
 	// ID would normally be set when listing FS within the parent instance.
 	fs.SetTestID("/instance/fs")
@@ -124,7 +135,7 @@ func (suite *fsTestSuite) TestFSRead() {
 	entry := suite.find(fs, "var/log/path1/a file")
 	suite.Equal("a file", plugin.Name(entry))
 
-	execResult := createResult("hello")
+	execResult := suite.createResult("hello")
 	exec.On("Exec", mock.Anything, "cat", []string{"/var/log/path1/a file"}, plugin.ExecOptions{Elevate: true}).Return(execResult, nil)
 	rdr, err := entry.(plugin.Readable).Open(context.Background())
 	suite.NoError(err)
@@ -145,7 +156,7 @@ type mockExecutor struct {
 	mock.Mock
 }
 
-func (m *mockExecutor) Exec(ctx context.Context, cmd string, args []string, opts plugin.ExecOptions) (plugin.ExecResult, error) {
+func (m *mockExecutor) Exec(ctx context.Context, cmd string, args []string, opts plugin.ExecOptions) (plugin.ExecCommand, error) {
 	arger := m.Called(ctx, cmd, args, opts)
-	return arger.Get(0).(plugin.ExecResult), arger.Error(1)
+	return arger.Get(0).(plugin.ExecCommand), arger.Error(1)
 }