支援 equals 相等的物件(可重複物件)作為 WeakHashMap 的 Key

2022-11-18 18:01:06

原文地址

程式碼地址

問題

長連結場景下通常有一個類似 Map<String, Set<Long>> 的結構,用來查詢一個邏輯組內的哪些使用者,String 型別的 Entry.key 是邏輯組 key,Set<Long> 型別的 Entry.value 存放邏輯組內的使用者 Id,那麼這個 Map 顯然要在邏輯組內使用者為 0 時刪除這個 Entry,以避免記憶體漏失。

刪除 Map 的 value 很容易聯想到 remove,但並行的處理很複雜,還要單獨開一個執行緒,如果可以自動刪除就好了,而 WeakHashMap 就可以自動刪除 value,前提它是 Entry.key 不存在參照時刪除 Entry.value,那麼只要將使用者的生命週期和 Entry.key 關聯上即可,以 Netty 的 Channel 為例就是將該 Entry.key 放到 Channel.attr 中。

上面稍微一看就有問題,Entry.key 是一個 String 型別的變數,字串存在常數池(字串其實挺好的),Channel 就算銷燬了也不會丟失對 WeakHashMap Entry.value 的參照,如果每次都 new 一個物件呢?問題更大,此時只有第一個使用者強參照 WeakHashMap 的 Entry.value(即 new Set 再 add),其他使用者僅僅是獲取到了(此時 Entry.key 是第一個使用者的,而不是當前使用者的),這樣第一個使用者下線時,這個 Set 就會被 GC。顯而易見問題是 Entry.Key 參照不一致導致的,只要給使用者返回永遠相同的 Entry.key 即可。

如何返回永遠相同的物件呢?感覺又回到了原點,因為返回一樣的物件顯然是 Map<String, Object>,但這個 Map 同樣不能記憶體漏失,不過情況略有不同,區別在於查詢 Set 變成了一個巢狀的查詢(String -> Object -> Set<Long>),而使用者強參照的 Entry.key 變成了 Object,即 Object 物件的生命週期跟隨使用者走即可( WeakHashMap<Object, Set<Long>> 負責 GC Set),也就是 WeakHashMap<Object, WeakReference<Object>>。

解決

下面給出程式碼:

package io.github.hligaty.util;

import java.lang.ref.WeakReference;
import java.util.Objects;
import java.util.WeakHashMap;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * Recreatable key objects.
 * With recreatable key objects,
 * the automatic removal of WeakHashMap entries whose keys have been discarded may prove to be confusing,
 * but WeakKey will not.
 *
 * @param <K> the type of keys maintained
 * @author hligaty
 * @see java.util.WeakHashMap
 */
public class WeakKey<K> {
    private static final WeakHashMap<WeakKey<?>, WeakReference<WeakKey<?>>> cache = new WeakHashMap<>();
    private static final ReadWriteLock cacheLock = new ReentrantReadWriteLock();
    private static final WeakHashMap<Thread, WeakKey<?>> shadowCache = new WeakHashMap<>();
    private static final ReadWriteLock shadowCacheLock = new ReentrantReadWriteLock();
    
    private K key;

    private WeakKey() {
    }
    
    @SuppressWarnings("unchecked")
    public static <T> WeakKey<T> wrap(T key) {
        WeakKey<T> shadow = (WeakKey<T>) getShadow();
        shadow.key = key;
        cacheLock.readLock().lock();
        try {
            WeakReference<WeakKey<?>> ref = cache.get(shadow);
            if (ref != null) {
                shadow.key = null;
                return (WeakKey<T>) ref.get();
            }
        } finally {
            cacheLock.readLock().unlock();
        }
        cacheLock.writeLock().lock();
        try {
            WeakReference<WeakKey<?>> newRef = cache.get(shadow);
            shadow.key = null;
            if (newRef == null) {
                WeakKey<T> weakKey = new WeakKey<>();
                weakKey.key = key;
                newRef = new WeakReference<>(weakKey);
                cache.put(weakKey, newRef);
                return weakKey;
            }
            return (WeakKey<T>) newRef.get();
        } finally {
            cacheLock.writeLock().unlock();
        }
    }

    private static WeakKey<?> getShadow() {
        Thread thread = Thread.currentThread();
        shadowCacheLock.readLock().lock();
        WeakKey<?> shadow;
        try {
            shadow = shadowCache.get(thread);
            if (shadow != null) {
                return shadow;
            }
        } finally {
            shadowCacheLock.readLock().unlock();
        }
        shadowCacheLock.writeLock().lock();
        try {
            shadow = shadowCache.get(thread);
            if (shadow == null) {
                shadow = new WeakKey<>();
                shadowCache.put(thread, shadow);
                return shadow;
            }
            return shadow;
        } finally {
            shadowCacheLock.writeLock().unlock();
        }
    }

    public K unwrap() {
        return key;
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;
        WeakKey<?> weakKey = (WeakKey<?>) o;
        return Objects.equals(key, weakKey.key);
    }

    @Override
    public int hashCode() {
        return Objects.hash(key);
    }

    @Override
    public String toString() {
        return "WeakKey{" +
                "attr=" + key +
                '}';
    }
}

WeakKey 是前面說的 Object,使用時將需要釋放的資料 Data 放到以 WeakKey 為 key 的 WeakHashMap(WeakHashMap<WeakKey, Data>),這樣當全部使用者釋放 WeakKey 參照時就可以完成 WeakHashMap Entry 的 GC(包括 WeakKey 和 Data)。

WeakKey 的主要工作是將使用者傳入的 key 封裝一下再返回,保證全域性唯一和記憶體安全,核心結構是 WeakHashMap<WeakKey<?>, WeakReference<WeakKey<?>>> cache,Entry.key 是對使用者 key 封裝的 WeakKey,Entry.value 是 Entry.key 外層巢狀的 WeakReference,作用是避免 value 對 key 強參照而無法對 Entry GC。因此 cache 只要沒人強參照裡面的 WeakKey,這個 map 在 GC 後就是空的,這樣就完成了目標,其餘的就是優化了。

如果想在 cache 裡查到 WeakKey,那麼首先要新建一個 WeakKey,再把 key 賦值到 WeakKey 中,再通過這個新建的 WeakKey 查詢,像下面一樣:

String key = "key";
WeakKey<String> weakKey = new WeakKey<>();
weakKey.key = key;
WeakReference<WeakKey<?>> ref = cache.get(weakKey);

每次查詢都新建物件,有點沙雕,這裡使用快取物件賦值再查詢就可以,另外要保證執行緒安全,threadLocal 沒大問題(ThreadLocal.withInitial(WeakKey::new)),只是不能在 finally 裡 remove(remove 的話下次還得新建),線上程池裡使用問題不大,不過還有另一種辦法,就是 WeakHashMap<Thread, WeakKey<?>>,它可以保證這個快取中的「影子」物件在這個執行緒只建立一次,當執行緒被 GC 的同時刪除「影子」物件,與 threadLocal 的區別只是犧牲了一些加讀鎖的時間。

測試

下面的 WeakHashMap put 了 Arrays.asList(705, 630, 818) 和 Collections.singletonList(705630818) 兩個資料,只有後面的 key 被方法參照了,因此在 GC 後 前一個 key 在 map 中找不到 value,而後一個 key 能獲取到 value。

package io.github.hligaty.util;

import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;

import java.util.*;

class WeakKeyTest {
    
    @Test
    public void testWeakKey() throws InterruptedException {
        WeakHashMap<WeakKey<List<Integer>>, Object> map = new WeakHashMap<>();
        map.put(WeakKey.wrap(Arrays.asList(705, 630, 818)), new Object());
        WeakKey<List<Integer>> weakKey = WeakKey.wrap(Collections.singletonList(705630818));
        map.put(weakKey, new Object());
        System.gc();
        Thread.sleep(5000L);
        Assertions.assertNull(map.get(WeakKey.wrap(Arrays.asList(705, 630, 818))));
        Assertions.assertNotNull(map.get(WeakKey.wrap(Collections.singletonList(705630818))));
    }
}

其他

如果你想使用 null,那 WeakKey 是支援的,但需要注意一點,如果你有兩個不同型別的 key 使用了 WeakKey,而兩者都允許 WeakKey.wrap(null),那麼當有一個型別的使用者持有 WeakKey.wrap(null),另一個型別的 WeakKey.wrap(null) 是不會被釋放的,因為顯然 Objects.equals(null, null) 為 true。