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