gosh: handle same WriteCloser passed to Add{Stdout,Stderr}Writer
Also, a few minor cleanups:
- Update all tests to set sh.Opts.Fatalf to nil as the
canonical way to disable fatal'ing
- Tweak TestAddWritersUnwrappedStdoutStderr to avoid loops,
since the loop-free impl is still short and will give more
useful output on failure
- Expand comment for Cmd.Args to make it clear that it
doesn't include the path
Note, Cmd.Args not including the path is different from
exec.Cmd, but it seems more sensible given that otherwise
there's an awkward overlap between Cmd.Path and Cmd.Args.
(Also, exec.Cmd special-cases for Cmd.Args==nil, setting it
to []string{Cmd.Path} in that case, which seems a bit
inelegant.)
Change-Id: Icadc8bb892129b67fc060857e4efa0545ea5ce01
diff --git a/gosh/cmd.go b/gosh/cmd.go
index f113530..effbac9 100644
--- a/gosh/cmd.go
+++ b/gosh/cmd.go
@@ -36,7 +36,7 @@
Path string
// Vars is the map of env vars for this Cmd.
Vars map[string]string
- // Args is the list of args for this Cmd.
+ // Args is the list of args for this Cmd, not including the path.
Args []string
// PropagateOutput is inherited from Shell.Opts.PropagateChildOutput.
PropagateOutput bool
@@ -106,22 +106,30 @@
}
// AddStdoutWriter configures this Cmd to tee the child's stdout to the given
-// wc, which will be closed when the process exits.
+// WriteCloser, which will be closed when the process exits.
//
-// Use NopWriteCloser to extend an io.Writer to io.WriteCloser, or to prevent an
-// existing io.WriteCloser from being closed. It is an error to pass in
-// os.Stdout or os.Stderr, since they shouldn't be closed.
+// If the same WriteCloser is passed to both AddStdoutWriter and
+// AddStderrWriter, Cmd will ensure that its methods are never called
+// concurrently and that Close is only called once.
+//
+// Use NopWriteCloser to extend a Writer to a WriteCloser, or to prevent an
+// existing WriteCloser from being closed. It is an error to pass in os.Stdout
+// or os.Stderr, since they shouldn't be closed.
func (c *Cmd) AddStdoutWriter(wc io.WriteCloser) {
c.sh.Ok()
c.handleError(c.addStdoutWriter(wc))
}
// AddStderrWriter configures this Cmd to tee the child's stderr to the given
-// wc, which will be closed when the process exits.
+// WriteCloser, which will be closed when the process exits.
//
-// Use NopWriteCloser to extend an io.Writer to io.WriteCloser, or to prevent an
-// existing io.WriteCloser from being closed. It is an error to pass in
-// os.Stdout or os.Stderr, since they shouldn't be closed.
+// If the same WriteCloser is passed to both AddStdoutWriter and
+// AddStderrWriter, Cmd will ensure that its methods are never called
+// concurrently and that Close is only called once.
+//
+// Use NopWriteCloser to extend a Writer to a WriteCloser, or to prevent an
+// existing WriteCloser from being closed. It is an error to pass in os.Stdout
+// or os.Stderr, since they shouldn't be closed.
func (c *Cmd) AddStderrWriter(wc io.WriteCloser) {
c.sh.Ok()
c.handleError(c.addStderrWriter(wc))
@@ -261,8 +269,14 @@
}
func (c *Cmd) closeClosers() {
+ // If the same WriteCloser was passed to both AddStdoutWriter and
+ // AddStderrWriter, we should only close it once.
+ cm := map[io.Closer]bool{}
for _, c := range c.closers {
- c.Close() // best-effort; ignore returned error
+ if !cm[c] {
+ cm[c] = true
+ c.Close() // best-effort; ignore returned error
+ }
}
}
@@ -332,7 +346,7 @@
if c.PropagateOutput {
c.addWriter(writers, std)
// Don't add std to c.closers, since we don't want to close os.Stdout or
- // os.Stderr for the entire address space when c finishes.
+ // os.Stderr for the entire address space when c exits.
}
if c.OutputDir != "" {
suffix := "stderr"
@@ -431,6 +445,17 @@
// TODO(sadovsky): Maybe wrap every child process with a "supervisor" process
// that calls WatchParent().
+type threadSafeWriter struct {
+ mu sync.Mutex
+ w io.Writer
+}
+
+func (w *threadSafeWriter) Write(p []byte) (int, error) {
+ w.mu.Lock()
+ defer w.mu.Unlock()
+ return w.w.Write(p)
+}
+
func (c *Cmd) start() error {
if c.calledStart {
return errAlreadyCalledStart
@@ -443,6 +468,29 @@
if c.sh.calledCleanup {
return errAlreadyCalledCleanup
}
+ // Wrap Writers in threadSafeWriters as needed so that if the same WriteCloser
+ // was passed to both AddStdoutWriter and AddStderrWriter, the writes are
+ // serialized.
+ isStdoutWriter := map[io.Writer]bool{}
+ for _, w := range c.stdoutWriters {
+ isStdoutWriter[w] = true
+ }
+ safe := map[io.Writer]*threadSafeWriter{}
+ for i, w := range c.stderrWriters {
+ if isStdoutWriter[w] {
+ if safe[w] == nil {
+ safe[w] = &threadSafeWriter{w: w}
+ }
+ c.stderrWriters[i] = safe[w]
+ }
+ }
+ if len(safe) > 0 {
+ for i, w := range c.stdoutWriters {
+ if s := safe[w]; s != nil {
+ c.stdoutWriters[i] = s
+ }
+ }
+ }
// Configure the command.
c.c.Path = c.Path
c.c.Env = mapToSlice(c.Vars)
diff --git a/gosh/shell.go b/gosh/shell.go
index 192b795..02ee145 100644
--- a/gosh/shell.go
+++ b/gosh/shell.go
@@ -568,7 +568,7 @@
}
// NopWriteCloser returns a WriteCloser with a no-op Close method wrapping the
-// provided Writer w.
+// provided Writer.
func NopWriteCloser(w io.Writer) io.WriteCloser {
return nopWriteCloser{w}
}
diff --git a/gosh/shell_test.go b/gosh/shell_test.go
index c96b406..b1c9c12 100644
--- a/gosh/shell_test.go
+++ b/gosh/shell_test.go
@@ -13,6 +13,7 @@
import (
"bufio"
+ "bytes"
"errors"
"fmt"
"io"
@@ -150,7 +151,7 @@
ok(t, err)
eq(t, cwd, startDir)
// The next sh.Popd() will fail.
- sh.Opts.Fatalf = func(string, ...interface{}) {}
+ sh.Opts.Fatalf = nil
sh.Popd()
nok(t, sh.Err)
}
@@ -285,7 +286,7 @@
eq(t, sh.Fn(printfIntsFn, "%v %v", 1, 2).Stdout(), "1 2")
// Error cases.
- sh.Opts.Fatalf = func(string, ...interface{}) {}
+ sh.Opts.Fatalf = nil
reset := func() {
nok(t, sh.Err)
sh.Err = nil
@@ -368,7 +369,7 @@
c.Stdin = "foo"
c.StdinPipe().Write([]byte("bar"))
c.StdinPipe().Close()
- sh.Opts.Fatalf = func(string, ...interface{}) {}
+ sh.Opts.Fatalf = nil
c.Start()
nok(t, sh.Err)
}
@@ -443,7 +444,7 @@
})
// Tests that it's safe to add wrapped os.Stdout and os.Stderr as writers.
-func TestAddWriters(t *testing.T) {
+func TestAddWritersWrappedStdoutStderr(t *testing.T) {
sh := gosh.NewShell(gosh.Opts{Fatalf: makeFatalf(t), Logf: t.Logf})
defer sh.Cleanup()
@@ -453,20 +454,73 @@
}
// Tests that adding non-wrapped os.Stdout or os.Stderr fails.
-func TestAddWritersUnwrappedStdoutStderr(t *testing.T) {
+func TestAddWritersNonWrappedStdoutStderr(t *testing.T) {
sh := gosh.NewShell(gosh.Opts{Fatalf: makeFatalf(t), Logf: t.Logf})
defer sh.Cleanup()
- for _, addFn := range []func(*gosh.Cmd, io.WriteCloser){(*gosh.Cmd).AddStdoutWriter, (*gosh.Cmd).AddStderrWriter} {
- for _, std := range []io.WriteCloser{os.Stdout, os.Stderr} {
- c := sh.Fn(writeMoreFn)
- sh.Opts.Fatalf = nil
- addFn(c, std)
- nok(t, sh.Err)
- sh.Err = nil
- sh.Opts.Fatalf = makeFatalf(t)
- }
- }
+ c := sh.Fn(writeMoreFn)
+ sh.Opts.Fatalf = nil
+ c.AddStdoutWriter(os.Stdout)
+ nok(t, sh.Err)
+ sh.Err = nil
+ c.AddStdoutWriter(os.Stderr)
+ nok(t, sh.Err)
+ sh.Err = nil
+ c.AddStderrWriter(os.Stdout)
+ nok(t, sh.Err)
+ sh.Err = nil
+ c.AddStderrWriter(os.Stderr)
+ nok(t, sh.Err)
+ sh.Err = nil
+}
+
+// Demonstrates how to capture combined stdout and stderr, a la
+// exec.Cmd.CombinedOutput.
+func TestCombinedStdoutStderr(t *testing.T) {
+ sh := gosh.NewShell(gosh.Opts{Fatalf: makeFatalf(t), Logf: t.Logf})
+ defer sh.Cleanup()
+
+ c := sh.Fn(writeFn, true, true)
+ buf := &bytes.Buffer{}
+ // Note, we must share a single NopWriteCloser so that Cmd detects that we
+ // passed the same WriteCloser to both AddStdoutWriter and AddStderrWriter.
+ wc := gosh.NopWriteCloser(buf)
+ c.AddStdoutWriter(wc)
+ c.AddStderrWriter(wc)
+ c.Run()
+ // Note, we can't assume any particular ordering of stdout and stderr, so we
+ // simply check the length of the combined output.
+ eq(t, len(buf.String()), 4)
+}
+
+type countingWriteCloser struct {
+ io.Writer
+ count int
+}
+
+func (wc *countingWriteCloser) Close() error {
+ wc.count++
+ return nil
+}
+
+// Tests that Close is called exactly once on a given WriteCloser, even if that
+// WriteCloser is passed to Add{Stdout,Stderr}Writer multiple times.
+func TestAddWritersCloseOnce(t *testing.T) {
+ sh := gosh.NewShell(gosh.Opts{Fatalf: makeFatalf(t), Logf: t.Logf})
+ defer sh.Cleanup()
+
+ c := sh.Fn(writeFn, true, true)
+ buf := &bytes.Buffer{}
+ wc := &countingWriteCloser{Writer: buf}
+ c.AddStdoutWriter(wc)
+ c.AddStdoutWriter(wc)
+ c.AddStderrWriter(wc)
+ c.AddStderrWriter(wc)
+ c.Run()
+ // Note, we can't assume any particular ordering of stdout and stderr, so we
+ // simply check the length of the combined output.
+ eq(t, len(buf.String()), 8)
+ eq(t, wc.count, 1)
}
// Tests piping from one Cmd's stdout/stderr to another's stdin. It should be
@@ -532,7 +586,7 @@
// Configure the "failed to start" command.
c1.StdinPipe()
c1.Stdin = "foo"
- sh.Opts.Fatalf = func(string, ...interface{}) {}
+ sh.Opts.Fatalf = nil
c1.Start()
nok(t, sh.Err)
sh.Err = nil
@@ -577,7 +631,7 @@
// If ExitErrorIsOk is false, exit code 1 triggers sh.HandleError.
c = sh.Fn(exitFn, 1)
- sh.Opts.Fatalf = func(string, ...interface{}) {}
+ sh.Opts.Fatalf = nil
c.Run()
nok(t, c.Err)
nok(t, sh.Err)
