refactored task scavenging mechanism

Author: Jano

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

@@ -23,7 +23,6 @@ of this software and associated documentation files (the "Software"), to deal
 */
 package com.jano7.executor;
 
-import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.LinkedList;
 import java.util.concurrent.Executor;
@@ -33,18 +32,39 @@ public final class KeySequentialRunner<Key> {
 
     private final class KeyRunner {
 
-        private final LinkedList<Runnable> tasks = new LinkedList<>();
-        private boolean active = false;
+        private final Key key;
+        private LinkedList<Runnable> tasks;
 
-        public synchronized void run(Runnable task) {
-            if (active) {
-                tasks.addFirst(task);
-            } else {
-                active = true;
+        KeyRunner(Key key) {
+            this.key = key;
+        }
+
+        synchronized void run(Runnable task) {
+            if (tasks == null) {
+                tasks = new LinkedList<>();
                 runTask(task);
+            } else {
+                tasks.add(task);
             }
         }
 
+        private synchronized Runnable pollTask() {
+            return tasks.poll();
+        }
+
+        private Runnable nextTask() {
+            Runnable runnable = pollTask();
+            if (runnable == null) {
+                synchronized (KeySequentialRunner.this) {
+                    runnable = pollTask();
+                    if (runnable == null) {
+                        keyRunners.remove(key);
+                    }
+                }
+            }
+            return runnable;
+        }
+
         private void runTask(Runnable task) {
             underlyingExecutor.execute(() -> {
                 task.run();
@@ -60,18 +80,6 @@ private void runTask(Runnable task) {
                 }
             });
         }
-
-        private synchronized Runnable nextTask() {
-            Runnable runnable = tasks.pollLast();
-            if (runnable == null) {
-                active = false;
-            }
-            return runnable;
-        }
-
-        public synchronized boolean isActive() {
-            return active;
-        }
     }
 
     private final Executor underlyingExecutor;
@@ -90,26 +98,12 @@ public KeySequentialRunner(Executor underlyingExecutor, TaskExceptionHandler exc
     }
 
     public synchronized void run(Key key, Runnable task) {
-        KeyRunner runner = keyRunners.get(key);
-        if (runner == null) {
-            runner = new KeyRunner();
-            keyRunners.put(key, runner);
-        }
-        runner.run(() -> {
+        keyRunners.computeIfAbsent(key, KeyRunner::new).run(() -> {
             try {
                 task.run();
             } catch (Throwable t) {
                 exceptionHandler.handleTaskException(t);
             }
         });
-        scavengeInactiveRunners();
-    }
-
-    private void scavengeInactiveRunners() {
-        for (Key key : new ArrayList<>(keyRunners.keySet())) {
-            if (!keyRunners.get(key).isActive()) {
-                keyRunners.remove(key);
-            }
-        }
     }
 }

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

@@ -27,7 +27,10 @@ of this software and associated documentation files (the "Software"), to deal
 
 import java.lang.reflect.Field;
 import java.util.Map;
-import java.util.concurrent.*;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -38,15 +41,8 @@ public class KeySequentialRunnerTest {
     public void noMemoryLeak() throws InterruptedException, IllegalAccessException, NoSuchFieldException {
         final int threadCount = 10;
         ExecutorService underlyingExecutor = Executors.newFixedThreadPool(threadCount);
-        Executor testExecutor = (Runnable runnable) -> {
-            try {
-                underlyingExecutor.execute(runnable);
-            } catch (RejectedExecutionException ignored) {
-                runnable.run();
-            }
-        };
 
-        KeySequentialRunner<String> runner = new KeySequentialRunner<>(testExecutor);
+        KeySequentialRunner<String> runner = new KeySequentialRunner<>(underlyingExecutor);
 
         Field keyRunners = KeySequentialRunner.class.getDeclaredField("keyRunners");
         keyRunners.setAccessible(true);
@@ -70,9 +66,6 @@ public void noMemoryLeak() throws InterruptedException, IllegalAccessException,
         underlyingExecutor.shutdown();
         underlyingExecutor.awaitTermination(Long.MAX_VALUE, TimeUnit.SECONDS);
 
-        runner.run("key", () -> {
-        });
-
         assertEquals(0, ((Map<?, ?>) keyRunners.get(runner)).size());
     }
 }