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

[zeroshade]

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 <zotthewiz...@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 <zotthewiz...@gmail.com>
    Signed-off-by: Matt Topol <zotthewiz...@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
 }