Fixed auto-read problem with result summary

All result stream should be buffered/consumed when result summary is
requested or when session is closed. This means underlying channel
should read everything that's available. It did not do so when summary
was requested after number of buffered records exceeded high watermark.

This commit makes fetching of summary or failure enable auto-read on
the underlying channel. So SUCCESS and FAILURE messages will be read
without waiting for all records to be consumed.

Author: lutovich

driver/src/main/java/org/neo4j/driver/internal/handlers/PullAllResponseHandler.java

@@ -170,6 +170,10 @@ else if ( summary != null )
         {
             if ( summaryFuture == null )
             {
+                // neither SUCCESS nor FAILURE message has arrived, register future to be notified when it arrives
+                // future will be completed with summary on SUCCESS and completed exceptionally on FAILURE
+                // enable auto-read, otherwise we might not read SUCCESS/FAILURE if records are not consumed
+                connection.enableAutoRead();
                 summaryFuture = new CompletableFuture<>();
             }
             return summaryFuture;
@@ -190,6 +194,10 @@ else if ( finished )
         {
             if ( failureFuture == null )
             {
+                // neither SUCCESS nor FAILURE message has arrived, register future to be notified when it arrives
+                // future will be completed with null on SUCCESS and completed with Throwable on FAILURE
+                // enable auto-read, otherwise we might not read SUCCESS/FAILURE if records are not consumed
+                connection.enableAutoRead();
                 failureFuture = new CompletableFuture<>();
             }
             return failureFuture;

driver/src/test/java/org/neo4j/driver/internal/handlers/PullAllResponseHandlerTest.java

@@ -708,7 +708,6 @@ public void shouldNotDisableAutoReadWhenSummaryRequested()
             handler.onRecord( values( "a", "b" ) );
         }
 
-        verify( connection, never() ).enableAutoRead();
         verify( connection, never() ).disableAutoRead();
 
         handler.onSuccess( emptyMap() );
@@ -745,7 +744,6 @@ public void shouldNotDisableAutoReadWhenFailureRequested()
             handler.onRecord( values( 123, 456 ) );
         }
 
-        verify( connection, never() ).enableAutoRead();
         verify( connection, never() ).disableAutoRead();
 
         IllegalStateException error = new IllegalStateException( "Wrong config" );
@@ -766,6 +764,61 @@ public void shouldNotDisableAutoReadWhenFailureRequested()
         assertNull( await( handler.nextAsync() ) );
     }
 
+    @Test
+    public void shouldEnableAutoReadOnConnectionWhenFailureRequestedButNotAvailable() throws Exception
+    {
+        Connection connection = connectionMock();
+        PullAllResponseHandler handler = newHandler( asList( "key1", "key2" ), connection );
+
+        handler.onRecord( values( 1, 2 ) );
+        handler.onRecord( values( 3, 4 ) );
+
+        verify( connection, never() ).enableAutoRead();
+        verify( connection, never() ).disableAutoRead();
+
+        CompletableFuture<Throwable> failureFuture = handler.failureAsync().toCompletableFuture();
+        assertFalse( failureFuture.isDone() );
+
+        verify( connection ).enableAutoRead();
+        verify( connection, never() ).disableAutoRead();
+
+        assertNotNull( await( handler.nextAsync() ) );
+        assertNotNull( await( handler.nextAsync() ) );
+
+        RuntimeException error = new RuntimeException( "Oh my!" );
+        handler.onFailure( error );
+
+        assertTrue( failureFuture.isDone() );
+        assertEquals( error, failureFuture.get() );
+    }
+
+    @Test
+    public void shouldEnableAutoReadOnConnectionWhenSummaryRequestedButNotAvailable() throws Exception
+    {
+        Connection connection = connectionMock();
+        PullAllResponseHandler handler = newHandler( asList( "key1", "key2", "key3" ), connection );
+
+        handler.onRecord( values( 1, 2, 3 ) );
+        handler.onRecord( values( 4, 5, 6 ) );
+
+        verify( connection, never() ).enableAutoRead();
+        verify( connection, never() ).disableAutoRead();
+
+        CompletableFuture<ResultSummary> summaryFuture = handler.summaryAsync().toCompletableFuture();
+        assertFalse( summaryFuture.isDone() );
+
+        verify( connection ).enableAutoRead();
+        verify( connection, never() ).disableAutoRead();
+
+        assertNotNull( await( handler.nextAsync() ) );
+        assertNotNull( await( handler.nextAsync() ) );
+
+        handler.onSuccess( emptyMap() );
+
+        assertTrue( summaryFuture.isDone() );
+        assertNotNull( summaryFuture.get() );
+    }
+
     private static PullAllResponseHandler newHandler()
     {
         return newHandler( new Statement( "RETURN 1" ) );

driver/src/test/java/org/neo4j/driver/v1/integration/SessionIT.java

@@ -1405,6 +1405,52 @@ public void shouldAllowAccessingRecordsAfterSessionClosed()
         }
     }
 
+    @Test
+    public void shouldAllowToConsumeRecordsSlowlyAndCloseSession() throws InterruptedException
+    {
+        Session session = neo4j.driver().session();
+
+        StatementResult result = session.run( "UNWIND range(10000, 0, -1) AS x RETURN 10 / x" );
+
+        // consume couple records slowly with a sleep in-between
+        for ( int i = 0; i < 10; i++ )
+        {
+            assertTrue( result.hasNext() );
+            assertNotNull( result.next() );
+            Thread.sleep( 50 );
+        }
+
+        try
+        {
+            session.close();
+            fail( "Exception expected" );
+        }
+        catch ( ClientException e )
+        {
+            assertThat( e, is( arithmeticError() ) );
+        }
+    }
+
+    @Test
+    public void shouldAllowToConsumeRecordsSlowlyAndRetrieveSummary() throws InterruptedException
+    {
+        try ( Session session = neo4j.driver().session() )
+        {
+            StatementResult result = session.run( "UNWIND range(8000, 1, -1) AS x RETURN 42 / x" );
+
+            // consume couple records slowly with a sleep in-between
+            for ( int i = 0; i < 12; i++ )
+            {
+                assertTrue( result.hasNext() );
+                assertNotNull( result.next() );
+                Thread.sleep( 50 );
+            }
+
+            ResultSummary summary = result.summary();
+            assertNotNull( summary );
+        }
+    }
+
     private void assumeServerIs31OrLater()
     {
         ServerVersion serverVersion = ServerVersion.version( neo4j.driver() );