walk/git: split ls-files and stat paths concurrently

walk/git.go

@@ -8,15 +8,22 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"runtime"
 	"strconv"
 	"strings"
+	"sync"
 
 	"github.com/charmbracelet/log"
 	"github.com/numtide/treefmt/v2/git"
 	"github.com/numtide/treefmt/v2/stats"
 	"golang.org/x/sync/errgroup"
 )
 
+type gitEntry struct {
+	relative string
+	gitlink  bool // mode 160000, i.e. a submodule
+}
+
 type GitReader struct {
 	root string
 	path string
@@ -25,101 +32,29 @@ type GitReader struct {
 	stats *stats.Stats
 
 	eg      *errgroup.Group
-	scanner *bufio.Scanner
+	filesCh chan *File
 }
 
 func (g *GitReader) Read(ctx context.Context, files []*File) (n int, err error) {
-	// ensure we record how many files we traversed
 	defer func() {
 		g.stats.Add(stats.Traversed, n)
 	}()
 
-	nextFile := func() (string, error) {
-		for line := g.scanner.Text(); len(line) > 0; line = g.scanner.Text() {
-			lineSplit := strings.Split(line, "\t")
-
-			var stage, file string
-			// Untracked files just show as `<filename>`, while tracked files show as `<mode> <object> <stage><TAB><file>`
-			if len(lineSplit) == 1 {
-				stage, file = "", lineSplit[0]
-			} else {
-				stage, file = lineSplit[0], lineSplit[1]
-			}
-
-			// 160000 is the mode for submodules, skip them because they are separate projects that may have their own
-			// formatting rules
-			if strings.HasPrefix(stage, "160000") {
-				g.scanner.Scan()
-
-				continue
-			}
-
-			if file[0] != '"' {
-				return file, nil
-			}
-
-			unquoted, err := strconv.Unquote(file)
-			if err != nil {
-				return "", fmt.Errorf("failed to unquote file %s: %w", file, err)
-			}
-
-			return unquoted, nil
-		}
-
-		return "", io.EOF
-	}
-
 LOOP:
-
 	for n < len(files) {
 		select {
-		// exit early if the context was cancelled
 		case <-ctx.Done():
 			return n, ctx.Err() //nolint:wrapcheck
 
-		default:
-			// read the next file
-			if g.scanner.Scan() {
-				entry, err := nextFile()
-				if err != nil {
-					return n, err
-				}
-
-				path := filepath.Join(g.root, g.path, entry)
-
-				g.log.Debugf("processing file: %s", path)
-
-				info, err := os.Lstat(path)
-
-				switch {
-				case os.IsNotExist(err):
-					// the underlying file might have been removed
-					g.log.Warnf(
-						"Path %s is in the worktree but appears to have been removed from the filesystem", path,
-					)
-
-					continue
-				case err != nil:
-					return n, fmt.Errorf("failed to stat %s: %w", path, err)
-				case info.Mode()&os.ModeSymlink == os.ModeSymlink:
-					// we skip reporting symlinks stored in Git, they should
-					// point to local files which we would list anyway.
-					continue
-				}
-
-				files[n] = &File{
-					Path:    path,
-					RelPath: filepath.Join(g.path, entry),
-					Info:    info,
-				}
-
-				n++
-			} else {
-				// nothing more to read
+		case file, ok := <-g.filesCh:
+			if !ok {
 				err = io.EOF
 
 				break LOOP
 			}
+
+			files[n] = file
+			n++
 		}
 	}
 
@@ -135,12 +70,114 @@ func (g *GitReader) Close() error {
 	return nil
 }
 
+func (g *GitReader) stat(entry gitEntry) {
+	if entry.gitlink {
+		// submodules are separate projects with their own formatting rules
+		return
+	}
+
+	path := filepath.Join(g.root, entry.relative)
+
+	g.log.Debugf("processing file: %s", path)
+
+	info, err := os.Lstat(path)
+
+	switch {
+	case os.IsNotExist(err):
+		g.log.Warnf(
+			"Path %s is in the worktree but appears to have been removed from the filesystem", path,
+		)
+
+		return
+	case err != nil:
+		g.log.Errorf("failed to stat %s: %v", path, err)
+
+		return
+	case info.Mode()&os.ModeSymlink == os.ModeSymlink:
+		// symlinks point at files we list anyway
+		return
+	}
+
+	g.filesCh <- &File{
+		Path:    path,
+		RelPath: entry.relative,
+		Info:    info,
+	}
+}
+
+func lsFiles(dir string, staged bool, prefix string, out chan<- gitEntry, args ...string) error {
+	//nolint:gosec // args are fixed flag sets assembled in NewGitReader, not user input.
+	cmd := exec.CommandContext(context.Background(), "git", append([]string{"ls-files"}, args...)...)
+	cmd.Dir = dir
+
+	stdout, err := cmd.StdoutPipe()
+	if err != nil {
+		return fmt.Errorf("failed to create stdout pipe: %w", err)
+	}
+
+	if err := cmd.Start(); err != nil {
+		return fmt.Errorf("failed to start git ls-files: %w", err)
+	}
+
+	scanErr := scanLsFiles(stdout, staged, prefix, out)
+
+	// Always reap the child. If scanning aborted early git may be blocked on a
+	// full pipe, so kill it first to guarantee Wait returns.
+	if scanErr != nil {
+		_ = cmd.Process.Kill()
+	}
+
+	if err := cmd.Wait(); err != nil && scanErr == nil {
+		return fmt.Errorf("git ls-files failed: %w", err)
+	}
+
+	return scanErr
+}
+
+func scanLsFiles(r io.Reader, staged bool, prefix string, out chan<- gitEntry) error {
+	scanner := bufio.NewScanner(r)
+	for scanner.Scan() {
+		line := scanner.Text()
+
+		var gitlink bool
+
+		path := line
+		if staged {
+			// <mode> <object> <stage>\t<file>
+			if mode, file, ok := strings.Cut(line, "\t"); ok {
+				gitlink = strings.HasPrefix(mode, "160000")
+				path = file
+			}
+		}
+
+		if path == "" {
+			continue
+		}
+
+		if path[0] == '"' {
+			unquoted, err := strconv.Unquote(path)
+			if err != nil {
+				return fmt.Errorf("failed to unquote file %s: %w", path, err)
+			}
+
+			path = unquoted
+		}
+
+		out <- gitEntry{relative: filepath.Join(prefix, path), gitlink: gitlink}
+	}
+
+	if err := scanner.Err(); err != nil {
+		return fmt.Errorf("failed to read git ls-files output: %w", err)
+	}
+
+	return nil
+}
+
 func NewGitReader(
 	root string,
 	path string,
 	statz *stats.Stats,
 ) (*GitReader, error) {
-	// check if the root is a git repository
 	isGit, err := git.IsInsideWorktree(root)
 	if err != nil {
 		return nil, fmt.Errorf("failed to check if %s is a git repository: %w", root, err)
@@ -150,34 +187,69 @@ func NewGitReader(
 		return nil, fmt.Errorf("%s is not a git repository", root)
 	}
 
-	// create an errgroup for executing in the background
-	eg := &errgroup.Group{}
+	dir := filepath.Join(root, path)
+
+	g := &GitReader{
+		root:    root,
+		path:    path,
+		stats:   statz,
+		log:     log.WithPrefix("walk | git"),
+		eg:      &errgroup.Group{},
+		filesCh: make(chan *File, BatchSize*runtime.NumCPU()),
+	}
+
+	entries := make(chan gitEntry, BatchSize)
+
+	// `--cached` and `--others` are queried separately because git buffers all
+	// output until the untracked scan finishes when both are combined; the
+	// index-only query streams immediately so formatters start without waiting.
+	var producers sync.WaitGroup
 
-	// create a pipe to capture the command output
-	r, w := io.Pipe()
+	producers.Add(2)
 
-	// create a command which will execute from the specified sub path within root
-	cmd := exec.CommandContext(
-		context.Background(),
-		"git", "ls-files", "--cached", "--others", "--exclude-standard", "--stage",
-	)
-	cmd.Dir = filepath.Join(root, path)
-	cmd.Stdout = w
+	g.eg.Go(func() error {
+		defer producers.Done()
 
-	// execute the command in the background
-	eg.Go(func() error {
-		return w.CloseWithError(cmd.Run())
+		return lsFiles(dir, true, path, entries, "--cached", "--stage")
 	})
 
-	// create a new scanner for reading the output
-	scanner := bufio.NewScanner(r)
+	g.eg.Go(func() error {
+		defer producers.Done()
 
-	return &GitReader{
-		eg:      eg,
-		root:    root,
-		path:    path,
-		stats:   statz,
-		scanner: scanner,
-		log:     log.WithPrefix("walk | git"),
-	}, nil
+		return lsFiles(dir, false, path, entries, "--others", "--exclude-standard")
+	})
+
+	g.eg.Go(func() error {
+		producers.Wait()
+		close(entries)
+
+		return nil
+	})
+
+	var workers sync.WaitGroup
+
+	statWorkers := runtime.GOMAXPROCS(0)
+
+	workers.Add(statWorkers)
+
+	for range statWorkers {
+		g.eg.Go(func() error {
+			defer workers.Done()
+
+			for e := range entries {
+				g.stat(e)
+			}
+
+			return nil
+		})
+	}
+
+	g.eg.Go(func() error {
+		workers.Wait()
+		close(g.filesCh)
+
+		return nil
+	})
+
+	return g, nil
 }