From 1372d63d78d653a08ad8c8bad5b6abeb0d3155c1 Mon Sep 17 00:00:00 2001
From: Rod Vagg <rod@vagg.org>
Date: Mon, 6 Jan 2025 21:12:06 +1100
Subject: [PATCH] fixup! feat(shed): actor state diff stats

---
 cmd/lotus-shed/state-stats.go | 53 ++++++++++++++++++++++++++---------
 1 file changed, 39 insertions(+), 14 deletions(-)

diff --git a/cmd/lotus-shed/state-stats.go b/cmd/lotus-shed/state-stats.go
index 4f6c69cdd47..e186ec3baa9 100644
--- a/cmd/lotus-shed/state-stats.go
+++ b/cmd/lotus-shed/state-stats.go
@@ -50,22 +50,26 @@ type actorStats struct {
 	Actor   *types.Actor
 	Fields  []fieldItem
 	Stats   api.ObjStat
+	Blocks  []cid.Cid `json:",omitempty"`
 }
 
 type fieldItem struct {
-	Name  string
-	Cid   cid.Cid
-	Stats api.ObjStat
+	Name   string
+	Cid    cid.Cid
+	Stats  api.ObjStat
+	Blocks []cid.Cid `json:",omitempty"`
 }
 
 type job struct {
 	c   cid.Cid
 	key string // prefix path for the region being recorded i.e. "/state/mineractor"
 }
+
 type cidCall struct {
 	c    cid.Cid
 	resp chan bool
 }
+
 type result struct {
 	key   string
 	stats api.ObjStat
@@ -122,14 +126,16 @@ func (cng *cacheNodeGetter) GetMany(ctx context.Context, list []cid.Cid) <-chan
 }
 
 type dagStatCollector struct {
-	ds   format.NodeGetter
-	walk func(format.Node) ([]*format.Link, error)
+	ds          format.NodeGetter
+	walk        func(format.Node) ([]*format.Link, error)
+	collectCids bool
 
 	statsLk sync.Mutex
 	stats   api.ObjStat
+	cids    []cid.Cid
 }
 
-func (dsc *dagStatCollector) record(ctx context.Context, nd format.Node) error {
+func (dsc *dagStatCollector) record(c cid.Cid, nd format.Node) error {
 	size, err := nd.Size()
 	if err != nil {
 		return err
@@ -140,6 +146,9 @@ func (dsc *dagStatCollector) record(ctx context.Context, nd format.Node) error {
 
 	dsc.stats.Size = dsc.stats.Size + size
 	dsc.stats.Links = dsc.stats.Links + 1
+	if dsc.collectCids {
+		dsc.cids = append(dsc.cids, c)
+	}
 
 	return nil
 }
@@ -150,7 +159,7 @@ func (dsc *dagStatCollector) walkLinks(ctx context.Context, c cid.Cid) ([]*forma
 		return nil, err
 	}
 
-	if err := dsc.record(ctx, nd); err != nil {
+	if err := dsc.record(c, nd); err != nil {
 		return nil, err
 	}
 
@@ -632,6 +641,10 @@ the total state of the actor in either tipset.
 			Name:  "diff-tipset",
 			Usage: "specify tipset to diff against, stat output will include only the mutated state between the two tipsets (pass comma separated array of cids or @<height> to specify tipset by height)",
 		},
+		&cli.BoolFlag{
+			Name:  "list-blocks",
+			Usage: "list the CIDs of blocks in the stat set being processed, in the case of a diff-tipset this will be the blocks that are different between the two tipsets",
+		},
 		&cli.IntFlag{
 			Name:  "workers",
 			Usage: "number of workers to use when processing",
@@ -695,6 +708,7 @@ the total state of the actor in either tipset.
 
 		numWorkers := cctx.Int("workers")
 		dagCacheSize := cctx.Int("dag-cache-size")
+		listBlocks := cctx.Bool("list-blocks")
 
 		jobs := make(chan address.Address, numWorkers)
 		results := make(chan actorStats, numWorkers)
@@ -728,7 +742,7 @@ the total state of the actor in either tipset.
 						}
 					}
 
-					actStats, err := sc.collectStats(ctx, addr, actor, dag)
+					actStats, err := sc.collectStats(ctx, addr, actor, dag, listBlocks)
 					if err != nil {
 						return err
 					}
@@ -927,7 +941,7 @@ type statCollector struct {
 	fieldCidSets map[string]*cid.Set
 }
 
-func (sc *statCollector) collectStats(ctx context.Context, addr address.Address, actor *types.Actor, dag format.NodeGetter) (actorStats, error) {
+func (sc *statCollector) collectStats(ctx context.Context, addr address.Address, actor *types.Actor, dag format.NodeGetter, collectCids bool) (actorStats, error) {
 	log.Infow("actor", "addr", addr, "code", actor.Code, "name", builtin.ActorNameByCode(actor.Code))
 
 	nd, err := dag.Get(ctx, actor.Head)
@@ -978,8 +992,9 @@ func (sc *statCollector) collectStats(ctx context.Context, addr address.Address,
 	}
 
 	dsc := &dagStatCollector{
-		ds:   dag,
-		walk: carWalkFunc,
+		ds:          dag,
+		walk:        carWalkFunc,
+		collectCids: collectCids,
 	}
 
 	if err := merkledag.Walk(ctx, dsc.walkLinks, actor.Head, sc.rootCidSet.Visit, merkledag.Concurrent()); err != nil {
@@ -987,18 +1002,28 @@ func (sc *statCollector) collectStats(ctx context.Context, addr address.Address,
 	}
 
 	actStats.Stats = dsc.stats
+	if collectCids {
+		actStats.Blocks = dsc.cids
+	}
 
 	for _, field := range fields {
 		dsc := &dagStatCollector{
-			ds:   dag,
-			walk: carWalkFunc,
+			ds:          dag,
+			walk:        carWalkFunc,
+			collectCids: collectCids,
 		}
 
-		if err := merkledag.Walk(ctx, dsc.walkLinks, field.Cid, sc.fieldCidSets[field.Name].Visit, merkledag.Concurrent()); err != nil {
+		if err := merkledag.Walk(ctx, func(ctx context.Context, c cid.Cid) ([]*format.Link, error) {
+			links, err := dsc.walkLinks(ctx, c)
+			return links, err
+		}, field.Cid, sc.fieldCidSets[field.Name].Visit, merkledag.Concurrent()); err != nil {
 			return actorStats{}, err
 		}
 
 		field.Stats = dsc.stats
+		if collectCids {
+			field.Blocks = dsc.cids
+		}
 
 		actStats.Fields = append(actStats.Fields, field)
 	}