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