From c6ad8fdaedf9b243457972bcf902ee0836f455ee Mon Sep 17 00:00:00 2001
From: vyzo <vyzo@hackzen.org>
Date: Tue, 6 Jul 2021 01:08:44 +0300
Subject: [PATCH] use walkObjectRaw for computing object weights

cids that come out of the hotstore with ForEach are raw.
---
 blockstore/splitstore/splitstore.go | 57 ++++++++++++++++++++++++++++-
 1 file changed, 55 insertions(+), 2 deletions(-)

diff --git a/blockstore/splitstore/splitstore.go b/blockstore/splitstore/splitstore.go
index 591fa479a..f58f318be 100644
--- a/blockstore/splitstore/splitstore.go
+++ b/blockstore/splitstore/splitstore.go
@@ -1143,6 +1143,57 @@ func (s *SplitStore) walkObject(c cid.Cid, walked *cid.Set, f func(cid.Cid) erro
 	return nil
 }
 
+// like walkObject, but treats Raw root cids as potential DAGs; this is necessary for walking cids
+// that come out of the blockstore itself.
+func (s *SplitStore) walkObjectRaw(c cid.Cid, walked *cid.Set, f func(cid.Cid) error) error {
+	if !walked.Visit(c) {
+		return nil
+	}
+
+	if err := f(c); err != nil {
+		if err == errStopWalk {
+			return nil
+		}
+
+		return err
+	}
+
+	switch c.Prefix().Codec {
+	case cid.DagCBOR, cid.Raw:
+	default:
+		return nil
+	}
+
+	var links []cid.Cid
+	err := s.view(c, func(data []byte) error {
+		return cbg.ScanForLinks(bytes.NewReader(data), func(c cid.Cid) {
+			links = append(links, c)
+		})
+	})
+
+	if err != nil {
+		// don't fail if the scan fails
+		log.Warnf("error scanning linked block (cid: %s): %s", c, err)
+		return nil
+	}
+
+	for _, c := range links {
+		// these are internal references and should no longer be raw, so we recurse with walkObject
+		err := s.walkObject(c, walked, f)
+		if err != nil {
+			if xerrors.Is(err, bstore.ErrNotFound) {
+				// potential false positive
+				log.Warnf("error walking link (cid: %s): %s", c, err)
+				continue
+			}
+
+			return xerrors.Errorf("error walking link (cid: %s): %w", c, err)
+		}
+	}
+
+	return nil
+}
+
 // like walkObject, but the object may be potentially incomplete (references missing from the hotstore)
 func (s *SplitStore) walkObjectIncomplete(c cid.Cid, walked *cid.Set, f, missing func(cid.Cid) error) error {
 	if !walked.Visit(c) {
@@ -1274,12 +1325,14 @@ func (s *SplitStore) moveColdBlocks(cold []cid.Cid) error {
 func (s *SplitStore) sortObjects(cids []cid.Cid) error {
 	weight := make(map[cid.Cid]int)
 	for _, c := range cids {
-		if c.Prefix().Codec != cid.DagCBOR {
+		switch c.Prefix().Codec {
+		case cid.DagCBOR, cid.Raw:
+		default:
 			continue
 		}
 
 		w := 0
-		err := s.walkObject(c, cid.NewSet(),
+		err := s.walkObjectRaw(c, cid.NewSet(),
 			func(c cid.Cid) error {
 				wc, ok := weight[c]
 				if ok {