Add a wrapper command to detect timeouts

Add a command that can be used to wrap actions with a timeout, and
optionally run an extra debugging command on timeout.

Bug: 181095653
Test: run_with_timeout_test.go
Change-Id: I91df5c3fb5277968717815a4ad4612113766dab1

cmd/run_with_timeout/Android.bp

@@ -0,0 +1,27 @@
+// Copyright 2021 Google Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+blueprint_go_binary {
+    name: "run_with_timeout",
+    srcs: [
+        "run_with_timeout.go",
+    ],
+    testSrcs: [
+        "run_with_timeout_test.go",
+    ],
+}

cmd/run_with_timeout/run_with_timeout.go

@@ -0,0 +1,110 @@
+// Copyright 2021 Google Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// run_with_timeout is a utility that can kill a wrapped command after a configurable timeout,
+// optionally running a command to collect debugging information first.
+
+package main
+
+import (
+	"flag"
+	"fmt"
+	"io"
+	"os"
+	"os/exec"
+	"syscall"
+	"time"
+)
+
+var (
+	timeout      = flag.Duration("timeout", 0, "time after which to kill command (example: 60s)")
+	onTimeoutCmd = flag.String("on_timeout", "", "command to run with `PID=<pid> sh -c` after timeout.")
+)
+
+func usage() {
+	fmt.Fprintf(os.Stderr, "usage: %s [--timeout N] [--on_timeout CMD] -- command [args...]\n", os.Args[0])
+	flag.PrintDefaults()
+	fmt.Fprintln(os.Stderr, "run_with_timeout is a utility that can kill a wrapped command after a configurable timeout,")
+	fmt.Fprintln(os.Stderr, "optionally running a command to collect debugging information first.")
+
+	os.Exit(2)
+}
+
+func main() {
+	flag.Usage = usage
+	flag.Parse()
+
+	if flag.NArg() < 1 {
+		fmt.Fprintln(os.Stderr, "command is required")
+		usage()
+	}
+
+	err := runWithTimeout(flag.Arg(0), flag.Args()[1:], *timeout, *onTimeoutCmd,
+		os.Stdin, os.Stdout, os.Stderr)
+	if err != nil {
+		if exitErr, ok := err.(*exec.ExitError); ok {
+			fmt.Fprintln(os.Stderr, "process exited with error:", exitErr.Error())
+		} else {
+			fmt.Fprintln(os.Stderr, "error:", err.Error())
+		}
+		os.Exit(1)
+	}
+}
+
+func runWithTimeout(command string, args []string, timeout time.Duration, onTimeoutCmdStr string,
+	stdin io.Reader, stdout, stderr io.Writer) error {
+	cmd := exec.Command(command, args...)
+	cmd.Stdin, cmd.Stdout, cmd.Stderr = stdin, stdout, stderr
+	err := cmd.Start()
+	if err != nil {
+		return err
+	}
+
+	// waitCh will signal the subprocess exited.
+	waitCh := make(chan error)
+	go func() {
+		waitCh <- cmd.Wait()
+	}()
+
+	// timeoutCh will signal the subprocess timed out if timeout was set.
+	var timeoutCh <-chan time.Time = make(chan time.Time)
+	if timeout > 0 {
+		timeoutCh = time.After(timeout)
+	}
+
+	select {
+	case err := <-waitCh:
+		if exitErr, ok := err.(*exec.ExitError); ok {
+			return fmt.Errorf("process exited with error: %w", exitErr)
+		}
+		return err
+	case <-timeoutCh:
+		// Continue below.
+	}
+
+	// Process timed out before exiting.
+	defer cmd.Process.Signal(syscall.SIGKILL)
+
+	if onTimeoutCmdStr != "" {
+		onTimeoutCmd := exec.Command("sh", "-c", onTimeoutCmdStr)
+		onTimeoutCmd.Stdin, onTimeoutCmd.Stdout, onTimeoutCmd.Stderr = stdin, stdout, stderr
+		onTimeoutCmd.Env = append(os.Environ(), fmt.Sprintf("PID=%d", cmd.Process.Pid))
+		err := onTimeoutCmd.Run()
+		if err != nil {
+			return fmt.Errorf("on_timeout command %q exited with error: %w", onTimeoutCmdStr, err)
+		}
+	}
+
+	return fmt.Errorf("timed out after %s", timeout.String())
+}

cmd/run_with_timeout/run_with_timeout_test.go

@@ -0,0 +1,94 @@
+// Copyright 2021 Google Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"bytes"
+	"io"
+	"testing"
+	"time"
+)
+
+func Test_runWithTimeout(t *testing.T) {
+	type args struct {
+		command      string
+		args         []string
+		timeout      time.Duration
+		onTimeoutCmd string
+		stdin        io.Reader
+	}
+	tests := []struct {
+		name       string
+		args       args
+		wantStdout string
+		wantStderr string
+		wantErr    bool
+	}{
+		{
+			name: "no timeout",
+			args: args{
+				command: "echo",
+				args:    []string{"foo"},
+			},
+			wantStdout: "foo\n",
+		},
+		{
+			name: "timeout not reached",
+			args: args{
+				command: "echo",
+				args:    []string{"foo"},
+				timeout: 1 * time.Second,
+			},
+			wantStdout: "foo\n",
+		},
+		{
+			name: "timed out",
+			args: args{
+				command: "sh",
+				args:    []string{"-c", "sleep 1 && echo foo"},
+				timeout: 1 * time.Millisecond,
+			},
+			wantErr: true,
+		},
+		{
+			name: "on_timeout command",
+			args: args{
+				command:      "sh",
+				args:         []string{"-c", "sleep 1 && echo foo"},
+				timeout:      1 * time.Millisecond,
+				onTimeoutCmd: "echo bar",
+			},
+			wantStdout: "bar\n",
+			wantErr:    true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			stdout := &bytes.Buffer{}
+			stderr := &bytes.Buffer{}
+			err := runWithTimeout(tt.args.command, tt.args.args, tt.args.timeout, tt.args.onTimeoutCmd, tt.args.stdin, stdout, stderr)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("runWithTimeout() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if gotStdout := stdout.String(); gotStdout != tt.wantStdout {
+				t.Errorf("runWithTimeout() gotStdout = %v, want %v", gotStdout, tt.wantStdout)
+			}
+			if gotStderr := stderr.String(); gotStderr != tt.wantStderr {
+				t.Errorf("runWithTimeout() gotStderr = %v, want %v", gotStderr, tt.wantStderr)
+			}
+		})
+	}
+}