JDK1.8+ 中 ConcurrentHashMap#computeIfAbsent 源码解析与使用建议

最近发现网上讲解 ConcurrentHashMap 源码的文章对于computeIfAbsent方法的讲解比较少，这里来分析一下，读者也可以自行分析compute方法，原理类似。本文不考虑JDK1.7，分析的源码版本为JDK25。

前置知识

适用于延迟初始化和缓存场景。
ConcurrentHashMap 中锁粒度很细，为数组节点（链表头结点 | 红黑树根节点 | 特殊用法节点）。
computeIfAbsent为原子操作，必须加锁。
这个操作不支持递归，会抛异常。意味着缓存查询不能嵌套。
传入的函数不能修改 ConcurrentHashMap 本身。
不支持 null，可以用来表示没存。

底层实现的前置知识

拉链法（数组+链表），hash 冲突多时退化成红黑树。链表采用尾插法，保证桶位局部有序，提升空间局部性。
扰动处理hash，减少冲突，简单的位操作，对性能影响小。
底层数组延迟初始化，在putAll、put时初始化，避免可能的扩容开销。
tabAt 使用unsafe实现数组中元素 volatile 语义。
转发节点 ForwardingNode 通过 f.hash == MOVED 计算，避免使用instanceof
检测到其他线程正在扩容时，当前线程辅助扩容
哈希均匀分布时，某个节点链表长度为2的概率为1/8。
除非受到攻击或者错误编写hashcode方法，几乎不可能出现红黑树，出现红黑树会影响性能。

代码

Java 复制代码

/**
 * 如果指定的key还没有对应的value或者value为null，则使用给定的函数计算value并存储
 * 这是一个原子操作，适用于需要延迟初始化或缓存的场景
 */
public V computeIfAbsent(K key, Function<? super K, ? extends V> mappingFunction) {
    // 参数校验：key和mappingFunction都不能为null
    if (key == null || mappingFunction == null)
        throw new NullPointerException();
    
    // 对key的hashcode进行扰动处理，使hash分布更均匀
    int h = spread(key.hashCode());
    V val = null;
    int binCount = 0;
    
    // 无限循环，直到成功完成操作
    for (Node<K,V>[] tab = table;;) {
        Node<K,V> f; int n, i, fh; K fk; V fv;
        
        // 1. 如果table未初始化，则初始化table
        if (tab == null || (n = tab.length) == 0)
            tab = initTable();
        
        // 2. 如果目标桶为空，则尝试插入新节点
        else if ((f = tabAt(tab, i = (n - 1) & h)) == null) {
            // 创建一个ReservationNode作为占位符，防止其他线程同时操作同一个桶
            Node<K,V> r = new ReservationNode<K,V>();
            synchronized (r) {  // 对占位符加锁，这里加锁是防御性的
                // 使用CAS操作将占位符放入桶中
                if (casTabAt(tab, i, null, r)) {
                    binCount = 1;
                    Node<K,V> node = null;
                    try {
                        // 调用映射函数计算value
                        if ((val = mappingFunction.apply(key)) != null)
                            // 如果计算结果不为null，则创建新节点
                            node = new Node<K,V>(h, key, val);
                    } finally {
                        // 无论计算是否成功，都用计算结果替换占位符
                        setTabAt(tab, i, node);
                    }
                }
            }
            // 如果成功插入了节点，则跳出循环
            if (binCount != 0)
                break;
        }
        
        // 3. 如果遇到转发节点（表示正在扩容），则帮助完成扩容
        else if ((fh = f.hash) == MOVED)
            tab = helpTransfer(tab, f);
        
        // 4. 如果第一个节点就是目标节点（无需加锁快速检查）
        else if (fh == h    // hash值相同
                 && ((fk = f.key) == key || (fk != null && key.equals(fk)))  // key相同
                 && (fv = f.val) != null)  // value不为null
            return fv;  // 直接返回已存在的value
        
        // 5. 其他情况：需要加锁进行详细检查
        else {
            boolean added = false;
            // 对桶的第一个节点加锁，保证线程安全
            synchronized (f) {
                // 再次检查桶中的节点是否还是原来的节点（防止在获取锁期间被其他线程修改）
                if (tabAt(tab, i) == f) {
                    
                    // 5.1 如果是链表节点
                    if (fh >= 0) {
                        binCount = 1;
                        // 遍历链表查找目标key
                        for (Node<K,V> e = f;; ++binCount) {
                            K ek;
                            // 如果找到相同的key，直接返回对应的value
                            if (e.hash == h &&
                                ((ek = e.key) == key ||
                                 (ek != null && key.equals(ek)))) {
                                val = e.val;
                                break;
                            }
                            
                            Node<K,V> pred = e;
                            // 如果遍历到链表末尾仍未找到，则添加新节点
                            if ((e = e.next) == null) {
                                // 调用映射函数计算value
                                if ((val = mappingFunction.apply(key)) != null) {
                                    // 检查是否有递归更新的情况
                                    if (pred.next != null)
                                        throw new IllegalStateException("Recursive update");
                                    added = true;
                                    // 将新节点添加到链表末尾
                                    pred.next = new Node<K,V>(h, key, val);
                                }
                                break;
                            }
                        }
                    }
                    
                    // 5.2 如果是树节点
                    else if (f instanceof TreeBin) {
                        binCount = 2;  // 标记为树结构
                        TreeBin<K,V> t = (TreeBin<K,V>)f;
                        TreeNode<K,V> r, p;
                        // 在树中查找目标key
                        if ((r = t.root) != null &&
                            (p = r.findTreeNode(h, key, null)) != null)
                            val = p.val;  // 找到则返回value
                        else if ((val = mappingFunction.apply(key)) != null) {
                            // 未找到且计算结果不为null，则添加到树中
                            added = true;
                            t.putTreeVal(h, key, val);
                        }
                    }
                    
                    // 5.3 如果是ReservationNode（不应该出现的情况）
                    else if (f instanceof ReservationNode)
                        throw new IllegalStateException("Recursive update");
                }
            }
            
            // 检查节点数
            if (binCount != 0) {
                // 如果链表长度达到阈值，则转换为红黑树
                if (binCount >= TREEIFY_THRESHOLD)
                    treeifyBin(tab, i);
                
                // 如果没有添加新节点（说明key已存在），则返回已存在的value
                if (!added)
                    return val;
                break;  // 操作完成，跳出循环
            }
        }
    }
    
    // 如果添加了新节点，则更新计数器
    if (val != null)
        addCount(1L, binCount);
    
    return val;  // 返回计算得到的value或已存在的value
}

这里需要特别注意的是：目标桶为空，则尝试插入新节点时，这里的锁是防御性的，是为了防止重入修改。

且看如下代码：

java 复制代码

Map<String, Integer> map = new ConcurrentHashMap<>(16);
map.computeIfAbsent(
    "AaAa",
    key -> {
        return map.computeIfAbsent(
            "BBBB",
            key2 -> 42);
    }
);

笔者本地运行结果如下：

less 复制代码

Exception in thread "main" java.lang.IllegalStateException: Recursive update
	at java.base/java.util.concurrent.ConcurrentHashMap.computeIfAbsent(ConcurrentHashMap.java:1763)
	at com.example.parfun.Main.lambda$main$1(Main.java:28)
	at java.base/java.util.concurrent.ConcurrentHashMap.computeIfAbsent(ConcurrentHashMap.java:1708)
	at com.example.parfun.Main.main(Main.java:25)

执行外部computeIfAbsent时锁住的是特殊用法节点 ReservationNode，这个节点表示正在修改，内部调用computeIfAbsent时，代码锁住的是刚刚创建的头结点，执行到5.3抛出异常。

这个技巧在很多源码中都有使用。

一些问题分析

bugs.openjdk.org/browse/JDK-... 中提到computeIfAbsent可能导致死循环，此外老bug，已修复。
性能有坑 | 慎用 Java 8 ConcurrentHashMap 的 computeIfAbsent 中提到此操作会有性能问题，原因在于如果使用JDK1.8低版本，访问头结点时，需要加锁访问。使用新版本JDK（如本文源码）不会有这个问题。一般情况下使用影响不大，一方面，哈希表通常只有头结点（链表长度为1），另一方面，初始化加锁很少成为性能瓶颈。
工作十几年,第一次在线上遇到死锁问题中提到嵌套使用computeIfAbsent操作会造成死锁，如：map1.computeIfAbsent(key, k -> map2.computeIfAbsent(key2, k2 -> v2))。进一步说，只要参数map函数可能获取锁，都有可能造成死锁。所以，map函数应该尽量使用纯函数，有时甚至可以选择放弃原子操作的要求。
使用 Caffeine 中 AsyncCache，可以实现缓存 Map 的递归实现

java 复制代码

AsyncCache<Integer, Integer> cache = Caffeine.newBuilder().buildAsync();

int factorial(int x) {
  var future = new CompletableFuture<Integer>();
  var prior = cache.asMap().putIfAbsent(x, future);
  if (prior != null) {
    return prior.join();
  }
  int result = (x == 1) ? 1 : (x * factorial(x - 1));
  future.complete(result);
  return result;
}