package io.funtom.util.concurrent;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

public final class SynchronizedPerKeyExecutor<K> {

	private final Map<K, Lock> locks = new HashMap<>();
	private final Map<K, Integer> keyUsersCount = new HashMap<>();

	public void execute(K key, Runnable task) {
		Lock lock = getLockForKey(key);
		lock.lock();
		try {
			task.run();
		} finally {
			lock.unlock();
			freeLockForKey(key);
		}
	}

	public <R> R submit(K key, Callable<R> task) throws Exception {
		Lock lock = getLockForKey(key);
		lock.lock();
		try {
			return task.call();
		} finally {
			lock.unlock();
			freeLockForKey(key);
		}
	}

	public <R> R submitUnchecked(K key, Callable<R> task) {
		try {
			return submit(key, task);
		} catch (Exception e) {
			throw new UncheckedExecutionException(e);
		}
	}

	private synchronized Lock getLockForKey(K key) {
		Lock result;
		Integer currentUsers = keyUsersCount.get(key);
		if (currentUsers == null) {
			keyUsersCount.put(key, 1);
			result = new ReentrantLock();
			locks.put(key, result);
		} else {
			keyUsersCount.put(key, currentUsers + 1);
			result = locks.get(key);
		}
		return result;
	}

	private synchronized void freeLockForKey(K key) {
		int currentUsers = keyUsersCount.get(key);
		if (currentUsers == 1) {
			keyUsersCount.remove(key);
			locks.remove(key);
		} else {
			keyUsersCount.put(key, currentUsers - 1);
		}
	}
}

@@ -119,48 +119,8 @@
 le;%0A
-import java.util.concurrent.locks.Lock;%0A
 impo
@@ -247,14 +247,26 @@
 Lock
-%3E lock
+Executor%3E executor
 s = 
@@ -402,509 +402,539 @@
 Lock
- lock = getLockForKey(key);%0A%09%09lock.lock();%0A%09%09try %7B%0A%09%09%09task.run();%0A%09%09%7D finally %7B%0A%09%09%09lock.unlock();%0A%09%09%09freeLockForKey(key);%0A%09%09%7D%0A%09%7D%0A%0A%09public %3CR%3E R submit(K key, Callable%3CR%3E task) throws Exception %7B%0A%09%09Lock lock = getLockForKey(key);%0A%09%09lock.lock();%0A%09%09try %7B%0A%09%09%09return task.call();%0A%09%09%7D finally %7B%0A%09%09%09lock.unlock();%0A%09%09%09freeLockForKey(key);%0A%09%09%7D%0A%09%7D%0A%0A%09public %3CR%3E R submitUnchecked(K key, Callable%3CR%3E task) %7B%0A%09%09try %7B%0A%09%09%09return submit(key, task);%0A%09%09%7D catch (Exception e) %7B%0A%09%09%09throw new UncheckedExecutionException(e
+Executor executor = getExecutorForKey(key);%0A%09%09try %7B%0A%09%09%09executor.execute(task);%0A%09%09%7D finally %7B%0A%09%09%09freeExecutorForKey(key);%0A%09%09%7D%0A%09%7D%0A%0A%09public %3CR%3E R submit(K key, Callable%3CR%3E task) throws Exception %7B%0A%09%09LockExecutor executor = getExecutorForKey(key);%0A%09%09try %7B%0A%09%09%09return executor.submit(task);%0A%09%09%7D finally %7B%0A%09%09%09freeExecutorForKey(key);%0A%09%09%7D%0A%09%7D%0A%0A%09public %3CR%3E R submitUnchecked(K key, Callable%3CR%3E task) %7B%0A%09%09LockExecutor executor = getExecutorForKey(key);%0A%09%09try %7B%0A%09%09%09return executor.submitUnchecked(task);%0A%09%09%7D finally %7B%0A%09%09%09freeExecutorForKey(key
 );%0A%09
@@ -970,16 +970,28 @@
 Lock
- getLock
+Executor getExecutor
 ForK
@@ -1008,16 +1008,24 @@
 %7B%0A%09%09Lock
+Executor
  result;
@@ -1150,16 +1150,33 @@
 t = new 
+LockExecutor(new 
 Reentran
@@ -1182,25 +1182,30 @@
 ntLock()
+)
 ;%0A%09%09%09
-lock
+executor
 s.put(ke
@@ -1284,20 +1284,24 @@
 esult = 
-lock
+executor
 s.get(ke
@@ -1360,20 +1360,24 @@
 oid free
-Lock
+Executor
 ForKey(K
@@ -1489,20 +1489,24 @@
 ey);%0A%09%09%09
-lock
+executor
 s.remove

Improve SyncronizedPerKeyExecutor - uses LockExecutor instead of raw executor