Esempio n. 1
0
func (r *runner) Map(mr *simplemr.MR, key string, val interface{}) error {
	mi := val.(*mapInput)
	output := &mapOutput{
		key: key,
		mi:  mi}
	jirix := mi.jirix
	path := os.Getenv("SHELL")
	if path == "" {
		path = "sh"
	}
	var wg sync.WaitGroup
	cmd := exec.Command(path, "-c", strings.Join(r.args, " "))
	cmd.Env = envvar.MapToSlice(jirix.Env())
	cmd.Dir = mi.ProjectState.Project.Path
	cmd.Stdin = mi.jirix.Stdin()
	var stdoutCloser, stderrCloser io.Closer
	if runpFlags.interactive {
		cmd.Stdout = jirix.Stdout()
		cmd.Stderr = jirix.Stderr()
	} else {
		var stdout io.Writer
		stderr := r.serializedWriter(jirix.Stderr())
		var cleanup func()
		if runpFlags.collateOutput {
			// Write standard output to a file, stderr
			// is not collated.
			f, err := ioutil.TempFile("", mi.ProjectState.Project.Name+"-")
			if err != nil {
				return err
			}
			stdout = f
			output.outputFilename = f.Name()
			cleanup = func() {
				os.Remove(output.outputFilename)
			}
			// The child process will have exited by the
			// time this method returns so it's safe to close the file
			// here.
			defer f.Close()
		} else {
			stdout = r.serializedWriter(jirix.Stdout())
			cleanup = func() {}
		}
		if !runpFlags.showNamePrefix && !runpFlags.showKeyPrefix {
			// write directly to stdout, stderr if there's no prefix
			cmd.Stdout = stdout
			cmd.Stderr = stderr
		} else {
			stdoutReader, stdoutWriter, err := os.Pipe()
			if err != nil {
				cleanup()
				return err
			}
			stderrReader, stderrWriter, err := os.Pipe()
			if err != nil {
				cleanup()
				stdoutReader.Close()
				stdoutWriter.Close()
				return err
			}
			cmd.Stdout = stdoutWriter
			cmd.Stderr = stderrWriter
			// Record the write end of the pipe so that it can be closed
			// after the child has exited, this ensures that all goroutines
			// will finish.
			stdoutCloser = stdoutWriter
			stderrCloser = stderrWriter
			prefix := key
			if runpFlags.showNamePrefix {
				prefix = mi.ProjectState.Project.Name
			}
			wg.Add(2)
			go func() { copyWithPrefix(prefix, stdout, stdoutReader); wg.Done() }()
			go func() { copyWithPrefix(prefix, stderr, stderrReader); wg.Done() }()

		}
	}
	if err := cmd.Start(); err != nil {
		mi.result = err
	}
	done := make(chan error)
	go func() {
		done <- cmd.Wait()
	}()
	select {
	case output.err = <-done:
		if output.err != nil && runpFlags.exitOnError {
			mr.Cancel()
		}
	case <-mr.CancelCh():
		output.err = cmd.Process.Kill()
	}
	for _, closer := range []io.Closer{stdoutCloser, stderrCloser} {
		if closer != nil {
			closer.Close()
		}
	}
	wg.Wait()
	mr.MapOut(key, output)
	return nil
}