Ultra-fast walk in Golang.

walk.go

package walk

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"sync"
	"time"
	"unsafe"

	"golang.org/x/sys/unix"
)

const dentBufSize = 8 << 20 // 8 MiB

// SkipDir is a special error that can be returned by Handler to skip the current directory.
var SkipDir = errors.New("skip this directory")

// Handler is the type of the function called by Walk to visit each file or directory.
// The path argument is the absolute path of the current file.
// The fi argument is a FileInfo struct about this file.
type Handler func(path string, fi *FileInfo) error

// FileInfo is the exposed file info struct.
type FileInfo struct {
	name    string
	size    int64
	mode    os.FileMode
	modTime time.Time
	sys     *unix.Stat_t
}

func (fi *FileInfo) Name() string {
	return fi.name
}

func (fi *FileInfo) Size() int64 {
	return fi.size
}

func (fi *FileInfo) Mode() os.FileMode {
	return fi.mode
}

func (fi *FileInfo) ModTime() time.Time {
	return fi.modTime
}

func (fi *FileInfo) IsDir() bool {
	return fi.mode&os.ModeDir != 0
}

func (fi *FileInfo) Sys() interface{} {
	return fi.sys
}

type dirent struct {
	ino  uint64
	name string
	typ  byte
}

// Walk is a walking function similar to filepath.Walk, but differs in these aspects:
//  1. It takes Handler instead of filepath.WalkFunc.
//  2. The error that occurred during walking is directly returned, without passing to Handler.
//  3. Lots of magic targeting *nix systems. See the comments for details.
func Walk(p string, walkFunc Handler) error {
	// Memory allocations are expensive. Use a pool to reuse buffers.
	dentBufPool := &sync.Pool{
		New: func() interface{} {
			return make([]byte, dentBufSize)
		},
	}
	p = filepath.Clean(p)
	fi, err := Stat(p)
	if err != nil {
		return fmt.Errorf("failed to stat path: %w", err)
	}
	if err := walkFunc(p, fi); err != nil {
		if errors.Is(err, SkipDir) {
			return nil
		}
		return err
	}
	if !fi.IsDir() {
		// It's a file, do not go deeper.
		return nil
	}
	// As we only need the file descriptors afterwards, we directly use the syscall.
	// We should avoid os.OpenFile, because an os.File is closed by Go runtime on GC.
	dirFd, err := unix.Open(p, os.O_RDONLY|unix.O_DIRECTORY, 0)
	if err != nil {
		return fmt.Errorf("failed to open the walk path: %w", err)
	}
	return walk(p, dirFd, dentBufPool, walkFunc)
}

// walk does the real walking stuff. It receives an opened directory and iterates the items in it.
func walk(dirName string, dirFd int, dentBufPool *sync.Pool, walkFunc Handler) error {
	buf := dentBufPool.Get().([]byte)
	defer func() {
		// The directory is closed when this function ends, and the dent buffer is returned to the pool.
		_ = unix.Close(dirFd)
		dentBufPool.Put(buf)
	}()
	for {
		// Use a large buffer to get directory entries.
		// The buffer size in os.ReadDir is 8 KiB and is too small, causing many unnecessary syscall operations.
		n, err := unix.ReadDirent(dirFd, buf)
		if err != nil {
			return fmt.Errorf("failed to call getdents64 on %s: %w", dirName, err)
		}
		if n == 0 {
			return nil
		}
		dirents := parseDirentBuf(buf[:n])
		for _, dent := range dirents {
			// Use fstatat with the fd of the already opened parent directory to save time.
			// If we use the full path directly, the kernel has to walk through the full path and do heavy checks
			// like the permission.
			fi, err := StatAt(dirFd, dent.name)
			if err != nil {
				return err
			}
			p := dirName + string(os.PathSeparator) + dent.name
			if err := walkFunc(p, fi); err != nil {
				if errors.Is(err, SkipDir) {
					continue
				}
				return err
			}
			if fi.IsDir() {
				// Walk the subdirectories recursively.
				// Also, use openat with the already opened parent directory to save time.
				nextDirFd, err := unix.Openat(dirFd, dent.name, unix.O_RDONLY|unix.O_DIRECTORY, 0)
				if err != nil {
					return fmt.Errorf("failed to open directory %s in %s: %w", dent.name, dirName, err)
				}
				if err := walk(p, nextDirFd, dentBufPool, walkFunc); err != nil {
					return err
				}
			}
		}
	}
}

// parseDirentBuf parses the dir entries returned by the syscall.
func parseDirentBuf(buf []byte) []*dirent {
	dirents := make([]*dirent, 0, len(buf)>>5) // Divided by 32, a reasonable guess to avoid capacity growth.
	for len(buf) > 0 {
		unixDirent := (*unix.Dirent)(unsafe.Pointer(&buf[0]))
		buf = buf[unixDirent.Reclen:]
		name := getDirentName(unixDirent)
		if name == "." || name == ".." {
			// Ignore the useless parent entries.
			continue
		}
		dent := &dirent{
			ino:  unixDirent.Ino,
			name: name,
			typ:  unixDirent.Type,
		}
		dirents = append(dirents, dent)
	}
	// Reading files in ascending inode order significantly improves read performance when page cache is absent.
	// It applies to EXT4 and XFS at least. Other filesystems haven't been tested yet.
	// When the length is less than 3, we skip sort.Slice to avoid reflection overheads.
	switch len(dirents) {
	case 0, 1:
	case 2:
		if dirents[0].ino > dirents[1].ino {
			dirents[0], dirents[1] = dirents[1], dirents[0]
		}
	default:
		sort.Slice(dirents, func(i, j int) bool {
			return dirents[i].ino < dirents[j].ino
		})
	}
	return dirents
}

// getDirentName returns the name field of a unix.Dirent.
func getDirentName(dirent *unix.Dirent) string {
	name := make([]byte, len(dirent.Name))
	var i int
	for ; i < len(dirent.Name); i++ {
		if dirent.Name[i] == 0 {
			break
		}
		name[i] = byte(dirent.Name[i])
	}
	return string(name[:i])
}

// Stat stats a file and returns an Entry.
func Stat(path string) (*FileInfo, error) {
	var stat unix.Stat_t
	err := unix.Lstat(path, &stat)
	if err != nil {
		return nil, fmt.Errorf("failed to stat %s: %w", path, err)
	}
	e := parseStat(path, &stat)
	return e, nil
}

// StatAt stats a file in an opened directory and returns an Entry.
func StatAt(dirFd int, name string) (*FileInfo, error) {
	var stat unix.Stat_t
	err := unix.Fstatat(dirFd, name, &stat, unix.AT_SYMLINK_NOFOLLOW)
	if err != nil {
		return nil, fmt.Errorf("failed to stat %s at dir fd %d: %w", name, dirFd, err)
	}
	e := parseStat(name, &stat)
	return e, nil
}

// parseStat parses a unix.Stat_t struct and returns an Entry.
func parseStat(name string, t *unix.Stat_t) *FileInfo {
	mode := os.FileMode(t.Mode & 0o777)
	switch t.Mode & unix.S_IFMT {
	case unix.S_IFBLK:
		mode |= os.ModeDevice
	case unix.S_IFCHR:
		mode |= os.ModeDevice | os.ModeCharDevice
	case unix.S_IFDIR:
		mode |= os.ModeDir
	case unix.S_IFIFO:
		mode |= os.ModeNamedPipe
	case unix.S_IFLNK:
		mode |= os.ModeSymlink
	case unix.S_IFSOCK:
		mode |= os.ModeSocket
	}
	if t.Mode&unix.S_ISGID != 0 {
		mode |= os.ModeSetgid
	}
	if t.Mode&unix.S_ISUID != 0 {
		mode |= os.ModeSetuid
	}
	if t.Mode&unix.S_ISVTX != 0 {
		mode |= os.ModeSticky
	}
	modTime := time.Unix(t.Mtim.Sec, t.Mtim.Nsec)
	return &FileInfo{
		name:    name,
		size:    t.Size,
		mode:    mode,
		modTime: modTime,
		sys:     t,
	}
}