How can I perform a thread-safe get then remove with ConcurrentHashMap?

Question

In an interview, I was asked to check whether following code works as intended.

ConcurrentHashMap<Integer, Integer> chm = new ConcurrentHashMap<>();

if (chm.get(key) != null) {
    chm.get(key).doSomething();
    chm.remove(key);
}

According to JavaDocs, get returns value of last completed update operation. So if thread 1 already called chm.remove(key) and if thread 2 came inside the if statement and is about to call get method then we might get an exception. Is this correct?

How can I make this thread-safe?

Answer 1

Map.remove(key) returns the value if it has been removed. This is a very good trick in many situations, including yours:

Object value = chm.remove(key)
if(value != null)
{
     value.doSomething();
}

You can't work safely with a get then a remove because if two threads call your method at the same time, there's always a risk they call doSomething two or more times, before the key has been removed.

This is not possible if you remove it first. The code above is Threadsafe, also simpler.

Answer 2

You are correct. If this Map can be modified by multiple threads, it's possible that the first call to chm.get(key) will return a non-null value and the second call will return null (due to the removal of the key from the Map done by another thread), and thus chm.get(key).doSomething() will throw a NullPointerException.

You can make this code thread safe by using the local variable to store the result of chm.get(key):

ConcurrentHashMap<Integer, Integer> chm = new ConcurrentHashMap<Integer, Integer>();
Integer value = chm.get(key);

if(value != null) {
    value.doSomething(); // P.S. Integer class doesn't have a doSomething() method
                         // but I guess this is just an example of calling some arbitrary 
                         // instance method
    chm.remove(key);
}

BTW, even if that Map wasn't a ConcurentHashMap and only one thread had access to it, I'd still use a local variable, since it's more efficient than calling the get() method twice.

EDIT :

As commented below, this fix will not prevent doSomething() from being called multiple times for the same key/value by different threads. Whether or not that's the desired behavior is not clear.

If you wish to prevent the possiblity of doSomething() being called by multiple threads for the same key/value, you can use chm.remove(key) for both removing the key and obtaining the value at the same step.

This however runs the risk that doSomething() will not be executed at all for some key/value, since if the first call to doSomething() resulted in an exception, there won't be another call to doSomething() by another thread, since the key/value pair will no longer be in the Map. On the other hand, if you remove the key/value pair from the Map only after doSomething() executes succesfully, you guarantee that doSomething() is executed successfuly at least once for all the key/value pairs that were reomved from the Map.