github.com/vanadium-archive/go.jiri@v0.0.0-20160715023856-abfb8b131290/cmd/jiri/runp.go

// 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 main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"regexp"
	"sort"
	"strings"
	"sync"

	"v.io/jiri"
	"v.io/jiri/gitutil"
	"v.io/jiri/profiles/profilescmdline"
	"v.io/jiri/profiles/profilesreader"
	"v.io/jiri/project"
	"v.io/jiri/tool"
	"v.io/x/lib/cmdline"
	"v.io/x/lib/envvar"
	"v.io/x/lib/simplemr"
)

var (
	cmdRunP   *cmdline.Command
	runpFlags runpFlagValues
)

func newRunP() *cmdline.Command {
	return &cmdline.Command{
		Runner: jiri.RunnerFunc(runRunp),
		Name:   "runp",
		Short:  "Run a command in parallel across jiri projects",
		Long: `
Run a command in parallel across one or more jiri projects using the specified
profile target's environment. Commands are run using the shell specified by the
users $SHELL environment variable, or "sh" if that's not set. Thus commands
are run as $SHELL -c "args..."
 `,
		ArgsName: "<command line>",
		ArgsLong: `
	A command line to be run in each project specified by the supplied command
line flags. Any environment variables intended to be evaluated when the
command line is run must be quoted to avoid expansion before being passed to
runp by the shell.
`,
	}
}

type runpFlagValues struct {
	profilescmdline.ReaderFlagValues
	projectKeys      string
	verbose          bool
	interactive      bool
	hasUncommitted   bool
	hasUntracked     bool
	hasGerritMessage bool
	showNamePrefix   bool
	showKeyPrefix    bool
	exitOnError      bool
	collateOutput    bool
	editMessage      bool
	hasBranch        string
}

func registerCommonFlags(flags *flag.FlagSet, values *runpFlagValues) {
	profilescmdline.RegisterReaderFlags(flags, &values.ReaderFlagValues, "", jiri.ProfilesDBDir)
	flags.BoolVar(&values.verbose, "v", false, "Print verbose logging information")
	flags.StringVar(&values.projectKeys, "projects", "", "A Regular expression specifying project keys to run commands in. By default, runp will use projects that have the same branch checked as the current project unless it is run from outside of a project in which case it will default to using all projects.")
	flags.BoolVar(&values.hasUncommitted, "has-uncommitted", false, "If specified, match projects that have, or have no, uncommitted changes")
	flags.BoolVar(&values.hasUntracked, "has-untracked", false, "If specified, match projects that have, or have no, untracked files")
	flags.BoolVar(&values.hasGerritMessage, "has-gerrit-message", false, "If specified, match branches that have, or have no, gerrit message")
	flags.BoolVar(&values.interactive, "interactive", true, "If set, the command to be run is interactive and should not have its stdout/stderr manipulated. This flag cannot be used with -show-name-prefix, -show-key-prefix or -collate-stdout.")
	flags.BoolVar(&values.showNamePrefix, "show-name-prefix", false, "If set, each line of output from each project will begin with the name of the project followed by a colon. This is intended for use with long running commands where the output needs to be streamed. Stdout and stderr are spliced apart. This flag cannot be used with -interactive, -show-key-prefix or -collate-stdout.")
	flags.BoolVar(&values.showKeyPrefix, "show-key-prefix", false, "If set, each line of output from each project will begin with the key of the project followed by a colon. This is intended for use with long running commands where the output needs to be streamed. Stdout and stderr are spliced apart. This flag cannot be used with -interactive, -show-name-prefix or -collate-stdout")
	flags.BoolVar(&values.collateOutput, "collate-stdout", true, "Collate all stdout output from each parallel invocation and display it as if had been generated sequentially. This flag cannot be used with -show-name-prefix, -show-key-prefix or -interactive.")
	flags.BoolVar(&values.exitOnError, "exit-on-error", false, "If set, all commands will killed as soon as one reports an error, otherwise, each will run to completion.")
	flags.StringVar(&values.hasBranch, "has-branch", "", "A regular expression specifying branch names to use in matching projects. A project will match if the specified branch exists, even if it is not checked out.")
}

func init() {
	// Avoid an intialization loop between cmdline.Command.Runner which
	// refers to cmdRunP and runRunp referring back to cmdRunP.ParsedFlags.
	cmdRunP = newRunP()
	cmdRoot.Children = append(cmdRoot.Children, cmdRunP)
	registerCommonFlags(&cmdRunP.Flags, &runpFlags)
}

type mapInput struct {
	*project.ProjectState
	key          project.ProjectKey
	jirix        *jiri.X
	index, total int
	result       error
}

func newmapInput(jirix *jiri.X, state *project.ProjectState, key project.ProjectKey, index, total int) *mapInput {
	return &mapInput{
		ProjectState: state,
		key:          key,
		jirix:        jirix.Clone(tool.ContextOpts{}),
		index:        index,
		total:        total,
	}
}

func stateNames(states map[project.ProjectKey]*mapInput) []string {
	n := []string{}
	for _, state := range states {
		n = append(n, state.Project.Name)
	}
	sort.Strings(n)
	return n
}

func stateKeys(states map[project.ProjectKey]*mapInput) []string {
	n := []string{}
	for key := range states {
		n = append(n, string(key))
	}
	sort.Strings(n)
	return n
}

type runner struct {
	args                 []string
	reader               *profilesreader.Reader
	serializedWriterLock sync.Mutex
	collatedOutputLock   sync.Mutex
}

func (r *runner) serializedWriter(w io.Writer) io.Writer {
	return &sharedLockWriter{&r.serializedWriterLock, w}
}

type sharedLockWriter struct {
	mu *sync.Mutex
	f  io.Writer
}

func (lw *sharedLockWriter) Write(d []byte) (int, error) {
	lw.mu.Lock()
	defer lw.mu.Unlock()
	return lw.f.Write(d)
}

func copyWithPrefix(prefix string, w io.Writer, r io.Reader) {
	reader := bufio.NewReader(r)
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			if line != "" {
				fmt.Fprintf(w, "%v: %v\n", prefix, line)
			}
			break
		}
		fmt.Fprintf(w, "%v: %v", prefix, line)
	}
}

type mapOutput struct {
	mi             *mapInput
	outputFilename string
	key            string
	err            error
}

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
}

func (r *runner) Reduce(mr *simplemr.MR, key string, values []interface{}) error {
	for _, v := range values {
		mo := v.(*mapOutput)
		jirix := mo.mi.jirix
		if mo.err != nil {
			fmt.Fprintf(jirix.Stdout(), "FAILED: %v: %s %v\n", mo.key, strings.Join(r.args, " "), mo.err)
			return mo.err
		} else {
			if runpFlags.collateOutput {
				r.collatedOutputLock.Lock()
				defer r.collatedOutputLock.Unlock()
				defer os.Remove(mo.outputFilename)
				if fi, err := os.Open(mo.outputFilename); err == nil {
					io.Copy(jirix.Stdout(), fi)
					fi.Close()
				} else {
					return err
				}
			}
		}
	}
	return nil
}

func runp(jirix *jiri.X, cmd *cmdline.Command, args []string) error {
	hasUntrackedSet := profilescmdline.IsFlagSet(cmd.ParsedFlags, "has-untracked")
	hasUncommitedSet := profilescmdline.IsFlagSet(cmd.ParsedFlags, "has-uncommitted")
	hasGerritSet := profilescmdline.IsFlagSet(cmd.ParsedFlags, "has-gerrit-message")

	if runpFlags.interactive {
		runpFlags.collateOutput = false
	}

	var keysRE, branchRE *regexp.Regexp
	var err error

	if profilescmdline.IsFlagSet(cmd.ParsedFlags, "projects") {
		re := ""
		for _, pre := range strings.Split(runpFlags.projectKeys, ",") {
			re += pre + "|"
		}
		re = strings.TrimRight(re, "|")
		keysRE, err = regexp.Compile(re)
		if err != nil {
			return fmt.Errorf("failed to compile projects regexp: %q: %v", runpFlags.projectKeys, err)
		}
	}

	if profilescmdline.IsFlagSet(cmd.ParsedFlags, "has-branch") {
		branchRE, err = regexp.Compile(runpFlags.hasBranch)
		if err != nil {
			return fmt.Errorf("failed to compile has-branch regexp: %q: %v", runpFlags.hasBranch, err)
		}
	}

	for _, f := range []string{"show-key-prefix", "show-name-prefix"} {
		if profilescmdline.IsFlagSet(cmd.ParsedFlags, f) {
			if runpFlags.interactive && profilescmdline.IsFlagSet(cmd.ParsedFlags, "interactive") {
				fmt.Fprintf(jirix.Stderr(), "WARNING: interactive mode being disabled because %s was set\n", f)
			}
			runpFlags.interactive = false
			runpFlags.collateOutput = true
			break
		}
	}

	git := gitutil.New(jirix.NewSeq())
	homeBranch, err := git.CurrentBranchName()
	if err != nil {
		// jiri was run from outside of a project. Let's assume we'll
		// use all projects if none have been specified via the projects flag.
		if keysRE == nil {
			keysRE = regexp.MustCompile(".*")
		}
	}

	dirty := false
	if hasUntrackedSet || hasUncommitedSet {
		dirty = true
	}
	states, err := project.GetProjectStates(jirix, dirty)
	if err != nil {
		return err
	}
	mapInputs := map[project.ProjectKey]*mapInput{}
	var keys project.ProjectKeys
	for key, state := range states {
		if keysRE != nil {
			if !keysRE.MatchString(string(key)) {
				continue
			}
		} else {
			if state.CurrentBranch != homeBranch {
				continue
			}
		}
		if branchRE != nil {
			found := false
			for _, br := range state.Branches {
				if branchRE.MatchString(br.Name) {
					found = true
					break
				}
			}
			if !found {
				continue
			}
		}
		if hasUntrackedSet && (state.HasUntracked != runpFlags.hasUntracked) {
			continue
		}
		if hasUncommitedSet && (state.HasUncommitted != runpFlags.hasUncommitted) {
			continue
		}
		if hasGerritSet {
			hasMsg := false
			for _, br := range state.Branches {
				if (state.CurrentBranch == br.Name) && br.HasGerritMessage {
					hasMsg = true
					break
				}
			}
			if hasMsg != runpFlags.hasGerritMessage {
				continue
			}
		}
		mapInputs[key] = &mapInput{
			ProjectState: state,
			jirix:        jirix,
			key:          key,
		}
		keys = append(keys, key)
	}

	total := len(mapInputs)
	index := 1
	for _, mi := range mapInputs {
		mi.index = index
		mi.total = total
		index++
	}

	if runpFlags.verbose {
		fmt.Fprintf(jirix.Stdout(), "Project Names: %s\n", strings.Join(stateNames(mapInputs), " "))
		fmt.Fprintf(jirix.Stdout(), "Project Keys: %s\n", strings.Join(stateKeys(mapInputs), " "))
	}

	reader, err := profilesreader.NewReader(jirix, runpFlags.ProfilesMode, runpFlags.DBFilename)
	runner := &runner{
		reader: reader,
		args:   args,
	}
	mr := simplemr.MR{}
	if runpFlags.interactive {
		// Run one mapper at a time.
		mr.NumMappers = 1
		sort.Sort(keys)
	}
	in, out := make(chan *simplemr.Record, len(mapInputs)), make(chan *simplemr.Record, len(mapInputs))
	sigch := make(chan os.Signal)
	signal.Notify(sigch, os.Interrupt)
	go func() { <-sigch; mr.Cancel() }()
	go mr.Run(in, out, runner, runner)
	for _, key := range keys {
		in <- &simplemr.Record{string(key), []interface{}{mapInputs[key]}}
	}
	close(in)
	<-out
	return mr.Error()
}

func runRunp(jirix *jiri.X, args []string) error {
	return runp(jirix, cmdRunP, args)
}