1 /*
2 * Copyright (c) 2023, 2025, Oracle and/or its affiliates. All rights reserved.
3 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
4 *
5 * This code is free software; you can redistribute it and/or modify it
6 * under the terms of the GNU General Public License version 2 only, as
7 * published by the Free Software Foundation. Oracle designates this
8 * particular file as subject to the "Classpath" exception as provided
9 * by Oracle in the LICENSE file that accompanied this code.
10 *
11 * This code is distributed in the hope that it will be useful, but WITHOUT
12 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
13 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
14 * version 2 for more details (a copy is included in the LICENSE file that
15 * accompanied this code).
16 *
17 * You should have received a copy of the GNU General Public License version
18 * 2 along with this work; if not, write to the Free Software Foundation,
19 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
20 *
21 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
22 * or visit www.oracle.com if you need additional information or have any
23 * questions.
24 */
25
26 package jdk.internal.util;
27
28 import java.lang.ref.Reference;
29 import java.lang.ref.ReferenceQueue;
30 import java.lang.ref.SoftReference;
31 import java.lang.ref.WeakReference;
32 import java.util.AbstractMap;
33 import java.util.Collection;
34 import java.util.HashMap;
35 import java.util.Objects;
36 import java.util.Map;
37 import java.util.Set;
38 import java.util.concurrent.ConcurrentHashMap;
39 import java.util.function.Supplier;
40 import java.util.function.UnaryOperator;
41 import java.util.stream.Collectors;
42 import java.util.stream.Stream;
43
44 import jdk.internal.access.SharedSecrets;
45 import jdk.internal.misc.CDS;
46
47 /**
48 * This class provides management of {@link Map maps} where it is desirable to
49 * remove entries automatically when the key is garbage collected. This is
50 * accomplished by using a backing map where the keys are either a
51 * {@link WeakReference} or a {@link SoftReference}.
52 * Keys must be {@linkplain Class#isIdentity() identity objects.}
53 * <p>
54 * To create a {@link ReferencedKeyMap} the user must provide a {@link Supplier}
55 * of the backing map and whether {@link WeakReference} or
56 * {@link SoftReference} is to be used.
57 *
58 * {@snippet :
59 * // Use HashMap and WeakReference
60 * Map<Long, String> map = ReferencedKeyMap.create(false, HashMap::new);
61 * map.put(10_000_000L, "a");
62 * map.put(10_000_001L, "b");
63 * map.put(10_000_002L, "c");
64 * map.put(10_000_003L, "d");
65 * map.put(10_000_004L, "e");
66 *
67 * // Use ConcurrentHashMap and SoftReference
68 * map = ReferencedKeyMap.create(true, ConcurrentHashMap::new);
69 * map.put(20_000_000L, "v");
70 * map.put(20_000_001L, "w");
71 * map.put(20_000_002L, "x");
72 * map.put(20_000_003L, "y");
73 * map.put(20_000_004L, "z");
74 * }
75 *
76 * @implNote Care must be given that the backing map does replacement by
77 * replacing the value in the map entry instead of deleting the old entry and
78 * adding a new entry, otherwise replaced entries may end up with a strongly
79 * referenced key. {@link HashMap} and {@link ConcurrentHashMap} are known
80 * to be safe.
81 *
82 * @param <K> the type of keys maintained by this map
83 * @param <V> the type of mapped values
84 *
85 * @since 21
86 */
87 public final class ReferencedKeyMap<K, V> implements Map<K, V> {
88 /**
89 * true if {@link SoftReference} keys are to be used,
90 * {@link WeakReference} otherwise.
91 */
92 private final boolean isSoft;
93
94 /**
95 * Backing {@link Map}.
96 */
97 private final Map<ReferenceKey<K>, V> map;
98
99 /**
100 * {@link ReferenceQueue} for cleaning up entries.
101 */
102 private final ReferenceQueue<K> stale;
103
104 /**
105 * @return a supplier to create a {@code ConcurrentHashMap} appropriate for use in the
106 * create methods.
107 * @param <K> the type of keys maintained by the new map
108 * @param <V> the type of mapped values
109 */
110 public static <K, V> Supplier<Map<ReferenceKey<K>, V>> concurrentHashMapSupplier() {
111 return new Supplier<>() {
112 @Override
113 public Map<ReferenceKey<K>, V> get() {
114 return new ConcurrentHashMap<>();
115 }
116 };
117 }
118
119 /**
120 * Private constructor.
121 *
122 * @param isSoft true if {@link SoftReference} keys are to
123 * be used, {@link WeakReference} otherwise.
124 * @param map backing map
125 * @param stale {@link ReferenceQueue} for cleaning up entries
126 */
127 private ReferencedKeyMap(boolean isSoft, Map<ReferenceKey<K>, V> map, ReferenceQueue<K> stale) {
128 this.isSoft = isSoft;
129 this.map = map;
130 this.stale = stale;
131 }
132
133 /**
134 * Create a new {@link ReferencedKeyMap} map.
135 *
136 * @param isSoft true if {@link SoftReference} keys are to
137 * be used, {@link WeakReference} otherwise.
138 * @param supplier {@link Supplier} of the backing map
139 *
140 * @return a new map with {@link Reference} keys
141 *
142 * @param <K> the type of keys maintained by the new map
143 * @param <V> the type of mapped values
144 */
145 public static <K, V> ReferencedKeyMap<K, V>
146 create(boolean isSoft, Supplier<Map<ReferenceKey<K>, V>> supplier) {
147 return new ReferencedKeyMap<K, V>(isSoft, supplier.get(), new ReferenceQueue<>());
148 }
149
150 /**
151 * {@return a key suitable for a map entry}
152 *
153 * @param key unwrapped key
154 */
155 @SuppressWarnings("unchecked")
156 private ReferenceKey<K> entryKey(Object key) {
157 if (isSoft) {
158 return new SoftReferenceKey<>((K)key, stale);
159 } else {
160 return new WeakReferenceKey<>((K)key, stale);
161 }
162 }
163
164 /**
165 * {@return a key suitable for lookup}
166 *
167 * @param key unwrapped key
168 */
169 @SuppressWarnings("unchecked")
170 private ReferenceKey<K> lookupKey(Object key) {
171 return new StrongReferenceKey<>((K)key);
172 }
173
174 @Override
175 public int size() {
176 removeStaleReferences();
177 return map.size();
178 }
179
180 @Override
181 public boolean isEmpty() {
182 removeStaleReferences();
183 return map.isEmpty();
184 }
185
186 @Override
187 public boolean containsKey(Object key) {
188 Objects.requireNonNull(key, "key must not be null");
189 removeStaleReferences();
190 return map.containsKey(lookupKey(key));
191 }
192
193 @Override
194 public boolean containsValue(Object value) {
195 Objects.requireNonNull(value, "value must not be null");
196 removeStaleReferences();
197 return map.containsValue(value);
198 }
199
200 @Override
201 public V get(Object key) {
202 removeStaleReferences();
203 return getNoCheckStale(key);
204 }
205
206 // Internal get(key) without removing stale references that would modify the keyset.
207 // Use when iterating or streaming over the keys to avoid ConcurrentModificationException.
208 private V getNoCheckStale(Object key) {
209 Objects.requireNonNull(key, "key must not be null");
210 return map.get(lookupKey(key));
211 }
212
213 @Override
214 public V put(K key, V newValue) {
215 Objects.requireNonNull(key, "key must not be null");
216 Objects.requireNonNull(newValue, "value must not be null");
217 removeStaleReferences();
218 ReferenceKey<K> entryKey = entryKey(key);
219 // If {@code put} returns non-null then was actually a {@code replace}
220 // and older key was used. In that case the new key was not used and the
221 // reference marked stale.
222 V oldValue = map.put(entryKey, newValue);
223 if (oldValue != null) {
224 entryKey.unused();
225 }
226 return oldValue;
227 }
228
229 @Override
230 public V remove(Object key) {
231 // Rely on gc to clean up old key.
232 return map.remove(lookupKey(key));
233 }
234
235 @Override
236 public void putAll(Map<? extends K, ? extends V> m) {
237 removeStaleReferences();
238 for (Entry<? extends K, ? extends V> entry : m.entrySet()) {
239 K key = entry.getKey();
240 V value = entry.getValue();
241 put(key, value);
242 }
243 }
244
245 @Override
246 public void clear() {
247 removeStaleReferences();
248 // Rely on gc to clean up old keys.
249 map.clear();
250 }
251
252 /**
253 * Common routine for collecting the current set of keys.
254 *
255 * @return {@link Stream} of valid keys (unwrapped)
256 */
257 private Stream<K> filterKeySet() {
258 return map.keySet()
259 .stream()
260 .map(ReferenceKey::get)
261 .filter(Objects::nonNull);
262 }
263
264 @Override
265 public Set<K> keySet() {
266 removeStaleReferences();
267 return filterKeySet().collect(Collectors.toSet());
268 }
269
270 @Override
271 public Collection<V> values() {
272 removeStaleReferences();
273 return map.values();
274 }
275
276 @Override
277 public Set<Entry<K, V>> entrySet() {
278 removeStaleReferences();
279 return filterKeySet()
280 .map(k -> new AbstractMap.SimpleEntry<>(k, getNoCheckStale(k)))
281 .collect(Collectors.toSet());
282 }
283
284 @Override
285 public V putIfAbsent(K key, V newValue) {
286 removeStaleReferences();
287 ReferenceKey<K> entryKey = entryKey(key);
288 // If {@code putIfAbsent} returns non-null then was actually a
289 // {@code replace} and older key was used. In that case the new key was
290 // not used and the reference marked stale.
291 V oldValue = map.putIfAbsent(entryKey, newValue);
292 if (oldValue != null) {
293 entryKey.unused();
294 }
295 return oldValue;
296 }
297
298 @Override
299 public boolean remove(Object key, Object value) {
300 // Rely on gc to clean up old key.
301 return map.remove(lookupKey(key), value);
302 }
303
304 @Override
305 public boolean replace(K key, V oldValue, V newValue) {
306 removeStaleReferences();
307 // If replace is successful then the older key will be used and the
308 // lookup key will suffice.
309 return map.replace(lookupKey(key), oldValue, newValue);
310 }
311
312 @Override
313 public V replace(K key, V value) {
314 removeStaleReferences();
315 // If replace is successful then the older key will be used and the
316 // lookup key will suffice.
317 return map.replace(lookupKey(key), value);
318 }
319
320 @Override
321 public String toString() {
322 removeStaleReferences();
323 return filterKeySet()
324 .map(k -> k + "=" + getNoCheckStale(k))
325 .collect(Collectors.joining(", ", "{", "}"));
326 }
327
328 /**
329 * Removes enqueued weak references from map.
330 */
331 public void removeStaleReferences() {
332 while (true) {
333 Object key = stale.poll();
334 if (key == null) {
335 break;
336 }
337 map.remove(key);
338 }
339 }
340
341 @SuppressWarnings("unchecked")
342 public void prepareForAOTCache() {
343 // We are running the AOT assembly phase. The JVM has a single Java thread, so
344 // we don't have any concurrent threads that may modify the map while we are
345 // iterating its keys.
346 //
347 // Also, the java.lang.ref.Reference$ReferenceHandler thread is not running,
348 // so even if the GC has put some of the keys on the pending ReferencePendingList,
349 // none of the keys would have been added to the stale queue yet.
350 assert CDS.isSingleThreadVM();
351
352 for (ReferenceKey<K> key : map.keySet()) {
353 Object referent = key.get();
354 if (referent == null) {
355 // We don't need this key anymore. Add to stale queue
356 ((Reference)key).enqueue();
357 } else {
358 // Make sure the referent cannot be collected. Otherwise, when
359 // the referent is collected, the GC may push the key onto
360 // Universe::reference_pending_list() at an unpredictable time,
361 // making it difficult to correctly serialize the key's
362 // state into the CDS archive.
363 //
364 // See aotReferenceObjSupport.cpp for more info.
365 CDS.keepAlive(referent);
366 }
367 Reference.reachabilityFence(referent);
368 }
369
370 // Remove all keys enqueued above
371 removeStaleReferences();
372 }
373
374
375 /**
376 * Puts an entry where the key and the value are the same. Used for
377 * interning values in a set.
378 *
379 * @implNote Requires a {@link ReferencedKeyMap} whose {@code V} type
380 * is a {@code ReferenceKey<K>}. Otherwise, a {@link ClassCastException} will
381 * be thrown.
382 *
383 * @param setMap {@link ReferencedKeyMap} where interning takes place
384 * @param key key to add
385 *
386 * @param <T> type of key
387 *
388 * @return the old key instance if found otherwise the new key instance
389 *
390 * @throws ClassCastException if {@code V} is not {@code EntryKey<T>}
391 */
392 static <T> T intern(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key) {
393 T value = existingKey(setMap, key);
394 if (value != null) {
395 return value;
396 }
397 return internKey(setMap, key);
398 }
399
400 /**
401 * Puts an entry where the key and the value are the same. Used for
402 * interning values in a set.
403 *
404 * @implNote Requires a {@link ReferencedKeyMap} whose {@code V} type
405 * is a {@code ReferenceKey<K>}. Otherwise, a {@link ClassCastException} will
406 * be thrown.
407 *
408 * @param setMap {@link ReferencedKeyMap} where interning takes place
409 * @param key key to add
410 * @param interner operation to apply to key before adding to map
411 *
412 * @param <T> type of key
413 *
414 * @return the old key instance if found otherwise the new key instance
415 *
416 * @throws ClassCastException if {@code V} is not {@code EntryKey<T>}
417 *
418 * @implNote This version of intern should not be called during phase1
419 * using a lambda. Use an UnaryOperator instance instead.
420 */
421 static <T> T intern(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key, UnaryOperator<T> interner) {
422 T value = existingKey(setMap, key);
423 if (value != null) {
424 return value;
425 }
426 key = interner.apply(key);
427 return internKey(setMap, key);
428 }
429
430 /**
431 * Check if the key already exists in the map.
432 *
433 * @param setMap {@link ReferencedKeyMap} where interning takes place
434 * @param key key to test
435 *
436 * @param <T> type of key
437 *
438 * @return key if found otherwise null
439 */
440 private static <T> T existingKey(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key) {
441 setMap.removeStaleReferences();
442 ReferenceKey<T> entryKey = setMap.get(setMap.lookupKey(key));
443 return entryKey != null ? entryKey.get() : null;
444 }
445
446 /**
447 * Attempt to add key to map.
448 *
449 * @param setMap {@link ReferencedKeyMap} where interning takes place
450 * @param key key to add
451 *
452 * @param <T> type of key
453 *
454 * @return the old key instance if found otherwise the new key instance
455 */
456 private static <T> T internKey(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key) {
457 ReferenceKey<T> entryKey = setMap.entryKey(key);
458 T interned;
459 do {
460 setMap.removeStaleReferences();
461 ReferenceKey<T> existing = setMap.map.putIfAbsent(entryKey, entryKey);
462 if (existing == null) {
463 return key;
464 } else {
465 // If {@code putIfAbsent} returns non-null then was actually a
466 // {@code replace} and older key was used. In that case the new
467 // key was not used and the reference marked stale.
468 interned = existing.get();
469 if (interned != null) {
470 entryKey.unused();
471 }
472 }
473 } while (interned == null);
474 return interned;
475 }
476
477
478 /**
479 * Attempt to add key to map if absent.
480 *
481 * @param setMap {@link ReferencedKeyMap} where interning takes place
482 * @param key key to add
483 *
484 * @param <T> type of key
485 *
486 * @return true if the key was added
487 */
488 static <T> boolean internAddKey(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key) {
489 ReferenceKey<T> entryKey = setMap.entryKey(key);
490 setMap.removeStaleReferences();
491 ReferenceKey<T> existing = setMap.map.putIfAbsent(entryKey, entryKey);
492 if (existing == null) {
493 return true;
494 } else {
495 // If {@code putIfAbsent} returns non-null then was actually a
496 // {@code replace} and older key was used. In that case the new
497 // key was not used and the reference marked stale.
498 entryKey.unused();
499 return false;
500 }
501 }
502
503 }