Merge pull request #584 from zhenlineo/2.0-stress-tests

Added stress tests for reactive sessions

Author: zhenlineo

driver/src/test/java/org/neo4j/driver/stress/AbstractRxQuery.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) 2002-2019 "Neo4j,"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.neo4j.driver.stress;
+
+import org.neo4j.driver.AccessMode;
+import org.neo4j.driver.Driver;
+import org.neo4j.driver.reactive.RxSession;
+
+public abstract class AbstractRxQuery<C extends AbstractContext> implements RxCommand<C>
+{
+    protected final Driver driver;
+    protected final boolean useBookmark;
+
+    public AbstractRxQuery( Driver driver, boolean useBookmark )
+    {
+        this.driver = driver;
+        this.useBookmark = useBookmark;
+    }
+
+    public RxSession newSession( AccessMode mode, C context )
+    {
+        if ( useBookmark )
+        {
+            return driver.rxSession( t -> t.withDefaultAccessMode( mode ).withBookmarks( context.getBookmark() ) );
+        }
+        return driver.rxSession( t -> t.withDefaultAccessMode( mode ) );
+    }
+}

driver/src/test/java/org/neo4j/driver/stress/AbstractStressTestBase.java

@@ -22,12 +22,16 @@
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
 
 import java.lang.management.ManagementFactory;
 import java.lang.management.OperatingSystemMXBean;
 import java.lang.reflect.Method;
 import java.net.URI;
+import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -61,8 +65,12 @@
 import org.neo4j.driver.async.StatementResultCursor;
 import org.neo4j.driver.internal.InternalDriver;
 import org.neo4j.driver.internal.logging.DevNullLogger;
+import org.neo4j.driver.internal.util.EnabledOnNeo4jWith;
 import org.neo4j.driver.internal.util.Futures;
 import org.neo4j.driver.internal.util.Iterables;
+import org.neo4j.driver.internal.util.Neo4jFeature;
+import org.neo4j.driver.reactive.RxSession;
+import org.neo4j.driver.reactive.RxTransaction;
 import org.neo4j.driver.types.Node;
 import org.neo4j.driver.util.DaemonThreadFactory;
 
@@ -88,6 +96,7 @@ abstract class AbstractStressTestBase<C extends AbstractContext>
 
     private static final int BIG_DATA_TEST_NODE_COUNT = Integer.getInteger( "bigDataTestNodeCount", 30_000 );
     private static final int BIG_DATA_TEST_BATCH_SIZE = Integer.getInteger( "bigDataTestBatchSize", 10_000 );
+    private static final Duration DEFAULT_BLOCKING_TIME_OUT = Duration.ofMinutes( 5 );
 
     private LoggerNameTrackingLogging logging;
     private ExecutorService executor;
@@ -115,7 +124,6 @@ void setUp()
     @AfterEach
     void tearDown()
     {
-        System.out.println( driver.metrics() );
         executor.shutdownNow();
         if ( driver != null )
         {
@@ -136,6 +144,13 @@ void asyncApiStressTest() throws Throwable
         runStressTest( this::launchAsyncWorkerThreads );
     }
 
+    @Test
+    @EnabledOnNeo4jWith( Neo4jFeature.BOLT_V4 )
+    void rxApiStressTest() throws Throwable
+    {
+        runStressTest( this::launchRxWorkerThreads );
+    }
+
     @Test
     void blockingApiBigDataTest()
     {
@@ -150,6 +165,14 @@ void asyncApiBigDataTest() throws Throwable
         readNodesAsync( driver, bookmark, BIG_DATA_TEST_NODE_COUNT );
     }
 
+    @Test
+    @EnabledOnNeo4jWith( Neo4jFeature.BOLT_V4 )
+    void rxApiBigDataTest() throws Throwable
+    {
+        String bookmark = createNodesRx( bigDataTestBatchCount(), BIG_DATA_TEST_BATCH_SIZE, driver );
+        readNodesRx( driver, bookmark, BIG_DATA_TEST_NODE_COUNT );
+    }
+
     private void runStressTest( Function<C,List<Future<?>>> threadLauncher ) throws Throwable
     {
         C context = createContext();
@@ -252,6 +275,71 @@ private Future<Void> launchBlockingWorkerThread( ExecutorService executor, List<
         } );
     }
 
+    private List<Future<?>> launchRxWorkerThreads( C context )
+    {
+        List<RxCommand<C>> commands = createRxCommands();
+        List<Future<?>> futures = new ArrayList<>();
+
+        for ( int i = 0; i < THREAD_COUNT; i++ )
+        {
+            Future<Void> future = launchRxWorkerThread( executor, commands, context );
+            futures.add( future );
+        }
+        return futures;
+    }
+
+    private List<RxCommand<C>> createRxCommands()
+    {
+        return Arrays.asList(
+                new RxReadQuery<>( driver, false ),
+                new RxReadQuery<>( driver, true ),
+
+                new RxWriteQuery<>( this, driver, false ),
+                new RxWriteQuery<>( this, driver, true ),
+
+                new RxReadQueryInTx<>( driver, false ),
+                new RxReadQueryInTx<>( driver, true ),
+
+                new RxWriteQueryInTx<>( this, driver, false ),
+                new RxWriteQueryInTx<>( this, driver, true ),
+
+                new RxReadQueryWithRetries<>( driver, false ),
+                new RxReadQueryWithRetries<>( driver, false ),
+
+                new RxWriteQueryWithRetries<>( this, driver, false ),
+                new RxWriteQueryWithRetries<>( this, driver, true ),
+
+                new RxFailingQuery<>( driver ),
+                new RxFailingQueryInTx<>( driver ),
+                new RxFailingQueryWithRetries<>( driver )
+        );
+    }
+
+    private Future<Void> launchRxWorkerThread( ExecutorService executor, List<RxCommand<C>> commands, C context )
+    {
+        return executor.submit( () ->
+        {
+            while ( !context.isStopped() )
+            {
+                CompletableFuture<Void> allCommands = executeRxCommands( context, commands, ASYNC_BATCH_SIZE );
+                assertNull( allCommands.get() );
+            }
+            return null;
+        } );
+    }
+
+    private CompletableFuture<Void> executeRxCommands( C context, List<RxCommand<C>> commands, int count )
+    {
+        CompletableFuture<Void>[] executions = new CompletableFuture[count];
+        for ( int i = 0; i < count; i++ )
+        {
+            RxCommand<C> command = randomOf( commands );
+            CompletionStage<Void> execution = command.execute( context );
+            executions[i] = execution.toCompletableFuture();
+        }
+        return CompletableFuture.allOf( executions );
+    }
+
     private List<Future<?>> launchAsyncWorkerThreads( C context )
     {
         List<AsyncCommand<C>> commands = createAsyncCommands();
@@ -529,6 +617,57 @@ private static void readNodesAsync( Driver driver, String bookmark, int expected
         System.out.println( "Reading nodes with async API took: " + NANOSECONDS.toMillis( end - start ) + "ms" );
     }
 
+    private String createNodesRx( int batchCount, int batchSize, InternalDriver driver )
+    {
+        long start = System.nanoTime();
+
+        RxSession session = driver.rxSession();
+
+        Flux.concat( Flux.range( 0, batchCount ).map( batchIndex ->
+            session.writeTransaction( tx -> createNodesInTxRx( tx, batchIndex, batchSize ) )
+        ) ).blockLast( DEFAULT_BLOCKING_TIME_OUT ); // throw any error if happened
+
+        long end = System.nanoTime();
+        System.out.println( "Node creation with reactive API took: " + NANOSECONDS.toMillis( end - start ) + "ms" );
+
+        return session.lastBookmark();
+    }
+
+    private Publisher<Void> createNodesInTxRx( RxTransaction tx, int batchIndex, int batchSize )
+    {
+        return Flux.concat( Flux.range( 0, batchSize ).map( index -> batchIndex * batchSize + index ).map( nodeIndex -> {
+            Statement statement = createNodeInTxStatement( nodeIndex );
+            return Flux.from( tx.run( statement ).summary() ).then(); // As long as there is no error
+        } ) );
+    }
+
+    private void readNodesRx( InternalDriver driver, String bookmark, int expectedNodeCount )
+    {
+        long start = System.nanoTime();
+
+        RxSession session = driver.rxSession( t -> t.withBookmarks( bookmark ) );
+        AtomicInteger nodesSeen = new AtomicInteger();
+
+        Publisher<Void> readQuery = session.readTransaction( tx -> Flux.from( tx.run( "MATCH (n:Node) RETURN n" ).records() ).doOnNext( record -> {
+            Node node = record.get( 0 ).asNode();
+            nodesSeen.incrementAndGet();
+
+            List<String> labels = Iterables.asList( node.labels() );
+            assertEquals( 2, labels.size() );
+            assertTrue( labels.contains( "Test" ) );
+            assertTrue( labels.contains( "Node" ) );
+
+            verifyNodeProperties( node );
+        } ).then() );
+
+        Flux.from( readQuery ).blockLast( DEFAULT_BLOCKING_TIME_OUT );
+
+        assertEquals( expectedNodeCount, nodesSeen.get() );
+
+        long end = System.nanoTime();
+        System.out.println( "Reading nodes with async API took: " + NANOSECONDS.toMillis( end - start ) + "ms" );
+    }
+
     private static Void createNodesInTx( Transaction tx, int batchIndex, int batchSize )
     {
         for ( int index = 0; index < batchSize; index++ )

driver/src/test/java/org/neo4j/driver/stress/RxCommand.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright (c) 2002-2019 "Neo4j,"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.neo4j.driver.stress;
+
+import java.util.concurrent.CompletionStage;
+
+public interface RxCommand<C extends AbstractContext>
+{
+    CompletionStage<Void> execute( C context );
+}

driver/src/test/java/org/neo4j/driver/stress/RxFailingQuery.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright (c) 2002-2019 "Neo4j,"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.neo4j.driver.stress;
+
+import reactor.core.publisher.Flux;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+
+import org.neo4j.driver.AccessMode;
+import org.neo4j.driver.Driver;
+import org.neo4j.driver.internal.util.Futures;
+import org.neo4j.driver.reactive.RxSession;
+
+import static org.hamcrest.Matchers.either;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.junit.MatcherAssert.assertThat;
+import static org.neo4j.driver.internal.util.Matchers.arithmeticError;
+
+public class RxFailingQuery<C extends AbstractContext> extends AbstractRxQuery<C>
+{
+    public RxFailingQuery( Driver driver )
+    {
+        super( driver, false );
+    }
+
+    @Override
+    public CompletionStage<Void> execute( C context )
+    {
+        CompletableFuture<Void> queryFinished = new CompletableFuture<>();
+        Flux.using( () -> newSession( AccessMode.READ, context ),
+                session -> session.run( "UNWIND [10, 5, 0] AS x RETURN 10 / x" ).records(),
+                RxSession::close )
+                .subscribe( record -> {
+                    assertThat( record.get( 0 ).asInt(), either( equalTo( 1 ) ).or( equalTo( 2 ) ) );
+                    queryFinished.complete( null );
+                }, error -> {
+                    Throwable cause = Futures.completionExceptionCause( error );
+                    assertThat( cause, is( arithmeticError() ) );
+                    queryFinished.complete( null );
+                });
+        return queryFinished;
+    }
+}