Fix memory leak

Author: ennuite

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

@@ -21,6 +21,8 @@
 import io.netty.util.concurrent.DefaultThreadFactory;
 import java.sql.SQLException;
 import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Properties;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -42,6 +44,8 @@ public final class ArrowFlightConnection extends AvaticaConnection {
   private final ArrowFlightSqlClientHandler clientHandler;
   private final ArrowFlightConnectionConfigImpl config;
   private ExecutorService executorService;
+  int metadataResultSetCount;
+  Map<Integer, ArrowFlightJdbcFlightStreamResultSet> metadataResultSetMap = new HashMap<>();
 
   /**
    * Creates a new {@link ArrowFlightConnection}.
@@ -66,6 +70,7 @@ private ArrowFlightConnection(
     this.config = Preconditions.checkNotNull(config, "Config cannot be null.");
     this.allocator = Preconditions.checkNotNull(allocator, "Allocator cannot be null.");
     this.clientHandler = Preconditions.checkNotNull(clientHandler, "Handler cannot be null.");
+    this.metadataResultSetCount = 0;
   }
 
   /**
@@ -190,7 +195,9 @@ public void close() throws SQLException {
     } catch (final Exception e) {
       topLevelException = e;
     }
+    // copies of the collections are used to avoid concurrent modification problems
     ArrayList<AutoCloseable> closeables = new ArrayList<>(statementMap.values());
+    closeables.addAll(new ArrayList<>(metadataResultSetMap.values()));
     closeables.add(clientHandler);
     closeables.addAll(allocator.getChildAllocators());
     closeables.add(allocator);

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

@@ -54,6 +54,7 @@ public final class ArrowFlightJdbcFlightStreamResultSet
   private VectorSchemaRoot currentVectorSchemaRoot;
 
   private Schema schema;
+  private Integer id = null; // used for metadata result sets only
 
   /** Public constructor used by ArrowFlightJdbcFactory. */
   ArrowFlightJdbcFlightStreamResultSet(
@@ -82,6 +83,8 @@ private ArrowFlightJdbcFlightStreamResultSet(
     super(null, state, signature, resultSetMetaData, timeZone, firstFrame);
     this.connection = connection;
     this.flightInfo = flightInfo;
+    this.id = connection.metadataResultSetCount++;
+    connection.metadataResultSetMap.put(id, this);
   }
 
   /**
@@ -234,7 +237,14 @@ protected void cancel() {
 
   @Override
   public synchronized void close() {
+
     try {
+      if (isClosed()) {
+        return;
+      }
+      if (id != null) { // id is only set for metadata result sets
+        this.connection.metadataResultSetMap.remove(id);
+      }
       if (flightEndpointDataQueue != null) {
         // flightStreamQueue should close currentFlightStream internally
         flightEndpointDataQueue.close();