This is similar to the golang.org/x/tools/internal/fastwalk I'd
previously written but not recursive and using mem.RO.
The metrics package already had some Linux-specific directory reading
code in it. Move that out to a new general package that can be reused
by portlist too, which helps its scanning of all /proc files:
name old time/op new time/op delta
FindProcessNames-8 2.79ms ± 6% 2.45ms ± 7% -12.11% (p=0.000 n=10+10)
name old alloc/op new alloc/op delta
FindProcessNames-8 62.9kB ± 0% 33.5kB ± 0% -46.76% (p=0.000 n=9+10)
name old allocs/op new allocs/op delta
FindProcessNames-8 2.25k ± 0% 0.38k ± 0% -82.98% (p=0.000 n=9+10)
Change-Id: I75db393032c328f12d95c39f71c9742c375f207a
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
main
Brad Fitzpatrick
3 years ago
committed by
Brad Fitzpatrick
parent
21ef7e5c35
commit
db2cc393af
@ -0,0 +1,54 @@ |
||||
// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
|
||||
// Use of this source code is governed by a BSD-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
// Package dirwalk contains code to walk a directory.
|
||||
package dirwalk |
||||
|
||||
import ( |
||||
"io" |
||||
"io/fs" |
||||
"os" |
||||
|
||||
"go4.org/mem" |
||||
) |
||||
|
||||
var osWalkShallow func(name mem.RO, fn WalkFunc) error |
||||
|
||||
// WalkFunc is the callback type used with WalkShallow.
|
||||
//
|
||||
// The name and de are only valid for the duration of func's call
|
||||
// and should not be retained.
|
||||
type WalkFunc func(name mem.RO, de fs.DirEntry) error |
||||
|
||||
// WalkShallow reads the entries in the named directory and calls fn for each.
|
||||
// It does not recurse into subdirectories.
|
||||
//
|
||||
// If fn returns an error, iteration stops and WalkShallow returns that value.
|
||||
//
|
||||
// On Linux, WalkShallow does not allocate, so long as certain methods on the
|
||||
// WalkFunc's DirEntry are not called which necessarily allocate.
|
||||
func WalkShallow(dirName mem.RO, fn WalkFunc) error { |
||||
if f := osWalkShallow; f != nil { |
||||
return f(dirName, fn) |
||||
} |
||||
of, err := os.Open(dirName.StringCopy()) |
||||
if err != nil { |
||||
return err |
||||
} |
||||
defer of.Close() |
||||
for { |
||||
fis, err := of.ReadDir(100) |
||||
for _, de := range fis { |
||||
if err := fn(mem.S(de.Name()), de); err != nil { |
||||
return err |
||||
} |
||||
} |
||||
if err != nil { |
||||
if err == io.EOF { |
||||
return nil |
||||
} |
||||
return err |
||||
} |
||||
} |
||||
} |
||||
@ -0,0 +1,168 @@ |
||||
// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
|
||||
// Use of this source code is governed by a BSD-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package dirwalk |
||||
|
||||
import ( |
||||
"fmt" |
||||
"io/fs" |
||||
"os" |
||||
"path/filepath" |
||||
"sync" |
||||
"syscall" |
||||
"unsafe" |
||||
|
||||
"go4.org/mem" |
||||
"golang.org/x/sys/unix" |
||||
) |
||||
|
||||
func init() { |
||||
osWalkShallow = linuxWalkShallow |
||||
} |
||||
|
||||
var dirEntPool = &sync.Pool{New: func() any { return new(linuxDirEnt) }} |
||||
|
||||
func linuxWalkShallow(dirName mem.RO, fn WalkFunc) error { |
||||
const blockSize = 8 << 10 |
||||
buf := make([]byte, blockSize) // stack-allocated; doesn't escape
|
||||
|
||||
nameb := mem.Append(buf[:0], dirName) |
||||
nameb = append(nameb, 0) |
||||
|
||||
fd, err := sysOpen(nameb) |
||||
if err != nil { |
||||
return err |
||||
} |
||||
defer syscall.Close(fd) |
||||
|
||||
bufp := 0 // starting read position in buf
|
||||
nbuf := 0 // end valid data in buf
|
||||
|
||||
de := dirEntPool.Get().(*linuxDirEnt) |
||||
defer de.cleanAndPutInPool() |
||||
de.root = dirName |
||||
|
||||
for { |
||||
if bufp >= nbuf { |
||||
bufp = 0 |
||||
nbuf, err = readDirent(fd, buf) |
||||
if err != nil { |
||||
return err |
||||
} |
||||
if nbuf <= 0 { |
||||
return nil |
||||
} |
||||
} |
||||
consumed, name := parseDirEnt(&de.d, buf[bufp:nbuf]) |
||||
bufp += consumed |
||||
if len(name) == 0 || string(name) == "." || string(name) == ".." { |
||||
continue |
||||
} |
||||
de.name = mem.B(name) |
||||
if err := fn(de.name, de); err != nil { |
||||
return err |
||||
} |
||||
} |
||||
} |
||||
|
||||
type linuxDirEnt struct { |
||||
root mem.RO |
||||
d syscall.Dirent |
||||
name mem.RO |
||||
} |
||||
|
||||
func (de *linuxDirEnt) cleanAndPutInPool() { |
||||
de.root = mem.RO{} |
||||
de.name = mem.RO{} |
||||
dirEntPool.Put(de) |
||||
} |
||||
|
||||
func (de *linuxDirEnt) Name() string { return de.name.StringCopy() } |
||||
func (de *linuxDirEnt) Info() (fs.FileInfo, error) { |
||||
return os.Lstat(filepath.Join(de.root.StringCopy(), de.name.StringCopy())) |
||||
} |
||||
func (de *linuxDirEnt) IsDir() bool { |
||||
return de.d.Type == syscall.DT_DIR |
||||
} |
||||
func (de *linuxDirEnt) Type() fs.FileMode { |
||||
switch de.d.Type { |
||||
case syscall.DT_BLK: |
||||
return fs.ModeDevice // shrug
|
||||
case syscall.DT_CHR: |
||||
return fs.ModeCharDevice |
||||
case syscall.DT_DIR: |
||||
return fs.ModeDir |
||||
case syscall.DT_FIFO: |
||||
return fs.ModeNamedPipe |
||||
case syscall.DT_LNK: |
||||
return fs.ModeSymlink |
||||
case syscall.DT_REG: |
||||
return 0 |
||||
case syscall.DT_SOCK: |
||||
return fs.ModeSocket |
||||
default: |
||||
return fs.ModeIrregular // shrug
|
||||
} |
||||
} |
||||
|
||||
func direntNamlen(dirent *syscall.Dirent) int { |
||||
const fixedHdr = uint16(unsafe.Offsetof(syscall.Dirent{}.Name)) |
||||
limit := dirent.Reclen - fixedHdr |
||||
const dirNameLen = 256 // sizeof syscall.Dirent.Name
|
||||
if limit > dirNameLen { |
||||
limit = dirNameLen |
||||
} |
||||
for i := uint16(0); i < limit; i++ { |
||||
if dirent.Name[i] == 0 { |
||||
return int(i) |
||||
} |
||||
} |
||||
panic("failed to find terminating 0 byte in dirent") |
||||
} |
||||
|
||||
func parseDirEnt(dirent *syscall.Dirent, buf []byte) (consumed int, name []byte) { |
||||
// golang.org/issue/37269
|
||||
copy(unsafe.Slice((*byte)(unsafe.Pointer(dirent)), unsafe.Sizeof(syscall.Dirent{})), buf) |
||||
if v := unsafe.Offsetof(dirent.Reclen) + unsafe.Sizeof(dirent.Reclen); uintptr(len(buf)) < v { |
||||
panic(fmt.Sprintf("buf size of %d smaller than dirent header size %d", len(buf), v)) |
||||
} |
||||
if len(buf) < int(dirent.Reclen) { |
||||
panic(fmt.Sprintf("buf size %d < record length %d", len(buf), dirent.Reclen)) |
||||
} |
||||
consumed = int(dirent.Reclen) |
||||
if dirent.Ino == 0 { // File absent in directory.
|
||||
return |
||||
} |
||||
name = unsafe.Slice((*byte)(unsafe.Pointer(&dirent.Name[0])), direntNamlen(dirent)) |
||||
return |
||||
} |
||||
|
||||
func sysOpen(name []byte) (fd int, err error) { |
||||
if len(name) == 0 || name[len(name)-1] != 0 { |
||||
return 0, syscall.EINVAL |
||||
} |
||||
var dirfd int = unix.AT_FDCWD |
||||
for { |
||||
r0, _, e1 := syscall.Syscall(unix.SYS_OPENAT, uintptr(dirfd), |
||||
uintptr(unsafe.Pointer(&name[0])), 0) |
||||
if e1 == 0 { |
||||
return int(r0), nil |
||||
} |
||||
if e1 == syscall.EINTR { |
||||
// Since https://golang.org/doc/go1.14#runtime we
|
||||
// need to loop on EINTR on more places.
|
||||
continue |
||||
} |
||||
return 0, syscall.Errno(e1) |
||||
} |
||||
} |
||||
|
||||
func readDirent(fd int, buf []byte) (n int, err error) { |
||||
for { |
||||
nbuf, err := syscall.ReadDirent(fd, buf) |
||||
if err != syscall.EINTR { |
||||
return nbuf, err |
||||
} |
||||
} |
||||
} |
||||
@ -0,0 +1,93 @@ |
||||
// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
|
||||
// Use of this source code is governed by a BSD-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package dirwalk |
||||
|
||||
import ( |
||||
"fmt" |
||||
"os" |
||||
"path/filepath" |
||||
"reflect" |
||||
"runtime" |
||||
"sort" |
||||
"testing" |
||||
|
||||
"go4.org/mem" |
||||
) |
||||
|
||||
func TestWalkShallowOSSpecific(t *testing.T) { |
||||
if osWalkShallow == nil { |
||||
t.Skip("no OS-specific implementation") |
||||
} |
||||
testWalkShallow(t, false) |
||||
} |
||||
|
||||
func TestWalkShallowPortable(t *testing.T) { |
||||
testWalkShallow(t, true) |
||||
} |
||||
|
||||
func testWalkShallow(t *testing.T, portable bool) { |
||||
if portable { |
||||
old := osWalkShallow |
||||
defer func() { osWalkShallow = old }() |
||||
osWalkShallow = nil |
||||
} |
||||
d := t.TempDir() |
||||
|
||||
t.Run("basics", func(t *testing.T) { |
||||
if err := os.WriteFile(filepath.Join(d, "foo"), []byte("1"), 0600); err != nil { |
||||
t.Fatal(err) |
||||
} |
||||
if err := os.WriteFile(filepath.Join(d, "bar"), []byte("22"), 0400); err != nil { |
||||
t.Fatal(err) |
||||
} |
||||
if err := os.Mkdir(filepath.Join(d, "baz"), 0777); err != nil { |
||||
t.Fatal(err) |
||||
} |
||||
|
||||
var got []string |
||||
if err := WalkShallow(mem.S(d), func(name mem.RO, de os.DirEntry) error { |
||||
var size int64 |
||||
if fi, err := de.Info(); err != nil { |
||||
t.Errorf("Info stat error on %q: %v", de.Name(), err) |
||||
} else if !fi.IsDir() { |
||||
size = fi.Size() |
||||
} |
||||
got = append(got, fmt.Sprintf("%q %q dir=%v type=%d size=%v", name.StringCopy(), de.Name(), de.IsDir(), de.Type(), size)) |
||||
return nil |
||||
}); err != nil { |
||||
t.Fatal(err) |
||||
} |
||||
sort.Strings(got) |
||||
want := []string{ |
||||
`"bar" "bar" dir=false type=0 size=2`, |
||||
`"baz" "baz" dir=true type=2147483648 size=0`, |
||||
`"foo" "foo" dir=false type=0 size=1`, |
||||
} |
||||
if !reflect.DeepEqual(got, want) { |
||||
t.Errorf("mismatch:\n got %#q\nwant %#q", got, want) |
||||
} |
||||
}) |
||||
|
||||
t.Run("err_not_exist", func(t *testing.T) { |
||||
err := WalkShallow(mem.S(filepath.Join(d, "not_exist")), func(name mem.RO, de os.DirEntry) error { |
||||
return nil |
||||
}) |
||||
if !os.IsNotExist(err) { |
||||
t.Errorf("unexpected error: %v", err) |
||||
} |
||||
}) |
||||
|
||||
t.Run("allocs", func(t *testing.T) { |
||||
allocs := int(testing.AllocsPerRun(1000, func() { |
||||
if err := WalkShallow(mem.S(d), func(name mem.RO, de os.DirEntry) error { return nil }); err != nil { |
||||
t.Fatal(err) |
||||
} |
||||
})) |
||||
t.Logf("allocs = %v", allocs) |
||||
if !portable && runtime.GOOS == "linux" && allocs != 0 { |
||||
t.Errorf("unexpected allocs: got %v, want 0", allocs) |
||||
} |
||||
}) |
||||
} |
||||
Loading…
Reference in new issue