ReferenceIter: use a pipeline

Allow iteration to be canceled using a context.

Author: mhagger

git/ref_iter.go

@@ -2,71 +2,89 @@ package git
 
 import (
 	"bufio"
+	"context"
+	"fmt"
 	"io"
-	"os"
-	"os/exec"
+
+	"github.com/github/git-sizer/internal/pipe"
 )
 
 // ReferenceIter is an iterator that interates over references.
 type ReferenceIter struct {
-	cmd     *exec.Cmd
-	out     io.ReadCloser
-	f       *bufio.Reader
-	errChan <-chan error
+	refCh chan Reference
+	errCh chan error
 }
 
 // NewReferenceIter returns an iterator that iterates over all of the
 // references in `repo`.
-func (repo *Repository) NewReferenceIter() (*ReferenceIter, error) {
-	cmd := repo.GitCommand(
-		"for-each-ref", "--format=%(objectname) %(objecttype) %(objectsize) %(refname)",
-	)
-
-	out, err := cmd.StdoutPipe()
-	if err != nil {
-		return nil, err
+func (repo *Repository) NewReferenceIter(ctx context.Context) (*ReferenceIter, error) {
+	iter := ReferenceIter{
+		refCh: make(chan Reference),
+		errCh: make(chan error),
 	}
 
-	cmd.Stderr = os.Stderr
+	p := pipe.New()
+	p.Add(
+		// Output all references and their values:
+		pipe.CommandStage(
+			"git-for-each-ref",
+			repo.GitCommand(
+				"for-each-ref",
+				"--format=%(objectname) %(objecttype) %(objectsize) %(refname)",
+			),
+		),
+
+		// Read the references and send them to `iter.refCh`, then close
+		// the channel.
+		pipe.Function(
+			"parse-refs",
+			func(ctx context.Context, env pipe.Env, stdin io.Reader, stdout io.Writer) error {
+				defer close(iter.refCh)
+
+				in := bufio.NewReader(stdin)
+				for {
+					line, err := in.ReadBytes('\n')
+					if err != nil {
+						if err == io.EOF {
+							return nil
+						}
+						return fmt.Errorf("reading 'git for-each-ref' output: %w", err)
+					}
+
+					ref, err := ParseReference(string(line[:len(line)-1]))
+					if err != nil {
+						return fmt.Errorf("parsing 'git for-each-ref' output: %w", err)
+					}
+					select {
+					case iter.refCh <- ref:
+					case <-ctx.Done():
+						return ctx.Err()
+					}
+				}
+			},
+		),
+	)
 
-	err = cmd.Start()
+	err := p.Start(ctx)
 	if err != nil {
 		return nil, err
 	}
 
-	return &ReferenceIter{
-		cmd:     cmd,
-		out:     out,
-		f:       bufio.NewReader(out),
-		errChan: make(chan error, 1),
-	}, nil
+	go func() {
+		iter.errCh <- p.Wait()
+	}()
+
+	return &iter, nil
 }
 
 // Next returns either the next reference or a boolean `false` value
 // indicating that the iteration is over. On errors, return an error
 // (in this case, the caller must still call `Close()`).
 func (iter *ReferenceIter) Next() (Reference, bool, error) {
-	line, err := iter.f.ReadString('\n')
-	if err != nil {
-		if err != io.EOF {
-			return Reference{}, false, err
-		}
-		return Reference{}, false, nil
-	}
-	ref, err := ParseReference(line[:len(line)-1])
-	if err != nil {
-		return ref, false, err
+	ref, ok := <-iter.refCh
+	if !ok {
+		return Reference{}, false, <-iter.errCh
 	}
 
 	return ref, true, nil
 }
-
-// Close closes the iterator and frees up resources.
-func (iter *ReferenceIter) Close() error {
-	err := iter.out.Close()
-	err2 := iter.cmd.Wait()
-	if err == nil {
-		err = err2
-	}
-	return err
-}

sizes/graph.go

@@ -2,6 +2,7 @@ package sizes
 
 import (
 	"bufio"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -69,15 +70,13 @@ func ScanRepositoryUsingGraph(
 ) (HistorySize, error) {
 	graph := NewGraph(rg, nameStyle)
 
-	refIter, err := repo.NewReferenceIter()
+	ctx, cancel := context.WithCancel(context.TODO())
+	defer cancel()
+
+	refIter, err := repo.NewReferenceIter(ctx)
 	if err != nil {
 		return HistorySize{}, err
 	}
-	defer func() {
-		if refIter != nil {
-			refIter.Close()
-		}
-	}()
 
 	iter, in, err := repo.NewObjectIter("--stdin", "--date-order")
 	if err != nil {
@@ -134,8 +133,6 @@ func ScanRepositoryUsingGraph(
 				return
 			}
 		}
-		err := refIter.Close()
-		refIter = nil
 		errChan <- err
 	}()