feat: rewrite dag stat to be awesomerer

Fixes #8794
Fixes #8791

This adds:
- skip-raw-leaves option, that avoids downloading raw leaves when you
already know their size in the parrent block
- multithreaded download
- properly count blocks that show up multiple times with different
multicodec
- do not count size of inlined blocks

Speed:
Way way faster (it's hard to make an objective number because that depends a lot on how much multithreading helps you and how many raw-leaves block you skip), in my case: `>100x` faster.

Author: Jorropo

core/commands/dag/dag.go

@@ -17,10 +17,11 @@ import (
 )
 
 const (
-	pinRootsOptionName = "pin-roots"
-	progressOptionName = "progress"
-	silentOptionName   = "silent"
-	statsOptionName    = "stats"
+	pinRootsOptionName      = "pin-roots"
+	progressOptionName      = "progress"
+	skipRawleavesOptionName = "skip-raw-leaves"
+	silentOptionName        = "silent"
+	statsOptionName         = "stats"
 )
 
 // DagCmd provides a subset of commands for interacting with ipld dag objects
@@ -305,6 +306,7 @@ Note: This command skips duplicate blocks in reporting both size and the number
 	},
 	Options: []cmds.Option{
 		cmds.BoolOption(progressOptionName, "p", "Return progressive data while reading through the DAG").WithDefault(true),
+		cmds.BoolOption(skipRawleavesOptionName, "s", "Skip raw leaves block when size is already known, faster but can lead to underestimating the size with malformed DAGs.").WithDefault(true),
 	},
 	Run:  dagStat,
 	Type: DagStat{},

core/commands/dag/stat.go

@@ -1,21 +1,164 @@
 package dagcmd
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"os"
 
 	"github.com/ipfs/go-ipfs/core/commands/cmdenv"
 	"github.com/ipfs/go-ipfs/core/commands/e"
-	"github.com/ipfs/go-merkledag/traverse"
 	"github.com/ipfs/interface-go-ipfs-core/path"
 
+	cid "github.com/ipfs/go-cid"
 	cmds "github.com/ipfs/go-ipfs-cmds"
+	ipld "github.com/ipfs/go-ipld-format"
 	mdag "github.com/ipfs/go-merkledag"
+	mh "github.com/multiformats/go-multihash"
 )
 
+// Implement traversal here because there is a lot of adhoc logic to take care off
+// This is an async (multithreaded downloads but with a single contention point for processing) DFS implementation
+type statTraversal struct {
+	runners int
+	stats   DagStat
+
+	ctx           context.Context
+	cancel        context.CancelFunc
+	res           cmds.ResponseEmitter
+	getter        ipld.NodeGetter
+	nodes         chan *ipld.NodeOption
+	progressive   bool
+	skipRawleaves bool
+	stated        map[string]struct{}
+	seen          map[string]struct{}
+}
+
+func newStatTraversal(ctx context.Context, getter ipld.NodeGetter, res cmds.ResponseEmitter, progressive bool, skipRawleaves bool) *statTraversal {
+	ctx, cancel := context.WithCancel(ctx)
+	return &statTraversal{
+		ctx:           ctx,
+		cancel:        cancel,
+		getter:        getter,
+		res:           res,
+		nodes:         make(chan *ipld.NodeOption),
+		progressive:   progressive,
+		skipRawleaves: skipRawleaves,
+		// Use two different maps to correctly coiunt blocks with matching multihashes but different codecs
+		stated: make(map[string]struct{}),
+		seen:   make(map[string]struct{}),
+	}
+}
+
+func (t *statTraversal) pump() error {
+	defer t.cancel()
+	for {
+		select {
+		case <-t.ctx.Done():
+			return t.ctx.Err()
+		case n := <-t.nodes:
+			if n.Err != nil {
+				return n.Err
+			}
+			t.runners--
+			err := t.handleStating(n.Node.Cid(), uint64(len(n.Node.RawData())))
+			if err != nil {
+				return err
+			}
+			err = t.handleRecursion(n.Node)
+			if err != nil {
+				return err
+			}
+
+			if t.runners == 0 {
+				// FINISHED !
+				return nil
+			}
+		}
+	}
+}
+
+func (t *statTraversal) handleStating(c cid.Cid, nodeLen uint64) error {
+	k := string(c.Hash())
+	if _, alreadyCounted := t.stated[k]; alreadyCounted {
+		return nil
+	}
+	t.stated[k] = struct{}{}
+
+	if c.Prefix().MhType != mh.IDENTITY { // Do not count the size of inlined blocks
+		t.stats.Size += nodeLen
+	}
+	t.stats.NumBlocks++
+
+	if t.progressive {
+		if err := t.res.Emit(&t.stats); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (t *statTraversal) handleRecursion(node ipld.Node) error {
+	scan := make([]cid.Cid, 0, len(node.Links())) // Prealoc enough capacity
+	for _, l := range node.Links() {
+		k := l.Cid.KeyString()
+		if _, alreadySeen := t.seen[k]; alreadySeen {
+			continue
+		}
+		t.seen[k] = struct{}{}
+
+		if t.skipRawleaves {
+			prefix := l.Cid.Prefix()
+			if prefix.Codec == cid.Raw && l.Size != 0 /* still fetch links with likely missing size */ {
+				err := t.handleStating(l.Cid, l.Size)
+				if err != nil {
+					return err
+				}
+				continue
+			}
+		}
+
+		scan = append(scan, l.Cid)
+	}
+
+	t.runners += len(scan)
+	go func() {
+		c := t.getter.GetMany(t.ctx, scan)
+		for {
+			select {
+			case <-t.ctx.Done():
+				return
+			case v, ok := <-c:
+				if !ok {
+					return
+				}
+				select {
+				case <-t.ctx.Done():
+					return
+				case t.nodes <- v:
+				}
+			}
+		}
+	}()
+	return nil
+}
+
+func (t *statTraversal) traverse(c cid.Cid) error {
+	t.seen[c.KeyString()] = struct{}{}
+	t.runners = 1
+	go func() {
+		node, err := t.getter.Get(t.ctx, c)
+		select {
+		case <-t.ctx.Done():
+		case t.nodes <- &ipld.NodeOption{Node: node, Err: err}:
+		}
+	}()
+	return t.pump()
+}
+
 func dagStat(req *cmds.Request, res cmds.ResponseEmitter, env cmds.Environment) error {
 	progressive := req.Options[progressOptionName].(bool)
+	skipRawleaves := req.Options[skipRawleavesOptionName].(bool)
 
 	api, err := cmdenv.GetApi(env, req)
 	if err != nil {
@@ -32,35 +175,15 @@ func dagStat(req *cmds.Request, res cmds.ResponseEmitter, env cmds.Environment)
 	}
 
 	nodeGetter := mdag.NewSession(req.Context, api.Dag())
-	obj, err := nodeGetter.Get(req.Context, rp.Cid())
-	if err != nil {
-		return err
-	}
-
-	dagstats := &DagStat{}
-	err = traverse.Traverse(obj, traverse.Options{
-		DAG:   nodeGetter,
-		Order: traverse.DFSPre,
-		Func: func(current traverse.State) error {
-			dagstats.Size += uint64(len(current.Node.RawData()))
-			dagstats.NumBlocks++
+	t := newStatTraversal(req.Context, nodeGetter, res, progressive, skipRawleaves)
 
-			if progressive {
-				if err := res.Emit(dagstats); err != nil {
-					return err
-				}
-			}
-			return nil
-		},
-		ErrFunc:        nil,
-		SkipDuplicates: true,
-	})
+	err = t.traverse(rp.Cid())
 	if err != nil {
 		return fmt.Errorf("error traversing DAG: %w", err)
 	}
 
 	if !progressive {
-		if err := res.Emit(dagstats); err != nil {
+		if err := res.Emit(&t.stats); err != nil {
 			return err
 		}
 	}