perf: further improve PartSet (backport #2439) (#2469)

Various performance optimizations to `PartSet`, especially related to
`Encode` and `Decode` function.

Resolves: #2403 

---

You can checkout first commit to get benchmark "baseline".

```
 go-perftuner benchstat main.txt latest.txt
name                                                             old time/op    new time/op    delta
PartSetEncodeDecode/dataSize=130857600,partSize=65536/encode-16     633ms ± 3%     294ms ± 8%  -53.64%  (p=0.000 n=10+10)
PartSetEncodeDecode/dataSize=130857600,partSize=65536/decode-16    78.1ms ± 0%    78.3ms ± 2%        ~  (p=0.489 n=9+9)
PartSetEncodeDecode/dataSize=523430400,partSize=65536/encode-16     2.68s ± 2%     1.26s ± 5%  -52.81%  (p=0.000 n=10+10)
PartSetEncodeDecode/dataSize=523430400,partSize=65536/decode-16     309ms ± 0%     305ms ± 0%        ~  (p=0.000 n=9+8)

name                                                             old alloc/op   new alloc/op   delta
PartSetEncodeDecode/dataSize=130857600,partSize=65536/encode-16     681MB ± 0%     532MB ± 0%  -21.89%  (p=0.000 n=9+8)
PartSetEncodeDecode/dataSize=130857600,partSize=65536/decode-16    4.88MB ± 0%    3.06MB ± 0%  -37.32%  (p=0.000 n=10+10)
PartSetEncodeDecode/dataSize=523430400,partSize=65536/encode-16    2.72GB ± 0%    2.13GB ± 0%  -21.90%  (p=0.000 n=10+9)
PartSetEncodeDecode/dataSize=523430400,partSize=65536/decode-16    19.6MB ± 0%    12.8MB ± 0%  -34.57%  (p=0.000 n=10+10)

name                                                             old allocs/op  new allocs/op  delta
PartSetEncodeDecode/dataSize=130857600,partSize=65536/encode-16     62.0k ± 0%     24.1k ± 0%  -61.18%  (p=0.000 n=10+10)
PartSetEncodeDecode/dataSize=130857600,partSize=65536/decode-16     64.1k ± 0%     48.1k ± 0%  -24.93%  (p=0.000 n=10+8)
PartSetEncodeDecode/dataSize=523430400,partSize=65536/encode-16      264k ± 0%       96k ± 0%  -63.60%  (p=0.000 n=10+9)
PartSetEncodeDecode/dataSize=523430400,partSize=65536/decode-16      256k ± 0%      192k ± 0%  -24.98%  (p=0.000 n=10+10)
```<hr>This is an automatic backport of pull request #2439 done by
[Mergify](https://mergify.com).

Co-authored-by: Tomasz Zdybał <[email protected]>

Author: mergify[bot]

crypto/merkle/proof.go

@@ -215,7 +215,14 @@ type ProofNode struct {
 // starting from a leaf ProofNode.
 func (spn *ProofNode) FlattenAunts() [][]byte {
 	// Nonrecursive impl.
-	innerHashes := [][]byte{}
+
+	// count the length of the path from the root to the leaf.
+	cnt := 0
+	for p := spn; p != nil; p = p.Parent {
+		cnt++
+	}
+
+	innerHashes := make([][]byte, 0, cnt-1)
 	for spn != nil {
 		switch {
 		case spn.Left != nil:

types/part_set.go

@@ -255,42 +255,83 @@ func NewPartSetFromData(data []byte, partSize uint32) (ops *PartSet, err error)
 	ops.proofs = make([]merkle.Proof, total)
 	for i := uint32(0); i < total; i++ {
 		ops.proofs[i] = *proofs[i]
-		ops.partsBitArray.SetIndex(int(i), true)
 	}
+	ops.partsBitArray.Fill()
 	ops.count = total
 	ops.byteSize = int64(len(data))
 
 	return ops, nil
 }
 
+// newPartSetFromChunks creates a new PartSet from given data chunks, and other data.
+func newPartSetFromChunks(chunks [][]byte, root cmtbytes.HexBytes, proofs []*merkle.Proof, partSize int) (*PartSet, error) {
+	total := len(chunks)
+	if total != len(proofs) {
+		return nil, fmt.Errorf("chunks and proofs have different lengths: %d != %d", len(chunks), len(proofs))
+	}
+	if root == nil {
+		return nil, fmt.Errorf("root is nil")
+	}
+
+	// create a new partset using the new parity parts.
+	ps := NewPartSetFromHeader(PartSetHeader{
+		Total: uint32(total),
+		Hash:  root,
+	}, uint32(partSize))
+
+	// access ps directly, without mutex, because we know it is not used elsewhere
+	for i := 0; i < total; i++ {
+		start := i * partSize
+		end := start + len(chunks[i])
+
+		// Ensure we don't exceed buffer bounds
+		if end > len(ps.buffer) {
+			return nil, fmt.Errorf("part data exceeds buffer bounds")
+		}
+
+		copy(ps.buffer[start:end], chunks[i])
+		ps.proofs[i] = *proofs[i]
+	}
+	ps.partsBitArray.Fill()
+	ps.count = uint32(total)
+	ps.lastPartSize = len(chunks[total-1])
+	ps.byteSize = int64(len(ps.buffer))
+	return ps, nil
+}
+
 // Encode Extend erasure encodes the block parts. Only the original parts should be
 // provided. The parity data is formed into its own PartSet and returned
 // alongside the length of the last part. The length of the last part is
 // necessary because the last part may be padded with zeros after decoding. These zeros must be removed before computi
 func Encode(ops *PartSet, partSize uint32) (*PartSet, int, error) {
-	chunks := make([][]byte, 2*ops.Total())
-	for i := range chunks {
-		if i < int(ops.Total()) {
-			chunks[i] = ops.GetPart(i).Bytes.Bytes()
-			continue
-		}
-		chunks[i] = make([]byte, partSize)
+	total := int(ops.Total())
+	chunks := make([][]byte, 2*total)
+	ops.mtx.Lock()
+	for i := range total {
+		chunks[i] = ops.getPartBytes(i)
+	}
+	ops.mtx.Unlock()
+
+	ps := int(partSize)
+	parityBuffer := make([]byte, total*ps) // allocate once, only slice later
+	for i := 0; i < total; i++ {
+		chunks[total+i] = parityBuffer[i*ps : (i+1)*ps]
 	}
 
 	// pad ONLY the last chunk and not the part with zeros if necessary AFTER the root has been generated
-	lastLen := len(ops.GetPart(int(ops.Total() - 1)).Bytes.Bytes())
+	lastLen := len(ops.GetPartBytes(total - 1))
 	if lastLen < int(partSize) {
 		padded := make([]byte, partSize)
-		count := copy(padded, chunks[ops.Total()-1])
-		if count < len(chunks[ops.Total()-1]) {
+		count := copy(padded, chunks[total-1])
+		if count < len(chunks[total-1]) {
 			return nil, 0, fmt.Errorf("copy failed of unpadded part with index %d: %d < %d", ops.Total()-1, count, len(chunks[ops.Total()-1]))
 		}
-		chunks[ops.Total()-1] = padded
+		chunks[total-1] = padded
 	}
 
 	// init an encoder if it is not already initialized using the original
 	// number of parts.
-	enc, err := reedsolomon.New(int(ops.Total()), int(ops.Total()))
+	enc, err := reedsolomon.New(total, total)
 	if err != nil {
 		return nil, 0, err
 	}
@@ -302,26 +343,12 @@ func Encode(ops *PartSet, partSize uint32) (*PartSet, int, error) {
 	}
 
 	// only the parity data is needed for the new partset.
-	chunks = chunks[ops.Total():]
+	chunks = chunks[total:]
 	eroot, eproofs := merkle.ParallelProofsFromByteSlices(chunks)
 
-	// create a new partset using the new parity parts.
-	eps := NewPartSetFromHeader(PartSetHeader{
-		Total: ops.Total(),
-		Hash:  eroot,
-	}, partSize)
-	for i := uint32(0); i < ops.Total(); i++ {
-		added, err := eps.AddPart(&Part{
-			Index: i,
-			Bytes: chunks[i],
-			Proof: *eproofs[i],
-		})
-		if err != nil {
-			return nil, 0, err
-		}
-		if !added {
-			return nil, 0, fmt.Errorf("couldn't add parity part %d", i)
-		}
+	eps, err := newPartSetFromChunks(chunks, eroot, eproofs, ps)
+	if err != nil {
+		return nil, 0, err
 	}
 	return eps, lastLen, nil
 }
@@ -541,9 +568,6 @@ func (ps *PartSet) addPart(part *Part) (bool, error) {
 		return false, errors.New("nil part")
 	}
 
-	ps.mtx.Lock()
-	defer ps.mtx.Unlock()
-
 	// Invalid part index
 	if part.Index >= ps.total {
 		return false, ErrPartSetUnexpectedIndex
@@ -563,18 +587,22 @@ func (ps *PartSet) addPart(part *Part) (bool, error) {
 		return false, fmt.Errorf("part data exceeds buffer bounds")
 	}
 
-	copy(ps.buffer[start:end], part.Bytes)
+	ps.mtx.Lock()
 
+	copy(ps.buffer[start:end], part.Bytes)
 	ps.proofs[part.Index] = part.Proof
 
 	// Track last part size if this is the last part
 	if part.Index == ps.total-1 {
 		ps.lastPartSize = len(part.Bytes)
 	}
-
-	ps.partsBitArray.SetIndex(int(part.Index), true)
 	ps.count++
 	ps.byteSize += int64(len(part.Bytes))
+
+	ps.mtx.Unlock()
+
+	ps.partsBitArray.SetIndex(int(part.Index), true)
+
 	return true, nil
 }
 
@@ -632,8 +660,6 @@ func (ps *PartSet) GetPart(index int) *Part {
 }
 
 func (ps *PartSet) HasPart(index int) bool {
-	ps.mtx.Lock()
-	defer ps.mtx.Unlock()
 	return ps.partsBitArray.GetIndex(index)
 }
 

types/part_set_test.go

@@ -441,3 +441,37 @@ func benchPartSetFromData(b *testing.B, data []byte, partSize uint32) (ops *Part
 
 	return ops, nil
 }
+
+func BenchmarkPartSetEncodeDecode(b *testing.B) {
+	cases := []struct {
+		dataSize int
+		partSize uint32
+	}{
+		{MaxBlockSizeBytes, BlockPartSizeBytes},
+		{4 * MaxBlockSizeBytes, BlockPartSizeBytes},
+	}
+
+	for c := range cases {
+		b.Run(fmt.Sprintf("dataSize=%d,partSize=%d", cases[c].dataSize, cases[c].partSize), func(b *testing.B) {
+			data := cmtrand.Bytes(cases[c].dataSize)
+			ops, err := NewPartSetFromData(data, cases[c].partSize)
+			require.NoError(b, err)
+			var eps *PartSet
+			lastPartLen := 0
+			b.Run("encode", func(b *testing.B) {
+				for i := 0; i < b.N; i++ {
+					b.ReportAllocs()
+					eps, lastPartLen, err = Encode(ops, cases[c].partSize)
+					require.NoError(b, err)
+				}
+			})
+			b.Run("decode", func(b *testing.B) {
+				for i := 0; i < b.N; i++ {
+					b.ReportAllocs()
+					_, _, err := Decode(ops, eps, lastPartLen)
+					require.NoError(b, err)
+				}
+			})
+		})
+	}
+}