add egpool to limit number of live goroutines in errgroup

use in importbatch

Author: jaffee

egpool/egpool.go

@@ -0,0 +1,60 @@
+package egpool
+
+import (
+	"sync"
+
+	"golang.org/x/sync/errgroup"
+)
+
+type Group struct {
+	PoolSize int
+
+	jobs chan func() error
+
+	errMu    sync.Mutex
+	firstErr error
+	poolEG   errgroup.Group
+}
+
+func (eg *Group) Go(f func() error) {
+	if eg.PoolSize == 0 {
+		eg.PoolSize = 1
+	}
+
+	if eg.jobs == nil {
+		eg.jobs = make(chan func() error, eg.PoolSize*2)
+		eg.startProcessJobsPool()
+	}
+
+	eg.jobs <- f
+}
+
+func (eg *Group) startProcessJobsPool() {
+	eg.poolEG = errgroup.Group{}
+	for i := 0; i < eg.PoolSize; i++ {
+		eg.poolEG.Go(eg.processJobs)
+	}
+}
+
+func (eg *Group) processJobs() error {
+	for jobFn := range eg.jobs {
+		err := jobFn()
+		if err != nil {
+			eg.errMu.Lock()
+			if eg.firstErr == nil {
+				eg.firstErr = err
+			}
+			eg.errMu.Unlock()
+		}
+	}
+	return nil
+}
+
+func (eg *Group) Wait() error {
+	if eg.jobs == nil {
+		return nil
+	}
+	close(eg.jobs)
+	_ = eg.poolEG.Wait() // never returns err
+	return eg.firstErr
+}

egpool/egpool_test.go

@@ -0,0 +1,36 @@
+package egpool_test
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/pilosa/go-pilosa/egpool"
+)
+
+func TestEGPool(t *testing.T) {
+	eg := egpool.Group{}
+
+	a := make([]int, 10)
+
+	for i := 0; i < 10; i++ {
+		i := i
+		eg.Go(func() error {
+			a[i] = i
+			if i == 7 {
+				return errors.New("blah")
+			}
+			return nil
+		})
+	}
+
+	err := eg.Wait()
+	if err == nil || err.Error() != "blah" {
+		t.Errorf("expected err blah, got: %v", err)
+	}
+
+	for i := 0; i < 10; i++ {
+		if a[i] != i {
+			t.Errorf("expected a[%d] to be %d, but is %d", i, i, a[i])
+		}
+	}
+}

gpexp/importbatch.go

@@ -4,9 +4,9 @@ import (
 	"time"
 
 	"github.com/pilosa/go-pilosa"
+	"github.com/pilosa/go-pilosa/egpool"
 	"github.com/pilosa/pilosa/roaring"
 	"github.com/pkg/errors"
-	"golang.org/x/sync/errgroup"
 )
 
 // TODO if using column translation, column ids might get way out of
@@ -633,7 +633,7 @@ func (b *Batch) doTranslation() error {
 }
 
 func (b *Batch) doImport() error {
-	eg := errgroup.Group{}
+	eg := egpool.Group{PoolSize: 50}
 
 	frags, clearFrags, err := b.makeFragments()
 	if err != nil {
@@ -792,7 +792,7 @@ func (b *Batch) importValueData() error {
 	if shardWidth == 0 {
 		shardWidth = pilosa.DefaultShardWidth
 	}
-	eg := errgroup.Group{}
+	eg := egpool.Group{PoolSize: 50}
 	ids := make([]uint64, len(b.ids))
 	for field, values := range b.values {
 		// grow our temp ids slice to full length
@@ -864,7 +864,7 @@ func (b *Batch) importMutexData() error {
 	if shardWidth == 0 {
 		shardWidth = pilosa.DefaultShardWidth
 	}
-	eg := errgroup.Group{}
+	eg := egpool.Group{PoolSize: 50}
 	ids := make([]uint64, 0, len(b.ids))
 
 	for findex, rowIDs := range b.rowIDs {