drain with timeout

Author: Jan Gaspar

README.md

@@ -71,13 +71,12 @@ boundedExecutor.execute(aTask);
 
 // execute more tasks ... at most 10 will be scheduled
 
-// before shutting down you can call a drain() method which blocks until all submitted task have been executed
-boundedExecutor.drain();
+// before shutting down you can call a 'drain' method which blocks until all submitted task have been executed
+boundedExecutor.drain(aTimeout, TimeUnit.SECONDS); // returns true if drained; false if timeout elapses
 
-// newly submitted tasks will be rejected after calling drain()
+// newly submitted tasks will be rejected after calling 'drain'
 
-underlyingExecutor.shutdown();
-underlyingExecutor.awaitTermination(aTimeout, TimeUnit.SECONDS);
+underlyingExecutor.shutdownNow(); // safe to call 'shutdownNow' if drained as there should be no active tasks
 ```
 The source code of the examples can be found [here](src/test/java/com/jano7/executor/Examples.java).
 
@@ -88,6 +87,6 @@ from multiple threads without synchronization.
 <dependency>
   <groupId>com.jano7</groupId>
   <artifactId>executor</artifactId>
-  <version>1.0.3</version>
+  <version>1.0.4</version>
 </dependency>
 ```

src/main/java/com/jano7/executor/BoundedExecutor.java

@@ -26,6 +26,7 @@ of this software and associated documentation files (the "Software"), to deal
 import java.util.concurrent.Executor;
 import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
 
 public final class BoundedExecutor implements Executor {
 
@@ -74,11 +75,11 @@ public void execute(Runnable task) {
         }
     }
 
-    public void drain() {
-        synchronized (this) {
-            semaphore.acquireUninterruptibly(maxTasks);
+    public synchronized boolean drain(long timeout, TimeUnit unit) throws InterruptedException {
+        if (semaphore.tryAcquire(maxTasks, timeout, unit)) {
             drained = true;
+            semaphore.release(maxTasks);
         }
-        semaphore.release(maxTasks);
+        return drained;
     }
 }

src/test/java/com/jano7/executor/BoundedExecutorTest.java

@@ -117,33 +117,46 @@ public void execute(Runnable command) {
         assertTrue(thrown);
     }
 
-    @Test(timeout = 5000)
+    @Test(timeout = 10000)
     public void drain() throws InterruptedException {
-        ExecutorService underlying = Executors.newFixedThreadPool(5);
-
         for (int i = 0; i < 1000; ++i) {
+            ExecutorService underlying = Executors.newFixedThreadPool(5);
             BoundedExecutor bounded = new BoundedExecutor(10, underlying);
+            CountDownLatch latch = new CountDownLatch(1);
             AtomicInteger completed = new AtomicInteger(0);
+
             for (int j = 0; j < 10; ++j) {
                 bounded.execute(completed::incrementAndGet);
             }
-            bounded.drain();
-            assertEquals(10, completed.get());
+            bounded.execute(() -> {
+                try {
+                    latch.await();
+                    completed.incrementAndGet();
+                } catch (InterruptedException e) {
+                    throw new RuntimeException(e);
+                }
+            });
+
+            assertFalse(bounded.drain(1, TimeUnit.MILLISECONDS));
+
+            latch.countDown();
+
+            assertTrue(bounded.drain(Long.MAX_VALUE, TimeUnit.SECONDS));
+            assertEquals(11, completed.get());
+            assertTrue(underlying.shutdownNow().isEmpty());
         }
 
-        underlying.shutdown();
-        underlying.awaitTermination(Long.MAX_VALUE, TimeUnit.SECONDS);
     }
 
     @Test(timeout = 5000, expected = RejectedExecutionException.class)
-    public void rejectTasksAfterDrain() {
+    public void rejectTasksAfterDrain() throws InterruptedException {
         ExecutorService underlying = Executors.newCachedThreadPool();
         BoundedExecutor bounded = new BoundedExecutor(10, underlying);
 
         bounded.execute(() -> {
         });
 
-        bounded.drain();
+        bounded.drain(Long.MAX_VALUE, TimeUnit.SECONDS);
         try {
             bounded.execute(() -> {
             });
@@ -153,15 +166,15 @@ public void rejectTasksAfterDrain() {
     }
 
     @Test(timeout = 5000)
-    public void callDrainMultipleTime() {
+    public void callDrainMultipleTime() throws InterruptedException {
         ExecutorService underlying = Executors.newCachedThreadPool();
         BoundedExecutor bounded = new BoundedExecutor(10, underlying);
 
         bounded.execute(() -> {
         });
 
-        bounded.drain();
-        bounded.drain();
+        bounded.drain(Long.MAX_VALUE, TimeUnit.SECONDS);
+        bounded.drain(Long.MAX_VALUE, TimeUnit.SECONDS);
 
         underlying.shutdownNow();
     }

src/test/java/com/jano7/executor/Examples.java

@@ -76,13 +76,12 @@ public static void boundedExecutorExample() throws InterruptedException {
 
         // execute more tasks ... at most 10 will be scheduled
 
-        // before shutting down you can call a drain() method which blocks until all submitted task have been executed
-        boundedExecutor.drain();
+        // before shutting down you can call a 'drain' method which blocks until all submitted task have been executed
+        boundedExecutor.drain(aTimeout, TimeUnit.SECONDS); // returns true if drained; false if timeout elapses
 
-        // newly submitted tasks will be rejected after calling drain()
+        // newly submitted tasks will be rejected after calling 'drain'
 
-        underlyingExecutor.shutdown();
-        underlyingExecutor.awaitTermination(aTimeout, TimeUnit.SECONDS);
+        underlyingExecutor.shutdownNow(); // safe to call 'shutdownNow' if drained as there should be no active tasks
     }
 
     public static void main(String[] args) throws InterruptedException {