长链接场景下通常有一个类似 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) 是不会被释放的,因为显然 null == null