Repository: snowflakedb/gosnowflake
Pull request #1157 (2024-06-13 23:04:17)

Proposed goroutine wrapper solution
### Description

As discussed in slack, I would like to have the ability to wrap a panic in a go routine, rather than letting it crash our service. We want to simply use the recover function to grab the panic, and then send that panic to sentry (the service we use to track errors). For us this would make debugging easier and more productive. 

I have added a variable `GoroutineWrapper` which is a function that takes in a function. I call this variable around all of the goroutines that are spawned throughout the codebase. On init, a caller can provide a function which does what it wants to wrap goroutines. If nothing is specified, we fallback on `defaultDoesNothing` which just calls the goroutine with no wrapping/handling as is happening now. Please let me know your thoughts, I am happy to change anything! 

### Checklist
- [X] Created tests which fail without the change (if possible)
- [ ] Extended the README / documentation, if necessary - I will do this once we are in a state of knowing for sure how we want this to look/if this is something we want to do 



Patch:
diff --git a/async.go b/async.go
index a290f3e..ed870d4 100644
--- a/async.go
+++ b/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
 }
 
diff --git a/authexternalbrowser.go b/authexternalbrowser.go
index cc76c8c..8fd2db8 100644
--- a/authexternalbrowser.go
+++ b/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")
diff --git a/chunk_downloader.go b/chunk_downloader.go
index 77a04be..1f1881c 100644
--- a/chunk_downloader.go
+++ b/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
 }
 
diff --git a/function_wrappers.go b/function_wrappers.go
new file mode 100644
index 0000000..d60145b
--- /dev/null
+++ b/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


Test patch:
diff --git a/function_wrapper_test.go b/function_wrapper_test.go
new file mode 100644
index 0000000..585e6af
--- /dev/null
+++ b/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")
+	})
+}