Merge pull request golang#164 from erizocosmico/bugfix-84/hung-processes

Introduce monitoredCmd

Author: sdboyer

_testdata/cmd/echosleep.go

@@ -0,0 +1,17 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"time"
+)
+
+func main() {
+	n := flag.Int("n", 1, "number of iterations before stopping")
+	flag.Parse()
+
+	for i := 0; i < *n; i++ {
+		fmt.Println("foo")
+		time.Sleep(time.Duration(i) * 100 * time.Millisecond)
+	}
+}

cmd.go

@@ -0,0 +1,99 @@
+package gps
+
+import (
+	"bytes"
+	"fmt"
+	"os/exec"
+	"time"
+)
+
+// monitoredCmd wraps a cmd and will keep monitoring the process until it
+// finishes or a certain amount of time has passed and the command showed
+// no signs of activity.
+type monitoredCmd struct {
+	cmd     *exec.Cmd
+	timeout time.Duration
+	stdout  *activityBuffer
+	stderr  *activityBuffer
+}
+
+func newMonitoredCmd(cmd *exec.Cmd, timeout time.Duration) *monitoredCmd {
+	stdout := newActivityBuffer()
+	stderr := newActivityBuffer()
+	cmd.Stderr = stderr
+	cmd.Stdout = stdout
+	return &monitoredCmd{cmd, timeout, stdout, stderr}
+}
+
+// run will wait for the command to finish and return the error, if any. If the
+// command does not show any activity for more than the specified timeout the
+// process will be killed.
+func (c *monitoredCmd) run() error {
+	ticker := time.NewTicker(c.timeout)
+	done := make(chan error, 1)
+	defer ticker.Stop()
+	go func() { done <- c.cmd.Run() }()
+
+	for {
+		select {
+		case <-ticker.C:
+			if c.hasTimedOut() {
+				if err := c.cmd.Process.Kill(); err != nil {
+					return &killCmdError{err}
+				}
+
+				return &timeoutError{c.timeout}
+			}
+		case err := <-done:
+			return err
+		}
+	}
+}
+
+func (c *monitoredCmd) hasTimedOut() bool {
+	t := time.Now().Add(-c.timeout)
+	return c.stderr.lastActivity.Before(t) &&
+		c.stdout.lastActivity.Before(t)
+}
+
+func (c *monitoredCmd) combinedOutput() ([]byte, error) {
+	if err := c.run(); err != nil {
+		return c.stderr.buf.Bytes(), err
+	}
+
+	return c.stdout.buf.Bytes(), nil
+}
+
+// activityBuffer is a buffer that keeps track of the last time a Write
+// operation was performed on it.
+type activityBuffer struct {
+	buf          *bytes.Buffer
+	lastActivity time.Time
+}
+
+func newActivityBuffer() *activityBuffer {
+	return &activityBuffer{
+		buf: bytes.NewBuffer(nil),
+	}
+}
+
+func (b *activityBuffer) Write(p []byte) (int, error) {
+	b.lastActivity = time.Now()
+	return b.buf.Write(p)
+}
+
+type timeoutError struct {
+	timeout time.Duration
+}
+
+func (e timeoutError) Error() string {
+	return fmt.Sprintf("command killed after %s of no activity", e.timeout)
+}
+
+type killCmdError struct {
+	err error
+}
+
+func (e killCmdError) Error() string {
+	return fmt.Sprintf("error killing command after timeout: %s", e.err)
+}

cmd_test.go

@@ -0,0 +1,51 @@
+package gps
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"testing"
+	"time"
+)
+
+func mkTestCmd(iterations int) *monitoredCmd {
+	return newMonitoredCmd(
+		exec.Command("./echosleep", "-n", fmt.Sprint(iterations)),
+		200*time.Millisecond,
+	)
+}
+
+func TestMonitoredCmd(t *testing.T) {
+	err := exec.Command("go", "build", "./_testdata/cmd/echosleep.go").Run()
+	if err != nil {
+		t.Errorf("Unable to build echosleep binary: %s", err)
+	}
+	defer os.Remove("./echosleep")
+
+	cmd := mkTestCmd(2)
+	err = cmd.run()
+	if err != nil {
+		t.Errorf("Expected command not to fail: %s", err)
+	}
+
+	expectedOutput := "foo\nfoo\n"
+	if cmd.stdout.buf.String() != expectedOutput {
+		t.Errorf("Unexpected output:\n\t(GOT): %s\n\t(WNT): %s", cmd.stdout.buf.String(), expectedOutput)
+	}
+
+	cmd = mkTestCmd(10)
+	err = cmd.run()
+	if err == nil {
+		t.Error("Expected command to fail")
+	}
+
+	_, ok := err.(*timeoutError)
+	if !ok {
+		t.Errorf("Expected a timeout error, but got: %s", err)
+	}
+
+	expectedOutput = "foo\nfoo\nfoo\nfoo\n"
+	if cmd.stdout.buf.String() != expectedOutput {
+		t.Errorf("Unexpected output:\n\t(GOT): %s\n\t(WNT): %s", cmd.stdout.buf.String(), expectedOutput)
+	}
+}