Merge pull request #464 from celestiaorg/hlib/shares-batching

Author: Wondertan

go.mod

@@ -21,6 +21,7 @@ require (
 	github.com/ipfs/go-ds-badger2 v0.1.1
 	github.com/ipfs/go-ipfs-blockstore v0.1.6
 	github.com/ipfs/go-ipfs-exchange-interface v0.0.1
+	github.com/ipfs/go-ipfs-exchange-offline v0.0.1
 	github.com/ipfs/go-ipfs-routing v0.1.0
 	github.com/ipfs/go-ipld-format v0.2.0
 	github.com/ipfs/go-log/v2 v2.5.0

ipld/read_test.go

@@ -10,8 +10,14 @@ import (
 	"testing"
 	"time"
 
+	"github.com/ipfs/go-blockservice"
 	"github.com/ipfs/go-cid"
+	ds "github.com/ipfs/go-datastore"
+	dssync "github.com/ipfs/go-datastore/sync"
+	blockstore "github.com/ipfs/go-ipfs-blockstore"
+	offline "github.com/ipfs/go-ipfs-exchange-offline"
 	format "github.com/ipfs/go-ipld-format"
+	"github.com/ipfs/go-merkledag"
 	mdutils "github.com/ipfs/go-merkledag/test"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -365,6 +371,43 @@ func TestGetLeavesByNamespace_MultipleRowsContainingSameNamespaceId(t *testing.T
 
 }
 
+func TestBatchSize(t *testing.T) {
+	tests := []struct {
+		name      string
+		origWidth int
+	}{
+		{"2", 2},
+		{"4", 4},
+		{"8", 8},
+		{"16", 16},
+		{"32", 32},
+		// {"64", 64}, // test case too large for CI with race detector
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second*time.Duration(tt.origWidth))
+			defer cancel()
+
+			bs := blockstore.NewBlockstore(dssync.MutexWrap(ds.NewMapDatastore()))
+			dag := merkledag.NewDAGService(blockservice.New(bs, offline.Exchange(bs)))
+
+			eds := generateRandEDS(t, tt.origWidth)
+			_, err := PutData(ctx, ExtractODSShares(eds), dag)
+			require.NoError(t, err)
+
+			out, err := bs.AllKeysChan(ctx)
+			require.NoError(t, err)
+
+			var count int
+			for range out {
+				count++
+			}
+			extendedWidth := tt.origWidth * 2
+			assert.Equalf(t, count, batchSize(extendedWidth), "batchSize(%v)", extendedWidth)
+		})
+	}
+}
+
 func putErasuredDataToDag(t *testing.T, rawData [][]byte) (format.DAGService, da.DataAvailabilityHeader) {
 	// calc square size
 	squareSize := uint64(math.Sqrt(float64(len(rawData))))

ipld/write.go

@@ -13,23 +13,17 @@ import (
 	"github.com/tendermint/tendermint/pkg/wrapper"
 )
 
-// BatchSize defines an amount of IPLD Nodes to be buffered and written at once.
-// This configuration is very database backend specific and the current default(128) is not optimized for the version of
-// Badger we use. We set it to one to avoid test flakiness, as some tests may read for data that was not written yet.
-// TODO(@Wondertan): Find out the perfect value for Badger(e.g. ask PL folks) or migrate to go-car IPLD
-//  storage(preferred).
-const BatchSize = 1
-
 // PutData posts erasured block data to IPFS using the provided ipld.NodeAdder.
 func PutData(ctx context.Context, shares [][]byte, adder ipld.NodeAdder) (*rsmt2d.ExtendedDataSquare, error) {
 	if len(shares) == 0 {
 		return nil, fmt.Errorf("empty data") // empty block is not an empty Data
 	}
-	// create nmt adder wrapping batch adder
-	batchAdder := NewNmtNodeAdder(ctx, ipld.NewBatch(ctx, adder, ipld.MaxSizeBatchOption(BatchSize)))
+	squareSize := int(math.Sqrt(float64(len(shares))))
+	// create nmt adder wrapping batch adder with calculated size
+	bs := batchSize(squareSize * 2)
+	batchAdder := NewNmtNodeAdder(ctx, ipld.NewBatch(ctx, adder, ipld.MaxSizeBatchOption(bs)))
 	// create the nmt wrapper to generate row and col commitments
-	squareSize := uint64(math.Sqrt(float64(len(shares))))
-	tree := wrapper.NewErasuredNamespacedMerkleTree(squareSize, nmt.NodeVisitor(batchAdder.Visit))
+	tree := wrapper.NewErasuredNamespacedMerkleTree(uint64(squareSize), nmt.NodeVisitor(batchAdder.Visit))
 	// recompute the eds
 	eds, err := rsmt2d.ComputeExtendedDataSquare(shares, rsmt2d.NewRSGF8Codec(), tree.Constructor)
 	if err != nil {
@@ -55,3 +49,20 @@ func ExtractODSShares(eds *rsmt2d.ExtendedDataSquare) [][]byte {
 	}
 	return origShares
 }
+
+// batchSize calculates the amount of nodes that are generated from block of 'squareSizes'
+// to be batched in one write.
+func batchSize(squareSize int) int {
+	// (squareSize*2-1) - amount of nodes in a generated binary tree
+	// squareSize*2 - the total number of trees, both for rows and cols
+	// (squareSize*squareSize) - all the shares
+	//
+	// Note that while our IPLD tree looks like this:
+	// ---X
+	// -X---X
+	// X-X-X-X
+	// X-X-X-X
+	// here we count leaves only once: the CIDs are the same for columns and rows
+	// and for the last two layers as well:
+	return (squareSize*2-1)*squareSize*2 - (squareSize * squareSize)
+}