refactor chunk data packs storage

Author: zhangchiqing

Diff for: storage/badger/all.go

@@ -29,7 +29,6 @@ func InitAll(metrics module.CacheMetrics, db *badger.DB) *storage.All {
 	transactionResults := NewTransactionResults(metrics, db, 10000)
 	collections := NewCollections(db, transactions)
 	events := NewEvents(metrics, db)
-	chunkDataPacks := NewChunkDataPacks(metrics, db, collections, 1000)
 
 	return &storage.All{
 		Headers:                   headers,
@@ -46,7 +45,6 @@ func InitAll(metrics module.CacheMetrics, db *badger.DB) *storage.All {
 		VersionBeacons:            versionBeacons,
 		Results:                   results,
 		Receipts:                  receipts,
-		ChunkDataPacks:            chunkDataPacks,
 		Commits:                   commits,
 		Transactions:              transactions,
 		TransactionResults:        transactionResults,

Diff for: storage/badger/chunk_data_pack_test.go

@@ -21,7 +21,7 @@ type ChunkDataPacks interface {
 	// BatchRemove removes ChunkDataPack c keyed by its ChunkID in provided batch
 	// No errors are expected during normal operation, even if no entries are matched.
 	// If Badger unexpectedly fails to process the request, the error is wrapped in a generic error and returned.
-	BatchRemove(chunkID flow.Identifier, batch BatchStorage) error
+	BatchRemove(chunkID flow.Identifier, batch ReaderBatchWriter) error
 }
 
 // StoredChunkDataPack is an in-storage representation of chunk data pack.

Diff for: storage/badger/operation/chunkDataPacks.go

@@ -0,0 +1,24 @@
+package operation
+
+import (
+	"github.com/onflow/flow-go/model/flow"
+	"github.com/onflow/flow-go/storage"
+)
+
+// InsertChunkDataPack inserts a chunk data pack keyed by chunk ID.
+// any error are exceptions
+func InsertChunkDataPack(w storage.Writer, c *storage.StoredChunkDataPack) error {
+	return UpsertByKey(w, MakePrefix(codeChunkDataPack, c.ChunkID), c)
+}
+
+// RetrieveChunkDataPack retrieves a chunk data pack by chunk ID.
+// it returns storage.ErrNotFound if the chunk data pack is not found
+func RetrieveChunkDataPack(r storage.Reader, chunkID flow.Identifier, c *storage.StoredChunkDataPack) error {
+	return RetrieveByKey(r, MakePrefix(codeChunkDataPack, chunkID), c)
+}
+
+// RemoveChunkDataPack removes the chunk data pack with the given chunk ID.
+// any error are exceptions
+func RemoveChunkDataPack(w storage.Writer, chunkID flow.Identifier) error {
+	return RemoveByKey(w, MakePrefix(codeChunkDataPack, chunkID))
+}

Diff for: storage/chunkDataPacks.go renamed to storage/chunk_data_packs.go

@@ -1,18 +1,19 @@
-package operation
+package operation_test
 
 import (
 	"testing"
 
-	"github.com/dgraph-io/badger/v2"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/onflow/flow-go/storage"
+	"github.com/onflow/flow-go/storage/operation"
+	"github.com/onflow/flow-go/storage/operation/dbtest"
 	"github.com/onflow/flow-go/utils/unittest"
 )
 
 func TestChunkDataPack(t *testing.T) {
-	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
+	dbtest.RunWithDB(t, func(t *testing.T, db storage.DB) {
 		collectionID := unittest.IdentifierFixture()
 		expected := &storage.StoredChunkDataPack{
 			ChunkID:      unittest.IdentifierFixture(),
@@ -23,27 +24,31 @@ func TestChunkDataPack(t *testing.T) {
 
 		t.Run("Retrieve non-existent", func(t *testing.T) {
 			var actual storage.StoredChunkDataPack
-			err := db.View(RetrieveChunkDataPack(expected.ChunkID, &actual))
+			err := operation.RetrieveChunkDataPack(db.Reader(), expected.ChunkID, &actual)
 			assert.Error(t, err)
 		})
 
 		t.Run("Save", func(t *testing.T) {
-			err := db.Update(InsertChunkDataPack(expected))
+			err := db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
+				return operation.InsertChunkDataPack(rw.Writer(), expected)
+			})
 			require.NoError(t, err)
 
 			var actual storage.StoredChunkDataPack
-			err = db.View(RetrieveChunkDataPack(expected.ChunkID, &actual))
+			err = operation.RetrieveChunkDataPack(db.Reader(), expected.ChunkID, &actual)
 			assert.NoError(t, err)
 
 			assert.Equal(t, *expected, actual)
 		})
 
 		t.Run("Remove", func(t *testing.T) {
-			err := db.Update(RemoveChunkDataPack(expected.ChunkID))
+			err := db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
+				return operation.RemoveChunkDataPack(rw.Writer(), expected.ChunkID)
+			})
 			require.NoError(t, err)
 
 			var actual storage.StoredChunkDataPack
-			err = db.View(RetrieveChunkDataPack(expected.ChunkID, &actual))
+			err = operation.RetrieveChunkDataPack(db.Reader(), expected.ChunkID, &actual)
 			assert.Error(t, err)
 		})
 	})