chore: use errgroup

Author: rootulp

datasquare.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"math"
 	"sync"
+
+	"golang.org/x/sync/errgroup"
 )
 
 // ErrUnevenChunks is thrown when non-nil chunks are not all of equal size.
@@ -187,37 +189,41 @@ func (ds *dataSquare) resetRoots() {
 	}
 }
 
-func (ds *dataSquare) computeRoots() {
-	var wg sync.WaitGroup
+func (ds *dataSquare) computeRoots() error {
+	g := new(errgroup.Group)
 
 	rowRoots := make([][]byte, ds.width)
 	colRoots := make([][]byte, ds.width)
 
 	for i := uint(0); i < ds.width; i++ {
-		wg.Add(2)
-
-		go func(i uint) {
-			defer wg.Done()
+		i := i // https://go.dev/doc/faq#closures_and_goroutines
+		g.Go(func() error {
 			rowRoot, err := ds.getRowRoot(i)
 			if err != nil {
-				panic(err)
+				return err
 			}
 			rowRoots[i] = rowRoot
-		}(i)
+			return nil
+		})
 
-		go func(i uint) {
-			defer wg.Done()
+		g.Go(func() error {
 			colRoot, err := ds.getColRoot(i)
 			if err != nil {
-				panic(err)
+				return err
 			}
 			colRoots[i] = colRoot
-		}(i)
+			return nil
+		})
+	}
+
+	err := g.Wait()
+	if err != nil {
+		return err
 	}
 
-	wg.Wait()
 	ds.rowRoots = rowRoots
 	ds.colRoots = colRoots
+	return nil
 }
 
 // getRowRoots returns the Merkle roots of all the rows in the square.

datasquare_test.go

@@ -146,7 +146,8 @@ func TestLazyRootGeneration(t *testing.T) {
 		colRoots = append(colRoots, colRoot)
 	}
 
-	square.computeRoots()
+	err = square.computeRoots()
+	assert.NoError(t, err)
 
 	if !reflect.DeepEqual(square.rowRoots, rowRoots) && !reflect.DeepEqual(square.colRoots, colRoots) {
 		t.Error("getRowRoot or getColRoot did not produce identical roots to computeRoots")
@@ -213,7 +214,8 @@ func BenchmarkEDSRoots(b *testing.B) {
 			func(b *testing.B) {
 				for n := 0; n < b.N; n++ {
 					square.resetRoots()
-					square.computeRoots()
+					err := square.computeRoots()
+					assert.NoError(b, err)
 				}
 			},
 		)