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