Let transactions fail gracefully on shutdown

We do this by checking if the database was shutdown while we were waiting for the transaction lock.

community/kernel/src/main/java/org/neo4j/kernel/impl/api/KernelTransactions.java

@@ -29,6 +29,7 @@
 import org.neo4j.collection.pool.LinkedQueuePool;
 import org.neo4j.collection.pool.MarshlandPool;
 import org.neo4j.function.Factory;
+import org.neo4j.graphdb.DatabaseShutdownException;
 import org.neo4j.graphdb.TransactionFailureException;
 import org.neo4j.kernel.AvailabilityGuard;
 import org.neo4j.kernel.api.KernelTransaction;
@@ -300,6 +301,10 @@ Set<KernelTransactionImplementation> getAllTransactions()
 
     private void assertRunning()
     {
+        if ( availabilityGuard.isShutdown() )
+        {
+            throw new DatabaseShutdownException();
+        }
         if ( stopped )
         {
             throw new IllegalStateException( "Can't start new transaction with stopped " + getClass() );

community/kernel/src/test/java/org/neo4j/kernel/impl/api/KernelTransactionsTest.java

@@ -34,6 +34,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReferenceArray;
 
+import org.neo4j.graphdb.DatabaseShutdownException;
 import org.neo4j.graphdb.security.AuthorizationExpiredException;
 import org.neo4j.kernel.AvailabilityGuard;
 import org.neo4j.kernel.api.KernelTransaction;
@@ -72,6 +73,8 @@
 import org.neo4j.test.Race;
 import org.neo4j.test.rule.concurrent.OtherThreadRule;
 
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.util.concurrent.locks.LockSupport.parkNanos;
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.CoreMatchers.instanceOf;
 import static org.hamcrest.CoreMatchers.is;
@@ -91,10 +94,6 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
-
-import static java.util.concurrent.TimeUnit.MILLISECONDS;
-import static java.util.concurrent.locks.LockSupport.parkNanos;
-
 import static org.neo4j.helpers.collection.Iterators.asSet;
 import static org.neo4j.kernel.api.KernelTransaction.Type.explicit;
 import static org.neo4j.kernel.api.security.SecurityContext.AUTH_DISABLED;
@@ -411,22 +410,15 @@ public void shouldNotLeakTransactionOnSecurityContextFreezeFailure() throws Thro
     }
 
     @Test
-    public void shouldStartNewTransactionAfterGuardMarkedAsShutdown() throws Throwable
+    public void exceptionWhenStartingNewTransactionOnShutdownInstance() throws Throwable
     {
-        // This is because the first thing thast happens when shutting down is that AvailabilityGuard gets
-        // marked as shut down. After that, transactions that just made it passed some layers in db.beginTx()
-        // or similar will have time to complete (other parts of shutdown will take care of that).
-
         KernelTransactions kernelTransactions = newKernelTransactions();
         SecurityContext securityContext = mock( SecurityContext.class );
 
         availabilityGuard.shutdown();
 
-        try ( KernelTransaction transaction =
-                kernelTransactions.newInstance( KernelTransaction.Type.explicit, securityContext, 0L ) )
-        {
-            assertNotNull( transaction );
-        }
+        expectedException.expect( DatabaseShutdownException.class );
+        kernelTransactions.newInstance( KernelTransaction.Type.explicit, securityContext, 0L );
     }
 
     @Test

community/kernel/src/test/java/org/neo4j/kernel/impl/api/TransactionCompletionAndShutdownRaceIT.java

@@ -21,11 +21,13 @@
 
 import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 
 import java.io.File;
 import java.util.concurrent.Future;
 import java.util.concurrent.atomic.AtomicBoolean;
 
+import org.neo4j.graphdb.DatabaseShutdownException;
 import org.neo4j.graphdb.GraphDatabaseService;
 import org.neo4j.graphdb.Node;
 import org.neo4j.graphdb.Transaction;
@@ -43,8 +45,8 @@
 import org.neo4j.test.rule.TestDirectory;
 import org.neo4j.test.rule.concurrent.OtherThreadRule;
 
+import static org.hamcrest.core.Is.isA;
 import static org.junit.Assert.assertNotNull;
-
 import static org.neo4j.helpers.collection.Iterables.single;
 
 /**
@@ -65,6 +67,8 @@ public class TransactionCompletionAndShutdownRaceIT
     public final OtherThreadRule<Void> transactor = new OtherThreadRule<>( "Transactor" );
     @Rule
     public final OtherThreadRule<Void> shutter = new OtherThreadRule<>( "Shutter" );
+    @Rule
+    public ExpectedException expectedException = ExpectedException.none();
 
     @Test
     public void shouldAlwaysAwaitTransactionCompletionBeforeShuttingDown() throws Exception
@@ -81,7 +85,7 @@ public void shouldAlwaysAwaitTransactionCompletionBeforeShuttingDown() throws Ex
             barrierInstalled.set( true );
 
             // WHEN
-            Future<Object> transactionFuture = transactor.execute( state -> doTransaction( db ) );
+            Future<Object> transactionFuture = transactor.execute( state -> doTransaction( db, barrier ) );
             Future<Object> shutterFuture = shutter.execute( state -> doShutdown( db, barrier ) );
 
             // THEN
@@ -141,10 +145,6 @@ public void require( AvailabilityRequirement requirement )
                                     public void await( long millis ) throws UnavailableException
                                     {
                                         super.await( millis );
-                                        if ( barrierInstaller.get() )
-                                        {
-                                            barrier.reached();
-                                        }
                                     }
                                 };
                             }
@@ -162,6 +162,94 @@ private Object doShutdown( GraphDatabaseService db, Barrier.Control barrier )
         return null;
     }
 
+    private Object doTransaction( GraphDatabaseService db, Barrier.Control barrier )
+    {
+        try ( Transaction tx = db.beginTx() )
+        {
+            barrier.reached();
+            db.createNode();
+            tx.success();
+        }
+        return null;
+    }
+
+    @Test
+    public void shouldNotStartTransactoinOnDatabaseThatIsKnownToBeShuttingDown() throws Exception
+    {
+        // GIVEN
+        Barrier.Control barrier = new Barrier.Control();
+        AtomicBoolean barrierInstalled = new AtomicBoolean();
+        File storeDir = directory.absolutePath();
+        TestGraphDatabaseFactory dbFactory = dbFactoryWithBarrierControlledTransactionStats2( barrier,
+                barrierInstalled );
+
+        {
+            GraphDatabaseService db = dbFactory.newImpermanentDatabase( storeDir );
+            barrierInstalled.set( true );
+
+            // WHEN
+            Future<Object> transactionFuture = transactor.execute( state -> doTransaction( db ) );
+            Future<Object> shutterFuture = shutter.execute( state -> doShutdown( db, barrier ) );
+
+            // THEN
+            shutterFuture.get();
+            expectedException.expectCause( isA( DatabaseShutdownException.class ) );
+            transactionFuture.get();
+            barrierInstalled.set( false );
+        }
+    }
+
+    private TestGraphDatabaseFactory dbFactoryWithBarrierControlledTransactionStats2( Barrier.Control barrier,
+            AtomicBoolean barrierInstaller )
+    {
+        return new TestGraphDatabaseFactory()
+        {
+            @Override
+            protected GraphDatabaseFacadeFactory newTestGraphDatabaseFacadeFactory( File storeDir, Config config,
+                    TestGraphDatabaseFactoryState state )
+            {
+                return new TestGraphDatabaseFacadeFactory( state, true )
+                {
+                    @Override
+                    protected PlatformModule createPlatform( File storeDir, Config config, Dependencies dependencies,
+                            GraphDatabaseFacade graphDatabaseFacade )
+                    {
+                        return new TestGraphDatabaseFacadeFactory.TestDatabasePlatformModule( storeDir, config,
+                                databaseInfo, dependencies, graphDatabaseFacade )
+                        {
+                            @Override
+                            protected AvailabilityGuard createAvailabilityGuard()
+                            {
+                                return new AvailabilityGuard( clock, NullLog.getInstance() )
+                                {
+                                    @Override
+                                    public void require( AvailabilityRequirement requirement )
+                                    {
+                                        super.require( requirement );
+                                        if ( barrierInstaller.get() )
+                                        {
+                                            barrier.release();
+                                        }
+                                    }
+
+                                    @Override
+                                    public void await( long millis ) throws UnavailableException
+                                    {
+                                        super.await( millis );
+                                        if ( barrierInstaller.get() )
+                                        {
+                                            barrier.reached();
+                                        }
+                                    }
+                                };
+                            }
+                        };
+                    }
+                };
+            }
+        };
+    }
+
     private Object doTransaction( GraphDatabaseService db )
     {
         try ( Transaction tx = db.beginTx() )