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