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
017package com.google.common.util.concurrent;
018
019import static com.google.common.base.Preconditions.checkNotNull;
020
021import com.google.common.annotations.Beta;
022import com.google.common.annotations.GwtCompatible;
023import com.google.common.base.Function;
024import com.google.common.collect.Maps;
025
026import java.util.Collections;
027import java.util.Map;
028import java.util.concurrent.ConcurrentHashMap;
029import 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
056public 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}