1 /*
2 * Copyright (c) 2023, 2024, 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
46 /**
47 * This class provides management of {@link Map maps} where it is desirable to
48 * remove entries automatically when the key is garbage collected. This is
49 * accomplished by using a backing map where the keys are either a
50 * {@link WeakReference} or a {@link SoftReference}.
51 * <p>
52 * To create a {@link ReferencedKeyMap} the user must provide a {@link Supplier}
53 * of the backing map and whether {@link WeakReference} or
54 * {@link SoftReference} is to be used.
55 *
56 * {@snippet :
57 * // Use HashMap and WeakReference
58 * Map<Long, String> map = ReferencedKeyMap.create(false, HashMap::new);
59 * map.put(10_000_000L, "a");
60 * map.put(10_000_001L, "b");
61 * map.put(10_000_002L, "c");
62 * map.put(10_000_003L, "d");
63 * map.put(10_000_004L, "e");
64 *
65 * // Use ConcurrentHashMap and SoftReference
66 * map = ReferencedKeyMap.create(true, ConcurrentHashMap::new);
67 * map.put(20_000_000L, "v");
68 * map.put(20_000_001L, "w");
69 * map.put(20_000_002L, "x");
70 * map.put(20_000_003L, "y");
71 * map.put(20_000_004L, "z");
72 * }
73 *
74 * @implNote Care must be given that the backing map does replacement by
75 * replacing the value in the map entry instead of deleting the old entry and
76 * adding a new entry, otherwise replaced entries may end up with a strongly
77 * referenced key. {@link HashMap} and {@link ConcurrentHashMap} are known
78 * to be safe.
79 *
80 * @param <K> the type of keys maintained by this map
81 * @param <V> the type of mapped values
82 *
83 * @since 21
84 */
85 public final class ReferencedKeyMap<K, V> implements Map<K, V> {
86 /**
87 * true if {@link SoftReference} keys are to be used,
88 * {@link WeakReference} otherwise.
89 */
90 private final boolean isSoft;
91
92 /**
93 * Backing {@link Map}.
94 */
95 private final Map<ReferenceKey<K>, V> map;
96
97 /**
98 * {@link ReferenceQueue} for cleaning up entries.
99 */
100 private final ReferenceQueue<K> stale;
101
102 /**
103 * @return a supplier to create a {@code ConcurrentHashMap} appropriate for use in the
104 * create methods.
105 * @param <K> the type of keys maintained by the new map
106 * @param <V> the type of mapped values
107 */
108 public static <K, V> Supplier<Map<ReferenceKey<K>, V>> concurrentHashMapSupplier() {
109 return new Supplier<>() {
110 @Override
111 public Map<ReferenceKey<K>, V> get() {
112 return new ConcurrentHashMap<>();
113 }
114 };
115 }
116
117 /**
118 * Private constructor.
119 *
120 * @param isSoft true if {@link SoftReference} keys are to
121 * be used, {@link WeakReference} otherwise.
122 * @param map backing map
123 * @param stale {@link ReferenceQueue} for cleaning up entries
124 */
125 private ReferencedKeyMap(boolean isSoft, Map<ReferenceKey<K>, V> map, ReferenceQueue<K> stale) {
126 this.isSoft = isSoft;
127 this.map = map;
128 this.stale = stale;
129 }
130
131 /**
132 * Create a new {@link ReferencedKeyMap} map.
133 *
134 * @param isSoft true if {@link SoftReference} keys are to
135 * be used, {@link WeakReference} otherwise.
136 * @param supplier {@link Supplier} of the backing map
137 *
138 * @return a new map with {@link Reference} keys
139 *
140 * @param <K> the type of keys maintained by the new map
141 * @param <V> the type of mapped values
142 */
143 public static <K, V> ReferencedKeyMap<K, V>
144 create(boolean isSoft, Supplier<Map<ReferenceKey<K>, V>> supplier) {
145 return create(isSoft, false, supplier);
146 }
147
148 /**
149 * Create a new {@link ReferencedKeyMap} map.
150 *
151 * @param isSoft true if {@link SoftReference} keys are to
152 * be used, {@link WeakReference} otherwise.
153 * @param useNativeQueue true if uses NativeReferenceQueue
154 * otherwise use {@link ReferenceQueue}.
155 * @param supplier {@link Supplier} of the backing map
156 *
157 * @return a new map with {@link Reference} keys
158 *
159 * @param <K> the type of keys maintained by the new map
160 * @param <V> the type of mapped values
161 */
162 public static <K, V> ReferencedKeyMap<K, V>
163 create(boolean isSoft, boolean useNativeQueue, Supplier<Map<ReferenceKey<K>, V>> supplier) {
164 return new ReferencedKeyMap<K, V>(isSoft, supplier.get(),
165 useNativeQueue ? SharedSecrets.getJavaLangRefAccess().newNativeReferenceQueue()
166 : new ReferenceQueue<>()
167 );
168 }
169
170 /**
171 * {@return a key suitable for a map entry}
172 *
173 * @param key unwrapped key
174 */
175 @SuppressWarnings("unchecked")
176 private ReferenceKey<K> entryKey(Object key) {
177 if (isSoft) {
178 return new SoftReferenceKey<>((K)key, stale);
179 } else {
180 return new WeakReferenceKey<>((K)key, stale);
181 }
182 }
183
184 /**
185 * {@return a key suitable for lookup}
186 *
187 * @param key unwrapped key
188 */
189 @SuppressWarnings("unchecked")
190 private ReferenceKey<K> lookupKey(Object key) {
191 return new StrongReferenceKey<>((K)key);
192 }
193
194 @Override
195 public int size() {
196 removeStaleReferences();
197 return map.size();
198 }
199
200 @Override
201 public boolean isEmpty() {
202 removeStaleReferences();
203 return map.isEmpty();
204 }
205
206 @Override
207 public boolean containsKey(Object key) {
208 Objects.requireNonNull(key, "key must not be null");
209 removeStaleReferences();
210 return map.containsKey(lookupKey(key));
211 }
212
213 @Override
214 public boolean containsValue(Object value) {
215 Objects.requireNonNull(value, "value must not be null");
216 removeStaleReferences();
217 return map.containsValue(value);
218 }
219
220 @Override
221 public V get(Object key) {
222 removeStaleReferences();
223 return getNoCheckStale(key);
224 }
225
226 // Internal get(key) without removing stale references that would modify the keyset.
227 // Use when iterating or streaming over the keys to avoid ConcurrentModificationException.
228 private V getNoCheckStale(Object key) {
229 Objects.requireNonNull(key, "key must not be null");
230 return map.get(lookupKey(key));
231 }
232
233 @Override
234 public V put(K key, V newValue) {
235 Objects.requireNonNull(key, "key must not be null");
236 Objects.requireNonNull(newValue, "value must not be null");
237 removeStaleReferences();
238 ReferenceKey<K> entryKey = entryKey(key);
239 // If {@code put} returns non-null then was actually a {@code replace}
240 // and older key was used. In that case the new key was not used and the
241 // reference marked stale.
242 V oldValue = map.put(entryKey, newValue);
243 if (oldValue != null) {
244 entryKey.unused();
245 }
246 return oldValue;
247 }
248
249 @Override
250 public V remove(Object key) {
251 // Rely on gc to clean up old key.
252 return map.remove(lookupKey(key));
253 }
254
255 @Override
256 public void putAll(Map<? extends K, ? extends V> m) {
257 removeStaleReferences();
258 for (Entry<? extends K, ? extends V> entry : m.entrySet()) {
259 K key = entry.getKey();
260 V value = entry.getValue();
261 put(key, value);
262 }
263 }
264
265 @Override
266 public void clear() {
267 removeStaleReferences();
268 // Rely on gc to clean up old keys.
269 map.clear();
270 }
271
272 /**
273 * Common routine for collecting the current set of keys.
274 *
275 * @return {@link Stream} of valid keys (unwrapped)
276 */
277 private Stream<K> filterKeySet() {
278 return map.keySet()
279 .stream()
280 .map(ReferenceKey::get)
281 .filter(Objects::nonNull);
282 }
283
284 @Override
285 public Set<K> keySet() {
286 removeStaleReferences();
287 return filterKeySet().collect(Collectors.toSet());
288 }
289
290 @Override
291 public Collection<V> values() {
292 removeStaleReferences();
293 return map.values();
294 }
295
296 @Override
297 public Set<Entry<K, V>> entrySet() {
298 removeStaleReferences();
299 return filterKeySet()
300 .map(k -> new AbstractMap.SimpleEntry<>(k, getNoCheckStale(k)))
301 .collect(Collectors.toSet());
302 }
303
304 @Override
305 public V putIfAbsent(K key, V newValue) {
306 removeStaleReferences();
307 ReferenceKey<K> entryKey = entryKey(key);
308 // If {@code putIfAbsent} returns non-null then was actually a
309 // {@code replace} and older key was used. In that case the new key was
310 // not used and the reference marked stale.
311 V oldValue = map.putIfAbsent(entryKey, newValue);
312 if (oldValue != null) {
313 entryKey.unused();
314 }
315 return oldValue;
316 }
317
318 @Override
319 public boolean remove(Object key, Object value) {
320 // Rely on gc to clean up old key.
321 return map.remove(lookupKey(key), value);
322 }
323
324 @Override
325 public boolean replace(K key, V oldValue, V newValue) {
326 removeStaleReferences();
327 // If replace is successful then the older key will be used and the
328 // lookup key will suffice.
329 return map.replace(lookupKey(key), oldValue, newValue);
330 }
331
332 @Override
333 public V replace(K key, V value) {
334 removeStaleReferences();
335 // If replace is successful then the older key will be used and the
336 // lookup key will suffice.
337 return map.replace(lookupKey(key), value);
338 }
339
340 @Override
341 public String toString() {
342 removeStaleReferences();
343 return filterKeySet()
344 .map(k -> k + "=" + getNoCheckStale(k))
345 .collect(Collectors.joining(", ", "{", "}"));
346 }
347
348 /**
349 * Removes enqueued weak references from map.
350 */
351 public void removeStaleReferences() {
352 while (true) {
353 Object key = stale.poll();
354 if (key == null) {
355 break;
356 }
357 map.remove(key);
358 }
359 }
360
361 /**
362 * Puts an entry where the key and the value are the same. Used for
363 * interning values in a set.
364 *
365 * @implNote Requires a {@link ReferencedKeyMap} whose {@code V} type
366 * is a {@code ReferenceKey<K>}. Otherwise, a {@link ClassCastException} will
367 * be thrown.
368 *
369 * @param setMap {@link ReferencedKeyMap} where interning takes place
370 * @param key key to add
371 *
372 * @param <T> type of key
373 *
374 * @return the old key instance if found otherwise the new key instance
375 *
376 * @throws ClassCastException if {@code V} is not {@code EntryKey<T>}
377 */
378 static <T> T intern(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key) {
379 T value = existingKey(setMap, key);
380 if (value != null) {
381 return value;
382 }
383 return internKey(setMap, key);
384 }
385
386 /**
387 * Puts an entry where the key and the value are the same. Used for
388 * interning values in a set.
389 *
390 * @implNote Requires a {@link ReferencedKeyMap} whose {@code V} type
391 * is a {@code ReferenceKey<K>}. Otherwise, a {@link ClassCastException} will
392 * be thrown.
393 *
394 * @param setMap {@link ReferencedKeyMap} where interning takes place
395 * @param key key to add
396 * @param interner operation to apply to key before adding to map
397 *
398 * @param <T> type of key
399 *
400 * @return the old key instance if found otherwise the new key instance
401 *
402 * @throws ClassCastException if {@code V} is not {@code EntryKey<T>}
403 *
404 * @implNote This version of intern should not be called during phase1
405 * using a lambda. Use an UnaryOperator instance instead.
406 */
407 static <T> T intern(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key, UnaryOperator<T> interner) {
408 T value = existingKey(setMap, key);
409 if (value != null) {
410 return value;
411 }
412 key = interner.apply(key);
413 return internKey(setMap, key);
414 }
415
416 /**
417 * Check if the key already exists in the map.
418 *
419 * @param setMap {@link ReferencedKeyMap} where interning takes place
420 * @param key key to test
421 *
422 * @param <T> type of key
423 *
424 * @return key if found otherwise null
425 */
426 private static <T> T existingKey(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key) {
427 setMap.removeStaleReferences();
428 ReferenceKey<T> entryKey = setMap.get(setMap.lookupKey(key));
429 return entryKey != null ? entryKey.get() : null;
430 }
431
432 /**
433 * Attempt to add key to map.
434 *
435 * @param setMap {@link ReferencedKeyMap} where interning takes place
436 * @param key key to add
437 *
438 * @param <T> type of key
439 *
440 * @return the old key instance if found otherwise the new key instance
441 */
442 private static <T> T internKey(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key) {
443 ReferenceKey<T> entryKey = setMap.entryKey(key);
444 T interned;
445 do {
446 setMap.removeStaleReferences();
447 ReferenceKey<T> existing = setMap.map.putIfAbsent(entryKey, entryKey);
448 if (existing == null) {
449 return key;
450 } else {
451 // If {@code putIfAbsent} returns non-null then was actually a
452 // {@code replace} and older key was used. In that case the new
453 // key was not used and the reference marked stale.
454 interned = existing.get();
455 if (interned != null) {
456 entryKey.unused();
457 }
458 }
459 } while (interned == null);
460 return interned;
461 }
462
463
464 /**
465 * Attempt to add key to map if absent.
466 *
467 * @param setMap {@link ReferencedKeyMap} where interning takes place
468 * @param key key to add
469 *
470 * @param <T> type of key
471 *
472 * @return true if the key was added
473 */
474 static <T> boolean internAddKey(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key) {
475 ReferenceKey<T> entryKey = setMap.entryKey(key);
476 setMap.removeStaleReferences();
477 ReferenceKey<T> existing = setMap.map.putIfAbsent(entryKey, entryKey);
478 if (existing == null) {
479 return true;
480 } else {
481 // If {@code putIfAbsent} returns non-null then was actually a
482 // {@code replace} and older key was used. In that case the new
483 // key was not used and the reference marked stale.
484 entryKey.unused();
485 return false;
486 }
487 }
488
489 }