Hello, is there a coroutine-friendly alternative to

ConcurrentHashMap

? The idea is that an update to the map via

compute

might be a suspending function, and the map should lock out the hash bucket being updated until that function has completed. I have this naive version which locks out the entire map:

class CoroutineMap<K : Any, V : Any> {

    private val data: MutableMap<K, V> = mutableMapOf()
    private val dataLock = Mutex()
    
    suspend fun compute(key: K, update: suspend (V?) -> V?): V? = dataLock.withLock { 
        val updated = update(data[key])
        if (updated == null) data.remove(key) else data[key] = updated
        updated
    }
}

but a version with per-bucket mutexes would be better, I believe?

You can still use ConcurrentHashMap + getOrPut, the rest varies on exact usecase. For example, your case may be rewritten as

class Cell<V: Any>(var value: V? = null) {
   val lock = Mutex()
   
   fun update(update: suspend (V?) -> V?): V?): V? {
      lock.withLock { value = update(value)
   }
}

class CoroutineMap<K : Any, V : Any> {

private val map = ConcurrentHashMap<K, Cell<V>>()
    
    suspend fun compute(key: K, update: suspend (V?) -> V?): V? = 
        map.getOrPut(key, ::Cell).update(update)
    }
}

In this case it is lock per key

Having a lock per bucket is efficient for ConcurrentHashMap; would it be efficient in the coroutine case with Mutexes?

I may have misunderstood your code, but doesn't this create a new Cell (and therefore a different Mutex) for each call?

getOrPut will only create a new Cell if there isn't already one for that key