chore: free references to unused FlightStreams to help GC

Author: karel-rehor

src/main/java/com/influxdb/v3/client/internal/FlightSqlClient.java

@@ -30,10 +30,12 @@
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
+import java.util.ListIterator;
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Spliterator;
 import java.util.Spliterators;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 import javax.annotation.Nonnull;
@@ -72,6 +74,8 @@
 final class FlightSqlClient implements AutoCloseable {
 
     private static final Logger LOG = LoggerFactory.getLogger(FlightSqlClient.class);
+    private static int AUTOCLOSEABLE_CHECK_LIMIT = 10;
+    private static Map<AutoCloseable, Boolean> CLOSEABLE_CLOSED_LEDGER = new ConcurrentHashMap<>();
 
     private final FlightClient client;
 
@@ -134,12 +138,34 @@ Stream<VectorSchemaRoot> execute(@Nonnull final String query,
         Ticket ticket = new Ticket(json.getBytes(StandardCharsets.UTF_8));
         FlightStream stream = client.getStream(ticket, callOptionArray);
         FlightSqlIterator iterator = new FlightSqlIterator(stream);
-        autoCloseables.add(stream);
+        addToAutoCloseable(stream);
 
         Spliterator<VectorSchemaRoot> spliterator = Spliterators.spliteratorUnknownSize(iterator, Spliterator.NONNULL);
         return StreamSupport.stream(spliterator, false).onClose(iterator::close);
     }
 
+    private void addToAutoCloseable(AutoCloseable closeable) {
+        // need to occasionally clean up references to closed streams
+        // in order to ensure memory can get freed.
+        if(autoCloseables.size() > AUTOCLOSEABLE_CHECK_LIMIT) {
+            LOG.info("checking to cleanup stale flight streams from {} known streams", autoCloseables.size());
+
+            ListIterator<AutoCloseable> iter = autoCloseables.listIterator();
+            while(iter.hasNext()){
+                AutoCloseable autoCloseable = iter.next();
+                if(CLOSEABLE_CLOSED_LEDGER.get(autoCloseable)){
+                    LOG.info("removing closed stream {}", autoCloseable);
+                    CLOSEABLE_CLOSED_LEDGER.keySet().remove(autoCloseable);
+                    iter.remove();
+                }
+            }
+        }
+
+        autoCloseables.add(closeable);
+        CLOSEABLE_CLOSED_LEDGER.put(closeable, false);
+        LOG.debug("autoCloseables count {}, LEDGER count {}", autoCloseables.size(),  CLOSEABLE_CLOSED_LEDGER.size());
+    }
+
     @Override
     public void close() throws Exception {
         autoCloseables.add(client);
@@ -262,6 +288,7 @@ public boolean hasNext() {
                 // Nothing left to read - close the stream
                 try {
                     flightStream.close();
+                    CLOSEABLE_CLOSED_LEDGER.replace(flightStream, true);
                 } catch (Exception e) {
                     LOG.error("Error while closing FlightStream: ", e);
                 }

src/test/java/com/influxdb/v3/client/integration/E2ETest.java

@@ -435,6 +435,55 @@ public void testNoAllocatorMemoryLeak() {
 
     }
 
+    @EnabledIfEnvironmentVariable(named = "TESTING_INFLUXDB_URL", matches = ".*")
+    @EnabledIfEnvironmentVariable(named = "TESTING_INFLUXDB_TOKEN", matches = ".*")
+    @EnabledIfEnvironmentVariable(named = "TESTING_INFLUXDB_DATABASE", matches = ".*")
+    @Test
+    public void testMultipleQueries() throws Exception {
+
+        Instant now = Instant.now();
+        String measurement = "test_" + now.toEpochMilli() % 1000;
+
+        try (InfluxDBClient client = InfluxDBClient.getInstance(
+            System.getenv("TESTING_INFLUXDB_URL"),
+            System.getenv("TESTING_INFLUXDB_TOKEN").toCharArray(),
+            System.getenv("TESTING_INFLUXDB_DATABASE"),
+            null)) {
+
+            List<Point> points = List.of(
+                Point.measurement(measurement)
+                    .setTag("type", "test")
+                    .setFloatField("rads", 3.14)
+                    .setIntegerField("life", 42)
+                    .setTimestamp(now.minus(2, ChronoUnit.SECONDS)),
+                Point.measurement(measurement)
+                    .setTag("type", "test")
+                    .setFloatField("rads", 3.14)
+                    .setIntegerField("life", 42)
+                    .setTimestamp(now.minus(1, ChronoUnit.SECONDS)),
+                Point.measurement(measurement)
+                    .setTag("type", "test")
+                    .setFloatField("rads", 3.14)
+                    .setIntegerField("life", 42)
+                    .setTimestamp(now));
+
+            client.writePoints(points);
+            String query = "SELECT * FROM " + measurement;
+
+            // TODO just checking FlightStream cleanup should be test in FlightSqlClient
+            for(int i = 0; i < 20; i++) {
+                try (Stream<PointValues> stream = client.queryPoints(query)) {
+                    stream.forEach(pointValues -> {
+                        Assertions.assertThat(pointValues.getField("life")).isEqualTo(42L);
+                        Assertions.assertThat(pointValues.getField("rads")).isEqualTo(3.14);
+                    });
+                }
+            }
+        }
+
+
+    }
+
 
     private void assertGetDataSuccess(@Nonnull final InfluxDBClient influxDBClient) {
         influxDBClient.writePoint(