feature: use bloom to reduce cpu usage

fs/du_unix.go

@@ -20,10 +20,13 @@ package fs
 
 import (
 	"context"
+	"encoding/binary"
 	"fmt"
 	"os"
 	"path/filepath"
 	"syscall"
+
+	"github.com/bits-and-blooms/bloom/v3"
 )
 
 // blocksUnitSize is the unit used by `st_blocks` in `stat` in bytes.
@@ -58,12 +61,46 @@ func fileInfoStat(path string, fi os.FileInfo) (*syscall.Stat_t, error) {
 	return stat, nil
 }
 
+type InodeStore struct {
+	filter *bloom.BloomFilter
+	real   map[inode]struct{}
+}
+
+func NewInodeStore(expectedItems uint, falsePositiveRate float64) *InodeStore {
+	return &InodeStore{
+		filter: bloom.NewWithEstimates(expectedItems, falsePositiveRate),
+		real:   make(map[inode]struct{}, expectedItems),
+	}
+}
+
+func (s *InodeStore) Add(i inode) {
+	b := i.MarshalToBuf()
+	s.filter.Add(b)
+	s.real[i] = struct{}{}
+}
+
+func (s *InodeStore) Exists(i inode) bool {
+	b := i.MarshalToBuf()
+	if !s.filter.Test(b) {
+		return false
+	}
+	_, ok := s.real[i]
+	return ok
+}
+
+func (i inode) MarshalToBuf() []byte {
+	var arr [16]byte
+	binary.BigEndian.PutUint64(arr[0:8], i.dev)
+	binary.BigEndian.PutUint64(arr[8:16], i.ino)
+	return arr[:]
+}
+
 func diskUsage(ctx context.Context, roots ...string) (Usage, error) {
 	var (
 		size   int64
 		inodes = map[inode]struct{}{} // expensive!
 	)
-
+	store := NewInodeStore(500000, 0.01)
 	for _, root := range roots {
 		if err := filepath.Walk(root, func(path string, fi os.FileInfo, err error) error {
 			if err != nil {
@@ -81,11 +118,10 @@ func diskUsage(ctx context.Context, roots ...string) (Usage, error) {
 				return err
 			}
 			inoKey := newInode(stat)
-			if _, ok := inodes[inoKey]; !ok {
-				inodes[inoKey] = struct{}{}
+			if ok := store.Exists(inoKey); !ok {
+				store.Add(inoKey)
 				size += stat.Blocks * blocksUnitSize
 			}
-
 			return nil
 		}); err != nil {
 			return Usage{}, err

go.mod

@@ -4,11 +4,15 @@ go 1.21
 
 require (
 	github.com/Microsoft/go-winio v0.6.2
+	github.com/bits-and-blooms/bloom/v3 v3.7.1
 	github.com/containerd/log v0.1.0
 	github.com/opencontainers/go-digest v1.0.0
 	golang.org/x/sync v0.8.0
 	golang.org/x/sys v0.26.0
 	google.golang.org/protobuf v1.35.1
 )
 
-require github.com/sirupsen/logrus v1.9.3 // indirect
+require (
+	github.com/bits-and-blooms/bitset v1.24.2 // indirect
+	github.com/sirupsen/logrus v1.9.3 // indirect
+)

go.sum

@@ -1,5 +1,9 @@
 github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY=
 github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU=
+github.com/bits-and-blooms/bitset v1.24.2 h1:M7/NzVbsytmtfHbumG+K2bremQPMJuqv1JD3vOaFxp0=
+github.com/bits-and-blooms/bitset v1.24.2/go.mod h1:7hO7Gc7Pp1vODcmWvKMRA9BNmbv6a/7QIWpPxHddWR8=
+github.com/bits-and-blooms/bloom/v3 v3.7.1 h1:WXovk4TRKZttAMJfoQx6K2DM0zNIt8w+c67UqO+etV0=
+github.com/bits-and-blooms/bloom/v3 v3.7.1/go.mod h1:rZzYLLje2dfzXfAkJNxQQHsKurAyK55KUnL43Euk0hU=
 github.com/containerd/log v0.1.0 h1:TCJt7ioM2cr/tfR8GPbGf9/VRAX8D2B4PjzCpfX540I=
 github.com/containerd/log v0.1.0/go.mod h1:VRRf09a7mHDIRezVKTRCrOq78v577GXq3bSa3EhrzVo=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -17,6 +21,8 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.8.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ8=
 github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
+github.com/twmb/murmur3 v1.1.8 h1:8Yt9taO/WN3l08xErzjeschgZU2QSrwm1kclYq+0aRg=
+github.com/twmb/murmur3 v1.1.8/go.mod h1:Qq/R7NUyOfr65zD+6Q5IHKsJLwP7exErjN6lyyq3OSQ=
 golang.org/x/sync v0.8.0 h1:3NFvSEYkUoMifnESzZl15y791HH1qU2xm6eCJU5ZPXQ=
 golang.org/x/sync v0.8.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=