lotus/node/repo/retrievalstoremgr/retrievalstoremgr_test.go

package retrievalstoremgr_test

import (
	"context"
	"math/rand"
	"testing"

	"github.com/ipfs/go-cid"
	"github.com/ipfs/go-datastore"
	"github.com/ipfs/go-datastore/query"
	dss "github.com/ipfs/go-datastore/sync"
	format "github.com/ipfs/go-ipld-format"
	dag "github.com/ipfs/go-merkledag"
	"github.com/stretchr/testify/require"

	"github.com/filecoin-project/go-multistore"

	"github.com/filecoin-project/lotus/lib/blockstore"
	"github.com/filecoin-project/lotus/node/repo/importmgr"
	"github.com/filecoin-project/lotus/node/repo/retrievalstoremgr"
)

func TestMultistoreRetrievalStoreManager(t *testing.T) {
	ctx := context.Background()
	ds := dss.MutexWrap(datastore.NewMapDatastore())
	multiDS, err := multistore.NewMultiDstore(ds)
	require.NoError(t, err)
	imgr := importmgr.New(multiDS, ds)
	retrievalStoreMgr := retrievalstoremgr.NewMultiStoreRetrievalStoreManager(imgr)

	var stores []retrievalstoremgr.RetrievalStore
	for i := 0; i < 5; i++ {
		store, err := retrievalStoreMgr.NewStore()
		require.NoError(t, err)
		stores = append(stores, store)
		nds := generateNodesOfSize(5, 100)
		err = store.DAGService().AddMany(ctx, nds)
		require.NoError(t, err)
	}

	t.Run("creates all keys", func(t *testing.T) {
		qres, err := ds.Query(query.Query{KeysOnly: true})
		require.NoError(t, err)
		all, err := qres.Rest()
		require.NoError(t, err)
		require.Len(t, all, 31)
	})

	t.Run("loads DAG services", func(t *testing.T) {
		for _, store := range stores {
			mstore, err := multiDS.Get(*store.StoreID())
			require.NoError(t, err)
			require.Equal(t, mstore.DAG, store.DAGService())
		}
	})

	t.Run("delete stores", func(t *testing.T) {
		err := retrievalStoreMgr.ReleaseStore(stores[4])
		require.NoError(t, err)
		storeIndexes := multiDS.List()
		require.Len(t, storeIndexes, 4)

		qres, err := ds.Query(query.Query{KeysOnly: true})
		require.NoError(t, err)
		all, err := qres.Rest()
		require.NoError(t, err)
		require.Len(t, all, 25)
	})
}

func TestBlockstoreRetrievalStoreManager(t *testing.T) {
	ctx := context.Background()
	ds := dss.MutexWrap(datastore.NewMapDatastore())
	bs := blockstore.NewBlockstore(ds)
	retrievalStoreMgr := retrievalstoremgr.NewBlockstoreRetrievalStoreManager(bs)
	var stores []retrievalstoremgr.RetrievalStore
	var cids []cid.Cid
	for i := 0; i < 5; i++ {
		store, err := retrievalStoreMgr.NewStore()
		require.NoError(t, err)
		stores = append(stores, store)
		nds := generateNodesOfSize(5, 100)
		err = store.DAGService().AddMany(ctx, nds)
		require.NoError(t, err)
		for _, nd := range nds {
			cids = append(cids, nd.Cid())
		}
	}

	t.Run("creates all keys", func(t *testing.T) {
		qres, err := ds.Query(query.Query{KeysOnly: true})
		require.NoError(t, err)
		all, err := qres.Rest()
		require.NoError(t, err)
		require.Len(t, all, 25)
	})

	t.Run("loads DAG services, all DAG has all nodes", func(t *testing.T) {
		for _, store := range stores {
			dagService := store.DAGService()
			for _, cid := range cids {
				_, err := dagService.Get(ctx, cid)
				require.NoError(t, err)
			}
		}
	})

	t.Run("release store has no effect", func(t *testing.T) {
		err := retrievalStoreMgr.ReleaseStore(stores[4])
		require.NoError(t, err)
		qres, err := ds.Query(query.Query{KeysOnly: true})
		require.NoError(t, err)
		all, err := qres.Rest()
		require.NoError(t, err)
		require.Len(t, all, 25)
	})
}

var seedSeq int64 = 0

func randomBytes(n int64) []byte {
	randBytes := make([]byte, n)
	r := rand.New(rand.NewSource(seedSeq))
	_, _ = r.Read(randBytes)
	seedSeq++
	return randBytes
}

func generateNodesOfSize(n int, size int64) []format.Node {
	generatedNodes := make([]format.Node, 0, n)
	for i := 0; i < n; i++ {
		b := dag.NewRawNode(randomBytes(size))
		generatedNodes = append(generatedNodes, b)

	}
	return generatedNodes
}