[arrow] branch main updated: GH-35328: [Go][FlightSQL] Fix flaky test for FlightSql driver (#38044)

[ze...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git


The following commit(s) were added to refs/heads/main by this push:
     new 1cad7a7878 GH-35328: [Go][FlightSQL] Fix flaky test for FlightSql driver (#38044)
1cad7a7878 is described below

commit 1cad7a787847add45261abe5acdd3ba080e54a75
Author: Matt Topol <zo...@gmail.com>
AuthorDate: Thu Oct 5 12:49:15 2023 -0400

    GH-35328: [Go][FlightSQL] Fix flaky test for FlightSql driver (#38044)
    
    
    
    ### Rationale for this change
    Fixing a flaky test that is *very* difficult to reproduce.
    
    ### What changes are included in this PR?
    Adding an explicit call to drain the remaining batches from the parameter batch reader in the `DoPutPreparedStatementQuery` of the `MockServer` in the tests. This ensures that the gRPC connection doesn't close on us in between the client writing the schema message and writing the record batch message producing an `io.EOF` error. This is an extremely rare occurence based on goroutine scheduling due to the client and mockserver both running in the same process for the tests, in local te [...]
    
    ### Are there any user-facing changes?
    No
    
    * Closes: #35328
    
    Authored-by: Matt Topol <zo...@gmail.com>
    Signed-off-by: Matt Topol <zo...@gmail.com>
---
 go/arrow/flight/flightsql/driver/driver_test.go | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/go/arrow/flight/flightsql/driver/driver_test.go b/go/arrow/flight/flightsql/driver/driver_test.go
index 3ec8533979..4449b39097 100644
--- a/go/arrow/flight/flightsql/driver/driver_test.go
+++ b/go/arrow/flight/flightsql/driver/driver_test.go
@@ -812,6 +812,17 @@ func (s *MockServer) DoPutPreparedStatementQuery(ctx context.Context, qry flight
 		return fmt.Errorf("parameter schema: %w", arrow.ErrInvalid)
 	}
 
+	// GH-35328: it's rare, but this function can complete execution and return
+	// closing the reader *after* the schema is written but *before* the parameter batch
+	// is written (race condition based on goroutine scheduling). In that situation,
+	// the client call to Write the parameter record batch will return an io.EOF because
+	// this end of the connection will have closed before it attempted to send the batch.
+	// This created a flaky test situation that was difficult to reproduce (1-4 failures
+	// in 5000 runs). We can avoid this flakiness by simply *explicitly* draining the
+	// record batch messages from the reader before returning.
+	for r.Next() {
+	}
+
 	return nil
 }