vpython: Add primary execution package.

vpython/run.go

Index: vpython/run.go
diff --git a/vpython/run.go b/vpython/run.go
new file mode 100644
index 0000000000000000000000000000000000000000..16f959738e10e1128d544105856ee366898db1f4
--- /dev/null
+++ b/vpython/run.go
@@ -0,0 +1,138 @@
+// Copyright 2017 The LUCI Authors. All rights reserved.
+// Use of this source code is governed under the Apache License, Version 2.0
+// that can be found in the LICENSE file.
+
+package vpython
+
+import (
+	"os"
+	"os/exec"
+	"os/signal"
+
+	"github.com/luci/luci-go/common/errors"
+	"github.com/luci/luci-go/common/logging"
+	"github.com/luci/luci-go/common/system/environ"
+	"github.com/luci/luci-go/vpython/venv"
+
+	"golang.org/x/net/context"
+)
+
+var (
+	// EnvSpecPath is the exported enviornment variable for the specification path.
+	//
+	// This is added to the bootstrap enviornment used by Run to allow subprocess
+	// "vpython" invocations to automatically inherit the same environment.
+	EnvSpecPath = "VPYTHON_VENV_SPEC_PATH"
+)
+
+// Run sets up a Python VirtualEnv and executes the supplied Options.
+//
+// Run returns nil if if the Python environment was successfully set-up and the
+// Python interpreter was successfully run with a zero return code. If the
+// Python interpreter returns a non-zero return code, a PythonError (potentially
+// wrapped) will be returned.
+//
+// A generalized return code to return for an error value can be obtained via
+// ReturnCode.
+//
+// Run consists of:
+//
+//	- Identify the target Python script to run (if there is one).
+//	- Identifying the Python interpreter to use.
+//	- Composing the environment specification.
+//	- Constructing the virtual environment (download, install).
+//	- Execute the Python process with the supplied arguments.
+//
+// The Python subprocess is bound to the lifetime of ctx, and will be terminated
+// if ctx is cancelled.
+func Run(c context.Context, opts Options) error {
+	// Resolve our Options.
+	if err := opts.resolve(c); err != nil {
+		return errors.Annotate(err).Reason("could not resolve options").Err()
+	}
+
+	// Create a local cancellation option (signal handling).
+	c, cancelFunc := context.WithCancel(c)
+	defer cancelFunc()
+
+	// Create our virtual enviornment root directory.
+	err := venv.With(c, opts.EnvConfig, opts.WaitForEnv, func(c context.Context, ve *venv.Env) error {
+		// Build the augmented environment variables.
+		e := opts.Environ
+		if e.Len() == 0 {
+			// If no environment was supplied, use the system environment.
+			e = environ.System()
+		}
+
+		e.Set("VIRTUAL_ENV", ve.Root) // Set by VirtualEnv script.
+		if ve.SpecPath != "" {
+			e.Set(EnvSpecPath, ve.SpecPath)
+		}
+
+		// Run our bootstrapped Python command.
+		cmd := ve.InterpreterCommand()
+		cmd.WorkDir = opts.WorkDir
+		cmd.Isolated = true
+		cmd.Env = e.Sorted()
+
+		pythonCmd, err := cmd.Prepare(c, opts.Args...)
+		if err != nil {
+			return errors.Annotate(err).Reason("failed to prepare command").Err()
+		}
+		pythonCmd.Stdin = os.Stdin
+
+		if err := runAndForwardSignals(c, pythonCmd, cancelFunc); err != nil {
+			return errors.Annotate(err).Reason("failed to execute bootstrapped Python").Err()
+		}
+		return nil
+	})
+	if err != nil {
+		return errors.Annotate(err).Err()
+	}
+	return nil
+}
+
+func runAndForwardSignals(c context.Context, cmd *exec.Cmd, cancelFunc context.CancelFunc) error {
+	signalC := make(chan os.Signal, 1)
+	signalDoneC := make(chan struct{})
+	signal.Notify(signalC, forwardedSignals...)
+	defer func() {
+		signal.Stop(signalC)
+
+		close(signalC)
+		<-signalDoneC
+	}()
+
+	if err := cmd.Start(); err != nil {
+		return errors.Annotate(err).Reason("failed to start process").Err()
+	}
+
+	logging.Fields{
+		"pid": cmd.Process.Pid,
+	}.Debugf(c, "Python subprocess has started!")
+
+	// Start our signal forwarding goroutine, now that the process is running.
+	go func() {
+		defer func() {
+			close(signalDoneC)
+		}()
+
+		for sig := range signalC {
+			logging.Debugf(c, "Forwarding signal: %v", sig)
+			if err := cmd.Process.Signal(sig); err != nil {
+				logging.Fields{
+					logging.ErrorKey: err,
+					"signal":         sig,
+				}.Errorf(c, "Failed to forward signal; terminating immediately.")
+				cancelFunc()
+			}
+		}
+	}()
+
+	err := cmd.Wait()
+	logging.Debugf(c, "Python subprocess has terminated: %v", err)
+	if err != nil {
+		return errors.Annotate(err).Err()
+	}
+	return nil
+}