1 /*
2  *  Licensed to the Apache Software Foundation (ASF) under one or more
3  *  contributor license agreements.  See the NOTICE file distributed with
4  *  this work for additional information regarding copyright ownership.
5  *  The ASF licenses this file to You under the Apache License, Version 2.0
6  *  (the "License"); you may not use this file except in compliance with
7  *  the License.  You may obtain a copy of the License at
8  *
9  *      http://www.apache.org/licenses/LICENSE-2.0
10  *
11  *  Unless required by applicable law or agreed to in writing, software
12  *  distributed under the License is distributed on an "AS IS" BASIS,
13  *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14  *  See the License for the specific language governing permissions and
15  *  limitations under the License.
16  */
17 package org.apache.tomcat.util.collections;
18
19 import java.lang.ref.Reference;
20 import java.lang.ref.ReferenceQueue;
21 import java.lang.ref.WeakReference;
22 import java.util.AbstractMap;
23 import java.util.AbstractSet;
24 import java.util.Collection;
25 import java.util.Iterator;
26 import java.util.Map;
27 import java.util.Objects;
28 import java.util.Set;
29 import java.util.concurrent.ConcurrentHashMap;
30 import java.util.concurrent.ConcurrentMap;
31
32 /**
33  * Concurrent hash map that holds its keys via weak references. Unlike
34  * <code>WeakHashMap</code> this class does not handle dead keys during common
35  * access operations, but expects you to call its {@link #maintain()} method
36  * periodically. Both keys and values are expected to be not-<code>null</code>.
37  *
38  * @param <K> The type of keys used with the Map instance
39  * @param <V> The type of values used with the Map instance
40  */
41 public class ManagedConcurrentWeakHashMap<K, V> extends AbstractMap<K, V> implements
42         ConcurrentMap<K, V> {
43
44     private final ConcurrentMap<Key, V> map = new ConcurrentHashMap<>();
45     private final ReferenceQueue<Object> queue = new ReferenceQueue<>();
46
47     /**
48      * Method, that has to be invoked periodically to clean dead keys from the
49      * map.
50      */
51     public void maintain() {
52         Key key;
53         while ((key = (Key) queue.poll()) != null) {
54             if (key.isDead()) {
55                 // No need to lookup if the key is not in the map
56                 continue;
57             }
58             key.ackDeath();
59             map.remove(key);
60         }
61     }
62
63     private static class Key extends WeakReference<Object> {
64         private final int hash;
65         private boolean dead;
66
67         public Key(Object key, ReferenceQueue<Object> queue) {
68             super(key, queue);
69             hash = key.hashCode();
70         }
71
72         @Override
73         public int hashCode() {
74             return hash;
75         }
76
77         @Override
78         public boolean equals(Object obj) {
79             if (this == obj) {
80                 return true;
81             }
82             if (dead) {
83                 // Post-mortem cleanup looks for this specific Reference
84                 // instance
85                 return false;
86             }
87             if (!(obj instanceof Reference<?>)) {
88                 return false;
89             }
90             Object oA = get();
91             Object oB = ((Reference<?>) obj).get();
92             if (oA == oB) {
93                 return true;
94             }
95             if (oA == null || oB == null) {
96                 return false;
97             }
98             return oA.equals(oB);
99         }
100
101         public void ackDeath() {
102             this.dead = true;
103         }
104
105         public boolean isDead() {
106             return dead;
107         }
108     }
109
110     /**
111      * Creates Key instance to be used to store values in the map. It is
112      * registered with the ReferenceQueue.
113      */
114     private Key createStoreKey(Object key) {
115         return new Key(key, queue);
116     }
117
118     /**
119      * Creates Key instance to be used only to lookup values in the map. It is
120      * not registered with the ReferenceQueue.
121      */
122     private Key createLookupKey(Object key) {
123         return new Key(key, null);
124     }
125
126     @Override
127     public int size() {
128         return map.size();
129     }
130
131     @Override
132     public boolean isEmpty() {
133         return map.isEmpty();
134     }
135
136     @Override
137     public boolean containsValue(Object value) {
138         if (value == null) {
139             return false;
140         }
141         return map.containsValue(value);
142     }
143
144     @Override
145     public boolean containsKey(Object key) {
146         if (key == null) {
147             return false;
148         }
149         return map.containsKey(createLookupKey(key));
150     }
151
152     @Override
153     public V get(Object key) {
154         if (key == null) {
155             return null;
156         }
157         return map.get(createLookupKey(key));
158     }
159
160     @Override
161     public V put(K key, V value) {
162         Objects.requireNonNull(value);
163         return map.put(createStoreKey(key), value);
164     }
165
166     @Override
167     public V remove(Object key) {
168         return map.remove(createLookupKey(key));
169     }
170
171     @Override
172     public void clear() {
173         map.clear();
174         // maintain() clears the queue, though it is not 100% reliable,
175         // as queue is populated by GC asynchronously
176         maintain();
177     }
178
179     @Override
180     public V putIfAbsent(K key, V value) {
181         Objects.requireNonNull(value);
182         Key storeKey = createStoreKey(key);
183         V oldValue = map.putIfAbsent(storeKey, value);
184         if (oldValue != null) { // ack that key has not been stored
185             storeKey.ackDeath();
186         }
187         return oldValue;
188     }
189
190     @Override
191     public boolean remove(Object key, Object value) {
192         if (value == null) {
193             return false;
194         }
195         return map.remove(createLookupKey(key), value);
196     }
197
198     @Override
199     public boolean replace(K key, V oldValue, V newValue) {
200         Objects.requireNonNull(newValue);
201         return map.replace(createLookupKey(key), oldValue, newValue);
202     }
203
204     @Override
205     public V replace(K key, V value) {
206         Objects.requireNonNull(value);
207         return map.replace(createLookupKey(key), value);
208     }
209
210     @Override
211     public Collection<V> values() {
212         return map.values();
213     }
214
215     @Override
216     public Set<Map.Entry<K, V>> entrySet() {
217         return new AbstractSet<Map.Entry<K, V>>() {
218             @Override
219             public boolean isEmpty() {
220                 return map.isEmpty();
221             }
222
223             @Override
224             public int size() {
225                 return map.size();
226             }
227
228             @Override
229             public Iterator<Map.Entry<K, V>> iterator() {
230                 return new Iterator<Map.Entry<K, V>>() {
231                     private final Iterator<Map.Entry<Key, V>> it = map
232                             .entrySet().iterator();
233
234                     @Override
235                     public boolean hasNext() {
236                         return it.hasNext();
237                     }
238
239                     @Override
240                     public Map.Entry<K, V> next() {
241                         return new Map.Entry<K, V>() {
242                             private final Map.Entry<Key, V> en = it.next();
243
244                             @SuppressWarnings("unchecked")
245                             @Override
246                             public K getKey() {
247                                 return (K) en.getKey().get();
248                             }
249
250                             @Override
251                             public V getValue() {
252                                 return en.getValue();
253                             }
254
255                             @Override
256                             public V setValue(V value) {
257                                 Objects.requireNonNull(value);
258                                 return en.setValue(value);
259                             }
260                         };
261                     }
262
263                     @Override
264                     public void remove() {
265                         it.remove();
266                     }
267                 };
268             }
269         };
270     }
271 }
272