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