Wait for processing thread to finish in LonelyProcessingStepTest

neo4j · May 10, 2018 · c7aef06 · c7aef06
1 parent 265e579
commit c7aef06
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 2 deletions.
diff --git a/.../kernel/src/main/java/org/neo4j/unsafe/impl/batchimport/staging/LonelyProcessingStep.java b/.../kernel/src/main/java/org/neo4j/unsafe/impl/batchimport/staging/LonelyProcessingStep.java
@@ -43,7 +43,7 @@ public LonelyProcessingStep( StageControl control, String name, Configuration co
     @Override
     public long receive( long ticket, Void nothing )
     {
-        new Thread( () -> {
+        startProcessing( () -> {
             assertHealthy();
             try
             {
@@ -75,10 +75,15 @@ public long receive( long ticket, Void nothing )
                     throw e;
                 }
             }
-        } ).start();
+        } );
         return 0;
     }
 
+    protected void startProcessing( Runnable runnable )
+    {
+        new Thread( runnable ).start();
+    }
+
     /**
      * Called once and signals the start of this step. Responsible for calling {@link #progress(long)}
      * at least now and then.

diff --git a/...nel/src/test/java/org/neo4j/unsafe/impl/batchimport/staging/LonelyProcessingStepTest.java b/...nel/src/test/java/org/neo4j/unsafe/impl/batchimport/staging/LonelyProcessingStepTest.java
@@ -19,11 +19,15 @@
  */
 package org.neo4j.unsafe.impl.batchimport.staging;
 
+import org.junit.After;
 import org.junit.ClassRule;
 import org.junit.Test;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 
 import org.neo4j.concurrent.BinaryLatch;
 import org.neo4j.test.rule.SuppressOutput;
@@ -37,6 +41,15 @@ public class LonelyProcessingStepTest
     @ClassRule
     public static SuppressOutput mute = SuppressOutput.suppressAll();
 
+    private final ExecutorService executors = Executors.newCachedThreadPool();
+
+    @After
+    public void tearDown() throws Exception
+    {
+        executors.shutdown();
+        executors.awaitTermination( 1, TimeUnit.MINUTES );
+    }
+
     @Test
     public void issuePanicBeforeCompletionOnError()
     {
@@ -75,6 +88,12 @@ protected void process()
             throw new RuntimeException( "Process exception" );
         }
 
+        @Override
+        protected void startProcessing( Runnable runnable )
+        {
+            executors.submit( runnable );
+        }
+
         @Override
         public void endOfUpstream()
         {