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  * Keys must be {@linkplain Class#isIdentity() identity objects.}
 52  * <p>
 53  * To create a {@link ReferencedKeyMap} the user must provide a {@link Supplier}
 54  * of the backing map and whether {@link WeakReference} or
 55  * {@link SoftReference} is to be used.
 56  *
 57  * {@snippet :
 58  * // Use HashMap and WeakReference
 59  * Map<Long, String> map = ReferencedKeyMap.create(false, HashMap::new);
 60  * map.put(10_000_000L, "a");
 61  * map.put(10_000_001L, "b");
 62  * map.put(10_000_002L, "c");
 63  * map.put(10_000_003L, "d");
 64  * map.put(10_000_004L, "e");
 65  *
 66  * // Use ConcurrentHashMap and SoftReference
 67  * map = ReferencedKeyMap.create(true, ConcurrentHashMap::new);
 68  * map.put(20_000_000L, "v");
 69  * map.put(20_000_001L, "w");
 70  * map.put(20_000_002L, "x");
 71  * map.put(20_000_003L, "y");
 72  * map.put(20_000_004L, "z");
 73  * }
 74  *
 75  * @implNote Care must be given that the backing map does replacement by
 76  * replacing the value in the map entry instead of deleting the old entry and
 77  * adding a new entry, otherwise replaced entries may end up with a strongly
 78  * referenced key. {@link HashMap} and {@link ConcurrentHashMap} are known
 79  * to be safe.
 80  *
 81  * @param <K> the type of keys maintained by this map
 82  * @param <V> the type of mapped values
 83  *
 84  * @since 21
 85  */
 86 public final class ReferencedKeyMap<K, V> implements Map<K, V> {
 87     /**
 88      * true if {@link SoftReference} keys are to be used,
 89      * {@link WeakReference} otherwise.
 90      */
 91     private final boolean isSoft;
 92 
 93     /**
 94      * Backing {@link Map}.
 95      */
 96     private final Map<ReferenceKey<K>, V> map;
 97 
 98     /**
 99      * {@link ReferenceQueue} for cleaning up entries.
100      */
101     private final ReferenceQueue<K> stale;
102 
103     /**
104      * Private constructor.
105      *
106      * @param isSoft          true if {@link SoftReference} keys are to
107      *                        be used, {@link WeakReference} otherwise.
108      * @param map             backing map
109      * @param stale           {@link ReferenceQueue} for cleaning up entries
110      */
111     private ReferencedKeyMap(boolean isSoft, Map<ReferenceKey<K>, V> map, ReferenceQueue<K> stale) {
112         this.isSoft = isSoft;
113         this.map = map;
114         this.stale = stale;
115     }
116 
117     /**
118      * Create a new {@link ReferencedKeyMap} map.
119      *
120      * @param isSoft          true if {@link SoftReference} keys are to
121      *                        be used, {@link WeakReference} otherwise.
122      * @param supplier        {@link Supplier} of the backing map
123      *
124      * @return a new map with {@link Reference} keys
125      *
126      * @param <K> the type of keys maintained by the new map
127      * @param <V> the type of mapped values
128      */
129     public static <K, V> ReferencedKeyMap<K, V>
130     create(boolean isSoft, Supplier<Map<ReferenceKey<K>, V>> supplier) {
131         return create(isSoft, false, supplier);
132     }
133 
134     /**
135      * Create a new {@link ReferencedKeyMap} map.
136      *
137      * @param isSoft          true if {@link SoftReference} keys are to
138      *                        be used, {@link WeakReference} otherwise.
139      * @param useNativeQueue  true if uses NativeReferenceQueue
140      *                        otherwise use {@link ReferenceQueue}.
141      * @param supplier        {@link Supplier} of the backing map
142      *
143      * @return a new map with {@link Reference} keys
144      *
145      * @param <K> the type of keys maintained by the new map
146      * @param <V> the type of mapped values
147      */
148     public static <K, V> ReferencedKeyMap<K, V>
149     create(boolean isSoft, boolean useNativeQueue, Supplier<Map<ReferenceKey<K>, V>> supplier) {
150         return new ReferencedKeyMap<K, V>(isSoft, supplier.get(),
151                 useNativeQueue ? SharedSecrets.getJavaLangRefAccess().newNativeReferenceQueue()
152                                : new ReferenceQueue<>()
153                 );
154     }
155 
156     /**
157      * {@return a key suitable for a map entry}
158      *
159      * @param key unwrapped key
160      */
161     @SuppressWarnings("unchecked")
162     private ReferenceKey<K> entryKey(Object key) {
163         if (isSoft) {
164             return new SoftReferenceKey<>((K)key, stale);
165         } else {
166             return new WeakReferenceKey<>((K)key, stale);
167         }
168     }
169 
170     /**
171      * {@return a key suitable for lookup}
172      *
173      * @param key unwrapped key
174      */
175     @SuppressWarnings("unchecked")
176     private ReferenceKey<K> lookupKey(Object key) {
177         return new StrongReferenceKey<>((K)key);
178     }
179 
180     @Override
181     public int size() {
182         removeStaleReferences();
183         return map.size();
184     }
185 
186     @Override
187     public boolean isEmpty() {
188         removeStaleReferences();
189         return map.isEmpty();
190     }
191 
192     @Override
193     public boolean containsKey(Object key) {
194         Objects.requireNonNull(key, "key must not be null");
195         removeStaleReferences();
196         return map.containsKey(lookupKey(key));
197     }
198 
199     @Override
200     public boolean containsValue(Object value) {
201         Objects.requireNonNull(value, "value must not be null");
202         removeStaleReferences();
203         return map.containsValue(value);
204     }
205 
206     @Override
207     public V get(Object key) {
208         Objects.requireNonNull(key, "key must not be null");
209         removeStaleReferences();
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, get(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 + "=" + get(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     /**
342      * Puts an entry where the key and the value are the same. Used for
343      * interning values in a set.
344      *
345      * @implNote Requires a {@link ReferencedKeyMap} whose {@code V} type
346      * is a {@code ReferenceKey<K>}. Otherwise, a {@link ClassCastException} will
347      * be thrown.
348      *
349      * @param setMap    {@link ReferencedKeyMap} where interning takes place
350      * @param key       key to add
351      *
352      * @param <T> type of key
353      *
354      * @return the old key instance if found otherwise the new key instance
355      *
356      * @throws ClassCastException if {@code V} is not {@code EntryKey<T>}
357      */
358     static <T> T intern(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key) {
359         T value = existingKey(setMap, key);
360         if (value != null) {
361             return value;
362         }
363         return internKey(setMap, key);
364     }
365 
366     /**
367      * Puts an entry where the key and the value are the same. Used for
368      * interning values in a set.
369      *
370      * @implNote Requires a {@link ReferencedKeyMap} whose {@code V} type
371      * is a {@code ReferenceKey<K>}. Otherwise, a {@link ClassCastException} will
372      * be thrown.
373      *
374      * @param setMap    {@link ReferencedKeyMap} where interning takes place
375      * @param key       key to add
376      * @param interner  operation to apply to key before adding to map
377      *
378      * @param <T> type of key
379      *
380      * @return the old key instance if found otherwise the new key instance
381      *
382      * @throws ClassCastException if {@code V} is not {@code EntryKey<T>}
383      *
384      * @implNote This version of intern should not be called during phase1
385      * using a lambda. Use an UnaryOperator instance instead.
386      */
387     static <T> T intern(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key, UnaryOperator<T> interner) {
388         T value = existingKey(setMap, key);
389         if (value != null) {
390             return value;
391         }
392         key = interner.apply(key);
393         return internKey(setMap, key);
394     }
395 
396     /**
397      * Check if the key already exists in the map.
398      *
399      * @param setMap    {@link ReferencedKeyMap} where interning takes place
400      * @param key       key to test
401      *
402      * @param <T> type of key
403      *
404      * @return key if found otherwise null
405      */
406     private static <T> T existingKey(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key) {
407         setMap.removeStaleReferences();
408         ReferenceKey<T> entryKey = setMap.get(setMap.lookupKey(key));
409         return entryKey != null ? entryKey.get() : null;
410     }
411 
412     /**
413      * Attempt to add key to map.
414      *
415      * @param setMap    {@link ReferencedKeyMap} where interning takes place
416      * @param key       key to add
417      *
418      * @param <T> type of key
419      *
420      * @return the old key instance if found otherwise the new key instance
421      */
422     private static <T> T internKey(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key) {
423         ReferenceKey<T> entryKey = setMap.entryKey(key);
424         T interned;
425         do {
426             setMap.removeStaleReferences();
427             ReferenceKey<T> existing = setMap.map.putIfAbsent(entryKey, entryKey);
428             if (existing == null) {
429                 return key;
430             } else {
431                 // If {@code putIfAbsent} returns non-null then was actually a
432                 // {@code replace} and older key was used. In that case the new
433                 // key was not used and the reference marked stale.
434                 interned = existing.get();
435                 if (interned != null) {
436                     entryKey.unused();
437                 }
438             }
439         } while (interned == null);
440         return interned;
441     }
442 
443 
444     /**
445      * Attempt to add key to map if absent.
446      *
447      * @param setMap    {@link ReferencedKeyMap} where interning takes place
448      * @param key       key to add
449      *
450      * @param <T> type of key
451      *
452      * @return true if the key was added
453      */
454     static <T> boolean internAddKey(ReferencedKeyMap<T, ReferenceKey<T>> setMap, T key) {
455         ReferenceKey<T> entryKey = setMap.entryKey(key);
456         setMap.removeStaleReferences();
457         ReferenceKey<T> existing = setMap.map.putIfAbsent(entryKey, entryKey);
458         if (existing == null) {
459             return true;
460         } else {
461             // If {@code putIfAbsent} returns non-null then was actually a
462             // {@code replace} and older key was used. In that case the new
463             // key was not used and the reference marked stale.
464             entryKey.unused();
465             return false;
466         }
467      }
468 
469 }