add: cleaning up the map at cost of synchronization

Drozdov Artyom · WonderCsabo · commit ce449c84b2b4 · 2015-05-06T09:38:21.000+02:00
diff --git a/AndroidAnnotations/androidannotations-api/src/main/java/org/androidannotations/api/UiThreadExecutor.java b/AndroidAnnotations/androidannotations-api/src/main/java/org/androidannotations/api/UiThreadExecutor.java
@@ -15,10 +15,12 @@
  */
 package org.androidannotations.api;
 
-import java.util.concurrent.ConcurrentHashMap;
+import java.util.HashMap;
+import java.util.Map;
 
 import android.os.Handler;
 import android.os.Looper;
+import android.os.Message;
 import android.os.SystemClock;
 
 /**
@@ -27,9 +29,20 @@
  */
 public class UiThreadExecutor {
 
-	private static final Handler HANDLER = new Handler(Looper.getMainLooper());
+	private static final Handler HANDLER = new Handler(Looper.getMainLooper()) {
+		@Override
+		public void handleMessage(Message msg) {
+			Runnable callback = msg.getCallback();
+			if (callback != null) {
+				callback.run();
+				decrementToken((Token) msg.obj);
+			} else {
+				super.handleMessage(msg);
+			}
+		}
+	};
 
-	private static final ConcurrentHashMap<String, Object> TOKENS = new ConcurrentHashMap<String, Object>();
+	private static final Map<String, Token> TOKENS = new HashMap<String, Token>();
 
 	private UiThreadExecutor() {
 		// should not be instantiated
@@ -52,20 +65,33 @@ public static void runTask(String id, Runnable task, long delay) {
 			return;
 		}
 		long time = SystemClock.uptimeMillis() + delay;
-		HANDLER.postAtTime(task, getToken(id), time);
+		HANDLER.postAtTime(task, nextToken(id), time);
 	}
 
-	private static Object getToken(String id) {
-		Object token = TOKENS.get(id);
-		if (token == null) {
-			token = new Object();
-			Object oldObject = TOKENS.putIfAbsent(id, token);
-			// if a concurrent thread was faster
-			if (oldObject != null) {
-				token = oldObject;
+	private static Token nextToken(String id) {
+		synchronized (TOKENS) {
+			Token token = TOKENS.get(id);
+			if (token == null) {
+				token = new Token(id);
+				TOKENS.put(id, token);
+			}
+			token.runnablesCount++;
+			return token;
+		}
+	}
+
+	private static void decrementToken(Token token) {
+		synchronized (TOKENS) {
+			if (--token.runnablesCount == 0) {
+				String id = token.id;
+				Token old = TOKENS.remove(id);
+				if (old != token) {
+					// a runnable finished after cancelling, we just removed a
+					// wrong token, lets put it back
+					TOKENS.put(id, old);
+				}
 			}
 		}
-		return token;
 	}
 
 	/**
@@ -75,12 +101,24 @@ private static Object getToken(String id) {
 	 *            the cancellation identifier
 	 */
 	public static void cancelAll(String id) {
-		Object token = TOKENS.get(id);
+		Token token;
+		synchronized (TOKENS) {
+			token = TOKENS.remove(id);
+		}
 		if (token == null) {
 			// nothing to cancel
 			return;
 		}
 		HANDLER.removeCallbacksAndMessages(token);
 	}
 
+	private static class Token {
+		int runnablesCount = 0;
+		final String id;
+
+		private Token(String id) {
+			this.id = id;
+		}
+	}
+
 }
diff --git a/AndroidAnnotations/functional-test-1-5/src/test/java/org/androidannotations/test15/UiThreadExecutorTest.java b/AndroidAnnotations/functional-test-1-5/src/test/java/org/androidannotations/test15/UiThreadExecutorTest.java
@@ -15,6 +15,7 @@
  */
 package org.androidannotations.test15;
 
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
 import java.util.concurrent.CountDownLatch;
@@ -38,11 +39,25 @@ public void oneTaskTest() throws Exception {
 			public void run() {
 				done.set(true);
 			}
-		}, 0);
+		}, 10);
 		Robolectric.runUiThreadTasksIncludingDelayedTasks();
 		assertTrue("Task is still under execution", done.get());
 	}
 
+	@Test
+	public void oneTaskCancelTest() throws Exception {
+		final AtomicBoolean done = new AtomicBoolean(false);
+		UiThreadExecutor.runTask("test", new Runnable() {
+			@Override
+			public void run() {
+				done.set(true);
+			}
+		}, 10);
+		UiThreadExecutor.cancelAll("test");
+		Robolectric.runUiThreadTasksIncludingDelayedTasks();
+		assertFalse("Task is not cancelled", done.get());
+	}
+
 	@Test
 	public void oneTaskInThreadTest() throws Exception {
 		final CountDownLatch taskStartedLatch = new CountDownLatch(1);
@@ -56,7 +71,7 @@ public void run() {
 						await(taskStartedLatch);
 						taskFinishedLatch.countDown();
 					}
-				}, 0);
+				}, 10);
 				taskStartedLatch.countDown();
 				Robolectric.runUiThreadTasksIncludingDelayedTasks();
 			}
