Merge pull request #645 from zhenlineo/4.0-fix-failing-rx-tck

Ensure subscriber references are freed.

Author: zhenlineo

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

@@ -30,6 +30,9 @@
 import org.neo4j.driver.internal.handlers.pulln.PullResponseHandler;
 import org.neo4j.driver.summary.ResultSummary;
 
+import static org.neo4j.driver.internal.cursor.RxStatementResultCursorImpl.RecordConsumerStatus.DISCARD_INSTALLED;
+import static org.neo4j.driver.internal.cursor.RxStatementResultCursorImpl.RecordConsumerStatus.INSTALLED;
+import static org.neo4j.driver.internal.cursor.RxStatementResultCursorImpl.RecordConsumerStatus.NOT_INSTALLED;
 import static org.neo4j.driver.internal.util.ErrorUtil.newResultConsumedError;
 
 public class RxStatementResultCursorImpl implements RxStatementResultCursor
@@ -39,8 +42,8 @@ public class RxStatementResultCursorImpl implements RxStatementResultCursor
     private final PullResponseHandler pullHandler;
     private final Throwable runResponseError;
     private final CompletableFuture<ResultSummary> summaryFuture = new CompletableFuture<>();
-    private BiConsumer<Record,Throwable> recordConsumer;
     private boolean resultConsumed;
+    private RecordConsumerStatus consumerStatus = NOT_INSTALLED;
 
     public RxStatementResultCursorImpl( RunResponseHandler runHandler, PullResponseHandler pullHandler )
     {
@@ -72,20 +75,17 @@ public void installRecordConsumer( BiConsumer<Record,Throwable> recordConsumer )
         {
             throw newResultConsumedError();
         }
-        if ( isRecordConsumerInstalled() )
+
+        if ( consumerStatus.isInstalled() )
         {
             return;
         }
-        this.recordConsumer = recordConsumer;
-        pullHandler.installRecordConsumer( this.recordConsumer );
+        consumerStatus = recordConsumer == DISCARD_RECORD_CONSUMER ?
+                         DISCARD_INSTALLED : INSTALLED;
+        pullHandler.installRecordConsumer( recordConsumer );
         assertRunCompletedSuccessfully();
     }
 
-    private boolean isRecordConsumerInstalled()
-    {
-        return this.recordConsumer != null;
-    }
-
     @Override
     public void request( long n )
     {
@@ -108,7 +108,7 @@ public CompletionStage<Throwable> discardAllFailureAsync()
     @Override
     public CompletionStage<Throwable> pullAllFailureAsync()
     {
-        if ( isRecordConsumerInstalled() && !isDone() )
+        if ( consumerStatus.isInstalled() && !isDone() )
         {
             return CompletableFuture.completedFuture( new TransactionNestingException(
                     "You cannot run another query or begin a new transaction in the same session before you've fully consumed the previous run result." ) );
@@ -146,7 +146,7 @@ private void assertRunCompletedSuccessfully()
     private void installSummaryConsumer()
     {
         pullHandler.installSummaryConsumer( ( summary, error ) -> {
-            if ( error != null && recordConsumer == DISCARD_RECORD_CONSUMER )
+            if ( error != null && consumerStatus.isDiscardConsumer() )
             {
                 // We will only report the error to summary if there is no user record consumer installed
                 // When a user record consumer is installed, the error will be reported to record consumer instead.
@@ -167,4 +167,30 @@ private void assertRunResponseArrived( RunResponseHandler runHandler )
             throw new IllegalStateException( "Should wait for response of RUN before allowing PULL." );
         }
     }
+
+    enum RecordConsumerStatus
+    {
+        NOT_INSTALLED( false, false ),
+        INSTALLED( true, false ),
+        DISCARD_INSTALLED( true, true );
+
+        private final boolean isInstalled;
+        private final boolean isDiscardConsumer;
+
+        RecordConsumerStatus( boolean isInstalled, boolean isDiscardConsumer )
+        {
+            this.isInstalled = isInstalled;
+            this.isDiscardConsumer = isDiscardConsumer;
+        }
+
+        boolean isInstalled()
+        {
+            return isInstalled;
+        }
+
+        boolean isDiscardConsumer()
+        {
+            return isDiscardConsumer;
+        }
+    }
 }

driver/src/main/java/org/neo4j/driver/internal/reactive/InternalRxStatementResult.java

@@ -20,10 +20,12 @@
 
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
+import reactor.core.publisher.FluxSink;
 import reactor.core.publisher.Mono;
 
 import java.util.List;
 import java.util.concurrent.CompletionStage;
+import java.util.function.BiConsumer;
 import java.util.function.Supplier;
 
 import org.neo4j.driver.Record;
@@ -63,20 +65,7 @@ public Publisher<Record> records()
                 }
                 else
                 {
-                    cursor.installRecordConsumer( ( r, e ) -> {
-                        if ( r != null )
-                        {
-                            sink.next( r );
-                        }
-                        else if ( e != null )
-                        {
-                            sink.error( e );
-                        }
-                        else
-                        {
-                            sink.complete();
-                        }
-                    } );
+                    cursor.installRecordConsumer( createRecordConsumer( sink ) );
                     sink.onCancel( cursor::cancel );
                     sink.onRequest( cursor::request );
                 }
@@ -89,6 +78,32 @@ else if ( e != null )
         } ), IGNORE );
     }
 
+    /**
+     * Defines how a subscriber shall consume records.
+     * A record consumer holds a reference to a subscriber.
+     * A publisher and/or a subscription who holds a reference to this consumer shall release the reference to this object
+     * after subscription is done or cancelled so that the subscriber can be garbage collected.
+     * @param sink the subscriber
+     * @return a record consumer.
+     */
+    private BiConsumer<Record,Throwable> createRecordConsumer( FluxSink<Record> sink )
+    {
+        return ( r, e ) -> {
+            if ( r != null )
+            {
+                sink.next( r );
+            }
+            else if ( e != null )
+            {
+                sink.error( e );
+            }
+            else
+            {
+                sink.complete();
+            }
+        };
+    }
+
     private CompletionStage<RxStatementResultCursor> getCursorFuture()
     {
         if ( cursorFuture != null )