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 new ReferencedKeyMap<K, V>(isSoft, supplier.get(), new ReferenceQueue<>()); 146 } 147 148 /** 149 * {@return a key suitable for a map entry} 150 * 151 * @param key unwrapped key 152 */ 153 @SuppressWarnings("unchecked") 154 private ReferenceKey<K> entryKey(Object key) { 155 if (isSoft) { 156 return new SoftReferenceKey<>((K)key, stale); 157 } else { 158 return new WeakReferenceKey<>((K)key, stale); 159 } 160 } 161 162 /** 163 * {@return a key suitable for lookup} 164 * 165 * @param key unwrapped key 166 */ 167 @SuppressWarnings("unchecked") 168 private ReferenceKey<K> lookupKey(Object key) { 169 return new StrongReferenceKey<>((K)key); 170 } 171 172 @Override 173 public int size() { 174 removeStaleReferences(); 175 return map.size(); 176 } 177 178 @Override 179 public boolean isEmpty() { 180 removeStaleReferences(); 181 return map.isEmpty(); 182 } 183 184 @Override 185 public boolean containsKey(Object key) { 186 Objects.requireNonNull(key, "key must not be null"); 187 removeStaleReferences(); 188 return map.containsKey(lookupKey(key)); 189 } 190 191 @Override 192 public boolean containsValue(Object value) { 193 Objects.requireNonNull(value, "value must not be null"); 194 removeStaleReferences(); 195 return map.containsValue(value); 196 } 197 198 @Override 199 public V get(Object key) { 200 removeStaleReferences(); 201 return getNoCheckStale(key); 202 } 203 204 // Internal get(key) without removing stale references that would modify the keyset. 205 // Use when iterating or streaming over the keys to avoid ConcurrentModificationException. 206 private V getNoCheckStale(Object key) { 207 Objects.requireNonNull(key, "key must not be null"); 208 return map.get(lookupKey(key)); 209 } 210 211 @Override 212 public V put(K key, V newValue) { 213 Objects.requireNonNull(key, "key must not be null"); 214 Objects.requireNonNull(newValue, "value must not be null"); 215 removeStaleReferences(); 216 ReferenceKey<K> entryKey = entryKey(key); 217 // If {@code put} returns non-null then was actually a {@code replace} 218 // and older key was used. In that case the new key was not used and the 219 // reference marked stale. 220 V oldValue = map.put(entryKey, newValue); 221 if (oldValue != null) { 222 entryKey.unused(); 223 } 224 return oldValue; 225 } 226 227 @Override 228 public V remove(Object key) { 229 // Rely on gc to clean up old key. 230 return map.remove(lookupKey(key)); 231 } 232 233 @Override 234 public void putAll(Map<? extends K, ? extends V> m) { 235 removeStaleReferences(); 236 for (Entry<? extends K, ? extends V> entry : m.entrySet()) { 237 K key = entry.getKey(); 238 V value = entry.getValue(); 239 put(key, value); 240 } 241 } 242 243 @Override 244 public void clear() { 245 removeStaleReferences(); 246 // Rely on gc to clean up old keys. 247 map.clear(); 248 } 249 250 /** 251 * Common routine for collecting the current set of keys. 252 * 253 * @return {@link Stream} of valid keys (unwrapped) 254 */ 255 private Stream<K> filterKeySet() { 256 return map.keySet() 257 .stream() 258 .map(ReferenceKey::get) 259 .filter(Objects::nonNull); 260 } 261 262 @Override 263 public Set<K> keySet() { 264 removeStaleReferences(); 265 return filterKeySet().collect(Collectors.toSet()); 266 } 267 268 @Override 269 public Collection<V> values() { 270 removeStaleReferences(); 271 return map.values(); 272 } 273 274 @Override 275 public Set<Entry<K, V>> entrySet() { 276 removeStaleReferences(); 277 return filterKeySet() 278 .map(k -> new AbstractMap.SimpleEntry<>(k, getNoCheckStale(k))) 279 .collect(Collectors.toSet()); 280 } 281 282 @Override 283 public V putIfAbsent(K key, V newValue) { 284 removeStaleReferences(); 285 ReferenceKey<K> entryKey = entryKey(key); 286 // If {@code putIfAbsent} returns non-null then was actually a 287 // {@code replace} and older key was used. In that case the new key was 288 // not used and the reference marked stale. 289 V oldValue = map.putIfAbsent(entryKey, newValue); 290 if (oldValue != null) { 291 entryKey.unused(); 292 } 293 return oldValue; 294 } 295 296 @Override 297 public boolean remove(Object key, Object value) { 298 // Rely on gc to clean up old key. 299 return map.remove(lookupKey(key), value); 300 } 301 302 @Override 303 public boolean replace(K key, V oldValue, V newValue) { 304 removeStaleReferences(); 305 // If replace is successful then the older key will be used and the 306 // lookup key will suffice. 307 return map.replace(lookupKey(key), oldValue, newValue); 308 } 309 310 @Override 311 public V replace(K key, V value) { 312 removeStaleReferences(); 313 // If replace is successful then the older key will be used and the 314 // lookup key will suffice. 315 return map.replace(lookupKey(key), value); 316 } 317 318 @Override 319 public String toString() { 320 removeStaleReferences(); 321 return filterKeySet() 322 .map(k -> k + "=" + getNoCheckStale(k)) 323 .collect(Collectors.joining(", ", "{", "}")); 324 } 325 326 /** 327 * Removes enqueued weak references from map. 328 */ 329 public void removeStaleReferences() { 330 while (true) { 331 Object key = stale.poll(); 332 if (key == null) { 333 break; 334 } 335 map.remove(key); 336 } 337 } 338 339 /** 340 * Puts an entry where the key and the value are the same. Used for 341 * interning values in a set. 342 * 343 * @implNote Requires a {@link ReferencedKeyMap} whose {@code V} type 344 * is a {@code ReferenceKey<K>}. Otherwise, a {@link ClassCastException} will 345 * be thrown. 346 * 347 * @param setMap {@link ReferencedKeyMap} where interning takes place 348 * @param key key to add 349 * 350 * @param <T> type of key 351 * 352 * @return the old key instance if found otherwise the new key instance 353 * 354 * @throws ClassCastException if {@code V} is not {@code EntryKey<T>} 355 */ 356 static <T> T intern(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key) { 357 T value = existingKey(setMap, key); 358 if (value != null) { 359 return value; 360 } 361 return internKey(setMap, key); 362 } 363 364 /** 365 * Puts an entry where the key and the value are the same. Used for 366 * interning values in a set. 367 * 368 * @implNote Requires a {@link ReferencedKeyMap} whose {@code V} type 369 * is a {@code ReferenceKey<K>}. Otherwise, a {@link ClassCastException} will 370 * be thrown. 371 * 372 * @param setMap {@link ReferencedKeyMap} where interning takes place 373 * @param key key to add 374 * @param interner operation to apply to key before adding to map 375 * 376 * @param <T> type of key 377 * 378 * @return the old key instance if found otherwise the new key instance 379 * 380 * @throws ClassCastException if {@code V} is not {@code EntryKey<T>} 381 * 382 * @implNote This version of intern should not be called during phase1 383 * using a lambda. Use an UnaryOperator instance instead. 384 */ 385 static <T> T intern(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key, UnaryOperator<T> interner) { 386 T value = existingKey(setMap, key); 387 if (value != null) { 388 return value; 389 } 390 key = interner.apply(key); 391 return internKey(setMap, key); 392 } 393 394 /** 395 * Check if the key already exists in the map. 396 * 397 * @param setMap {@link ReferencedKeyMap} where interning takes place 398 * @param key key to test 399 * 400 * @param <T> type of key 401 * 402 * @return key if found otherwise null 403 */ 404 private static <T> T existingKey(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key) { 405 setMap.removeStaleReferences(); 406 ReferenceKey<T> entryKey = setMap.get(setMap.lookupKey(key)); 407 return entryKey != null ? entryKey.get() : null; 408 } 409 410 /** 411 * Attempt to add key to map. 412 * 413 * @param setMap {@link ReferencedKeyMap} where interning takes place 414 * @param key key to add 415 * 416 * @param <T> type of key 417 * 418 * @return the old key instance if found otherwise the new key instance 419 */ 420 private static <T> T internKey(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key) { 421 ReferenceKey<T> entryKey = setMap.entryKey(key); 422 T interned; 423 do { 424 setMap.removeStaleReferences(); 425 ReferenceKey<T> existing = setMap.map.putIfAbsent(entryKey, entryKey); 426 if (existing == null) { 427 return key; 428 } else { 429 // If {@code putIfAbsent} returns non-null then was actually a 430 // {@code replace} and older key was used. In that case the new 431 // key was not used and the reference marked stale. 432 interned = existing.get(); 433 if (interned != null) { 434 entryKey.unused(); 435 } 436 } 437 } while (interned == null); 438 return interned; 439 } 440 441 442 /** 443 * Attempt to add key to map if absent. 444 * 445 * @param setMap {@link ReferencedKeyMap} where interning takes place 446 * @param key key to add 447 * 448 * @param <T> type of key 449 * 450 * @return true if the key was added 451 */ 452 static <T> boolean internAddKey(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key) { 453 ReferenceKey<T> entryKey = setMap.entryKey(key); 454 setMap.removeStaleReferences(); 455 ReferenceKey<T> existing = setMap.map.putIfAbsent(entryKey, entryKey); 456 if (existing == null) { 457 return true; 458 } else { 459 // If {@code putIfAbsent} returns non-null then was actually a 460 // {@code replace} and older key was used. In that case the new 461 // key was not used and the reference marked stale. 462 entryKey.unused(); 463 return false; 464 } 465 } 466 467 }