Rewrite walk to use less RAM

go.mod

@@ -13,7 +13,6 @@ require (
 	github.com/smartystreets/goconvey v1.7.2
 	github.com/spf13/cobra v1.8.1
 	github.com/termie/go-shutil v0.0.0-20140729215957-bcacb06fecae
-	github.com/wtsi-hgi/godirwalk v1.18.1
 	github.com/wtsi-ssg/wr v0.5.9
 )
 

go.sum

@@ -284,8 +284,6 @@ github.com/tklauser/numcpus v0.9.0 h1:lmyCHtANi8aRUgkckBgoDk1nHCux3n2cgkJLXdQGPD
 github.com/tklauser/numcpus v0.9.0/go.mod h1:SN6Nq1O3VychhC1npsWostA+oW+VOQTxZrS604NSRyI=
 github.com/ugorji/go/codec v1.2.12 h1:9LC83zGrHhuUA9l16C9AHXAqEV/2wBQ4nkvumAE65EE=
 github.com/ugorji/go/codec v1.2.12/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg=
-github.com/wtsi-hgi/godirwalk v1.18.1 h1:t7eaGXYBfTtfIEGLizPCC9fzASTvZtdhKEEri8TyyJs=
-github.com/wtsi-hgi/godirwalk v1.18.1/go.mod h1:rLa4FlI9kdT7o67jwFos8qgaX3K2sMC6XI4FXJ1iVyk=
 github.com/wtsi-ssg/wr v0.5.9 h1:lJWNuJfVvhTpXQqxRN5RbffhvK3HMog0fFpUFznvoz8=
 github.com/wtsi-ssg/wr v0.5.9/go.mod h1:njSdCX+xv1xzzw3Oy3Smid6s/IyIQEvLsKbRwaq4fC8=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=

walk/dirent.go

@@ -27,149 +27,240 @@
 package walk
 
 import (
+	"bytes"
 	"io/fs"
 	"os"
 	"sync"
-	"unsafe"
-
-	"github.com/wtsi-hgi/godirwalk"
 )
 
+func newDirentPool(size int) *sync.Pool {
+	return &sync.Pool{
+		New: func() any {
+			return &Dirent{
+				name:   make([]byte, 0, size),
+				parent: nullDirEnt,
+				next:   nullDirEnt,
+			}
+		},
+	}
+}
+
 var (
-	filePathPool64   = sync.Pool{New: func() any { x := make(FilePath, 0, 64); return &x }}   //nolint:gochecknoglobals,mnd,nlreturn,lll
-	filePathPool128  = sync.Pool{New: func() any { x := make(FilePath, 0, 128); return &x }}  //nolint:gochecknoglobals,mnd,nlreturn,lll
-	filePathPool256  = sync.Pool{New: func() any { x := make(FilePath, 0, 256); return &x }}  //nolint:gochecknoglobals,mnd,nlreturn,lll
-	filePathPool512  = sync.Pool{New: func() any { x := make(FilePath, 0, 512); return &x }}  //nolint:gochecknoglobals,mnd,nlreturn,lll
-	filePathPool1024 = sync.Pool{New: func() any { x := make(FilePath, 0, 1024); return &x }} //nolint:gochecknoglobals,mnd,nlreturn,lll
-	filePathPool2048 = sync.Pool{New: func() any { x := make(FilePath, 0, 2048); return &x }} //nolint:gochecknoglobals,mnd,nlreturn,lll
-	filePathPool4096 = sync.Pool{New: func() any { x := make(FilePath, 0, 4096); return &x }} //nolint:gochecknoglobals,mnd,nlreturn,lll
+	direntPool0   = newDirentPool(0)   //nolint:gochecknoglobals
+	direntPool32  = newDirentPool(32)  //nolint:gochecknoglobals,mnd
+	direntPool64  = newDirentPool(64)  //nolint:gochecknoglobals,mnd
+	direntPool128 = newDirentPool(128) //nolint:gochecknoglobals,mnd
+	dirEntPool256 = newDirentPool(257) //nolint:gochecknoglobals,mnd
+
+	nullDirEnt = new(Dirent) //nolint:gochecknoglobals
 )
 
-// FilePath is a byte-slice of a path, utilising object pools to reduce memory
-// allocations.
-//
-// It is the clients responsibility to call the Done method once it is no longer
-// needed.
-type FilePath []byte
+func init() { //nolint:gochecknoinits
+	nullDirEnt.parent = nullDirEnt
+	nullDirEnt.next = nullDirEnt
+}
 
-func newFilePathSize(size int) *FilePath {
+func getDirent(size int) *Dirent {
 	switch {
+	case size == 0:
+		return direntPool0.Get().(*Dirent) //nolint:forcetypeassert,errcheck
+	case size <= 32: //nolint:mnd
+		return direntPool32.Get().(*Dirent) //nolint:forcetypeassert,errcheck
 	case size <= 64: //nolint:mnd
-		return filePathPool64.Get().(*FilePath) //nolint:forcetypeassert
+		return direntPool64.Get().(*Dirent) //nolint:forcetypeassert,errcheck
 	case size <= 128: //nolint:mnd
-		return filePathPool128.Get().(*FilePath) //nolint:forcetypeassert
-	case size <= 256: //nolint:mnd
-		return filePathPool256.Get().(*FilePath) //nolint:forcetypeassert
-	case size <= 512: //nolint:mnd
-		return filePathPool512.Get().(*FilePath) //nolint:forcetypeassert
-	case size <= 1024: //nolint:mnd
-		return filePathPool1024.Get().(*FilePath) //nolint:forcetypeassert
-	case size <= 2048: //nolint:mnd
-		return filePathPool2048.Get().(*FilePath) //nolint:forcetypeassert
+		return direntPool128.Get().(*Dirent) //nolint:forcetypeassert,errcheck
 	}
 
-	return filePathPool4096.Get().(*FilePath) //nolint:forcetypeassert
+	return dirEntPool256.Get().(*Dirent) //nolint:forcetypeassert,errcheck
+}
+
+func putDirent(d *Dirent) {
+	d.name = d.name[:0]
+	d.parent = nullDirEnt
+	d.next = nullDirEnt
+	d.depth = 0
+
+	switch cap(d.name) {
+	case 0:
+		direntPool0.Put(d)
+	case 32: //nolint:mnd
+		direntPool32.Put(d)
+	case 64: //nolint:mnd
+		direntPool64.Put(d)
+	case 128: //nolint:mnd
+		direntPool128.Put(d)
+	default:
+		dirEntPool256.Put(d)
+	}
 }
 
-// NewFilePath creates a new FilePath, setting the value to the given string.
-func NewFilePath(path string) *FilePath {
-	c := newFilePathSize(len(path))
-	c.writeString(path)
+// Dirent represents a file system directory entry (a file or a directory),
+// providing information about the entry's path, type and inode.
+type Dirent struct {
+	parent *Dirent // left
+	name   []byte
+	depth  int16
+
+	// Type is the type bits of the file mode of this entry.
+	Type fs.FileMode
+
+	// Inode is the file system inode number for this entry.
+	Inode uint64
 
-	return c
+	next  *Dirent // right
+	ready sync.Mutex
 }
 
-func (f *FilePath) writeString(str string) {
-	*f = append(*f, str...)
+// IsDir returns true if we are a directory.
+func (d *Dirent) IsDir() bool {
+	return d.Type.IsDir()
 }
 
-func (f *FilePath) writeBytes(p []byte) {
-	*f = append(*f, p...)
+// IsRegular returns true if we are a regular file.
+func (d *Dirent) IsRegular() bool {
+	return d.Type.IsRegular()
 }
 
-// Done deallocates the underlying byte-slice; any uses of the Bytes method are
-// now invalid and may change.
-func (f *FilePath) Done() { //nolint:gocyclo
-	*f = (*f)[:0]
+// IsSymlink returns true if we are a symlink.
+func (d *Dirent) IsSymlink() bool {
+	return d.Type&os.ModeSymlink != 0
+}
 
-	switch cap(*f) {
-	case 64: //nolint:mnd
-		filePathPool64.Put(f)
-	case 128: //nolint:mnd
-		filePathPool128.Put(f)
-	case 256: //nolint:mnd
-		filePathPool256.Put(f)
-	case 512: //nolint:mnd
-		filePathPool512.Put(f)
-	case 1024: //nolint:mnd
-		filePathPool1024.Put(f)
-	case 2048: //nolint:mnd
-		filePathPool2048.Put(f)
-	case 4096: //nolint:mnd
-		filePathPool4096.Put(f)
+func (d *Dirent) appendTo(p []byte) []byte {
+	if d.parent != nil {
+		p = d.parent.appendTo(p)
 	}
+
+	return append(p, d.name...)
 }
 
-func (f *FilePath) sub(d *godirwalk.Dirent) *FilePath {
-	name := d.Name()
-	size := len(*f) + len(name)
+// Bytes returns the FilePath as a literal byte-slice.
+func (d *Dirent) Bytes() []byte {
+	return d.appendTo(nil)
+}
 
-	if d.IsDir() {
-		size++
+func (d *Dirent) compare(e *Dirent) int {
+	if d.depth < e.depth {
+		e = e.getDepth(d.depth)
+	} else if d.depth > e.depth {
+		d = d.getDepth(e.depth)
 	}
 
-	c := newFilePathSize(size)
-
-	c.writeBytes(*f)
-	c.writeString(name)
+	return e.compareTo(d)
+}
 
-	if d.IsDir() {
-		c.writeString("/")
+func (d *Dirent) getDepth(n int16) *Dirent {
+	for d.depth != n {
+		d = d.parent
 	}
 
-	return c
+	return d
 }
 
-// Bytes returns the FilePath as a literal byte-slice.
-func (f *FilePath) Bytes() []byte {
-	return *f
+func (d *Dirent) compareTo(e *Dirent) int {
+	if d == e {
+		return 0
+	}
+
+	cmp := d.parent.compareTo(e.parent)
+
+	if cmp == 0 {
+		return bytes.Compare(d.name, e.name)
+	}
+
+	return cmp
 }
 
-func (f *FilePath) string() string {
-	return unsafe.String(&(*f)[0], len(*f))
+func (d *Dirent) done() *Dirent {
+	next := d.next
+	d.next = nullDirEnt
+
+	if len(d.name) == 0 {
+		putDirent(d.parent)
+	}
+
+	if !d.IsDir() {
+		putDirent(d)
+	}
+
+	return next
 }
 
-// Dirent represents a file system directory entry (a file or a directory),
-// providing information about the entry's path, type and inode.
-type Dirent struct {
-	// Path is the complete path to the directory entry (including both
-	// directory and basename)
-	Path *FilePath
+func (d *Dirent) insert(e *Dirent) *Dirent { //nolint:gocyclo
+	if d == nullDirEnt {
+		return e
+	}
 
-	// Type is the type bits of the file mode of this entry.
-	Type os.FileMode
+	switch bytes.Compare(d.name, e.name) {
+	case 1:
+		d.parent = d.parent.insert(e)
+	case -1:
+		d.next = d.next.insert(e)
+	}
 
-	// Inode is the file system inode number for this entry.
-	Inode uint64
+	d.setDepth()
+
+	switch d.parent.depth - d.next.depth {
+	case -2:
+		if d.next.parent.depth > d.next.next.depth {
+			d.next = d.next.rotateRight()
+		}
+
+		return d.rotateLeft()
+	case 2: //nolint:mnd
+		if d.parent.next.depth > d.parent.parent.depth {
+			d.parent = d.parent.rotateLeft()
+		}
+
+		return d.rotateRight()
+	}
+
+	return d
 }
 
-// newDirentForDirectoryPath returns a Dirent for the given directory, with
-// a Type for directories and no Inode.
-func newDirentForDirectoryPath(dir string) Dirent {
-	return Dirent{Path: NewFilePath(dir), Type: fs.ModeDir}
+func (d *Dirent) setDepth() {
+	if d == nullDirEnt {
+		return
+	}
+
+	if d.parent.depth > d.next.depth {
+		d.depth = d.parent.depth + 1
+	} else {
+		d.depth = d.next.depth + 1
+	}
 }
 
-// IsDir returns true if we are a directory.
-func (d *Dirent) IsDir() bool {
-	return d.Type.IsDir()
+func (d *Dirent) rotateLeft() *Dirent {
+	n := d.next
+	d.next = n.parent
+	n.parent = d
+
+	d.setDepth()
+	n.setDepth()
+
+	return n
 }
 
-// IsRegular returns true if we are a regular file.
-func (d *Dirent) IsRegular() bool {
-	return d.Type.IsRegular()
+func (d *Dirent) rotateRight() *Dirent {
+	n := d.parent
+	d.parent = n.next
+	n.next = d
+
+	d.setDepth()
+	n.setDepth()
+
+	return n
 }
 
-// IsSymlink returns true if we are a symlink.
-func (d *Dirent) IsSymlink() bool {
-	return d.Type&os.ModeSymlink != 0
+func (d *Dirent) flatten(parent, prev *Dirent, depth int16) *Dirent {
+	if d == nullDirEnt {
+		return prev
+	}
+
+	d.parent.flatten(parent, prev, depth).next = d
+	d.parent = parent
+	d.depth = depth
+
+	return d.next.flatten(parent, d, depth)
 }

walk/dirent_test.go

@@ -71,11 +71,4 @@ func TestDirent(t *testing.T) {
 		So(d.IsRegular(), ShouldBeFalse)
 		So(d.IsSymlink(), ShouldBeTrue)
 	})
-
-	Convey("You can make a fake Direct for directories", t, func() {
-		d := newDirentForDirectoryPath("/a/dir")
-		So(d.IsDir(), ShouldBeTrue)
-		So(d.IsRegular(), ShouldBeFalse)
-		So(d.IsSymlink(), ShouldBeFalse)
-	})
 }