From e85dfc7499789c555e49bc4bea5f7fbf5900dd70 Mon Sep 17 00:00:00 2001
From: whyrusleeping <why@ipfs.io>
Date: Thu, 23 Jul 2020 15:32:28 -0700
Subject: [PATCH] add command to print state size statistics of actors

---
 api/api_full.go                   |   6 +-
 cli/state.go                      |   8 +++
 cmd/lotus-shed/main.go            |   2 +-
 cmd/lotus-shed/stateroot-stats.go | 114 +++++++++++++++++++++++++++++-
 4 files changed, 126 insertions(+), 4 deletions(-)

diff --git a/api/api_full.go b/api/api_full.go
index c557cb456..e3a0d8fb4 100644
--- a/api/api_full.go
+++ b/api/api_full.go
@@ -67,7 +67,11 @@ type FullNode interface {
 
 	// ChainHasObj checks if a given CID exists in the chain blockstore.
 	ChainHasObj(context.Context, cid.Cid) (bool, error)
-	ChainStatObj(context.Context, cid.Cid, cid.Cid) (ObjStat, error)
+
+	// ChainStatObj returns statistics about the graph referenced by 'obj'.
+	// If 'base' is also specified, then the returned stat will be a diff
+	// between the two objects.
+	ChainStatObj(ctx context.Context, obj cid.Cid, base cid.Cid) (ObjStat, error)
 
 	// ChainSetHead forcefully sets current chain head. Use with caution.
 	ChainSetHead(context.Context, types.TipSetKey) error
diff --git a/cli/state.go b/cli/state.go
index 8c5d712d0..d1b9f9b71 100644
--- a/cli/state.go
+++ b/cli/state.go
@@ -131,7 +131,15 @@ func LoadTipSet(ctx context.Context, cctx *cli.Context, api api.FullNode) (*type
 		return nil, nil
 	}
 
+	return ParseTipSetRef(ctx, api, tss)
+}
+
+func ParseTipSetRef(ctx context.Context, api api.FullNode, tss string) (*types.TipSet, error) {
 	if tss[0] == '@' {
+		if tss == "@head" {
+			return api.ChainHead(ctx)
+		}
+
 		var h uint64
 		if _, err := fmt.Sscanf(tss, "@%d", &h); err != nil {
 			return nil, xerrors.Errorf("parsing height tipset ref: %w", err)
diff --git a/cmd/lotus-shed/main.go b/cmd/lotus-shed/main.go
index 15cd3cfb3..7a2189c73 100644
--- a/cmd/lotus-shed/main.go
+++ b/cmd/lotus-shed/main.go
@@ -21,7 +21,7 @@ func main() {
 		keyinfoCmd,
 		noncefix,
 		bigIntParseCmd,
-		staterootStatsCmd,
+		staterootCmd,
 		importCarCmd,
 		commpToCidCmd,
 		fetchParamCmd,
diff --git a/cmd/lotus-shed/stateroot-stats.go b/cmd/lotus-shed/stateroot-stats.go
index 0546e5315..c02e0202a 100644
--- a/cmd/lotus-shed/stateroot-stats.go
+++ b/cmd/lotus-shed/stateroot-stats.go
@@ -2,17 +2,29 @@ package main
 
 import (
 	"fmt"
+	"sort"
 
+	"github.com/multiformats/go-multihash"
 	"github.com/urfave/cli/v2"
 
 	"github.com/ipfs/go-cid"
 
+	"github.com/filecoin-project/go-address"
+	"github.com/filecoin-project/lotus/api"
 	"github.com/filecoin-project/lotus/chain/types"
 	lcli "github.com/filecoin-project/lotus/cli"
 )
 
-var staterootStatsCmd = &cli.Command{
-	Name:        "stateroot-stats",
+var staterootCmd = &cli.Command{
+	Name: "stateroot",
+	Subcommands: []*cli.Command{
+		staterootDiffsCmd,
+		staterootStatCmd,
+	},
+}
+
+var staterootDiffsCmd = &cli.Command{
+	Name:        "diffs",
 	Description: "Walk down the chain and collect stats-obj changes between tipsets",
 	Flags: []cli.Flag{
 		&cli.StringFlag{
@@ -92,3 +104,101 @@ var staterootStatsCmd = &cli.Command{
 		return nil
 	},
 }
+
+type statItem struct {
+	Addr  address.Address
+	Actor *types.Actor
+	Stat  api.ObjStat
+}
+
+var staterootStatCmd = &cli.Command{
+	Name:  "stat",
+	Usage: "print statistics for the stateroot of a given block",
+	Flags: []cli.Flag{
+		&cli.StringFlag{
+			Name:  "tipset",
+			Usage: "specify tipset to start from",
+		},
+	},
+	Action: func(cctx *cli.Context) error {
+		api, closer, err := lcli.GetFullNodeAPI(cctx)
+		if err != nil {
+			return err
+		}
+
+		defer closer()
+		ctx := lcli.ReqContext(cctx)
+
+		ts, err := lcli.LoadTipSet(ctx, cctx, api)
+		if err != nil {
+			return err
+		}
+
+		if ts == nil {
+			ts, err = api.ChainHead(ctx)
+			if err != nil {
+				return err
+			}
+		}
+
+		var addrs []address.Address
+
+		for _, inp := range cctx.Args().Slice() {
+			a, err := address.NewFromString(inp)
+			if err != nil {
+				return err
+			}
+			addrs = append(addrs, a)
+		}
+
+		if len(addrs) == 0 {
+			allActors, err := api.StateListActors(ctx, ts.Key())
+			if err != nil {
+				return err
+			}
+			addrs = allActors
+		}
+
+		var infos []statItem
+		for _, a := range addrs {
+			act, err := api.StateGetActor(ctx, a, ts.Key())
+			if err != nil {
+				return err
+			}
+
+			stat, err := api.ChainStatObj(ctx, act.Head, cid.Undef)
+			if err != nil {
+				return err
+			}
+
+			infos = append(infos, statItem{
+				Addr:  a,
+				Actor: act,
+				Stat:  stat,
+			})
+		}
+
+		sort.Slice(infos, func(i, j int) bool {
+			return infos[i].Stat.Size > infos[j].Stat.Size
+		})
+
+		outcap := 10
+		if cctx.Args().Len() > outcap {
+			outcap = cctx.Args().Len()
+		}
+		if len(infos) < outcap {
+			outcap = len(infos)
+		}
+
+		fmt.Print("Addr\tType\tSize\n")
+		for _, inf := range infos[:outcap] {
+			cmh, err := multihash.Decode(inf.Actor.Code.Hash())
+			if err != nil {
+				return err
+			}
+
+			fmt.Printf("%s\t%s\t%d\n", inf.Addr, string(cmh.Digest), inf.Stat.Size)
+		}
+		return nil
+	},
+}