001 // Copyright 2011 Google Inc. All Rights Reserved.
002
003 package com.google.common.util.concurrent;
004
005 import static com.google.common.base.Preconditions.checkNotNull;
006
007 import com.google.common.annotations.Beta;
008 import com.google.common.base.Function;
009 import com.google.common.collect.Maps;
010
011 import java.util.Collections;
012 import java.util.Map;
013 import java.util.concurrent.ConcurrentHashMap;
014 import java.util.concurrent.atomic.AtomicLong;
015
016 /**
017 * A map containing {@code long} values that can be atomically updated. While writes to a
018 * traditional {@code Map} rely on {@code put(K, V)}, the typical mechanism for writing to this map
019 * is {@code addAndGet(K, long)}, which adds a {@code long} to the value currently associated with
020 * {@code K}. If a key has not yet been associated with a value, its implicit value is zero.
021 *
022 * <p>Most methods in this class treat absent values and zero values identically, as individually
023 * documented. Exceptions to this are {@link #containsKey}, {@link #size}, {@link #isEmpty},
024 * {@link #asMap}, and {@link #toString}.
025 *
026 * <p>Instances of this class may be used by multiple threads concurrently. All operations are
027 * atomic unless otherwise noted.
028 *
029 * <p><b>Note:</b> If your values are always positive and less than 2^31, you may wish to use a
030 * {@link com.google.common.collect.Multiset} such as
031 * {@link com.google.common.collect.ConcurrentHashMultiset} instead.
032 *
033 * <b>Warning:</b> Unlike {@code Multiset}, entries whose values are zero are not automatically
034 * removed from the map. Instead they must be removed manually with {@link #removeAllZeros}.
035 *
036 * @author Charles Fry
037 * @since 11.0
038 */
039 @Beta
040 public final class AtomicLongMap<K> {
041 private final ConcurrentHashMap<K, AtomicLong> map;
042
043 private AtomicLongMap(ConcurrentHashMap<K, AtomicLong> map) {
044 this.map = checkNotNull(map);
045 }
046
047 /**
048 * Creates an {@code AtomicLongMap}.
049 */
050 public static <K> AtomicLongMap<K> create() {
051 return new AtomicLongMap<K>(new ConcurrentHashMap<K, AtomicLong>());
052 }
053
054 /**
055 * Creates an {@code AtomicLongMap} with the same mappings as the specified {@code Map}.
056 */
057 public static <K> AtomicLongMap<K> create(Map<? extends K, ? extends Long> m) {
058 AtomicLongMap<K> result = create();
059 result.putAll(m);
060 return result;
061 }
062
063 /**
064 * Returns the value associated with {@code key}, or zero if there is no value associated with
065 * {@code key}.
066 */
067 public long get(K key) {
068 AtomicLong atomic = map.get(key);
069 return atomic == null ? 0L : atomic.get();
070 }
071
072 /**
073 * Increments by one the value currently associated with {@code key}, and returns the new value.
074 */
075 public long incrementAndGet(K key) {
076 return addAndGet(key, 1);
077 }
078
079 /**
080 * Decrements by one the value currently associated with {@code key}, and returns the new value.
081 */
082 public long decrementAndGet(K key) {
083 return addAndGet(key, -1);
084 }
085
086 /**
087 * Adds {@code delta} to the value currently associated with {@code key}, and returns the new
088 * value.
089 */
090 public long addAndGet(K key, long delta) {
091 outer: for (;;) {
092 AtomicLong atomic = map.get(key);
093 if (atomic == null) {
094 atomic = map.putIfAbsent(key, new AtomicLong(delta));
095 if (atomic == null) {
096 return delta;
097 }
098 // atomic is now non-null; fall through
099 }
100
101 for (;;) {
102 long oldValue = atomic.get();
103 if (oldValue == 0L) {
104 // don't compareAndSet a zero
105 if (map.replace(key, atomic, new AtomicLong(delta))) {
106 return delta;
107 }
108 // atomic replaced
109 continue outer;
110 }
111
112 long newValue = oldValue + delta;
113 if (atomic.compareAndSet(oldValue, newValue)) {
114 return newValue;
115 }
116 // value changed
117 }
118 }
119 }
120
121 /**
122 * Increments by one the value currently associated with {@code key}, and returns the old value.
123 */
124 public long getAndIncrement(K key) {
125 return getAndAdd(key, 1);
126 }
127
128 /**
129 * Decrements by one the value currently associated with {@code key}, and returns the old value.
130 */
131 public long getAndDecrement(K key) {
132 return getAndAdd(key, -1);
133 }
134
135 /**
136 * Adds {@code delta} to the value currently associated with {@code key}, and returns the old
137 * value.
138 */
139 public long getAndAdd(K key, long delta) {
140 outer: for (;;) {
141 AtomicLong atomic = map.get(key);
142 if (atomic == null) {
143 atomic = map.putIfAbsent(key, new AtomicLong(delta));
144 if (atomic == null) {
145 return 0L;
146 }
147 // atomic is now non-null; fall through
148 }
149
150 for (;;) {
151 long oldValue = atomic.get();
152 if (oldValue == 0L) {
153 // don't compareAndSet a zero
154 if (map.replace(key, atomic, new AtomicLong(delta))) {
155 return 0L;
156 }
157 // atomic replaced
158 continue outer;
159 }
160
161 long newValue = oldValue + delta;
162 if (atomic.compareAndSet(oldValue, newValue)) {
163 return oldValue;
164 }
165 // value changed
166 }
167 }
168 }
169
170 /**
171 * Associates {@code newValue} with {@code key} in this map, and returns the value previously
172 * associated with {@code key}, or zero if there was no such value.
173 */
174 public long put(K key, long newValue) {
175 outer: for (;;) {
176 AtomicLong atomic = map.get(key);
177 if (atomic == null) {
178 atomic = map.putIfAbsent(key, new AtomicLong(newValue));
179 if (atomic == null) {
180 return 0L;
181 }
182 // atomic is now non-null; fall through
183 }
184
185 for (;;) {
186 long oldValue = atomic.get();
187 if (oldValue == 0L) {
188 // don't compareAndSet a zero
189 if (map.replace(key, atomic, new AtomicLong(newValue))) {
190 return 0L;
191 }
192 // atomic replaced
193 continue outer;
194 }
195
196 if (atomic.compareAndSet(oldValue, newValue)) {
197 return oldValue;
198 }
199 // value changed
200 }
201 }
202 }
203
204 /**
205 * Copies all of the mappings from the specified map to this map. The effect of this call is
206 * equivalent to that of calling {@code put(k, v)} on this map once for each mapping from key
207 * {@code k} to value {@code v} in the specified map. The behavior of this operation is undefined
208 * if the specified map is modified while the operation is in progress.
209 */
210 public void putAll(Map<? extends K, ? extends Long> m) {
211 for (Map.Entry<? extends K, ? extends Long> entry : m.entrySet()) {
212 put(entry.getKey(), entry.getValue());
213 }
214 }
215
216 /**
217 * Removes and returns the value associated with {@code key}. If {@code key} is not
218 * in the map, this method has no effect and returns zero.
219 */
220 public long remove(K key) {
221 AtomicLong atomic = map.get(key);
222 if (atomic == null) {
223 return 0L;
224 }
225
226 for (;;) {
227 long oldValue = atomic.get();
228 if (oldValue == 0L || atomic.compareAndSet(oldValue, 0L)) {
229 // only remove after setting to zero, to avoid concurrent updates
230 map.remove(key, atomic);
231 // succeed even if the remove fails, since the value was already adjusted
232 return oldValue;
233 }
234 }
235 }
236
237 /**
238 * Removes all mappings from this map whose values are zero.
239 *
240 * <p>This method is not atomic: the map may be visible in intermediate states, where some
241 * of the zero values have been removed and others have not.
242 */
243 public void removeAllZeros() {
244 for (K key : map.keySet()) {
245 AtomicLong atomic = map.get(key);
246 if (atomic != null && atomic.get() == 0L) {
247 map.remove(key, atomic);
248 }
249 }
250 }
251
252 /**
253 * Returns the sum of all values in this map.
254 *
255 * <p>This method is not atomic: the sum may or may not include other concurrent operations.
256 */
257 public long sum() {
258 long sum = 0L;
259 for (AtomicLong value : map.values()) {
260 sum = sum + value.get();
261 }
262 return sum;
263 }
264
265 private transient Map<K, Long> asMap;
266
267 /**
268 * Returns a live, read-only view of the map backing this {@code AtomicLongMap}.
269 */
270 public Map<K, Long> asMap() {
271 Map<K, Long> result = asMap;
272 return (result == null) ? asMap = createAsMap() : result;
273 }
274
275 private Map<K, Long> createAsMap() {
276 return Collections.unmodifiableMap(
277 Maps.transformValues(map, new Function<AtomicLong, Long>() {
278 @Override
279 public Long apply(AtomicLong atomic) {
280 return atomic.get();
281 }
282 }));
283 }
284
285 /**
286 * Returns true if this map contains a mapping for the specified key.
287 */
288 public boolean containsKey(Object key) {
289 return map.containsKey(key);
290 }
291
292 /**
293 * Returns the number of key-value mappings in this map. If the map contains more than
294 * {@code Integer.MAX_VALUE} elements, returns {@code Integer.MAX_VALUE}.
295 */
296 public int size() {
297 return map.size();
298 }
299
300 /**
301 * Returns {@code true} if this map contains no key-value mappings.
302 */
303 public boolean isEmpty() {
304 return map.isEmpty();
305 }
306
307 /**
308 * Removes all of the mappings from this map. The map will be empty after this call returns.
309 *
310 * <p>This method is not atomic: the map may not be empty after returning if there were concurrent
311 * writes.
312 */
313 public void clear() {
314 map.clear();
315 }
316
317 @Override
318 public String toString() {
319 return map.toString();
320 }
321
322 /*
323 * ConcurrentMap operations which we may eventually add.
324 *
325 * The problem with these is that remove(K, long) has to be done in two phases by definition ---
326 * first decrementing to zero, and then removing. putIfAbsent or replace could observe the
327 * intermediate zero-state. Ways we could deal with this are:
328 *
329 * - Don't define any of the ConcurrentMap operations. This is the current state of affairs.
330 *
331 * - Define putIfAbsent and replace as treating zero and absent identically (as currently
332 * implemented below). This is a bit surprising with putIfAbsent, which really becomes
333 * putIfZero.
334 *
335 * - Allow putIfAbsent and replace to distinguish between zero and absent, but don't implement
336 * remove(K, long). Without any two-phase operations it becomes feasible for all remaining
337 * operations to distinguish between zero and absent. If we do this, then perhaps we should add
338 * replace(key, long).
339 *
340 * - Introduce a special-value private static final AtomicLong that would have the meaning of
341 * removal-in-progress, and rework all operations to properly distinguish between zero and
342 * absent.
343 */
344
345 /**
346 * If {@code key} is not already associated with a value or if {@code key} is associated with
347 * zero, associate it with {@code newValue}. Returns the previous value associated with
348 * {@code key}, or zero if there was no mapping for {@code key}.
349 */
350 long putIfAbsent(K key, long newValue) {
351 for (;;) {
352 AtomicLong atomic = map.get(key);
353 if (atomic == null) {
354 atomic = map.putIfAbsent(key, new AtomicLong(newValue));
355 if (atomic == null) {
356 return 0L;
357 }
358 // atomic is now non-null; fall through
359 }
360
361 long oldValue = atomic.get();
362 if (oldValue == 0L) {
363 // don't compareAndSet a zero
364 if (map.replace(key, atomic, new AtomicLong(newValue))) {
365 return 0L;
366 }
367 // atomic replaced
368 continue;
369 }
370
371 return oldValue;
372 }
373 }
374
375 /**
376 * If {@code (key, expectedOldValue)} is currently in the map, this method replaces
377 * {@code expectedOldValue} with {@code newValue} and returns true; otherwise, this method
378 * returns false.
379 *
380 * <p>If {@code expectedOldValue} is zero, this method will succeed if {@code (key, zero)}
381 * is currently in the map, or if {@code key} is not in the map at all.
382 */
383 boolean replace(K key, long expectedOldValue, long newValue) {
384 if (expectedOldValue == 0L) {
385 return putIfAbsent(key, newValue) == 0L;
386 } else {
387 AtomicLong atomic = map.get(key);
388 return (atomic == null) ? false : atomic.compareAndSet(expectedOldValue, newValue);
389 }
390 }
391
392 /**
393 * If {@code (key, value)} is currently in the map, this method removes it and returns
394 * true; otherwise, this method returns false.
395 */
396 boolean remove(K key, long value) {
397 AtomicLong atomic = map.get(key);
398 if (atomic == null) {
399 return false;
400 }
401
402 long oldValue = atomic.get();
403 if (oldValue != value) {
404 return false;
405 }
406
407 if (oldValue == 0L || atomic.compareAndSet(oldValue, 0L)) {
408 // only remove after setting to zero, to avoid concurrent updates
409 map.remove(key, atomic);
410 // succeed even if the remove fails, since the value was already adjusted
411 return true;
412 }
413
414 // value changed
415 return false;
416 }
417
418 }