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