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);
		}
	}
}

package io.funtom.util.concurrent;

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

public final class SynchronizedPerKeyExecutor<K> {

	private final Map<K, LockExecutor> executors = new HashMap<>();
	private final Map<K, Integer> keyUsersCount = new HashMap<>();

	public void execute(K key, Runnable task) {
		LockExecutor executor = getExecutorForKey(key);
		try {
			executor.execute(task);
		} finally {
			freeExecutorForKey(key);
		}
	}

	public <R> R submit(K key, Callable<R> task) throws Exception {
		LockExecutor executor = getExecutorForKey(key);
		try {
			return executor.submit(task);
		} finally {
			freeExecutorForKey(key);
		}
	}

	public <R> R submitUnchecked(K key, Callable<R> task) {
		LockExecutor executor = getExecutorForKey(key);
		try {
			return executor.submitUnchecked(task);
		} finally {
			freeExecutorForKey(key);
		}
	}

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

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

Improve SyncronizedPerKeyExecutor - uses LockExecutor instead of raw executor