GH-62: fixing missing concurrent close

Author: xborder

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java

@@ -32,6 +32,7 @@
 import org.apache.arrow.util.AutoCloseables;
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.types.pojo.Schema;
+import org.apache.calcite.avatica.AvaticaConnection;
 import org.apache.calcite.avatica.AvaticaResultSet;
 import org.apache.calcite.avatica.AvaticaResultSetMetaData;
 import org.apache.calcite.avatica.AvaticaStatement;
@@ -209,6 +210,13 @@ public boolean next() throws SQLException {
         continue;
       }
 
+      // No more data. If the queue was closed concurrently (e.g. statement.cancel()
+      // racing with the reader past super.next()), surface as "Statement canceled"
+      // to match Avatica's cancellation semantics.
+      if (flightEndpointDataQueue.isClosed()) {
+        throw AvaticaConnection.HELPER.createException("Statement canceled");
+      }
+
       if (statement != null && statement.isCloseOnCompletion()) {
         statement.close();
       }

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java

@@ -177,7 +177,17 @@ public void enqueue(final Collection<CloseableEndpointStreamPair> endpointReques
   /** Adds given {@link FlightStream} to the queue. */
   public synchronized void enqueue(final CloseableEndpointStreamPair endpointRequest) {
     checkNotNull(endpointRequest);
-    checkOpen();
+    if (isClosed()) {
+      // Concurrent close() raced with enqueue (e.g. statement.cancel() while the reader
+      // was mid-loop). Treat as end-of-stream: close the endpoint and return silently.
+      // The reader detects cancellation via a subsequent isClosed() check.
+      try {
+        endpointRequest.close();
+      } catch (final Exception e) {
+        LOGGER.error("Failed to close endpoint after queue was closed.", e);
+      }
+      return;
+    }
     endpointsToClose.add(endpointRequest);
     futures.add(
         completionService.submit(

flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueueTest.java

@@ -21,6 +21,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 
 import java.util.concurrent.CompletionService;
 import org.apache.arrow.driver.jdbc.client.CloseableEndpointStreamPair;
@@ -54,10 +55,11 @@ public void testNextShouldReturnNullUponClose() throws Exception {
   }
 
   @Test
-  public void testEnqueueShouldThrowExceptionUponClose() throws Exception {
+  public void testEnqueueShouldCloseEndpointSilentlyAfterClose() throws Exception {
     queue.close();
-    ThrowableAssertionUtils.simpleAssertThrowableClass(
-        IllegalStateException.class, () -> queue.enqueue(mock(CloseableEndpointStreamPair.class)));
+    final CloseableEndpointStreamPair endpoint = mock(CloseableEndpointStreamPair.class);
+    assertDoesNotThrow(() -> queue.enqueue(endpoint));
+    verify(endpoint).close();
   }
 
   @Test