/*
 * Copyright (c) 2002-2017 "Neo Technology,"
 * Network Engine for Objects in Lund AB [http://neotechnology.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.v1.stress;

import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;

import java.util.List;
import java.util.Random;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

import org.neo4j.driver.v1.Config;
import org.neo4j.driver.v1.Driver;
import org.neo4j.driver.v1.Session;
import org.neo4j.driver.v1.StatementResult;
import org.neo4j.driver.v1.util.TestNeo4j;

import static java.util.Arrays.asList;
import static org.junit.Assert.assertTrue;
import static org.neo4j.driver.v1.GraphDatabase.driver;

public class SessionPoolingStressIT
{
    @Rule
    public TestNeo4j neo4j = new TestNeo4j();

    private static final int N_THREADS = 50;
    private static final int MAX_TIME = 10000;

    private static final List<String> QUERIES = asList(
            "RETURN 1295 + 42", "UNWIND range(1,10000) AS x CREATE (n {prop:x}) DELETE n " );

    private Driver driver;
    private ExecutorService executor;

    @Before
    public void setUp() throws Exception
    {
        executor = Executors.newFixedThreadPool( N_THREADS );
    }

    @After
    public void tearDown() throws Exception
    {
        if ( executor != null )
        {
            executor.shutdownNow();
        }

        if ( driver != null )
        {
            driver.close();
        }
    }

    @Test
    public void shouldWorkFine() throws Throwable
    {
        Config config = Config.build()
                .withEncryptionLevel( Config.EncryptionLevel.NONE )
                .toConfig();

        driver = driver( neo4j.uri(), config );

        AtomicBoolean stop = new AtomicBoolean();
        AtomicReference<Throwable> failureReference = new AtomicReference<>();

        doWork( stop, failureReference );

        Thread.sleep( MAX_TIME );

        stop.set( true );
        executor.shutdown();
        assertTrue( executor.awaitTermination( MAX_TIME, TimeUnit.MILLISECONDS ) );

        Throwable failure = failureReference.get();
        if ( failure != null )
        {
            throw new AssertionError( "Some workers have failed", failure );
        }
    }

    private void doWork( AtomicBoolean stop, AtomicReference<Throwable> failure )
    {
        for ( int i = 0; i < N_THREADS; i++ )
        {
            executor.execute( new Worker( driver, stop, failure ) );
        }
    }

    private class Worker implements Runnable
    {
        private final Random random = ThreadLocalRandom.current();
        private final Driver driver;
        private final AtomicBoolean stop;
        private final AtomicReference<Throwable> failureReference;

        Worker( Driver driver, AtomicBoolean stop, AtomicReference<Throwable> failureReference )
        {
            this.driver = driver;
            this.stop = stop;
            this.failureReference = failureReference;
        }

        @Override
        public void run()
        {
            try
            {
                while ( !stop.get() )
                {
                    for ( String query : QUERIES )
                    {
                        runQuery( query );
                    }
                }
            }
            catch ( Throwable failure )
            {
                if ( !failureReference.compareAndSet( null, failure ) )
                {
                    Throwable firstFailure = failureReference.get();
                    synchronized ( firstFailure )
                    {
                        firstFailure.addSuppressed( failure );
                    }
                }
            }
        }

        private void runQuery( String query ) throws InterruptedException
        {
            try ( Session session = driver.session() )
            {
                StatementResult run = session.run( query );
                Thread.sleep( random.nextInt( 100 ) );
                run.consume();
                Thread.sleep( random.nextInt( 100 ) );
            }
        }
    }
}

@@ -836,16 +836,89 @@
 it.Test;
+%0Aimport org.junit.rules.TestWatcher;%0Aimport org.junit.runner.Description;
 %0A%0Aimport
@@ -930,24 +930,46 @@
 .util.List;%0A
+import java.util.Map;%0A
 import java.
@@ -1645,16 +1645,22 @@
  public 
+final 
 TestNeo4
@@ -1683,24 +1683,933 @@
 stNeo4j();%0A%0A
+    @Rule%0A    public final TestWatcher testWatcher = new TestWatcher()%0A    %7B%0A        @Override%0A        protected void failed( Throwable e, Description description )%0A        %7B%0A            super.failed( e, description );%0A%0A            StringBuilder sb = new StringBuilder();%0A            Map%3CThread,StackTraceElement%5B%5D%3E allStackTraces = Thread.getAllStackTraces();%0A            for ( Map.Entry%3CThread,StackTraceElement%5B%5D%3E entry : allStackTraces.entrySet() )%0A            %7B%0A                Thread thread = entry.getKey();%0A                sb.append( thread ).append( %22 -- %22 ).append( thread.getState() ).append( System.lineSeparator() );%0A                for ( StackTraceElement element : entry.getValue() )%0A                %7B%0A                    sb.append( %22    %22 ).append( element ).append( System.lineSeparator() );%0A                %7D%0A            %7D%0A%0A            System.out.println( sb.toString() );%0A        %7D%0A    %7D;%0A%0A
     private 
@@ -2670,19 +2670,20 @@
 nal int 
-MAX
+TEST
 _TIME = 
@@ -3717,19 +3717,20 @@
 .sleep( 
-MAX
+TEST
 _TIME );
@@ -3837,16 +3837,10 @@
 on( 
-MAX_TIME
+20
 , Ti
@@ -3850,13 +3850,8 @@
 nit.
-MILLI
 SECO

Increase executor shutdown timeout