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