github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/module/executiondatasync/tracker/storage_test.go (about)

     1  package tracker
     2  
     3  import (
     4  	"crypto/rand"
     5  	"testing"
     6  
     7  	"github.com/dgraph-io/badger/v2"
     8  	"github.com/ipfs/go-cid"
     9  	"github.com/rs/zerolog"
    10  	"github.com/stretchr/testify/assert"
    11  	"github.com/stretchr/testify/require"
    12  
    13  	"github.com/onflow/flow-go/module/blobs"
    14  )
    15  
    16  func randomCid() cid.Cid {
    17  	data := make([]byte, 1024)
    18  	_, _ = rand.Read(data)
    19  	return blobs.NewBlob(data).Cid()
    20  }
    21  
    22  // TestPrune tests that when a height is pruned, all CIDs appearing at or below the pruned
    23  // height, and their associated tracking data, should be removed from the database.
    24  func TestPrune(t *testing.T) {
    25  	expectedPrunedCIDs := make(map[cid.Cid]struct{})
    26  	storageDir := t.TempDir()
    27  	storage, err := OpenStorage(storageDir, 0, zerolog.Nop(), WithPruneCallback(func(c cid.Cid) error {
    28  		_, ok := expectedPrunedCIDs[c]
    29  		assert.True(t, ok, "unexpected CID pruned: %s", c.String())
    30  		delete(expectedPrunedCIDs, c)
    31  		return nil
    32  	}))
    33  	require.NoError(t, err)
    34  
    35  	// c1 and c2 are for height 1, and c3 and c4 are for height 2
    36  	// after pruning up to height 1, only c1 and c2 should be pruned
    37  	c1 := randomCid()
    38  	expectedPrunedCIDs[c1] = struct{}{}
    39  	c2 := randomCid()
    40  	expectedPrunedCIDs[c2] = struct{}{}
    41  	c3 := randomCid()
    42  	c4 := randomCid()
    43  
    44  	require.NoError(t, storage.Update(func(tbf TrackBlobsFn) error {
    45  		require.NoError(t, tbf(1, c1, c2))
    46  		require.NoError(t, tbf(2, c3, c4))
    47  
    48  		return nil
    49  	}))
    50  	require.NoError(t, storage.PruneUpToHeight(1))
    51  
    52  	prunedHeight, err := storage.GetPrunedHeight()
    53  	require.NoError(t, err)
    54  	assert.Equal(t, uint64(1), prunedHeight)
    55  
    56  	assert.Len(t, expectedPrunedCIDs, 0)
    57  
    58  	err = storage.db.View(func(txn *badger.Txn) error {
    59  		_, err := txn.Get(makeBlobRecordKey(1, c1))
    60  		assert.ErrorIs(t, err, badger.ErrKeyNotFound)
    61  		_, err = txn.Get(makeLatestHeightKey(c1))
    62  		assert.ErrorIs(t, err, badger.ErrKeyNotFound)
    63  		_, err = txn.Get(makeBlobRecordKey(1, c2))
    64  		assert.ErrorIs(t, err, badger.ErrKeyNotFound)
    65  		_, err = txn.Get(makeLatestHeightKey(c2))
    66  		assert.ErrorIs(t, err, badger.ErrKeyNotFound)
    67  
    68  		_, err = txn.Get(makeBlobRecordKey(2, c3))
    69  		assert.NoError(t, err)
    70  		_, err = txn.Get(makeLatestHeightKey(c3))
    71  		assert.NoError(t, err)
    72  		_, err = txn.Get(makeBlobRecordKey(2, c4))
    73  		assert.NoError(t, err)
    74  		_, err = txn.Get(makeLatestHeightKey(c4))
    75  		assert.NoError(t, err)
    76  
    77  		return nil
    78  	})
    79  	require.NoError(t, err)
    80  }
    81  
    82  // TestPruneNonLatestHeight test that when pruning a height at which a CID exists,
    83  // if that CID also exists at another height above the pruned height, the CID should not be pruned.
    84  func TestPruneNonLatestHeight(t *testing.T) {
    85  	storageDir := t.TempDir()
    86  	storage, err := OpenStorage(storageDir, 0, zerolog.Nop(), WithPruneCallback(func(c cid.Cid) error {
    87  		assert.Fail(t, "unexpected CID pruned: %s", c.String())
    88  		return nil
    89  	}))
    90  	require.NoError(t, err)
    91  
    92  	// c1 and c2 appear both at height 1 and 2
    93  	// therefore, when pruning up to height 1, both c1 and c2 should be retained
    94  	c1 := randomCid()
    95  	c2 := randomCid()
    96  
    97  	require.NoError(t, storage.Update(func(tbf TrackBlobsFn) error {
    98  		require.NoError(t, tbf(1, c1, c2))
    99  		require.NoError(t, tbf(2, c1, c2))
   100  
   101  		return nil
   102  	}))
   103  	require.NoError(t, storage.PruneUpToHeight(1))
   104  
   105  	prunedHeight, err := storage.GetPrunedHeight()
   106  	require.NoError(t, err)
   107  	assert.Equal(t, uint64(1), prunedHeight)
   108  
   109  	err = storage.db.View(func(txn *badger.Txn) error {
   110  		_, err = txn.Get(makeBlobRecordKey(2, c1))
   111  		assert.NoError(t, err)
   112  		_, err = txn.Get(makeLatestHeightKey(c1))
   113  		assert.NoError(t, err)
   114  		_, err = txn.Get(makeBlobRecordKey(2, c2))
   115  		assert.NoError(t, err)
   116  		_, err = txn.Get(makeLatestHeightKey(c2))
   117  		assert.NoError(t, err)
   118  
   119  		return nil
   120  	})
   121  	require.NoError(t, err)
   122  }