lib/modules/exec.go

package modules

import (
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"sync"
	"time"

	"veyron.io/veyron/veyron2/vlog"

	vexec "veyron.io/veyron/veyron/lib/exec"
)

// execHandle implements both the command and Handle interfaces.
type execHandle struct {
	mu         sync.Mutex
	cmd        *exec.Cmd
	entryPoint string
	handle     *vexec.ParentHandle
	sh         *Shell
	stderr     *os.File
	stdout     io.ReadCloser
	stdin      io.WriteCloser
}

func testFlags() []string {
	var fl []string
	// pass logging flags to any subprocesses
	for fname, fval := range vlog.Log.ExplicitlySetFlags() {
		fl = append(fl, "--"+fname+"="+fval)
	}
	timeout := flag.Lookup("test.timeout")
	if timeout == nil {
		// not a go test binary
		return fl
	}
	// must be a go test binary
	fl = append(fl, "-test.run=TestHelperProcess")
	val := timeout.Value.(flag.Getter).Get().(time.Duration)
	if val.String() != timeout.DefValue {
		// use supplied command value for subprocesses
		fl = append(fl, "--test.timeout="+timeout.Value.String())
	} else {
		// translate default value into 1m for subproccesses
		fl = append(fl, "--test.timeout=1m")
	}
	return fl
}

// IsTestHelperProces returns true if it is called in via
// -run=TestHelperProcess which normally only ever happens for subprocesses
// run from tests.
func IsTestHelperProcess() bool {
	runFlag := flag.Lookup("test.run")
	if runFlag == nil {
		return false
	}
	return runFlag.Value.String() == "TestHelperProcess"
}

func newExecHandle(entryPoint string) command {
	return &execHandle{entryPoint: entryPoint}
}

func (eh *execHandle) Stdout() io.Reader {
	eh.mu.Lock()
	defer eh.mu.Unlock()
	return eh.stdout
}

func (eh *execHandle) Stderr() io.Reader {
	eh.mu.Lock()
	defer eh.mu.Unlock()
	return eh.stderr
}

func (eh *execHandle) Stdin() io.Writer {
	eh.mu.Lock()
	defer eh.mu.Unlock()
	return eh.stdin
}

func (eh *execHandle) CloseStdin() {
	eh.mu.Lock()
	eh.stdin.Close()
	eh.mu.Unlock()
}

func osEnvironMap() map[string]string {
	m := make(map[string]string)
	for _, osv := range os.Environ() {
		if len(osv) == 0 {
			continue
		}
		parts := strings.SplitN(osv, "=", 2)
		key := parts[0]
		if len(parts) == 2 {
			m[key] = parts[1]
		} else {
			m[key] = ""
		}
	}
	return m
}
func (sh *Shell) mergeOSEnv() map[string]string {
	merged := osEnvironMap()
	sh.mu.Lock()
	for k, v := range sh.env {
		merged[k] = v
	}
	sh.mu.Unlock()
	return merged
}

func (eh *execHandle) envelope(sh *Shell, args ...string) ([]string, []string) {
	newargs := []string{os.Args[0]}
	newargs = append(newargs, testFlags()...)
	newargs = append(newargs, args...)
	return newargs, append(sh.MergedEnv(), eh.entryPoint)
}

func (eh *execHandle) start(sh *Shell, args ...string) (Handle, error) {
	eh.mu.Lock()
	defer eh.mu.Unlock()
	eh.sh = sh
	// Take care to not pass the command line as an arg to the child
	// process since that'll prevent parsing any subsequent args by
	// the flag package.
	newargs := append(testFlags(), args[1:]...)
	cmd := exec.Command(os.Args[0], newargs...)
	cmd.Env = append(sh.MergedEnv(), eh.entryPoint)
	fname := strings.TrimPrefix(eh.entryPoint, ShellEntryPoint+"=")
	stderr, err := newLogfile(strings.TrimLeft(fname, "-\n\t "))
	if err != nil {
		return nil, err
	}
	cmd.Stderr = stderr
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}

	handle := vexec.NewParentHandle(cmd)
	eh.stdout = stdout
	eh.stderr = stderr
	eh.stdin = stdin
	eh.handle = handle
	eh.cmd = cmd
	if err := handle.Start(); err != nil {
		return nil, err
	}
	err = handle.WaitForReady(10 * time.Second)
	return eh, err
}

func (eh *execHandle) Pid() int {
	return eh.cmd.Process.Pid
}

func (eh *execHandle) Shutdown(stdout, stderr io.Writer) error {
	eh.mu.Lock()
	defer eh.mu.Unlock()
	eh.stdin.Close()
	logFile := eh.stderr.Name()
	defer eh.sh.Forget(eh)

	defer func() {
		os.Remove(logFile)
	}()

	if stdout == nil && stderr == nil {
		return eh.cmd.Wait()
	}
	// Read from stdin before waiting for the child process to ensure
	// that we get to read all of its output.
	readTo(eh.stdout, stdout)

	procErr := eh.cmd.Wait()

	// Stderr is buffered to a file, so we can safely read it after we
	// wait for the process.
	eh.stderr.Close()
	if stderr != nil {
		stderrFile, err := os.Open(logFile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "failed to open %q: %s\n", logFile, err)
			return procErr
		}
		readTo(stderrFile, stderr)
		stderrFile.Close()
	}
	return procErr
}

const ShellEntryPoint = "VEYRON_SHELL_HELPER_PROCESS_ENTRY_POINT"

func RegisterChild(name, help string, main Main) {
	child.Lock()
	defer child.Unlock()
	child.mains[name] = &childEntryPoint{main, help}
}

// DispatchInTest will execute the requested subproccess command from within
// a unit test run as a subprocess.
func DispatchInTest() {
	if !IsTestHelperProcess() {
		return
	}
	if err := child.dispatch(); err != nil {
		fmt.Fprintf(os.Stderr, "Failed: %s\n", err)
		os.Exit(1)
	}
	os.Exit(0)
}

// Dispatch will execute the requested subprocess command from a within a
// a subprocess that is not a unit test.
func Dispatch() error {
	if IsTestHelperProcess() {
		return fmt.Errorf("use DispatchInTest in unittests")
	}
	return child.dispatch()
}

func (child *childRegistrar) hasCommand(name string) bool {
	child.Lock()
	_, present := child.mains[name]
	child.Unlock()
	return present
}

func (child *childRegistrar) dispatch() error {
	ch, _ := vexec.GetChildHandle()
	// Only signal that the child is ready or failed if we successfully get
	// a child handle. We most likely failed to get a child handle
	// because the subprocess was run directly from the command line.

	command := os.Getenv(ShellEntryPoint)
	if len(command) == 0 {
		err := fmt.Errorf("Failed to find entrypoint %q", ShellEntryPoint)
		if ch != nil {
			ch.SetFailed(err)
		}
		return err
	}

	child.Lock()
	m := child.mains[command]
	child.Unlock()

	if m == nil {
		err := fmt.Errorf("Shell command %q not registered", command)
		if ch != nil {
			ch.SetFailed(err)
		}
		return err
	}

	if ch != nil {
		ch.SetReady()
	}

	go func(pid int) {
		for {
			_, err := os.FindProcess(pid)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Looks like our parent exited: %v", err)
				os.Exit(1)
			}
			time.Sleep(time.Second)
		}
	}(os.Getppid())

	args := append([]string{command}, flag.Args()...)
	return m.fn(os.Stdin, os.Stdout, os.Stderr, osEnvironMap(), args...)
}

func (child *childRegistrar) addSubprocesses(sh *Shell, pattern string) error {
	re, err := regexp.Compile(pattern)
	if err != nil {
		return err
	}
	child.Lock()
	defer child.Unlock()
	found := false
	for name, subproc := range child.mains {
		if re.MatchString(name) {
			sh.addSubprocess(name, subproc.help)
			found = true
		}
	}
	if !found {
		return fmt.Errorf("patterh %q failed to match any registered commands", pattern)
	}
	return nil
}

// AddRegisteredSubprocesses adds any commands that match the regexp pattern
// to the supplied shell.
func AddRegisteredSubprocesses(sh *Shell, pattern string) error {
	return child.addSubprocesses(sh, pattern)
}

// WaitForEof returns when a read on its io.Reader parameter returns io.EOF
func WaitForEOF(stdin io.Reader) {
	buf := [1024]byte{}
	for {
		if _, err := stdin.Read(buf[:]); err == io.EOF {
			return
		}
	}
}