recover trie for all allocations

Author: Hitenjain14

code/go/0chain.net/blobber/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"github.com/0chain/blobber/code/go/0chain.net/blobbercore/allocation"
 	"github.com/0chain/blobber/code/go/0chain.net/core/common"
 	"github.com/0chain/blobber/code/go/0chain.net/core/logging"
 	"github.com/0chain/blobber/code/go/0chain.net/core/node"
@@ -58,6 +59,11 @@ func main() {
 		panic(err)
 	}
 
+	if recoverTrie {
+		logging.Logger.Info("Recovering trie")
+		allocation.RecoverTrie()
+	}
+
 	// todo: activate this when gRPC functionalities are implemented
 	// go startGRPCServer()
 

code/go/0chain.net/blobbercore/allocation/entity.go

@@ -8,9 +8,13 @@ import (
 	"time"
 
 	"github.com/0chain/blobber/code/go/0chain.net/blobbercore/datastore"
+	"github.com/0chain/blobber/code/go/0chain.net/blobbercore/filestore"
+	"github.com/0chain/blobber/code/go/0chain.net/blobbercore/reference"
 	"github.com/0chain/blobber/code/go/0chain.net/core/common"
+	"github.com/0chain/blobber/code/go/0chain.net/core/logging"
 	"github.com/0chain/common/core/encryption"
 	"github.com/0chain/common/core/util/wmpt"
+	"go.uber.org/zap"
 	"gorm.io/gorm/clause"
 
 	"gorm.io/gorm"
@@ -357,3 +361,52 @@ type ReadPoolRedeem struct {
 	PoolID  string `json:"pool_id"` // read pool ID
 	Balance int64  `json:"balance"` // balance reduction
 }
+
+func (a *Allocation) recoverTrie() error {
+
+	trie := wmpt.New(nil, datastore.GetBlockStore())
+	//fetch all the files of the allocation to rebuild the trie
+	var offsetPath string
+	for {
+		var (
+			pRefs *[]reference.PaginatedRef
+			err   error
+		)
+		err = datastore.GetStore().WithNewTransaction(func(ctx context.Context) error {
+			pRefs, _, offsetPath, err = reference.GetRefs(ctx, a.ID, "/", offsetPath, reference.FILE, 0, 100)
+			return err
+		})
+		if err != nil {
+			logging.Logger.Error("recover_trie_fetch_refs", zap.Error(err))
+			return err
+		}
+		if len(*pRefs) == 0 {
+			break
+		}
+		trie.SaveRoot()
+		for _, ref := range *pRefs {
+			decodedKey, _ := hex.DecodeString(ref.LookupHash)
+			decodedValue, _ := hex.DecodeString(ref.FileMetaHash)
+			err = trie.Update(decodedKey, decodedValue, uint64(ref.NumBlocks))
+			if err != nil {
+				logging.Logger.Error("recover_trie_update", zap.Error(err))
+				return err
+			}
+		}
+		batcher, err := trie.Commit(filestore.COLLAPSE_DEPTH)
+		if err != nil {
+			logging.Logger.Error("recover_trie_commit", zap.Error(err))
+			return err
+		}
+		err = batcher.Commit(true)
+		if err != nil {
+			logging.Logger.Error("recover_trie_batcher_commit", zap.Error(err))
+			return err
+		}
+		_ = trie.DeleteNodes()
+		if len(*pRefs) < 100 {
+			break
+		}
+	}
+	return nil
+}

code/go/0chain.net/blobbercore/allocation/workers.go

@@ -348,3 +348,34 @@ func deleteAllocation(ctx context.Context, a *Allocation) (err error) {
 			a.ID).Error
 	return err
 }
+
+func RecoverTrie() {
+	var (
+		allocs []*Allocation
+		err    error
+		offset int64
+	)
+
+	for {
+		err = datastore.GetStore().WithNewTransaction(func(ctx context.Context) error {
+			allocs, err = Repo.GetAllocations(ctx, offset)
+			return err
+		})
+		if err != nil {
+			logging.Logger.Error("recover_trie_fetch_alloc", zap.Error(err))
+			return
+		}
+		if len(allocs) == 0 {
+			return
+		}
+		offset += int64(len(allocs))
+		// recover trie of each allocation
+		for _, a := range allocs {
+			logging.Logger.Info("recover_trie", zap.String("allocation_id", a.ID))
+			err = a.recoverTrie()
+			if err != nil {
+				logging.Logger.Error("recover_trie", zap.Error(err))
+			}
+		}
+	}
+}

code/go/0chain.net/blobbercore/datastore/store.go

@@ -106,10 +106,6 @@ func OpenBlockStore() error {
 	return nil
 }
 
-func RecoverTrie() {
-
-}
-
 func FromContext(ctx context.Context) Store {
 	store := ctx.Value(ContextKeyStore)
 	if store != nil {