Skip to content

Commit

Permalink
Buffer output files and avoid allocation during path quoting.
Browse files Browse the repository at this point in the history
  • Loading branch information
mjkw31 committed Nov 27, 2024
1 parent de37f57 commit 10fc1d5
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 17 deletions.
39 changes: 25 additions & 14 deletions walk/file.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
package walk

import (
"bufio"
"fmt"
"io"
"os"
Expand All @@ -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
Expand All @@ -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
}

Expand All @@ -100,28 +115,24 @@ 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++

if f.filesI == f.filesMax {
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}
}
Expand Down
9 changes: 6 additions & 3 deletions walk/walk_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 10fc1d5

Please sign in to comment.