Delete results from transaction results holder when fully consumed (#…

…1571)

driver/src/main/java/org/neo4j/driver/internal/FailableCursor.java

@@ -28,4 +28,6 @@ public interface FailableCursor {
  * Pulling all unconsumed records into memory and returning failure if there is any pull errors.
  */
  CompletionStage<Throwable> pullAllFailureAsync();
+
+ CompletionStage<Void> consumed();
 }

driver/src/main/java/org/neo4j/driver/internal/async/ResultCursorsHolder.java

@@ -20,30 +20,39 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import org.neo4j.driver.internal.FailableCursor;
 
 public class ResultCursorsHolder {
- private final List<CompletionStage<? extends FailableCursor>> cursorStages =
- Collections.synchronizedList(new ArrayList<>());
+ private final List<CompletionStage<? extends FailableCursor>> cursorStages = new ArrayList<>();
 
- public void add(CompletionStage<? extends FailableCursor> cursorStage) {
+ void add(CompletionStage<? extends FailableCursor> cursorStage) {
  Objects.requireNonNull(cursorStage);
- cursorStages.add(cursorStage);
+ synchronized (this) {
+ cursorStages.add(cursorStage);
+ }
+ cursorStage.thenCompose(FailableCursor::consumed).whenComplete((ignored, throwable) -> {
+ synchronized (this) {
+ cursorStages.remove(cursorStage);
+ }
+ });
  }
 
  CompletionStage<Throwable> retrieveNotConsumedError() {
- var failures = retrieveAllFailures();
-
+ List<CompletionStage<? extends FailableCursor>> cursorStages;
+ synchronized (this) {
+ cursorStages = List.copyOf(this.cursorStages);
+ }
+ var failures = retrieveAllFailures(cursorStages);
  return CompletableFuture.allOf(failures).thenApply(ignore -> findFirstFailure(failures));
  }
 
  @SuppressWarnings("unchecked")
- private CompletableFuture<Throwable>[] retrieveAllFailures() {
+ private static CompletableFuture<Throwable>[] retrieveAllFailures(
+ List<CompletionStage<? extends FailableCursor>> cursorStages) {
  return cursorStages.stream()
  .map(ResultCursorsHolder::retrieveFailure)
  .map(CompletionStage::toCompletableFuture)

driver/src/main/java/org/neo4j/driver/internal/cursor/AsyncResultCursorImpl.java

@@ -32,6 +32,7 @@ public class AsyncResultCursorImpl implements AsyncResultCursor {
  private final Throwable runError;
  private final RunResponseHandler runHandler;
  private final PullAllResponseHandler pullAllHandler;
+ private final CompletableFuture<Void> consumedFuture = new CompletableFuture<>();
 
  public AsyncResultCursorImpl(
  Throwable runError, RunResponseHandler runHandler, PullAllResponseHandler pullAllHandler) {
@@ -47,7 +48,18 @@ public List<String> keys() {
 
  @Override
  public CompletionStage<ResultSummary> consumeAsync() {
- return pullAllHandler.consumeAsync();
+ var summaryFuture = new CompletableFuture<ResultSummary>();
+ pullAllHandler.consumeAsync().whenComplete((summary, throwable) -> {
+ throwable = Futures.completionExceptionCause(throwable);
+ if (throwable != null) {
+ consumedFuture.completeExceptionally(throwable);
+ summaryFuture.completeExceptionally(throwable);
+ } else {
+ consumedFuture.complete(null);
+ summaryFuture.complete(summary);
+ }
+ });
+ return summaryFuture;
  }
 
  @Override
@@ -138,4 +150,9 @@ private void internalForEachAsync(Consumer<Record> action, CompletableFuture<Voi
  public CompletableFuture<AsyncResultCursor> mapSuccessfulRunCompletionAsync() {
  return runError != null ? Futures.failedFuture(runError) : CompletableFuture.completedFuture(this);
  }
+
+ @Override
+ public CompletableFuture<Void> consumed() {
+ return consumedFuture;
+ }
 }

driver/src/main/java/org/neo4j/driver/internal/cursor/DisposableAsyncResultCursor.java

@@ -96,6 +96,11 @@ public CompletionStage<Throwable> pullAllFailureAsync() {
  return delegate.pullAllFailureAsync();
  }
 
+ @Override
+ public CompletionStage<Void> consumed() {
+ return delegate.consumed();
+ }
+
  private <T> CompletableFuture<T> assertNotDisposed() {
  if (isDisposed) {
  return failedFuture(newResultConsumedError());

driver/src/main/java/org/neo4j/driver/internal/cursor/RxResultCursorImpl.java

@@ -31,6 +31,7 @@
 import org.neo4j.driver.exceptions.TransactionNestingException;
 import org.neo4j.driver.internal.handlers.RunResponseHandler;
 import org.neo4j.driver.internal.handlers.pulln.PullResponseHandler;
+import org.neo4j.driver.internal.util.Futures;
 import org.neo4j.driver.summary.ResultSummary;
 
 public class RxResultCursorImpl implements RxResultCursor {
@@ -46,6 +47,7 @@ public class RxResultCursorImpl implements RxResultCursor {
  private boolean summaryFutureExposed;
  private boolean resultConsumed;
  private RecordConsumerStatus consumerStatus = NOT_INSTALLED;
+ private final CompletableFuture<Void> consumedFuture = new CompletableFuture<>();
 
  // for testing only
  public RxResultCursorImpl(RunResponseHandler runHandler, PullResponseHandler pullHandler) {
@@ -119,10 +121,26 @@ public CompletionStage<Throwable> pullAllFailureAsync() {
  return discardAllFailureAsync();
  }
 
+ @Override
+ public CompletionStage<Void> consumed() {
+ return consumedFuture;
+ }
+
  @Override
  public CompletionStage<ResultSummary> summaryAsync() {
  summaryFutureExposed = true;
- return summaryStage();
+ var summaryFuture = new CompletableFuture<ResultSummary>();
+ summaryStage().whenComplete((summary, throwable) -> {
+ throwable = Futures.completionExceptionCause(throwable);
+ if (throwable != null) {
+ consumedFuture.completeExceptionally(throwable);
+ summaryFuture.completeExceptionally(throwable);
+ } else {
+ consumedFuture.complete(null);
+ summaryFuture.complete(summary);
+ }
+ });
+ return summaryFuture;
  }
 
  @Override

driver/src/test/java/org/neo4j/driver/internal/async/ResultCursorsHolderTest.java

@@ -22,6 +22,8 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.BDDMockito.given;
+import static org.mockito.BDDMockito.then;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 import static org.neo4j.driver.testutil.TestUtil.await;
@@ -30,7 +32,9 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.TimeoutException;
+import java.util.stream.IntStream;
 import org.junit.jupiter.api.Test;
+import org.neo4j.driver.internal.FailableCursor;
 import org.neo4j.driver.internal.cursor.AsyncResultCursorImpl;
 import org.neo4j.driver.internal.util.Futures;
 
@@ -124,6 +128,38 @@ void shouldWaitForAllFailuresToArrive() {
  assertEquals(error1, await(failureFuture));
  }
 
+ @Test
+ void shouldRemoveConsumedResults() {
+ var holder = new ResultCursorsHolder();
+ var list = IntStream.range(0, 100)
+ .mapToObj(i -> {
+ var cursor = mock(FailableCursor.class);
+ var consume = new CompletableFuture<Void>();
+ given(cursor.consumed()).willReturn(consume);
+ holder.add(CompletableFuture.completedFuture(cursor));
+ if (i % 2 == 0) {
+ consume.complete(null);
+ given(cursor.discardAllFailureAsync())
+ .willReturn(CompletableFuture.failedFuture(new RuntimeException()));
+ } else {
+ given(cursor.discardAllFailureAsync()).willReturn(CompletableFuture.completedStage(null));
+ }
+ return cursor;
+ })
+ .toList();
+
+ holder.retrieveNotConsumedError().toCompletableFuture().join();
+
+ for (var i = 0; i < list.size(); i++) {
+ var cursor = list.get(i);
+ then(cursor).should().consumed();
+ if (i % 2 == 1) {
+ then(cursor).should().discardAllFailureAsync();
+ }
+ then(cursor).shouldHaveNoMoreInteractions();
+ }
+ }
+
  private static CompletionStage<AsyncResultCursorImpl> cursorWithoutError() {
  return cursorWithError(null);
  }
@@ -134,6 +170,7 @@ private static CompletionStage<AsyncResultCursorImpl> cursorWithError(Throwable
 
  private static CompletionStage<AsyncResultCursorImpl> cursorWithFailureFuture(CompletableFuture<Throwable> future) {
  var cursor = mock(AsyncResultCursorImpl.class);
+ when(cursor.consumed()).thenReturn(new CompletableFuture<>());
  when(cursor.discardAllFailureAsync()).thenReturn(future);
  return completedFuture(cursor);
  }

driver/src/test/java/org/neo4j/driver/internal/async/UnmanagedTransactionTest.java

@@ -724,6 +724,7 @@ private static Connection connectionWithBegin(Consumer<ResponseHandler> beginBeh
  private ResultCursorsHolder mockResultCursorWith(ClientException clientException) {
  var resultCursorsHolder = new ResultCursorsHolder();
  var cursor = mock(FailableCursor.class);
+ given(cursor.consumed()).willReturn(new CompletableFuture<>());
  doReturn(completedFuture(clientException)).when(cursor).discardAllFailureAsync();
  resultCursorsHolder.add(completedFuture(cursor));
  return resultCursorsHolder;