Designing a key based lock (or lock map)

Question

I'm trying to design a key-based locking facility: something like a normal reentrant lock, but instead of lock() and unlock(), you lock(key) and unlock(key), with the contract that no-one will be able to lock(key1) simultaneously if key.equals(key1).

Will this code work? Are there more efficients solutions? I particularly don't like the while loop while trying to put the lock in the map...

package luca;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.locks.ReentrantLock;

public class KeyedReentrantLock<K> {
    private ConcurrentMap<K, ReentrantLock> lockMap = new ConcurrentHashMap<K, ReentrantLock>();

    public void lock(K key) {
        ReentrantLock oldLock = lockMap.get(key);
        if (oldLock != null && oldLock.isHeldByCurrentThread()){
            // increase lock count and return.
            oldLock.lock();
            return;
        }
        ReentrantLock newLock = new ReentrantLock();
        newLock.lock();
        while ((oldLock = lockMap.putIfAbsent(key, newLock)) != null){
            // wait for the old lock to be released;
            oldLock.lock();
            oldLock.unlock();
        }
        return;
    }

    public void unlock(K key){
        ReentrantLock lock = lockMap.get(key);
        if (lock == null) throw new IllegalMonitorStateException("There was no lock for this key!");
        if (lock.getHoldCount() == 1){
            lockMap.remove(key);
        }
        lock.unlock();
    }

}

Answer 1

Why don't just use simple striped locking, like:

/**
 * Striped locks holder, contains array of {@link java.util.concurrent.locks.ReentrantLock}, on which lock/unlock
 * operations are performed. Purpose of this is to decrease lock contention.
 * <p>When client requests lock, it gives an integer argument, from which target lock is derived as follows:
 * index of lock in array equals to <code>id & (locks.length - 1)</code>.
 * Since <code>locks.length</code> is the power of 2, <code>locks.length - 1</code> is string of '1' bits,
 * and this means that all lower bits of argument are taken into account.
 * <p>Number of locks it can hold is bounded: it can be from set {2, 4, 8, 16, 32, 64}.
  */
public class StripedLock {
    private final ReentrantLock[] locks;

    /**
     * Default ctor, creates 16 locks
     */
    public StripedLock() {
        this(4);
    }

    /**
     * Creates array of locks, size of array may be any from set {2, 4, 8, 16, 32, 64} 
     * @param storagePower size of array will be equal to <code>Math.pow(2, storagePower)</code>
     */
    public StripedLock(int storagePower) {
        if (!(storagePower >= 1 && storagePower <= 6)) { throw new IllegalArgumentException("storage power must be in [1..6]"); }

        int lockSize = (int) Math.pow(2, storagePower);
        locks = new ReentrantLock[lockSize];
        for (int i = 0; i < locks.length; i++)
            locks[i] = new ReentrantLock();
    }

    /**
     * Locks lock associated with given id.
     * @param id value, from which lock is derived
     */
    public void lock(int id) {
        getLock(id).lock();
    }

    /**
     * Unlocks lock associated with given id.
     * @param id value, from which lock is derived 
     */
    public void unlock(int id) {
        getLock(id).unlock();
    }

    /**
     * Map function between integer and lock from locks array
     * @param id argument
     * @return lock which is result of function 
     */
    private ReentrantLock getLock(int id) {
        return locks[id & (locks.length - 1)];
    }
}