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