From 91e59e4376240fc7b309d1a7582b58ddd7f420a7 Mon Sep 17 00:00:00 2001
From: souhailanoor <90720110+souhailanoor@users.noreply.github.com>
Date: Wed, 29 Jun 2022 16:22:37 -0500
Subject: [PATCH] Create a pool of goroutines bounded to hardware limits. This
 should resolve "connection refused" error from spawning too many goroutines
 at once. (#2138)

---
 api.go             | 16 ++++++++++++----
 ctl/import_test.go |  6 +++---
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/api.go b/api.go
index 1324206eb..66fdfb8f5 100644
--- a/api.go
+++ b/api.go
@@ -15,6 +15,7 @@ import (
 	"math"
 	"net/url"
 	"os"
+	"runtime"
 	"sort"
 	"strconv"
 	"strings"
@@ -1556,11 +1557,15 @@ func (api *API) ImportWithTx(ctx context.Context, qcx *Qcx, req *ImportRequest,
 		options.IgnoreKeyCheck = true
 
 		var eg errgroup.Group
+		guard := make(chan struct{}, runtime.NumCPU()) // only run as many goroutines as CPUs available
 		for _, subReq := range reqs {
 			// TODO: if local node owns this shard we don't need to go through the client
+			guard <- struct{}{} // would block if guard channel is already filled
 			subReq := subReq
 			eg.Go(func() error {
-				return api.server.defaultClient.Import(ctx, qcx, subReq, options)
+				err := api.server.defaultClient.Import(ctx, qcx, subReq, options)
+				<-guard
+				return err
 			})
 		}
 		return eg.Wait()
@@ -1851,7 +1856,8 @@ func (api *API) ImportValueWithTx(ctx context.Context, qcx *Qcx, req *ImportValu
 	options.IgnoreKeyCheck = true
 	start := 0
 	shard := req.ColumnIDs[0] / ShardWidth
-	var eg errgroup.Group // TODO make this a pooled errgroup
+	var eg errgroup.Group
+	guard := make(chan struct{}, runtime.NumCPU()) // only run as many goroutines as CPUs available
 	for i, colID := range req.ColumnIDs {
 		if colID/ShardWidth != shard {
 			subreq := &ImportValueRequest{
@@ -1865,9 +1871,11 @@ func (api *API) ImportValueWithTx(ctx context.Context, qcx *Qcx, req *ImportValu
 			} else if req.FloatValues != nil {
 				subreq.FloatValues = req.FloatValues[start:i]
 			}
-
+			guard <- struct{}{} // would block if guard channel is already filled
 			eg.Go(func() error {
-				return api.server.defaultClient.ImportValue(ctx, qcx, subreq, options)
+				err := api.server.defaultClient.ImportValue(ctx, qcx, subreq, options)
+				<-guard
+				return err
 			})
 			start = i
 			shard = colID / ShardWidth
diff --git a/ctl/import_test.go b/ctl/import_test.go
index a5ac14b08..625475593 100644
--- a/ctl/import_test.go
+++ b/ctl/import_test.go
@@ -249,8 +249,8 @@ func TestImportCommand_KeyReplication(t *testing.T) {
 	// create a large import file in order to test the
 	// translateStoreBufferSize growth logic.
 	keyBytes := []byte{}
-	for row := 0; row < 100; row++ {
-		for col := 0; col < 100; col++ {
+	for row := 0; row < 1000; row++ {
+		for col := 0; col < 1000; col++ {
 			x := fmt.Sprintf("foo%d,bar%d\n", row, col)
 			keyBytes = append(keyBytes, x...)
 		}
@@ -303,7 +303,7 @@ func TestImportCommand_KeyReplication(t *testing.T) {
 			}
 
 			// Read body and unmarshal response.
-			exp := `{"results":[100]}` + "\n"
+			exp := `{"results":[1000]}` + "\n"
 			if body, err := ioutil.ReadAll(resp.Body); err != nil {
 				return fmt.Errorf("reading: %s", err)
 			} else if !reflect.DeepEqual(body, []byte(exp)) {