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