Proposed goroutine wrapper solution (#1157)

Author: madisonchamberlain

async.go

@@ -37,7 +37,12 @@ func (sr *snowflakeRestful) processAsync(
 	}
 
 	// spawn goroutine to retrieve asynchronous results
-	go sr.getAsync(ctx, headers, sr.getFullURL(respd.Data.GetResultURL, nil), timeout, res, rows, cfg)
+	go GoroutineWrapper(
+		ctx,
+		func() {
+			sr.getAsync(ctx, headers, sr.getFullURL(respd.Data.GetResultURL, nil), timeout, res, rows, cfg)
+		},
+	)
 	return respd, nil
 }
 

authexternalbrowser.go

@@ -211,9 +211,12 @@ func authenticateByExternalBrowser(
 	disableConsoleLogin ConfigBool,
 ) ([]byte, []byte, error) {
 	resultChan := make(chan authenticateByExternalBrowserResult, 1)
-	go func() {
-		resultChan <- doAuthenticateByExternalBrowser(ctx, sr, authenticator, application, account, user, password, disableConsoleLogin)
-	}()
+	go GoroutineWrapper(
+		ctx,
+		func() {
+			resultChan <- doAuthenticateByExternalBrowser(ctx, sr, authenticator, application, account, user, password, disableConsoleLogin)
+		},
+	)
 	select {
 	case <-time.After(externalBrowserTimeout):
 		return nil, nil, errors.New("authentication timed out")

chunk_downloader.go

@@ -148,7 +148,12 @@ func (scd *snowflakeChunkDownloader) schedule() {
 	select {
 	case nextIdx := <-scd.ChunksChan:
 		logger.WithContext(scd.ctx).Infof("schedule chunk: %v", nextIdx+1)
-		go scd.FuncDownload(scd.ctx, scd, nextIdx)
+		go GoroutineWrapper(
+			scd.ctx,
+			func() {
+				scd.FuncDownload(scd.ctx, scd, nextIdx)
+			},
+		)
 	default:
 		// no more download
 		logger.WithContext(scd.ctx).Info("no more download")
@@ -162,7 +167,12 @@ func (scd *snowflakeChunkDownloader) checkErrorRetry() (err error) {
 			errc.Error != context.Canceled &&
 			errc.Error != context.DeadlineExceeded {
 			// add the index to the chunks channel so that the download will be retried.
-			go scd.FuncDownload(scd.ctx, scd, errc.Index)
+			go GoroutineWrapper(
+				scd.ctx,
+				func() {
+					scd.FuncDownload(scd.ctx, scd, errc.Index)
+				},
+			)
 			scd.ChunksErrorCounter++
 			logger.WithContext(scd.ctx).Warningf("chunk idx: %v, err: %v. retrying (%v/%v)...",
 				errc.Index, errc.Error, scd.ChunksErrorCounter, maxChunkDownloaderErrorCounter)
@@ -508,55 +518,58 @@ func (scd *streamChunkDownloader) nextResultSet() error {
 }
 
 func (scd *streamChunkDownloader) start() error {
-	go func() {
-		readErr := io.EOF
-
-		logger.WithContext(scd.ctx).Infof(
-			"start downloading. downloader id: %v, %v/%v rows, %v chunks",
-			scd.id, len(scd.RowSet.RowType), scd.Total, len(scd.ChunkMetas))
-		t := time.Now()
-
-		defer func() {
-			if readErr == io.EOF {
-				logger.WithContext(scd.ctx).Infof("downloading done. downloader id: %v", scd.id)
-			} else {
-				logger.WithContext(scd.ctx).Debugf("downloading error. downloader id: %v", scd.id)
-			}
-			scd.readErr = readErr
-			close(scd.rowStream)
-
-			if r := recover(); r != nil {
-				if err, ok := r.(error); ok {
-					readErr = err
+	go GoroutineWrapper(
+		scd.ctx,
+		func() {
+			readErr := io.EOF
+
+			logger.WithContext(scd.ctx).Infof(
+				"start downloading. downloader id: %v, %v/%v rows, %v chunks",
+				scd.id, len(scd.RowSet.RowType), scd.Total, len(scd.ChunkMetas))
+			t := time.Now()
+
+			defer func() {
+				if readErr == io.EOF {
+					logger.WithContext(scd.ctx).Infof("downloading done. downloader id: %v", scd.id)
 				} else {
-					readErr = fmt.Errorf("%v", r)
+					logger.WithContext(scd.ctx).Debugf("downloading error. downloader id: %v", scd.id)
+				}
+				scd.readErr = readErr
+				close(scd.rowStream)
+
+				if r := recover(); r != nil {
+					if err, ok := r.(error); ok {
+						readErr = err
+					} else {
+						readErr = fmt.Errorf("%v", r)
+					}
 				}
+			}()
+
+			logger.WithContext(scd.ctx).Infof("sending initial set of rows in %vms", time.Since(t).Microseconds())
+			t = time.Now()
+			for _, row := range scd.RowSet.JSON {
+				scd.rowStream <- row
 			}
-		}()
-
-		logger.WithContext(scd.ctx).Infof("sending initial set of rows in %vms", time.Since(t).Microseconds())
-		t = time.Now()
-		for _, row := range scd.RowSet.JSON {
-			scd.rowStream <- row
-		}
-		scd.RowSet.JSON = nil
-
-		// Download and parse one chunk at a time. The fetcher will send each
-		// parsed row to the row stream. When an error occurs, the fetcher will
-		// stop writing to the row stream so we can stop processing immediately
-		for i, chunk := range scd.ChunkMetas {
-			logger.WithContext(scd.ctx).Infof("starting chunk fetch %d (%d rows)", i, chunk.RowCount)
-			if err := scd.fetcher.fetch(chunk.URL, scd.rowStream); err != nil {
-				logger.WithContext(scd.ctx).Debugf(
-					"failed chunk fetch %d: %#v, downloader id: %v, %v/%v rows, %v chunks",
-					i, err, scd.id, len(scd.RowSet.RowType), scd.Total, len(scd.ChunkMetas))
-				readErr = fmt.Errorf("chunk fetch: %w", err)
-				break
+			scd.RowSet.JSON = nil
+
+			// Download and parse one chunk at a time. The fetcher will send each
+			// parsed row to the row stream. When an error occurs, the fetcher will
+			// stop writing to the row stream so we can stop processing immediately
+			for i, chunk := range scd.ChunkMetas {
+				logger.WithContext(scd.ctx).Infof("starting chunk fetch %d (%d rows)", i, chunk.RowCount)
+				if err := scd.fetcher.fetch(chunk.URL, scd.rowStream); err != nil {
+					logger.WithContext(scd.ctx).Debugf(
+						"failed chunk fetch %d: %#v, downloader id: %v, %v/%v rows, %v chunks",
+						i, err, scd.id, len(scd.RowSet.RowType), scd.Total, len(scd.ChunkMetas))
+					readErr = fmt.Errorf("chunk fetch: %w", err)
+					break
+				}
+				logger.WithContext(scd.ctx).Infof("fetched chunk %d (%d rows) in %vms", i, chunk.RowCount, time.Since(t).Microseconds())
+				t = time.Now()
 			}
-			logger.WithContext(scd.ctx).Infof("fetched chunk %d (%d rows) in %vms", i, chunk.RowCount, time.Since(t).Microseconds())
-			t = time.Now()
-		}
-	}()
+		},
+	)
 	return nil
 }
 

function_wrapper_test.go

@@ -0,0 +1,47 @@
+package gosnowflake
+
+import (
+	"context"
+	"sync"
+	"testing"
+)
+
+var (
+	goWrapperCalled          = false
+	testGoRoutineWrapperLock sync.Mutex
+)
+
+func setGoWrapperCalled(value bool) {
+	testGoRoutineWrapperLock.Lock()
+	defer testGoRoutineWrapperLock.Unlock()
+	goWrapperCalled = value
+}
+func getGoWrapperCalled() bool {
+	testGoRoutineWrapperLock.Lock()
+	defer testGoRoutineWrapperLock.Unlock()
+	return goWrapperCalled
+}
+
+// this is the go wrapper function we are going to pass into GoroutineWrapper.
+// we will know that this has been called if the channel is closed
+var closeGoWrapperCalledChannel = func(ctx context.Context, f func()) {
+	setGoWrapperCalled(true)
+	f()
+}
+
+func TestGoWrapper(t *testing.T) {
+	runDBTest(t, func(dbt *DBTest) {
+		oldGoroutineWrapper := GoroutineWrapper
+		t.Cleanup(func() {
+			GoroutineWrapper = oldGoroutineWrapper
+		})
+
+		GoroutineWrapper = closeGoWrapperCalledChannel
+
+		ctx := WithAsyncMode(context.Background())
+		rows := dbt.mustQueryContext(ctx, "SELECT 1")
+		defer rows.Close()
+
+		assertTrueF(t, getGoWrapperCalled(), "channel should be closed, indicating our wrapper worked")
+	})
+}

function_wrappers.go

@@ -0,0 +1,18 @@
+package gosnowflake
+
+import "context"
+
+// GoroutineWrapperFunc is used to wrap goroutines. This is useful if the caller wants
+// to recover panics, rather than letting panics cause a system crash. A suggestion would be to
+// use use the recover functionality, and log the panic as is most useful to you
+type GoroutineWrapperFunc func(ctx context.Context, f func())
+
+// The default GoroutineWrapperFunc; this does nothing. With this default wrapper
+// panics will take down binary as expected
+var noopGoroutineWrapper = func(_ context.Context, f func()) {
+	f()
+}
+
+// GoroutineWrapper is used to hold the GoroutineWrapperFunc set by the client, or to
+// store the default goroutine wrapper which does nothing
+var GoroutineWrapper GoroutineWrapperFunc = noopGoroutineWrapper