夢の終着駅

1
Cache<String, Value> cache = Caffeine.newBuilder().build();
2

3
Map<String, Value> getAll(Set<String> keys) {
4
    // 命中键值对
5
    Map<String, Value> present = new HashMap<>();
6
    // 未命中Key
7
    Set<String> missing = new HashSet<>();
8

9
    // 筛选
10
    for (String k : keys) {
11
        Value v = cache.getIfPresent(k);
12
        if (v != null) present.put(k, v);
13
        else missing.add(k);
14
    }
15

16
    // 批量加载
17
    if (!missing.isEmpty()) {
18
        Map<String, Value> loaded = db.batchLoad(missing);
19
        cache.putAll(loaded);
20
        present.putAll(loaded);
21
    }
22
    return present;
23
}

1
public class BatchingSingleflight<K, V> {
2
    // 只有首次触发加载的线程会创建 future，其他并发到达者会复用这个 future。
3
    private final ConcurrentHashMap<K, CompletableFuture<V>> inflight = new ConcurrentHashMap<>();
4

5
    // 用于收集待批量加载的 key，保持插入顺序
6
    private final ConcurrentLinkedQueue<K> pendingQueue = new ConcurrentLinkedQueue<>();
7

8
    // 与queue配合用于去重
9
    private final Set<K> pendingSet = ConcurrentHashMap.newKeySet();
10

11
    // 用于执行批量加载任务
12
    private final ScheduledExecutorService scheduler;
13

14
    // 用于执行实际的任务或立即触发drain的普通executor
15
    private final Executor executor;
16

17
    private final Cache<K,V> cache;
18

19
    // 批量加载函数
20
    private final Function<Set<K>, Map<K, V>> batchLoader;
21

22
    // 最大批次大小
23
    private final int maxBatchSize;
24

25
    // 时窗（ms）
26
    private final long windowMillis;
27

28
    // 表示是否已经调度了一个drain任务（防止重复调度）
29
    private final AtomicBoolean drainScheduled = new AtomicBoolean(false);
30

31
    // 用于判断当前Key是否已经正在加载中
32
    public CompletableFuture<V> get(K key) {
33
        V cached = cache == null ? null : cache.getIfPresent(key);
34
        if (cached != null){
35
            // 正在加载则共享回调结果
36
             return CompletableFuture.completedFuture(cached);
37
        }
38
        // 反之则开始加载
39
        CompletableFuture<V> fut = inflight.computeIfAbsent(key, k -> {
40
            CompletableFuture<V> newF = new CompletableFuture<>();
41
            // 去重
42
            if (pendingSet.add(k)) {
43
                收集
44
                pendingQueue.add(k);
45
            }
46
            // 开启时间窗口
47
            scheduleDrainIfNeeded();
48
            return newF;
49
        });
50

51
        // 在回调完成后释放资源，防止内存泄漏
52
        fut.whenComplete((v, ex) -> inflight.remove(key, fut));
53
        return fut;
54
    }
55

56
    // 判断是否需要时间窗口
57
    private void scheduleDrainIfNeeded() {
58
        // 当前Key数量是否已经达到阈值
59
        if (pendingSet.size() >= maxBatchSize) {
60
            // 当前是否已经调度过了，避免并发问题
61
            if (drainScheduled.compareAndSet(false, true)) {
62
                // 立即执行
63
                executor.execute(this::drain);
64
            }
65
        } else {
66
            if (drainScheduled.compareAndSet(false, true)) {
67
                // 开启时间窗口
68
                scheduler.schedule(this::drain, windowMillis, TimeUnit.MILLISECONDS);
69
            }
70
        }
71
    }
72

73
    private void drain() {
74
        try {
75
            List<K> batchKeys = new ArrayList<>(maxBatchSize);
76
            // 获取需要加载的Keys
77
            while (batchKeys.size() < maxBatchSize) {
78
                K k = pendingQueue.poll();
79
                if (k == null) break;
80
                if (pendingSet.remove(k)) {
81
                    // 仅当成功移除时才真正加入，防止重复
82
                    batchKeys.add(k);
83
                }
84
            }
85
            // 为空则无需加载
86
            if (batchKeys.isEmpty()) return;
87
            // 进一步去重
88
            Set<K> keySet = new HashSet<>(batchKeys);
89

90
            CompletableFuture
91
                    // 执行加载函数
92
                    .supplyAsync(() -> batchLoader.apply(keySet), executor)
93
                    .whenComplete((map, ex) -> {
94
                        if (ex != null) {
95
                            // 获取结果
96
                            for (K k : keySet) {
97
                                CompletableFuture<V> f = inflight.get(k);
98
                                if (f != null) f.completeExceptionally(ex);
99
                            }
100
                        } else {
101
                            for (K k : keySet) {
102
                                CompletableFuture<V> f = inflight.get(k);
103
                                V value = map.get(k);
104
                                if (value != null) {
105
                                    // 写回本地缓存
106
                                    if (cache != null) cache.put(k, value);
107
                                    if (f != null) f.complete(value);
108
                                } else {
109
                                     // 库中也没有则需要做缓存穿透的防御
110
                                }
111
                            }
112
                        }
113

114
                        // 这里再做一次检查并清理中已经完成的future
115
                        for (K k : keySet) {
116
                            CompletableFuture<V> f = inflight.get(k);
117
                            if (f != null && f.isDone()) {
118
                                inflight.remove(k, f);
119
                            }
120
                        }
121
                    });
122
        } finally {
123
            // 回滚状态
124
            drainScheduled.set(false);
125
            // 如果还有剩余Key则再次加载
126
            if (!pendingSet.isEmpty()) scheduleDrainIfNeeded();
127
        }
128
    }
129
}