fix(go): fix potential deadlocks in reader (#60)

## What's Changed

## Fix Critical Deadlocks and Race Conditions in Snowflake Record Reader

This PR addresses multiple critical concurrency issues in the Snowflake
driver's `recordReader` that could cause complete application hangs
under normal racing conditions.

### Issues Fixed

*1. Critical Deadlock: `Release()` Blocking Forever*

*Problem*: When `Release()` was called while producer goroutines were
blocked on channel sends, a permanent deadlock occurred:

* `Release()` cancels context and attempts to drain channels
* Producer goroutines blocked on `ch <- rec` cannot see the cancellation
* Channels never close because producers never exit
* `Release()` blocks forever on `for rec := range ch`

*Fix:* Added a `done` channel that signals when all producer goroutines
have completed. `Release()` now waits for this signal before attempting
to drain channels.

*2. Severe Deadlock: Non-Context-Aware Channel Sends*

*Problem:* Channel send operations at lines 694 and 732 checked context
before the send but not during:

```go
for rr.Next() && ctx.Err() == nil {  // Context checked here
    // ... 
    ch <- rec  // But send blocks here without checking context
}
```

*Fix:* Wrapped all channel sends in `select` statements with context
awareness:

```go
select {
case chs[0] <- rec:
    // Successfully sent
case <-ctx.Done():
    rec.Release()
    return ctx.Err()
}
```

*3. Critical Race Condition: Nil Channel Reads*

*Problem:* Channels were created asynchronously in goroutines after
`newRecordReader` returned. If `Next()` was called quickly after
creation, it could read from uninitialized (nil) channels, causing
infinite blocking.

*Fix:* Initialize all channels upfront before starting any goroutines:

```go
chs := make([]chan arrow.RecordBatch, len(batches))
for i := range chs {
    chs[i] = make(chan arrow.RecordBatch, bufferSize)
}
```

*4. Goroutine Leaks on Initialization Errors*

*Problem:* Error paths only cleaned up the first channel, potentially
leaking goroutines if initialization failed after starting concurrent
operations.

*Fix:* Moved all error-prone initialization (GetStream, NewReader)
before goroutine creation, and added proper cleanup on errors.

----------------------

#### Changes

* Added `done` channel to `reader` struct to signal goroutine completion
* Initialize all channels upfront to eliminate race conditions
* Use context-aware sends with `select` statements for all channel
operations
* Update `Release()` to wait on `done` channel before draining
* Reorganize initialization to handle errors before starting goroutines
* Signal completion by closing `done` channel after all producers finish

#### Reproduction Scenarios Prevented

*Deadlock:*

1. bufferSize = 1, producer generates 2 records quickly
2. Channel becomes full after first record
3. Producer blocks on send
4. Consumer calls Release() before Next()
5. Without fix: permanent deadlock
6. With fix: producer responds to cancellation, Release() completes

*Race Condition:*

1. Query returns 3 batches
2. First batch processes quickly
3. Next() advances to second channel
4. Without fix: reads from nil channel, blocks forever
5. With fix: channel already initialized, works correctly

Backport of apache/arrow-adbc#3870.

Author: zeroshade

go/record_reader.go

@@ -548,6 +548,7 @@ type reader struct {
 	err        error
 
 	cancelFn context.CancelFunc
+	done     chan struct{} // signals all producer goroutines have finished
 }
 
 func newRecordReader(ctx context.Context, alloc memory.Allocator, ld gosnowflake.ArrowStreamLoader, bufferSize, prefetchConcurrency int, useHighPrecision bool, maxTimestampPrecision MaxTimestampPrecision) (array.RecordReader, error) {
@@ -655,48 +656,59 @@ func newRecordReader(ctx context.Context, alloc memory.Allocator, ld gosnowflake
 		return array.NewRecordReader(schema, results)
 	}
 
-	ch := make(chan arrow.RecordBatch, bufferSize)
-	group, ctx := errgroup.WithContext(compute.WithAllocator(ctx, alloc))
-	ctx, cancelFn := context.WithCancel(ctx)
-	group.SetLimit(prefetchConcurrency)
-
-	defer func() {
-		if err != nil {
-			close(ch)
-			cancelFn()
-		}
-	}()
-
-	chs := make([]chan arrow.RecordBatch, len(batches))
-	rdr := &reader{
-		refCount: 1,
-		chs:      chs,
-		err:      nil,
-		cancelFn: cancelFn,
-	}
-
+	// Handle empty batches case early
 	if len(batches) == 0 {
 		schema, err := rowTypesToArrowSchema(ctx, ld, useHighPrecision, maxTimestampPrecision)
 		if err != nil {
 			return nil, err
 		}
+		_, cancelFn := context.WithCancel(ctx)
+		rdr := &reader{
+			refCount: 1,
+			chs:      nil,
+			err:      nil,
+			cancelFn: cancelFn,
+			done:     make(chan struct{}),
+		}
+		close(rdr.done) // No goroutines to wait for
 		rdr.schema, _ = getTransformer(schema, ld, useHighPrecision, maxTimestampPrecision)
 		return rdr, nil
 	}
 
+	// Do all error-prone initialization first, before starting goroutines
 	r, err := batches[0].GetStream(ctx)
 	if err != nil {
 		return nil, errToAdbcErr(adbc.StatusIO, err)
 	}
 
 	rr, err := ipc.NewReader(r, ipc.WithAllocator(alloc))
 	if err != nil {
+		_ = r.Close() // Clean up the stream
 		return nil, adbc.Error{
 			Msg:  err.Error(),
 			Code: adbc.StatusInvalidState,
 		}
 	}
 
+	// Now setup concurrency primitives after error-prone operations
+	group, ctx := errgroup.WithContext(compute.WithAllocator(ctx, alloc))
+	ctx, cancelFn := context.WithCancel(ctx)
+	group.SetLimit(prefetchConcurrency)
+
+	// Initialize all channels upfront to avoid race condition
+	chs := make([]chan arrow.RecordBatch, len(batches))
+	for i := range chs {
+		chs[i] = make(chan arrow.RecordBatch, bufferSize)
+	}
+
+	rdr := &reader{
+		refCount: 1,
+		chs:      chs,
+		err:      nil,
+		cancelFn: cancelFn,
+		done:     make(chan struct{}),
+	}
+
 	var recTransform recordTransformer
 	rdr.schema, recTransform = getTransformer(rr.Schema(), ld, useHighPrecision, maxTimestampPrecision)
 
@@ -706,7 +718,7 @@ func newRecordReader(ctx context.Context, alloc memory.Allocator, ld gosnowflake
 			err = errors.Join(err, r.Close())
 		}()
 		if len(batches) > 1 {
-			defer close(ch)
+			defer close(chs[0])
 		}
 
 		for rr.Next() && ctx.Err() == nil {
@@ -715,18 +727,25 @@ func newRecordReader(ctx context.Context, alloc memory.Allocator, ld gosnowflake
 			if err != nil {
 				return err
 			}
-			ch <- rec
+
+			// Use context-aware send to prevent deadlock
+			select {
+			case chs[0] <- rec:
+				// Successfully sent
+			case <-ctx.Done():
+				// Context cancelled, clean up and exit
+				rec.Release()
+				return ctx.Err()
+			}
 		}
 		return rr.Err()
 	})
 
-	chs[0] = ch
-
 	lastChannelIndex := len(chs) - 1
 	go func() {
 		for i, b := range batches[1:] {
 			batch, batchIdx := b, i+1
-			chs[batchIdx] = make(chan arrow.RecordBatch, bufferSize)
+			// Channels already initialized above, no need to create them here
 			group.Go(func() (err error) {
 				// close channels (except the last) so that Next can move on to the next channel properly
 				if batchIdx != lastChannelIndex {
@@ -753,7 +772,16 @@ func newRecordReader(ctx context.Context, alloc memory.Allocator, ld gosnowflake
 					if err != nil {
 						return err
 					}
-					chs[batchIdx] <- rec
+
+					// Use context-aware send to prevent deadlock
+					select {
+					case chs[batchIdx] <- rec:
+						// Successfully sent
+					case <-ctx.Done():
+						// Context cancelled, clean up and exit
+						rec.Release()
+						return ctx.Err()
+					}
 				}
 
 				return rr.Err()
@@ -768,6 +796,8 @@ func newRecordReader(ctx context.Context, alloc memory.Allocator, ld gosnowflake
 		// don't close the last channel until after the group is finished,
 		// so that Next() can only return after reader.err may have been set
 		close(chs[lastChannelIndex])
+		// Signal that all producer goroutines have finished
+		close(rdr.done)
 	}()
 
 	return rdr, nil
@@ -819,7 +849,17 @@ func (r *reader) Release() {
 			r.rec.Release()
 		}
 		r.cancelFn()
+
+		// Wait for all producer goroutines to finish before draining channels
+		// This prevents deadlock where producers are blocked on sends
+		<-r.done
+
+		// Now safely drain remaining data from channels
+		// All channels should be closed at this point
 		for _, ch := range r.chs {
+			if ch == nil {
+				continue
+			}
 			for rec := range ch {
 				rec.Release()
 			}