blob: faa1cc3d5fd04b11b7a997c3341151ae8789547f [file] [log] [blame]
// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package gosh provides facilities for running and managing processes: start
// them, wait for them to exit, capture their output streams, pipe messages
// between them, terminate them (e.g. on SIGINT), and so on.
//
// Gosh is meant to be used in situations where you might otherwise be tempted
// to write a shell script. (Oh my gosh, no more shell scripts!)
//
// For usage examples, see shell_test.go and internal/gosh_example/main.go.
package gosh
import (
"errors"
"fmt"
"io/ioutil"
"log"
"math/rand"
"os"
"os/signal"
"path"
"path/filepath"
"sync"
"syscall"
"time"
)
const (
envBinDir = "GOSH_BIN_DIR"
envChildOutputDir = "GOSH_CHILD_OUTPUT_DIR"
envInvocation = "GOSH_INVOCATION"
envSpawnedByShell = "GOSH_SPAWNED_BY_SHELL"
)
var (
errAlreadyCalledCleanup = errors.New("gosh: already called Shell.Cleanup")
errDidNotCallMaybeRunFnAndExit = errors.New("gosh: did not call Shell.MaybeRunFnAndExit")
errDidNotCallNewShell = errors.New("gosh: did not call gosh.NewShell")
errShellErrIsNotNil = errors.New("gosh: Shell.Err is not nil")
)
// Shell represents a shell. Not thread-safe.
type Shell struct {
// Err is the most recent error (may be nil).
Err error
// Opts is the Opts struct for this Shell, with default values filled in.
Opts Opts
// Vars is the map of env vars for this Shell.
Vars map[string]string
// Args is the list of args to append to subsequent command invocations.
Args []string
// Internal state.
calledNewShell bool
dirStack []string // for pushd/popd
cleanupMu sync.Mutex // protects the fields below; held during cleanup
calledCleanup bool
cmds []*Cmd
tempFiles []*os.File
tempDirs []string
cleanupFns []func()
}
// Opts configures Shell.
type Opts struct {
// Errorf is called whenever an error is encountered.
// If not specified, defaults to panic(fmt.Sprintf(format, v...)).
Errorf func(format string, v ...interface{})
// Logf is called to log things.
// If not specified, defaults to log.Printf(format, v...).
Logf func(format string, v ...interface{})
// Child stdout and stderr are propagated up to the parent's stdout and stderr
// iff SuppressChildOutput is false.
SuppressChildOutput bool
// If specified, each child's stdout and stderr streams are also piped to
// files in this directory.
// If not specified, defaults to GOSH_CHILD_OUTPUT_DIR.
ChildOutputDir string
// Directory where BuildGoPkg() writes compiled binaries.
// If not specified, defaults to GOSH_BIN_DIR.
BinDir string
}
// NewShell returns a new Shell.
func NewShell(opts Opts) *Shell {
sh, err := newShell(opts)
sh.HandleError(err)
return sh
}
// HandleError sets sh.Err. If err is not nil, it also calls sh.Opts.Errorf.
func (sh *Shell) HandleError(err error) {
sh.Ok()
sh.Err = err
if err != nil && sh.Opts.Errorf != nil {
sh.Opts.Errorf("%v", err)
}
}
// Cmd returns a Cmd for an invocation of the named program.
func (sh *Shell) Cmd(name string, args ...string) *Cmd {
sh.Ok()
res, err := sh.cmd(nil, name, args...)
sh.HandleError(err)
return res
}
// Fn returns a Cmd for an invocation of the given registered Fn.
func (sh *Shell) Fn(fn *Fn, args ...interface{}) *Cmd {
sh.Ok()
res, err := sh.fn(fn, args...)
sh.HandleError(err)
return res
}
// Main returns a Cmd for an invocation of the given registered main() function.
// Intended usage: Have your program's main() call RealMain, then write a parent
// program that uses Shell.Main to run RealMain in a child process. With this
// approach, RealMain can be compiled into the parent program's binary. Caveat:
// potential flag collisions.
func (sh *Shell) Main(fn *Fn, args ...string) *Cmd {
sh.Ok()
res, err := sh.main(fn, args...)
sh.HandleError(err)
return res
}
// Wait waits for all commands started by this Shell to exit.
func (sh *Shell) Wait() {
sh.Ok()
sh.HandleError(sh.wait())
}
// Rename renames (moves) a file. It's just like os.Rename, but retries once on
// error.
func (sh *Shell) Rename(oldpath, newpath string) {
sh.Ok()
sh.HandleError(sh.rename(oldpath, newpath))
}
// BuildGoPkg compiles a Go package using the "go build" command and writes the
// resulting binary to sh.Opts.BinDir. Returns the absolute path to the binary.
// Included in Shell for convenience, but could have just as easily been
// provided as a utility function.
func (sh *Shell) BuildGoPkg(pkg string, flags ...string) string {
sh.Ok()
res, err := sh.buildGoPkg(pkg, flags...)
sh.HandleError(err)
return res
}
// MakeTempFile creates a new temporary file in os.TempDir, opens the file for
// reading and writing, and returns the resulting *os.File.
func (sh *Shell) MakeTempFile() *os.File {
sh.Ok()
res, err := sh.makeTempFile()
sh.HandleError(err)
return res
}
// MakeTempDir creates a new temporary directory in os.TempDir and returns the
// path of the new directory.
func (sh *Shell) MakeTempDir() string {
sh.Ok()
res, err := sh.makeTempDir()
sh.HandleError(err)
return res
}
// Pushd behaves like Bash pushd.
func (sh *Shell) Pushd(dir string) {
sh.Ok()
sh.HandleError(sh.pushd(dir))
}
// Popd behaves like Bash popd.
func (sh *Shell) Popd() {
sh.Ok()
sh.HandleError(sh.popd())
}
// AddToCleanup registers the given function to be called by Shell.Cleanup().
func (sh *Shell) AddToCleanup(fn func()) {
sh.Ok()
sh.HandleError(sh.addToCleanup(fn))
}
// Cleanup cleans up all resources (child processes, temporary files and
// directories) associated with this Shell. It is safe (and recommended) to call
// Cleanup after a Shell error.
func (sh *Shell) Cleanup() {
if !sh.calledNewShell {
panic(errDidNotCallNewShell)
}
sh.cleanupMu.Lock()
defer sh.cleanupMu.Unlock()
if sh.calledCleanup {
panic(errAlreadyCalledCleanup)
} else {
sh.calledCleanup = true
sh.cleanup()
}
}
// Ok panics iff this Shell is in a state where it's invalid to call other
// methods. This method is public to facilitate Shell wrapping.
func (sh *Shell) Ok() {
if !sh.calledNewShell {
panic(errDidNotCallNewShell)
}
// Panic on incorrect usage of Shell.
if sh.Err != nil {
panic(errShellErrIsNotNil)
}
sh.cleanupMu.Lock()
defer sh.cleanupMu.Unlock()
if sh.calledCleanup {
panic(errAlreadyCalledCleanup)
}
}
////////////////////////////////////////
// Internals
// onTerminationSignal starts a goroutine that listens for various termination
// signals and calls the given function when such a signal is received.
func onTerminationSignal(fn func(os.Signal)) {
ch := make(chan os.Signal, 1)
signal.Notify(ch, syscall.SIGINT, syscall.SIGQUIT, syscall.SIGTERM)
go func() {
fn(<-ch)
}()
}
// Note: On error, newShell returns a *Shell with Opts.Errorf initialized to
// simplify things for the caller.
func newShell(opts Opts) (*Shell, error) {
if opts.Errorf == nil {
opts.Errorf = func(format string, v ...interface{}) {
panic(fmt.Sprintf(format, v...))
}
}
if opts.Logf == nil {
opts.Logf = func(format string, v ...interface{}) {
log.Printf(format, v...)
}
}
if opts.ChildOutputDir == "" {
opts.ChildOutputDir = os.Getenv(envChildOutputDir)
}
sh := &Shell{
Opts: opts,
Vars: map[string]string{},
calledNewShell: true,
}
if sh.Opts.BinDir == "" {
sh.Opts.BinDir = os.Getenv(envBinDir)
if sh.Opts.BinDir == "" {
var err error
if sh.Opts.BinDir, err = sh.makeTempDir(); err != nil {
// Note: Here and below, we keep sh.calledCleanup false so that clients
// with a non-fatal Errorf implementation can safely defer sh.Cleanup()
// before checking sh.Err.
sh.cleanup()
return sh, err
}
}
}
// Set this process's PGID to its PID so that its child processes can be
// identified reliably.
// http://man7.org/linux/man-pages/man2/setpgid.2.html
// TODO(sadovsky): Is there any way to reliably kill all spawned subprocesses
// without modifying external state?
if err := syscall.Setpgid(0, 0); err != nil {
sh.cleanup()
return sh, err
}
// Call sh.cleanup() if needed when a termination signal is received.
onTerminationSignal(func(sig os.Signal) {
sh.logf("Received signal: %v\n", sig)
sh.cleanupMu.Lock()
defer sh.cleanupMu.Unlock()
if !sh.calledCleanup {
sh.calledCleanup = true
sh.cleanup()
}
// Note: We hold cleanupMu during os.Exit(1) so that the main goroutine will
// not call Shell.ok() or Shell.Cleanup() and panic before we exit.
os.Exit(1)
})
return sh, nil
}
func (sh *Shell) logf(format string, v ...interface{}) {
if sh.Opts.Logf != nil {
sh.Opts.Logf(format, v...)
}
}
func (sh *Shell) cmd(vars map[string]string, name string, args ...string) (*Cmd, error) {
if vars == nil {
vars = make(map[string]string)
}
vars[envSpawnedByShell] = "1"
c, err := newCmd(sh, mergeMaps(sliceToMap(os.Environ()), sh.Vars, vars), name, append(args, sh.Args...)...)
if err != nil {
return nil, err
}
c.SuppressOutput = sh.Opts.SuppressChildOutput
c.OutputDir = sh.Opts.ChildOutputDir
return c, nil
}
func (sh *Shell) fn(fn *Fn, args ...interface{}) (*Cmd, error) {
// Safeguard against the developer forgetting to call MaybeRunFnAndExit, which
// could lead to infinite recursion.
if !calledMaybeRunFnAndExit {
return nil, errDidNotCallMaybeRunFnAndExit
}
b, err := encInvocation(fn.name, args...)
if err != nil {
return nil, err
}
vars := map[string]string{envInvocation: string(b)}
return sh.cmd(vars, os.Args[0])
}
func (sh *Shell) main(fn *Fn, args ...string) (*Cmd, error) {
// Safeguard against the developer forgetting to call MaybeRunFnAndExit, which
// could lead to infinite recursion.
if !calledMaybeRunFnAndExit {
return nil, errDidNotCallMaybeRunFnAndExit
}
// Check that fn has the required signature.
t := fn.value.Type()
if t.NumIn() != 0 || t.NumOut() != 0 {
return nil, errors.New("main function must have no input or output parameters")
}
b, err := encInvocation(fn.name)
if err != nil {
return nil, err
}
vars := map[string]string{envInvocation: string(b)}
return sh.cmd(vars, os.Args[0], args...)
}
func (sh *Shell) wait() error {
// Note: It is illegal to call newCmd() concurrently with Shell.wait(), so we
// need not hold cleanupMu when accessing sh.cmds below.
var res error
for _, c := range sh.cmds {
if !c.calledStart() || c.calledWait {
continue
}
if err := c.wait(); err != nil {
sh.logf("Cmd.Wait() failed: %v\n", err)
if res == nil {
res = err
}
}
}
return res
}
func (sh *Shell) rename(oldpath, newpath string) error {
if err := os.Rename(oldpath, newpath); err != nil {
// Concurrent, same-directory rename operations sometimes fail on certain
// filesystems, so we retry once after a random backoff.
time.Sleep(time.Duration(rand.Int63n(1000)) * time.Millisecond)
if err := os.Rename(oldpath, newpath); err != nil {
return err
}
}
return nil
}
func (sh *Shell) buildGoPkg(pkg string, flags ...string) (string, error) {
binPath := filepath.Join(sh.Opts.BinDir, path.Base(pkg))
// If this binary has already been built, don't rebuild it.
if _, err := os.Stat(binPath); err == nil {
return binPath, nil
} else if !os.IsNotExist(err) {
return "", err
}
// Build binary to tempBinPath, then move it to binPath.
tempDir, err := ioutil.TempDir(sh.Opts.BinDir, "")
if err != nil {
return "", err
}
defer os.RemoveAll(tempDir)
tempBinPath := filepath.Join(tempDir, path.Base(pkg))
args := []string{"build", "-x", "-o", tempBinPath}
args = append(args, flags...)
args = append(args, pkg)
c, err := sh.cmd(nil, "go", args...)
if err != nil {
return "", err
}
c.SuppressOutput = true
if err := c.run(); err != nil {
return "", err
}
if err := sh.rename(tempBinPath, binPath); err != nil {
return "", err
}
return binPath, nil
}
func (sh *Shell) makeTempFile() (*os.File, error) {
sh.cleanupMu.Lock()
defer sh.cleanupMu.Unlock()
if sh.calledCleanup {
return nil, errAlreadyCalledCleanup
}
f, err := ioutil.TempFile("", "")
if err != nil {
return nil, err
}
sh.tempFiles = append(sh.tempFiles, f)
return f, nil
}
func (sh *Shell) makeTempDir() (string, error) {
sh.cleanupMu.Lock()
defer sh.cleanupMu.Unlock()
if sh.calledCleanup {
return "", errAlreadyCalledCleanup
}
name, err := ioutil.TempDir("", "")
if err != nil {
return "", err
}
sh.tempDirs = append(sh.tempDirs, name)
return name, nil
}
func (sh *Shell) pushd(dir string) error {
cwd, err := os.Getwd()
if err != nil {
return err
}
if err := os.Chdir(dir); err != nil {
return err
}
sh.dirStack = append(sh.dirStack, cwd)
return nil
}
func (sh *Shell) popd() error {
if len(sh.dirStack) == 0 {
return errors.New("dir stack is empty")
}
dir := sh.dirStack[len(sh.dirStack)-1]
if err := os.Chdir(dir); err != nil {
return err
}
sh.dirStack = sh.dirStack[:len(sh.dirStack)-1]
return nil
}
func (sh *Shell) addToCleanup(fn func()) error {
sh.cleanupMu.Lock()
defer sh.cleanupMu.Unlock()
if sh.calledCleanup {
return errAlreadyCalledCleanup
}
sh.cleanupFns = append(sh.cleanupFns, fn)
return nil
}
// forEachRunningCmd applies fn to each running child process.
func (sh *Shell) forEachRunningCmd(fn func(*Cmd)) bool {
anyRunning := false
for _, c := range sh.cmds {
if !c.calledStart() || c.c.Process == nil {
continue // not started
}
if pgid, err := syscall.Getpgid(c.c.Process.Pid); err != nil || pgid != os.Getpid() {
continue // not our child
}
anyRunning = true
if fn != nil {
fn(c)
}
}
return anyRunning
}
// Note: It is safe to run Shell.terminateRunningCmds() concurrently with
// Cmd.wait(). In particular, Shell.terminateRunningCmds() only reads
// c.c.Process.{Pid,Path} and calls c.c.Process.{Signal,Kill}, all of which are
// thread-safe with Cmd.wait().
func (sh *Shell) terminateRunningCmds() {
// Try SIGINT first; if that doesn't work, use SIGKILL.
anyRunning := sh.forEachRunningCmd(func(c *Cmd) {
if err := c.c.Process.Signal(os.Interrupt); err != nil {
sh.logf("%d.Signal(os.Interrupt) failed: %v\n", c.c.Process.Pid, err)
}
})
// If any child is still running, wait for 50ms.
if anyRunning {
time.Sleep(50 * time.Millisecond)
anyRunning = sh.forEachRunningCmd(func(c *Cmd) {
sh.logf("%s (PID %d) did not die\n", c.c.Path, c.c.Process.Pid)
})
}
// If any child is still running, wait for another second, then send SIGKILL
// to all running children.
if anyRunning {
time.Sleep(time.Second)
sh.forEachRunningCmd(func(c *Cmd) {
if err := c.c.Process.Kill(); err != nil {
sh.logf("%d.Kill() failed: %v\n", c.c.Process.Pid, err)
}
})
sh.logf("Sent SIGKILL to all remaining child processes\n")
}
}
func (sh *Shell) cleanup() {
// Terminate all children that are still running. Note, newShell() calls
// syscall.Setpgid().
pgid, pid := syscall.Getpgrp(), os.Getpid()
if pgid != pid {
sh.logf("PGID (%d) != PID (%d); skipping subprocess termination\n", pgid, pid)
} else {
sh.terminateRunningCmds()
}
// Close and delete all temporary files.
for _, tempFile := range sh.tempFiles {
name := tempFile.Name()
if err := tempFile.Close(); err != nil {
sh.logf("%q.Close() failed: %v\n", name, err)
}
if err := os.RemoveAll(name); err != nil {
sh.logf("os.RemoveAll(%q) failed: %v\n", name, err)
}
}
// Delete all temporary directories.
for _, tempDir := range sh.tempDirs {
if err := os.RemoveAll(tempDir); err != nil {
sh.logf("os.RemoveAll(%q) failed: %v\n", tempDir, err)
}
}
// Call any registered cleanup functions in LIFO order.
for i := len(sh.cleanupFns) - 1; i >= 0; i-- {
sh.cleanupFns[i]()
}
}
////////////////////////////////////////
// Public utilities
var calledMaybeRunFnAndExit = false
// MaybeRunFnAndExit must be called first thing in main() or TestMain(), before
// flags are parsed. In the parent process, it returns immediately with no
// effect. In a child process for a Shell.Fn() or Shell.Main() command, it runs
// the specified function, then exits.
func MaybeRunFnAndExit() {
calledMaybeRunFnAndExit = true
s := os.Getenv(envInvocation)
if s == "" {
return
}
os.Unsetenv(envInvocation)
// Call MaybeWatchParent rather than WatchParent so that envSpawnedByShell
// gets cleared.
MaybeWatchParent()
name, args, err := decInvocation(s)
if err != nil {
log.Fatal(err)
}
if err := Call(name, args...); err != nil {
log.Fatal(err)
}
os.Exit(0)
}
// Run calls MaybeRunFnAndExit(), then returns run(). Exported so that TestMain
// functions can simply call os.Exit(gosh.Run(m.Run)).
func Run(run func() int) int {
MaybeRunFnAndExit()
return run()
}