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.errorprone.annotations.CanIgnoreReturnValue;
024import java.io.Serializable;
025import java.util.Collections;
026import java.util.Map;
027import java.util.concurrent.ConcurrentHashMap;
028import java.util.concurrent.atomic.AtomicBoolean;
029import java.util.concurrent.atomic.AtomicLong;
030import java.util.function.LongBinaryOperator;
031import java.util.function.LongUnaryOperator;
032import javax.annotation.CheckForNull;
033
034/**
035 * A map containing {@code long} values that can be atomically updated. While writes to a
036 * traditional {@code Map} rely on {@code put(K, V)}, the typical mechanism for writing to this map
037 * is {@code addAndGet(K, long)}, which adds a {@code long} to the value currently associated with
038 * {@code K}. If a key has not yet been associated with a value, its implicit value is zero.
039 *
040 * <p>Most methods in this class treat absent values and zero values identically, as individually
041 * documented. Exceptions to this are {@link #containsKey}, {@link #size}, {@link #isEmpty}, {@link
042 * #asMap}, and {@link #toString}.
043 *
044 * <p>Instances of this class may be used by multiple threads concurrently. All operations are
045 * atomic unless otherwise noted.
046 *
047 * <p><b>Note:</b> If your values are always positive and less than 2^31, you may wish to use a
048 * {@link com.google.common.collect.Multiset} such as {@link
049 * com.google.common.collect.ConcurrentHashMultiset} instead.
050 *
051 * <p><b>Warning:</b> Unlike {@code Multiset}, entries whose values are zero are not automatically
052 * removed from the map. Instead they must be removed manually with {@link #removeAllZeros}.
053 *
054 * @author Charles Fry
055 * @since 11.0
056 */
057@GwtCompatible
058@ElementTypesAreNonnullByDefault
059public final class AtomicLongMap<K> implements Serializable {
060  private final ConcurrentHashMap<K, Long> map;
061
062  private AtomicLongMap(ConcurrentHashMap<K, Long> map) {
063    this.map = checkNotNull(map);
064  }
065
066  /** Creates an {@code AtomicLongMap}. */
067  public static <K> AtomicLongMap<K> create() {
068    return new AtomicLongMap<K>(new ConcurrentHashMap<>());
069  }
070
071  /** Creates an {@code AtomicLongMap} with the same mappings as the specified {@code Map}. */
072  public static <K> AtomicLongMap<K> create(Map<? extends K, ? extends Long> m) {
073    AtomicLongMap<K> result = create();
074    result.putAll(m);
075    return result;
076  }
077
078  /**
079   * Returns the value associated with {@code key}, or zero if there is no value associated with
080   * {@code key}.
081   */
082  public long get(K key) {
083    return map.getOrDefault(key, 0L);
084  }
085
086  /**
087   * Increments by one the value currently associated with {@code key}, and returns the new value.
088   */
089  @CanIgnoreReturnValue
090  public long incrementAndGet(K key) {
091    return addAndGet(key, 1);
092  }
093
094  /**
095   * Decrements by one the value currently associated with {@code key}, and returns the new value.
096   */
097  @CanIgnoreReturnValue
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  @CanIgnoreReturnValue
107  public long addAndGet(K key, long delta) {
108    return accumulateAndGet(key, delta, Long::sum);
109  }
110
111  /**
112   * Increments by one the value currently associated with {@code key}, and returns the old value.
113   */
114  @CanIgnoreReturnValue
115  public long getAndIncrement(K key) {
116    return getAndAdd(key, 1);
117  }
118
119  /**
120   * Decrements by one the value currently associated with {@code key}, and returns the old value.
121   */
122  @CanIgnoreReturnValue
123  public long getAndDecrement(K key) {
124    return getAndAdd(key, -1);
125  }
126
127  /**
128   * Adds {@code delta} to the value currently associated with {@code key}, and returns the old
129   * value.
130   */
131  @CanIgnoreReturnValue
132  public long getAndAdd(K key, long delta) {
133    return getAndAccumulate(key, delta, Long::sum);
134  }
135
136  /**
137   * Updates the value currently associated with {@code key} with the specified function, and
138   * returns the new value. If there is not currently a value associated with {@code key}, the
139   * function is applied to {@code 0L}.
140   *
141   * @since 21.0
142   */
143  @CanIgnoreReturnValue
144  public long updateAndGet(K key, LongUnaryOperator updaterFunction) {
145    checkNotNull(updaterFunction);
146    return map.compute(
147        key, (k, value) -> updaterFunction.applyAsLong((value == null) ? 0L : value.longValue()));
148  }
149
150  /**
151   * Updates the value currently associated with {@code key} with the specified function, and
152   * returns the old value. If there is not currently a value associated with {@code key}, the
153   * function is applied to {@code 0L}.
154   *
155   * @since 21.0
156   */
157  @CanIgnoreReturnValue
158  public long getAndUpdate(K key, LongUnaryOperator updaterFunction) {
159    checkNotNull(updaterFunction);
160    AtomicLong holder = new AtomicLong();
161    map.compute(
162        key,
163        (k, value) -> {
164          long oldValue = (value == null) ? 0L : value.longValue();
165          holder.set(oldValue);
166          return updaterFunction.applyAsLong(oldValue);
167        });
168    return holder.get();
169  }
170
171  /**
172   * Updates the value currently associated with {@code key} by combining it with {@code x} via the
173   * specified accumulator function, returning the new value. The previous value associated with
174   * {@code key} (or zero, if there is none) is passed as the first argument to {@code
175   * accumulatorFunction}, and {@code x} is passed as the second argument.
176   *
177   * @since 21.0
178   */
179  @CanIgnoreReturnValue
180  public long accumulateAndGet(K key, long x, LongBinaryOperator accumulatorFunction) {
181    checkNotNull(accumulatorFunction);
182    return updateAndGet(key, oldValue -> accumulatorFunction.applyAsLong(oldValue, x));
183  }
184
185  /**
186   * Updates the value currently associated with {@code key} by combining it with {@code x} via the
187   * specified accumulator function, returning the old value. The previous value associated with
188   * {@code key} (or zero, if there is none) is passed as the first argument to {@code
189   * accumulatorFunction}, and {@code x} is passed as the second argument.
190   *
191   * @since 21.0
192   */
193  @CanIgnoreReturnValue
194  public long getAndAccumulate(K key, long x, LongBinaryOperator accumulatorFunction) {
195    checkNotNull(accumulatorFunction);
196    return getAndUpdate(key, oldValue -> accumulatorFunction.applyAsLong(oldValue, x));
197  }
198
199  /**
200   * Associates {@code newValue} with {@code key} in this map, and returns the value previously
201   * associated with {@code key}, or zero if there was no such value.
202   */
203  @CanIgnoreReturnValue
204  public long put(K key, long newValue) {
205    return getAndUpdate(key, x -> newValue);
206  }
207
208  /**
209   * Copies all of the mappings from the specified map to this map. The effect of this call is
210   * equivalent to that of calling {@code put(k, v)} on this map once for each mapping from key
211   * {@code k} to value {@code v} in the specified map. The behavior of this operation is undefined
212   * if the specified map is modified while the operation is in progress.
213   */
214  public void putAll(Map<? extends K, ? extends Long> m) {
215    m.forEach(this::put);
216  }
217
218  /**
219   * Removes and returns the value associated with {@code key}. If {@code key} is not in the map,
220   * this method has no effect and returns zero.
221   */
222  @CanIgnoreReturnValue
223  public long remove(K key) {
224    Long result = map.remove(key);
225    return (result == null) ? 0L : result.longValue();
226  }
227
228  /**
229   * If {@code (key, value)} is currently in the map, this method removes it and returns true;
230   * otherwise, this method returns false.
231   */
232  boolean remove(K key, long value) {
233    return map.remove(key, value);
234  }
235
236  /**
237   * Atomically remove {@code key} from the map iff its associated value is 0.
238   *
239   * @since 20.0
240   */
241  @Beta
242  @CanIgnoreReturnValue
243  public boolean removeIfZero(K key) {
244    return remove(key, 0);
245  }
246
247  /**
248   * Removes all mappings from this map whose values are zero.
249   *
250   * <p>This method is not atomic: the map may be visible in intermediate states, where some of the
251   * zero values have been removed and others have not.
252   */
253  public void removeAllZeros() {
254    map.values().removeIf(x -> x == 0);
255  }
256
257  /**
258   * Returns the sum of all values in this map.
259   *
260   * <p>This method is not atomic: the sum may or may not include other concurrent operations.
261   */
262  public long sum() {
263    return map.values().stream().mapToLong(Long::longValue).sum();
264  }
265
266  @CheckForNull private transient Map<K, Long> asMap;
267
268  /** Returns a live, read-only view of the map backing this {@code AtomicLongMap}. */
269  public Map<K, Long> asMap() {
270    Map<K, Long> result = asMap;
271    return (result == null) ? asMap = createAsMap() : result;
272  }
273
274  private Map<K, Long> createAsMap() {
275    return Collections.unmodifiableMap(map);
276  }
277
278  /** Returns true if this map contains a mapping for the specified key. */
279  public boolean containsKey(Object key) {
280    return map.containsKey(key);
281  }
282
283  /**
284   * Returns the number of key-value mappings in this map. If the map contains more than {@code
285   * Integer.MAX_VALUE} elements, returns {@code Integer.MAX_VALUE}.
286   */
287  public int size() {
288    return map.size();
289  }
290
291  /** Returns {@code true} if this map contains no key-value mappings. */
292  public boolean isEmpty() {
293    return map.isEmpty();
294  }
295
296  /**
297   * Removes all of the mappings from this map. The map will be empty after this call returns.
298   *
299   * <p>This method is not atomic: the map may not be empty after returning if there were concurrent
300   * writes.
301   */
302  public void clear() {
303    map.clear();
304  }
305
306  @Override
307  public String toString() {
308    return map.toString();
309  }
310
311  /**
312   * If {@code key} is not already associated with a value or if {@code key} is associated with
313   * zero, associate it with {@code newValue}. Returns the previous value associated with {@code
314   * key}, or zero if there was no mapping for {@code key}.
315   */
316  long putIfAbsent(K key, long newValue) {
317    AtomicBoolean noValue = new AtomicBoolean(false);
318    Long result =
319        map.compute(
320            key,
321            (k, oldValue) -> {
322              if (oldValue == null || oldValue == 0) {
323                noValue.set(true);
324                return newValue;
325              } else {
326                return oldValue;
327              }
328            });
329    return noValue.get() ? 0L : result.longValue();
330  }
331
332  /**
333   * If {@code (key, expectedOldValue)} is currently in the map, this method replaces {@code
334   * expectedOldValue} with {@code newValue} and returns true; otherwise, this method returns false.
335   *
336   * <p>If {@code expectedOldValue} is zero, this method will succeed if {@code (key, zero)} is
337   * currently in the map, or if {@code key} is not in the map at all.
338   */
339  boolean replace(K key, long expectedOldValue, long newValue) {
340    if (expectedOldValue == 0L) {
341      return putIfAbsent(key, newValue) == 0L;
342    } else {
343      return map.replace(key, expectedOldValue, newValue);
344    }
345  }
346}