Merge pull request #2 from yuya-takeyama/parallel

Add -p (--parallels) option to set parallel degree of execution

Author: yuya-takeyama

README.md

@@ -11,6 +11,12 @@ foo bar baz
 foo bar baz
 ```
 
+### Set parallel degree of execution (-p)
+
+```
+$ ntimes 10 -p 3 -- sh -c 'echo "Hi!"; sleep 1; echo "Bye"'
+```
+
 ## Author
 
 Yuya Takeyama

main.go

@@ -1,27 +1,30 @@
 package main
 
 import (
+	"bytes"
 	"fmt"
 	"io"
 	"os"
 	"os/exec"
 	"strconv"
+	"sync"
 
 	flags "github.com/jessevdk/go-flags"
 )
 
 const AppName = "ntimes"
 
-type Options struct {
+type options struct {
+	Parallels   int  `short:"p" long:"parallels" description:"Parallel degree of execution" default:"1"`
 	ShowVersion bool `short:"v" long:"version" description:"Show version"`
 }
 
-var opts Options
+var opts options
 
 func main() {
 	parser := flags.NewParser(&opts, flags.Default^flags.PrintErrors)
 	parser.Name = AppName
-	parser.Usage = "[OPTIONS] -- COMMAND"
+	parser.Usage = "N [OPTIONS] -- COMMAND"
 
 	args, err := parser.Parse()
 
@@ -39,23 +42,61 @@ func main() {
 	cmdName := args[1]
 	cmdArgs := args[2:]
 
+	stdoutCh := make(chan io.ReadWriter)
+	exitCh := make(chan bool)
+
 	if err != nil {
 		panic(err)
 	}
 
-	ntimes(cnt, cmdName, cmdArgs, os.Stdin, os.Stdout, os.Stderr)
+	go printer(os.Stdout, stdoutCh, exitCh)
+
+	ntimes(cnt, cmdName, cmdArgs, os.Stdin, os.Stdout, os.Stderr, stdoutCh, opts.Parallels)
+
+	exitCh <- true
 }
 
-func ntimes(cnt int, cmdName string, cmdArgs []string, stdin io.Reader, stdout io.Writer, stderr io.Writer) {
+func ntimes(cnt int, cmdName string, cmdArgs []string, stdin io.Reader, stdout io.Writer, stderr io.Writer, stdoutCh chan io.ReadWriter, parallels int) {
+	var wg sync.WaitGroup
+
+	sema := make(chan bool, parallels)
+
 	for i := 0; i < cnt; i++ {
-		cmd := exec.Command(cmdName, cmdArgs...)
-		cmd.Stdin = stdin
-		cmd.Stdout = stdout
-		cmd.Stderr = stderr
-		err := cmd.Run()
-
-		if err != nil {
-			panic(err)
+		wg.Add(1)
+		go func() {
+			sema <- true
+
+			defer func() {
+				wg.Done()
+				<-sema
+			}()
+
+			stdoutBuffer := new(bytes.Buffer)
+
+			cmd := exec.Command(cmdName, cmdArgs...)
+			cmd.Stdin = stdin
+			cmd.Stdout = stdoutBuffer
+			cmd.Stderr = stderr
+			err := cmd.Run()
+
+			if err != nil {
+				panic(err)
+			}
+
+			stdoutCh <- stdoutBuffer
+		}()
+	}
+
+	wg.Wait()
+}
+
+func printer(stdout io.Writer, stdoutCh chan io.ReadWriter, exitCh chan bool) {
+	for {
+		select {
+		case r := <-stdoutCh:
+			io.Copy(stdout, r)
+		case <-exitCh:
+			return
 		}
 	}
 }

main_test.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"bytes"
+	"io"
 	"testing"
 )
 
@@ -10,7 +11,14 @@ func TestNtimes(t *testing.T) {
 	stdout := new(bytes.Buffer)
 	stderr := new(bytes.Buffer)
 
-	ntimes(3, "echo", []string{"foo", "bar", "baz"}, stdin, stdout, stderr)
+	stdoutCh := make(chan io.ReadWriter)
+	exitCh := make(chan bool)
+
+	go printer(stdout, stdoutCh, exitCh)
+
+	ntimes(3, "echo", []string{"foo", "bar", "baz"}, stdin, stdout, stderr, stdoutCh, 1)
+
+	exitCh <- true
 
 	expected := "foo bar baz\nfoo bar baz\nfoo bar baz\n"
 	actual := stdout.String()