From 10fc1d52893bc3ffe477b675835048c99be85a6b Mon Sep 17 00:00:00 2001
From: Michael Woolnough <mw31@sanger.ac.uk>
Date: Wed, 27 Nov 2024 09:02:13 +0000
Subject: [PATCH] Buffer output files and avoid allocation during path quoting.

---
 walk/file.go      | 39 +++++++++++++++++++++++++--------------
 walk/walk_test.go |  9 ++++++---
 2 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/walk/file.go b/walk/file.go
index 397b4a00..37736076 100644
--- a/walk/file.go
+++ b/walk/file.go
@@ -26,6 +26,7 @@
 package walk
 
 import (
+	"bufio"
 	"fmt"
 	"io"
 	"os"
@@ -45,10 +46,23 @@ func (e *WriteError) Error() string { return e.Err.Error() }
 
 func (e *WriteError) Unwrap() error { return e.Err }
 
+type bufferedFile struct {
+	bufio.Writer
+	io.Closer
+}
+
+func (b *bufferedFile) Close() error {
+	if err := b.Writer.Flush(); err != nil {
+		return err
+	}
+
+	return b.Closer.Close()
+}
+
 // Files represents a collection of output files that can be written to in a
 // round-robin.
 type Files struct {
-	files    []*os.File
+	files    []bufferedFile
 	Paths    []string
 	filesI   int
 	filesMax int
@@ -70,19 +84,20 @@ func NewFiles(outDir string, n int) (*Files, error) {
 		return nil, err
 	}
 
-	files := make([]*os.File, n)
+	files := make([]bufferedFile, n)
 	outPaths := make([]string, n)
 
 	for i := range files {
-		var err error
-
 		path := filepath.Join(outDir, fmt.Sprintf("walk.%d", i+1))
 
-		files[i], err = os.Create(path)
+		file, err := os.Create(path)
 		if err != nil {
 			return nil, err
 		}
 
+		files[i].Reset(file)
+		files[i].Closer = file
+
 		outPaths[i] = path
 	}
 
@@ -100,15 +115,16 @@ func NewFiles(outDir string, n int) (*Files, error) {
 //
 // It will terminate the walk if writes to our output files fail.
 func (f *Files) WritePaths() PathCallback {
+	var quoted [10240]byte
+
 	return func(entry *Dirent) error {
-		return f.writePath(strconv.Quote(entry.Path))
+		return f.writePath(append(strconv.AppendQuote(quoted[:0], entry.Path), '\n'))
 	}
 }
 
 // writePath is a thread-safe way of writing the given path to our next output
 // file. Returns a WriteError on failure to write to an output file.
-func (f *Files) writePath(path string) error {
-	f.mu.Lock()
+func (f *Files) writePath(path []byte) error {
 	i := f.filesI
 	f.filesI++
 
@@ -116,12 +132,7 @@ func (f *Files) writePath(path string) error {
 		f.filesI = 0
 	}
 
-	f.mu.Unlock()
-
-	f.mus[i].Lock()
-	defer f.mus[i].Unlock()
-
-	_, err := io.WriteString(f.files[i], path+"\n")
+	_, err := f.files[i].Write(path)
 	if err != nil {
 		err = &WriteError{Err: err}
 	}
diff --git a/walk/walk_test.go b/walk/walk_test.go
index e19e5611..4d8af362 100644
--- a/walk/walk_test.go
+++ b/walk/walk_test.go
@@ -75,6 +75,9 @@ func TestWalk(t *testing.T) {
 			err = w.Walk(walkDir, cb)
 			So(err, ShouldBeNil)
 
+			err = files.Close()
+			So(err, ShouldBeNil)
+
 			splitExpected := make([][]string, n)
 			splitI := 0
 
@@ -105,9 +108,6 @@ func TestWalk(t *testing.T) {
 
 			So(len(walkErrors), ShouldEqual, 0)
 
-			err = files.Close()
-			So(err, ShouldBeNil)
-
 			err = files.files[0].Close()
 			So(err, ShouldNotBeNil)
 
@@ -334,6 +334,9 @@ func testOutputToFiles(includDirs, ignoreSymlinks bool, walkDir, outDir string,
 	err = w.Walk(walkDir, cb)
 	So(err, ShouldBeNil)
 
+	err = files.Close()
+	So(err, ShouldBeNil)
+
 	outPath := filepath.Join(outDir, "walk.1")
 	So(files.Paths[0], ShouldEqual, outPath)
 	content, err := os.ReadFile(outPath)