Fix -stdin being blocked after a restart until Enter is pressed.

Closes #22.

os.Stdin.Read() blocks until a newline `\n` is provided (due to the OS doing
line-by-line buffering https://groups.google.com/g/golang-nuts/c/NvD-pOTASIk).
Therefore, do not tie the lifecycle of a command with os.Stdin.Read(). Instead,
read from os.Stdin continuously in a background goroutine and feed its data
into the stdinPipe of whichever exec.Cmd is currently running.

Author: bokwoon95

wgo_cmd.go

@@ -111,8 +111,7 @@ type WgoCmd struct {
 	// EnableStdin controls whether the Stdin field is used.
 	EnableStdin bool
 
-	// Stdin is where the last command gets its stdin input from (EnableStdin
-	// must be true).
+	// Stdin is where commands get stdin input from (EnableStdin must be true).
 	Stdin io.Reader
 
 	// Stdout is where the commands write their stdout output.
@@ -425,6 +424,31 @@ func (wgoCmd *WgoCmd) Run() error {
 	}
 	defer timer.Stop()
 
+	// Start a background job that continuously drains data from os.Stdin and
+	// feeds it into stdinPipe (which connected to an exec.Cmd). stdinPipe can
+	// be swapped out anytime when the exec.Cmd changes, so access is guarded
+	// by a mutex.
+	var stdinPipe io.WriteCloser
+	var stdinPipeMutex sync.Mutex
+	if wgoCmd.EnableStdin {
+		go func() {
+			p := make([]byte, 4096)
+			for {
+				n, err := wgoCmd.Stdin.Read(p)
+				if n > 0 {
+					stdinPipeMutex.Lock()
+					if stdinPipe != nil {
+						_, _ = stdinPipe.Write(p[:n])
+					}
+					stdinPipeMutex.Unlock()
+				}
+				if err != nil {
+					break
+				}
+			}
+		}()
+	}
+
 	for restartCount := 0; ; restartCount++ {
 	CMD_CHAIN:
 		for i, args := range wgoCmd.ArgsList {
@@ -494,26 +518,19 @@ func (wgoCmd *WgoCmd) Run() error {
 				}
 			}
 			// If the user enabled it, feed wgoCmd.Stdin to the command's
-			// Stdin. Only the last command gets to read from Stdin -- if we
-			// give Stdin to every command in the middle it will prevent the
-			// next command from being executed if they don't consume Stdin.
+			// Stdin.
 			//
 			// We have to use cmd.StdinPipe() here instead of assigning
 			// cmd.Stdin directly, otherwise `wgo run ./testdata/stdin` doesn't
 			// work interactively (the tests will pass, but somehow it won't
 			// actually work if you run it in person. I don't know why).
-			var wg sync.WaitGroup
-			if wgoCmd.EnableStdin && i == len(wgoCmd.ArgsList)-1 {
-				stdinPipe, err := cmd.StdinPipe()
+			if wgoCmd.EnableStdin {
+				stdinPipeMutex.Lock()
+				stdinPipe, err = cmd.StdinPipe()
+				stdinPipeMutex.Unlock()
 				if err != nil {
 					return err
 				}
-				wg.Add(1)
-				go func() {
-					defer wg.Done()
-					defer stdinPipe.Close()
-					_, _ = io.Copy(stdinPipe, wgoCmd.Stdin)
-				}()
 			}
 
 			// Step 2: Run the command in the background.
@@ -525,9 +542,8 @@ func (wgoCmd *WgoCmd) Run() error {
 				return err
 			}
 			go func() {
-				wg.Wait()
+				defer close(waitDone)
 				cmdResult <- cmd.Wait()
-				close(waitDone)
 			}()
 
 			// Step 3: Wait for events in the event loop.

wgo_cmd_test.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"errors"
 	"flag"
+	"io"
 	"log"
 	"math/rand"
 	"os"
@@ -862,16 +863,37 @@ func TestStdin(t *testing.T) {
 	t.Parallel()
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
-	wgoCmd, err := WgoCommand(ctx, 0, []string{"run", "-exit", "-dir", "testdata/stdin", "-stdin", "./testdata/stdin"})
+	// Precompile ./testdata/stdin executable so that wgoCmd.Run() isn't
+	// affected by Go compile times. We need to command to run _immediately_ so
+	// that it starts reading from stdin immediately.
+	executablePath := filepath.Join(os.TempDir(), t.Name()+".exe")
+	err := exec.Command("go", "build", "-o", executablePath, "./testdata/stdin").Run()
 	if err != nil {
 		t.Fatal(err)
 	}
-	wgoCmd.Stdin = strings.NewReader("foo\nbar\nbaz")
+	defer os.Remove(executablePath)
+	wgoCmd, err := WgoCommand(ctx, 0, []string{"-dir", "testdata/stdin", "-stdin", executablePath})
+	if err != nil {
+		t.Fatal(err)
+	}
+	// This is so awfully flaky, but we probably need to wait like 5 seconds
+	// before writing to stdin in case the command hasn't started. Just in
+	// case, idk, the test runner is running on a pentium cpu or something and
+	// causes tests to fail because it's too slow.
+	pipeReader, pipeWriter := io.Pipe()
+	wgoCmd.Stdin = pipeReader
+	go func() {
+		time.Sleep(5 * time.Second)
+		pipeWriter.Write([]byte("foo\nbar\nbaz\n"))
+		pipeWriter.Close()
+	}()
 	buf := &Buffer{}
 	wgoCmd.Stderr = buf
 	err = wgoCmd.Run()
 	if err != nil {
-		t.Fatal(err)
+		if !errors.Is(err, context.DeadlineExceeded) {
+			t.Fatal(err)
+		}
 	}
 	got := strings.TrimSpace(buf.String())
 	want := "1: foo\n2: bar\n3: baz"