Repository: openllb/hlb
Pull request #305 (2022-02-17 22:51:00)

Refactor InputSteerer to pkg/steer and add tests
- Also relaxed requirement to `io.WriteCloser` but would cast to `*io.PipeWriter` and `CloseWithError` if available.


Patch:
diff --git a/cmd/hlb/command/run.go b/cmd/hlb/command/run.go
index 732e6b3..b2a44c9 100644
--- a/cmd/hlb/command/run.go
+++ b/cmd/hlb/command/run.go
@@ -23,6 +23,7 @@ import (
 	"github.com/openllb/hlb/local"
 	"github.com/openllb/hlb/parser"
 	"github.com/openllb/hlb/pkg/filebuffer"
+	"github.com/openllb/hlb/pkg/steer"
 	"github.com/openllb/hlb/solver"
 	cli "github.com/urfave/cli/v2"
 	"github.com/xlab/treeprint"
@@ -94,18 +95,18 @@ var runCommand = &cli.Command{
 			DefaultPlatform:  c.String("platform"),
 		}
 
-		var inputSteerer *codegen.InputSteerer
+		var inputSteerer *steer.InputSteerer
 		if c.Bool("debug") {
 			pr, pw := io.Pipe()
 			r := bufio.NewReader(pr)
-			inputSteerer = codegen.NewInputSteerer(os.Stdin, pw)
+			inputSteerer = steer.NewInputSteerer(os.Stdin, pw)
 
 			ri.Debugger = codegen.NewDebugger(cln, os.Stderr, inputSteerer, r)
 		}
 
 		if c.Bool("shell-on-error") {
 			if inputSteerer == nil {
-				inputSteerer = codegen.NewInputSteerer(os.Stdin)
+				inputSteerer = steer.NewInputSteerer(os.Stdin)
 			}
 
 			ri.SolveErrorHandler = func(ctx context.Context, c gateway.Client, err error) {
diff --git a/codegen/debug.go b/codegen/debug.go
index e10e573..5241aea 100644
--- a/codegen/debug.go
+++ b/codegen/debug.go
@@ -8,7 +8,6 @@ import (
 	"os"
 	"os/exec"
 	"strings"
-	"sync"
 
 	"github.com/docker/buildx/util/progress"
 	shellquote "github.com/kballard/go-shellquote"
@@ -24,6 +23,7 @@ import (
 	"github.com/openllb/hlb/parser"
 	"github.com/openllb/hlb/parser/ast"
 	"github.com/openllb/hlb/pkg/llbutil"
+	"github.com/openllb/hlb/pkg/steer"
 	"github.com/openllb/hlb/solver"
 	"github.com/pkg/errors"
 	"golang.org/x/crypto/ssh/terminal"
@@ -56,7 +56,7 @@ func (s snapshot) fs() (Filesystem, error) {
 	return s.val.Filesystem()
 }
 
-func NewDebugger(c *client.Client, w io.Writer, inputSteerer *InputSteerer, promptReader *bufio.Reader) Debugger {
+func NewDebugger(c *client.Client, w io.Writer, inputSteerer *steer.InputSteerer, promptReader *bufio.Reader) Debugger {
 	var (
 		fd                *ast.FuncDecl
 		next              *ast.FuncDecl
@@ -872,63 +872,3 @@ type nopWriteCloser struct {
 func (w *nopWriteCloser) Close() error {
 	return nil
 }
-
-// InputSteerer is a mechanism for directing input to one of a set of
-// Readers. This is used when the debugger runs an exec: we can't
-// interrupt a Read from the exec context, so if we naively passed the
-// primary reader into the exec, it would swallow the next debugger
-// command after the exec session ends. To work around this, have a
-// goroutine which continuously reads from the input, and steers data
-// into the appropriate pipe depending whether we have an exec session
-// active.
-type InputSteerer struct {
-	mu  sync.Mutex
-	pws []*io.PipeWriter
-}
-
-func NewInputSteerer(inputReader io.Reader, pws ...*io.PipeWriter) *InputSteerer {
-	is := &InputSteerer{
-		pws: pws,
-	}
-
-	go func() {
-		var p [4096]byte
-		for {
-			n, err := inputReader.Read(p[:])
-			var pw *io.PipeWriter
-			is.mu.Lock()
-			if len(is.pws) != 0 {
-				pw = is.pws[len(is.pws)-1]
-			}
-			is.mu.Unlock()
-			if n != 0 && pw != nil {
-				pw.Write(p[:n])
-			}
-			if err != nil {
-				is.mu.Lock()
-				defer is.mu.Unlock()
-				for _, pw := range is.pws {
-					pw.CloseWithError(err)
-				}
-				return
-			}
-		}
-	}()
-	return is
-}
-
-// Push pushes a new pipe to steer input to, until Pop is called to steer it
-// back to the previous pipe.
-func (is *InputSteerer) Push(pw *io.PipeWriter) {
-	is.mu.Lock()
-	defer is.mu.Unlock()
-	is.pws = append(is.pws, pw)
-}
-
-// Pop causes future input to be directed to the pipe where it was going before
-// the last call to Push.
-func (is *InputSteerer) Pop() {
-	is.mu.Lock()
-	defer is.mu.Unlock()
-	is.pws = is.pws[:len(is.pws)-1]
-}
diff --git a/pkg/steer/input_steerer.go b/pkg/steer/input_steerer.go
new file mode 100644
index 0000000..f61d809
--- /dev/null
+++ b/pkg/steer/input_steerer.go
@@ -0,0 +1,68 @@
+package steer
+
+import (
+	"io"
+	"sync"
+)
+
+// InputSteerer is a mechanism for directing input to one of a set of
+// Readers. This is used when the debugger runs an exec: we can't
+// interrupt a Read from the exec context, so if we naively passed the
+// primary reader into the exec, it would swallow the next debugger
+// command after the exec session ends. To work around this, have a
+// goroutine which continuously reads from the input, and steers data
+// into the appropriate writer depending whether we have an exec session
+// active.
+type InputSteerer struct {
+	mu sync.Mutex
+	ws []io.WriteCloser
+}
+
+func NewInputSteerer(inputReader io.Reader, ws ...io.WriteCloser) *InputSteerer {
+	is := &InputSteerer{ws: ws}
+
+	go func() {
+		var p [4096]byte
+		for {
+			n, err := inputReader.Read(p[:])
+			var w io.WriteCloser
+			is.mu.Lock()
+			if len(is.ws) != 0 {
+				w = is.ws[len(is.ws)-1]
+			}
+			is.mu.Unlock()
+			if n != 0 && w != nil {
+				w.Write(p[:n])
+			}
+			if err != nil {
+				is.mu.Lock()
+				defer is.mu.Unlock()
+				for _, w := range is.ws {
+					if pw, ok := w.(*io.PipeWriter); ok {
+						pw.CloseWithError(err)
+					} else {
+						w.Close()
+					}
+				}
+				return
+			}
+		}
+	}()
+	return is
+}
+
+// Push pushes a new writer to steer input to, until Pop is called to steer it
+// back to the previous writer.
+func (is *InputSteerer) Push(w io.WriteCloser) {
+	is.mu.Lock()
+	defer is.mu.Unlock()
+	is.ws = append(is.ws, w)
+}
+
+// Pop causes future input to be directed to the writer where it was going before
+// the last call to Push.
+func (is *InputSteerer) Pop() {
+	is.mu.Lock()
+	defer is.mu.Unlock()
+	is.ws = is.ws[:len(is.ws)-1]
+}


Test patch:
diff --git a/pkg/steer/input_steerer_test.go b/pkg/steer/input_steerer_test.go
new file mode 100644
index 0000000..9e2677b
--- /dev/null
+++ b/pkg/steer/input_steerer_test.go
@@ -0,0 +1,61 @@
+package steer
+
+import (
+	"io"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+type fixedReader struct {
+	io.Reader
+}
+
+func (fr *fixedReader) Read(p []byte) (n int, err error) {
+	b := make([]byte, 1)
+	n, err = fr.Reader.Read(b)
+	copy(p, b)
+	return
+}
+
+func TestInputSteerer(t *testing.T) {
+	r := &fixedReader{
+		Reader: strings.NewReader("abc"),
+	}
+
+	pr, pw := io.Pipe()
+	is := NewInputSteerer(r, pw)
+
+	p := make([]byte, 1)
+	n, err := pr.Read(p)
+	require.NoError(t, err)
+	require.Equal(t, 1, n)
+	require.Equal(t, "a", string(p[:n]))
+
+	pr2, pw2 := io.Pipe()
+	is.Push(pw2)
+
+	// A new pipe writer was pushed, so reading from the previous pipe reader
+	// should block until it is popped off.
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		p = make([]byte, 1)
+		n, err = pr.Read(p)
+		require.NoError(t, err)
+		require.Equal(t, 1, n)
+	}()
+
+	p2 := make([]byte, 1)
+	n, err = pr2.Read(p2)
+	require.NoError(t, err)
+	require.Equal(t, 1, n)
+	require.Equal(t, "b", string(p2[:n]))
+
+	// After popping, the value read should be after what the popped off pipe
+	// reader read.
+	is.Pop()
+	<-done
+	require.Equal(t, "c", string(p[:n]))
+}