veyron/lib/modules: add support for per-command environment and logging.
Change-Id: Ie54dd05c161b06303cee7b878772d5e58ba31e17
diff --git a/lib/modules/exec.go b/lib/modules/exec.go
index 44ef680..35d6512 100644
--- a/lib/modules/exec.go
+++ b/lib/modules/exec.go
@@ -11,15 +11,15 @@
"sync"
"time"
- "veyron.io/veyron/veyron2/vlog"
-
vexec "veyron.io/veyron/veyron/lib/exec"
+ "veyron.io/veyron/veyron2/vlog"
)
// execHandle implements both the command and Handle interfaces.
type execHandle struct {
mu sync.Mutex
cmd *exec.Cmd
+ name string
entryPoint string
handle *vexec.ParentHandle
sh *Shell
@@ -63,8 +63,8 @@
return runFlag.Value.String() == "TestHelperProcess"
}
-func newExecHandle(entryPoint string) command {
- return &execHandle{entryPoint: entryPoint}
+func newExecHandle(name string) command {
+ return &execHandle{name: name, entryPoint: ShellEntryPoint + "=" + name}
}
func (eh *execHandle) Stdout() io.Reader {
@@ -91,40 +91,14 @@
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) {
+func (eh *execHandle) envelope(sh *Shell, env []string, args ...string) ([]string, []string) {
newargs := []string{os.Args[0]}
newargs = append(newargs, testFlags()...)
newargs = append(newargs, args...)
- return newargs, append(sh.MergedEnv(), eh.entryPoint)
+ return newargs, append(env, eh.entryPoint)
}
-func (eh *execHandle) start(sh *Shell, args ...string) (Handle, error) {
+func (eh *execHandle) start(sh *Shell, env []string, args ...string) (Handle, error) {
eh.mu.Lock()
defer eh.mu.Unlock()
eh.sh = sh
@@ -133,7 +107,7 @@
// the flag package.
newargs := append(testFlags(), args[1:]...)
cmd := exec.Command(os.Args[0], newargs...)
- cmd.Env = append(sh.MergedEnv(), eh.entryPoint)
+ cmd.Env = append(env, eh.entryPoint)
fname := strings.TrimPrefix(eh.entryPoint, ShellEntryPoint+"=")
stderr, err := newLogfile(strings.TrimLeft(fname, "-\n\t "))
if err != nil {
@@ -155,10 +129,13 @@
eh.stdin = stdin
eh.handle = handle
eh.cmd = cmd
+ vlog.VI(1).Infof("Start: %q args: %v", eh.name, cmd.Args)
+ vlog.VI(2).Infof("Start: %q env: %v", eh.name, cmd.Env)
if err := handle.Start(); err != nil {
return nil, err
}
- err = handle.WaitForReady(10 * time.Second)
+ // TODO(cnicolaou): make this timeout configurable
+ err = handle.WaitForReady(time.Minute)
return eh, err
}
@@ -169,6 +146,7 @@
func (eh *execHandle) Shutdown(stdout, stderr io.Writer) error {
eh.mu.Lock()
defer eh.mu.Unlock()
+ vlog.VI(1).Infof("Shutdown: %q", eh.name)
eh.stdin.Close()
logFile := eh.stderr.Name()
defer eh.sh.Forget(eh)
@@ -177,14 +155,19 @@
os.Remove(logFile)
}()
+ // TODO(cnicolaou): make this configurable
+ timeout := 10 * time.Second
if stdout == nil && stderr == nil {
- return eh.cmd.Wait()
+ return eh.handle.Wait(timeout)
}
- // 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()
+ if stdout != nil {
+ // 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.handle.Wait(timeout)
// Stderr is buffered to a file, so we can safely read it after we
// wait for the process.
@@ -192,7 +175,7 @@
if stderr != nil {
stderrFile, err := os.Open(logFile)
if err != nil {
- fmt.Fprintf(os.Stderr, "failed to open %q: %s\n", logFile, err)
+ vlog.VI(1).Infof("failed to open %q: %s\n", logFile, err)
return procErr
}
readTo(stderrFile, stderr)
@@ -216,8 +199,7 @@
return
}
if err := child.dispatch(); err != nil {
- fmt.Fprintf(os.Stderr, "Failed: %s\n", err)
- os.Exit(1)
+ vlog.Fatalf("Failed: %s", err)
}
os.Exit(0)
}
@@ -273,15 +255,14 @@
for {
_, err := os.FindProcess(pid)
if err != nil {
- fmt.Fprintf(os.Stderr, "Looks like our parent exited: %v", err)
- os.Exit(1)
+ vlog.Fatalf("Looks like our parent exited: %v", err)
}
time.Sleep(time.Second)
}
}(os.Getppid())
args := append([]string{command}, flag.Args()...)
- return m.fn(os.Stdin, os.Stdout, os.Stderr, osEnvironMap(), args...)
+ return m.fn(os.Stdin, os.Stdout, os.Stderr, envSliceToMap(os.Environ()), args...)
}
func (child *childRegistrar) addSubprocesses(sh *Shell, pattern string) error {
diff --git a/lib/modules/func.go b/lib/modules/func.go
index 68d103b..54ce341 100644
--- a/lib/modules/func.go
+++ b/lib/modules/func.go
@@ -5,6 +5,8 @@
"io"
"os"
"sync"
+
+ "veyron.io/veyron/veyron2/vlog"
)
type pipe struct {
@@ -12,6 +14,7 @@
}
type functionHandle struct {
mu sync.Mutex
+ name string
main Main
stdin, stdout pipe
stderr *os.File
@@ -20,8 +23,8 @@
wg sync.WaitGroup
}
-func newFunctionHandle(main Main) command {
- return &functionHandle{main: main}
+func newFunctionHandle(name string, main Main) command {
+ return &functionHandle{name: name, main: main}
}
func (fh *functionHandle) Stdout() io.Reader {
@@ -48,11 +51,11 @@
fh.mu.Unlock()
}
-func (fh *functionHandle) envelope(sh *Shell, args ...string) ([]string, []string) {
- return args, sh.MergedEnv()
+func (fh *functionHandle) envelope(sh *Shell, env []string, args ...string) ([]string, []string) {
+ return args, env
}
-func (fh *functionHandle) start(sh *Shell, args ...string) (Handle, error) {
+func (fh *functionHandle) start(sh *Shell, env []string, args ...string) (Handle, error) {
fh.mu.Lock()
defer fh.mu.Unlock()
fh.sh = sh
@@ -69,7 +72,7 @@
fh.stderr = stderr
fh.wg.Add(1)
- go func() {
+ go func(env []string) {
fh.mu.Lock()
stdin := fh.stdin.r
stdout := fh.stdout.w
@@ -77,7 +80,10 @@
main := fh.main
fh.mu.Unlock()
- err := main(stdin, stdout, stderr, sh.mergeOSEnv(), args...)
+ cenv := envSliceToMap(env)
+ vlog.VI(1).Infof("Start: %q args: %v", fh.name, args)
+ vlog.VI(2).Infof("Start: %q env: %v", fh.name, cenv)
+ err := main(stdin, stdout, stderr, cenv, args...)
if err != nil {
fmt.Fprintf(stderr, "%s\n", err)
}
@@ -88,7 +94,7 @@
fh.err = err
fh.mu.Unlock()
fh.wg.Done()
- }()
+ }(env)
return fh, nil
}
@@ -98,6 +104,7 @@
func (fh *functionHandle) Shutdown(stdout_w, stderr_w io.Writer) error {
fh.mu.Lock()
+ vlog.VI(1).Infof("Shutdown: %q", fh.name)
fh.stdin.w.Close()
stdout := fh.stdout.r
stderr := fh.stderr
diff --git a/lib/modules/modules_test.go b/lib/modules/modules_test.go
index 6869816..278edab 100644
--- a/lib/modules/modules_test.go
+++ b/lib/modules/modules_test.go
@@ -12,9 +12,11 @@
"veyron.io/veyron/veyron/lib/exec"
"veyron.io/veyron/veyron/lib/modules"
+ "veyron.io/veyron/veyron/lib/testutil"
)
func init() {
+ testutil.Init()
modules.RegisterChild("envtest", "envtest: <variables to print>...", PrintFromEnv)
modules.RegisterChild("printenv", "printenv", PrintEnv)
modules.RegisterChild("echo", "[args]*", Echo)
@@ -49,7 +51,7 @@
fmt.Fprintf(stdout, "%s%s\n", printEnvArgPrefix, a)
}
for k, v := range env {
- fmt.Fprintln(stdout, k+"="+v)
+ fmt.Fprintf(stdout, k+"="+v+"\n")
}
return nil
}
@@ -97,7 +99,7 @@
}
h.CloseStdin()
if !waitForInput(scanner) {
- t.Errorf("timeout")
+ t.Fatalf("timeout")
return
}
if got, want := scanner.Text(), "done"; got != want {
@@ -144,12 +146,12 @@
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
- if got, want := h.Shutdown(nil, nil), "exit status 1"; got == nil || got.Error() != want {
+ if got, want := h.Shutdown(nil, nil), "exit status 255"; got == nil || got.Error() != want {
t.Errorf("got %q, want %q", got, want)
}
}
-func testShutdown(t *testing.T, sh *modules.Shell) {
+func testShutdown(t *testing.T, sh *modules.Shell, isfunc bool) {
result := ""
args := []string{"a", "b c", "ddd"}
if _, err := sh.Start("echo", args...); err != nil {
@@ -166,6 +168,9 @@
if got, want := stdoutBuf.String(), stdoutOutput+result; got != want {
t.Errorf("got %q want %q", got, want)
}
+ if !isfunc {
+ stderrBuf.ReadString('\n') // Skip past the random # generator output
+ }
if got, want := stderrBuf.String(), stderrOutput; got != want {
t.Errorf("got %q want %q", got, want)
}
@@ -173,13 +178,13 @@
func TestShutdownSubprocess(t *testing.T) {
sh := modules.NewShell("echo")
- testShutdown(t, sh)
+ testShutdown(t, sh, false)
}
func TestShutdownFunction(t *testing.T) {
sh := modules.NewShell()
sh.AddFunction("echo", Echo, "[args]*")
- testShutdown(t, sh)
+ testShutdown(t, sh, true)
}
func TestErrorFunc(t *testing.T) {
@@ -224,7 +229,7 @@
childEnv = append(childEnv, o)
}
}
- shArgs, shEnv := sh.CommandEnvelope("printenv", args...)
+ shArgs, shEnv := sh.CommandEnvelope("printenv", nil, args...)
for _, want := range args {
if !find(want, childArgs) {
t.Errorf("failed to find %q in %s", want, childArgs)
@@ -254,7 +259,10 @@
defer sh.Cleanup(nil, nil)
sh.SetVar("a", "1")
os.Setenv("a", "wrong, should be 1")
- h, err := sh.Start("printenv")
+ sh.SetVar("b", "2 also wrong")
+ os.Setenv("b", "wrong, should be 2")
+
+ h, err := sh.StartWithEnv("printenv", []string{"b=2"})
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
@@ -263,7 +271,12 @@
o := scanner.Text()
if strings.HasPrefix(o, "a=") {
if got, want := o, "a=1"; got != want {
- t.Errorf("got: %s, want %s", got, want)
+ t.Errorf("got: %q, want %q", got, want)
+ }
+ }
+ if strings.HasPrefix(o, "b=") {
+ if got, want := o, "b=2"; got != want {
+ t.Errorf("got: %q, want %q", got, want)
}
}
}
diff --git a/lib/modules/shell.go b/lib/modules/shell.go
index 69e633a..cb68856 100644
--- a/lib/modules/shell.go
+++ b/lib/modules/shell.go
@@ -108,6 +108,14 @@
// CreateAndUseNewCredentials setups a new credentials directory and then
// configures the shell and all of its children to use to it.
+//
+// TODO(cnicolaou): this should use the principal already setup
+// with the runtime if the runtime has been initialized, if not,
+// it should create a new principal. As of now, this approach only works
+// for child processes that talk to each other, but not to the parent
+// process that started them since it's running with a different set of
+// credentials setup elsewhere. When this change is made it should
+// be possible to remove creating credentials in many unit tests.
func (sh *Shell) CreateAndUseNewCredentials() error {
dir, err := ioutil.TempDir("", "veyron_credentials")
if err != nil {
@@ -132,9 +140,8 @@
}
func (sh *Shell) addSubprocess(name string, help string) {
- entryPoint := ShellEntryPoint + "=" + name
sh.mu.Lock()
- sh.cmds[name] = &commandDesc{func() command { return newExecHandle(entryPoint) }, help}
+ sh.cmds[name] = &commandDesc{func() command { return newExecHandle(name) }, help}
sh.mu.Unlock()
}
@@ -142,7 +149,7 @@
// within the current process.
func (sh *Shell) AddFunction(name string, main Main, help string) {
sh.mu.Lock()
- sh.cmds[name] = &commandDesc{func() command { return newFunctionHandle(main) }, help}
+ sh.cmds[name] = &commandDesc{func() command { return newFunctionHandle(name, main) }, help}
sh.mu.Unlock()
}
@@ -179,9 +186,16 @@
// the Cleanup method. If the non-registered subprocess command does not
// exist then the Start command will return an error.
func (sh *Shell) Start(name string, args ...string) (Handle, error) {
+ return sh.StartWithEnv(name, nil, args...)
+}
+
+// StartWithEnv is like Start except with a set of environment variables
+// that override those in the Shell and the OS' environment.
+func (sh *Shell) StartWithEnv(name string, env []string, args ...string) (Handle, error) {
+ cenv := sh.MergedEnv(env)
cmd := sh.getCommand(name)
expanded := append([]string{name}, sh.expand(args...)...)
- h, err := cmd.factory().start(sh, expanded...)
+ h, err := cmd.factory().start(sh, cenv, expanded...)
if err != nil {
return nil, err
}
@@ -196,8 +210,7 @@
cmd := sh.cmds[name]
sh.mu.Unlock()
if cmd == nil {
- entryPoint := ShellEntryPoint + "=" + name
- cmd = &commandDesc{func() command { return newExecHandle(entryPoint) }, ""}
+ cmd = &commandDesc{func() command { return newExecHandle(name) }, ""}
}
return cmd
}
@@ -205,8 +218,8 @@
// CommandEnvelope returns the command line and environment that would be
// used for running the subprocess or function if it were started with the
// specifed arguments.
-func (sh *Shell) CommandEnvelope(name string, args ...string) ([]string, []string) {
- return sh.getCommand(name).factory().envelope(sh, args...)
+func (sh *Shell) CommandEnvelope(name string, env []string, args ...string) ([]string, []string) {
+ return sh.getCommand(name).factory().envelope(sh, sh.MergedEnv(env), args...)
}
// Forget tells the Shell to stop tracking the supplied Handle. This is
@@ -249,6 +262,13 @@
sh.env[key] = value
}
+// ClearVar removes the speficied variable from the Shell's environment
+func (sh *Shell) ClearVar(key string) {
+ sh.mu.Lock()
+ defer sh.mu.Unlock()
+ delete(sh.env, key)
+}
+
// Env returns the entire set of environment variables associated with this
// Shell as a string slice.
func (sh *Shell) Env() []string {
@@ -263,8 +283,10 @@
// Cleanup calls Shutdown on all of the Handles currently being tracked
// by the Shell and writes to stdout and stderr as per the Shutdown
-// method in the Handle interface.
-func (sh *Shell) Cleanup(stdout, stderr io.Writer) {
+// method in the Handle interface. Cleanup returns the error from the
+// last Shutdown that returned a non-nil error. The order that the
+// Shutdown routines are executed is not defined.
+func (sh *Shell) Cleanup(stdout, stderr io.Writer) error {
sh.mu.Lock()
handles := make(map[Handle]struct{})
for k, v := range sh.handles {
@@ -272,24 +294,35 @@
}
sh.handles = make(map[Handle]struct{})
sh.mu.Unlock()
+ var err error
for k, _ := range handles {
- k.Shutdown(stdout, stderr)
+ cerr := k.Shutdown(stdout, stderr)
+ if cerr != nil {
+ err = cerr
+ }
}
if len(sh.credDir) > 0 {
os.RemoveAll(sh.credDir)
}
+ return err
}
// MergedEnv returns a slice that contains the merged set of environment
-// variables from the OS environment and those in this Shell, preferring
-// values in the Shell environment over those found in the OS environment.
-func (sh *Shell) MergedEnv() []string {
- merged := sh.mergeOSEnv()
- env := []string{}
- for k, v := range merged {
- env = append(env, k+"="+v)
+// variables from the OS environment, those in this Shell and those provided
+// as a parameter to it. It prefers values from its parameter over those
+// from the Shell, over those from the OS.
+func (sh *Shell) MergedEnv(env []string) []string {
+ osmap := envSliceToMap(os.Environ())
+ evmap := envSliceToMap(env)
+ sh.mu.Lock()
+ m1 := mergeMaps(osmap, sh.env)
+ defer sh.mu.Unlock()
+ m2 := mergeMaps(m1, evmap)
+ r := []string{}
+ for k, v := range m2 {
+ r = append(r, k+"="+v)
}
- return env
+ return r
}
// Handle represents a running command.
@@ -325,6 +358,6 @@
// command is used to abstract the implementations of inprocess and subprocess
// commands.
type command interface {
- envelope(sh *Shell, args ...string) ([]string, []string)
- start(sh *Shell, args ...string) (Handle, error)
+ envelope(sh *Shell, env []string, args ...string) ([]string, []string)
+ start(sh *Shell, env []string, args ...string) (Handle, error)
}
diff --git a/lib/modules/util.go b/lib/modules/util.go
index 21fc4c2..c35c74e 100644
--- a/lib/modules/util.go
+++ b/lib/modules/util.go
@@ -6,6 +6,7 @@
"io"
"io/ioutil"
"os"
+ "strings"
)
func newLogfile(prefix string) (*os.File, error) {
@@ -25,3 +26,40 @@
fmt.Fprintf(w, "%s\n", scanner.Text())
}
}
+
+// envSliceToMap returns a map representation of a string slive
+// of environment variables.
+func envSliceToMap(env []string) map[string]string {
+ m := make(map[string]string)
+ if env == nil {
+ return m
+ }
+ for _, osv := range env {
+ 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
+}
+
+// mergeMaps merges two maps, a & b, with b taking preference over a.
+func mergeMaps(a, b map[string]string) map[string]string {
+ if len(b) == 0 {
+ return a
+ }
+ merged := make(map[string]string)
+ for k, v := range a {
+ merged[k] = v
+ }
+ for k, v := range b {
+ merged[k] = v
+ }
+ return merged
+}
