    001    /*
    002     * Copyright (C) 2007 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.collect;
    018
    019    import static com.google.common.base.Preconditions.checkArgument;
    020    import static com.google.common.base.Preconditions.checkNotNull;
    021    import static com.google.common.base.Preconditions.checkState;
    022    import static com.google.common.collect.Multisets.checkNonnegative;
    023
    024    import com.google.common.annotations.Beta;
    025    import com.google.common.annotations.VisibleForTesting;
    026    import com.google.common.collect.Serialization.FieldSetter;
    027    import com.google.common.math.IntMath;
    028    import com.google.common.primitives.Ints;
    029
    030    import java.io.IOException;
    031    import java.io.ObjectInputStream;
    032    import java.io.ObjectOutputStream;
    033    import java.io.Serializable;
    034    import java.util.Collection;
    035    import java.util.Iterator;
    036    import java.util.List;
    037    import java.util.Map;
    038    import java.util.Set;
    039    import java.util.concurrent.ConcurrentHashMap;
    040    import java.util.concurrent.ConcurrentMap;
    041    import java.util.concurrent.atomic.AtomicInteger;
    042
    043    import javax.annotation.Nullable;
    044
    045    /**
    046     * A multiset that supports concurrent modifications and that provides atomic versions of most
    047     * {@code Multiset} operations (exceptions where noted). Null elements are not supported.
    048     *
    049     * <p>See the Guava User Guide article on <a href=
    050     * "http://code.google.com/p/guava-libraries/wiki/NewCollectionTypesExplained#Multiset">
    051     * {@code Multiset}</a>.
    052     *
    053     * @author Cliff L. Biffle
    054     * @author mike nonemacher
    055     * @since 2.0 (imported from Google Collections Library)
    056     */
    057    public final class ConcurrentHashMultiset<E> extends AbstractMultiset<E> implements Serializable {
    058
    059      /*
    060       * The ConcurrentHashMultiset's atomic operations are implemented primarily in terms of
    061       * AtomicInteger's atomic operations, with some help from ConcurrentMap's atomic operations on
    062       * creation and removal (including automatic removal of zeroes). If the modification of an
    063       * AtomicInteger results in zero, we compareAndSet the value to zero; if that succeeds, we remove
    064       * the entry from the Map. If another operation sees a zero in the map, it knows that the entry is
    065       * about to be removed, so this operation may remove it (often by replacing it with a new
    066       * AtomicInteger).
    067       */
    068
    069      /** The number of occurrences of each element. */
    070      private final transient ConcurrentMap<E, AtomicInteger> countMap;
    071
    072      // This constant allows the deserialization code to set a final field. This holder class
    073      // makes sure it is not initialized unless an instance is deserialized.
    074      private static class FieldSettersHolder {
    075        static final FieldSetter<ConcurrentHashMultiset> COUNT_MAP_FIELD_SETTER =
    076            Serialization.getFieldSetter(ConcurrentHashMultiset.class, "countMap");
    077      }
    078
    079      /**
    080       * Creates a new, empty {@code ConcurrentHashMultiset} using the default
    081       * initial capacity, load factor, and concurrency settings.
    082       */
    083      public static <E> ConcurrentHashMultiset<E> create() {
    084        // TODO(schmoe): provide a way to use this class with other (possibly arbitrary)
    085        // ConcurrentMap implementors. One possibility is to extract most of this class into
    086        // an AbstractConcurrentMapMultiset.
    087        return new ConcurrentHashMultiset<E>(new ConcurrentHashMap<E, AtomicInteger>());
    088      }
    089
    090      /**
    091       * Creates a new {@code ConcurrentHashMultiset} containing the specified elements, using
    092       * the default initial capacity, load factor, and concurrency settings.
    093       *
    094       * <p>This implementation is highly efficient when {@code elements} is itself a {@link Multiset}.
    095       *
    096       * @param elements the elements that the multiset should contain
    097       */
    098      public static <E> ConcurrentHashMultiset<E> create(Iterable<? extends E> elements) {
    099        ConcurrentHashMultiset<E> multiset = ConcurrentHashMultiset.create();
    100        Iterables.addAll(multiset, elements);
    101        return multiset;
    102      }
    103
    104      /**
    105       * Creates a new, empty {@code ConcurrentHashMultiset} using {@code mapMaker}
    106       * to construct the internal backing map.
    107       *
    108       * <p>If this {@link MapMaker} is configured to use entry eviction of any kind, this eviction
    109       * applies to all occurrences of a given element as a single unit. However, most updates to the
    110       * multiset do not count as map updates at all, since we're usually just mutating the value
    111       * stored in the map, so {@link MapMaker#expireAfterAccess} makes sense (evict the entry that
    112       * was queried or updated longest ago), but {@link MapMaker#expireAfterWrite} doesn't, because
    113       * the eviction time is measured from when we saw the first occurrence of the object.
    114       *
    115       * <p>The returned multiset is serializable but any serialization caveats
    116       * given in {@code MapMaker} apply.
    117       *
    118       * <p>Finally, soft/weak values can be used but are not very useful: the values are created
    119       * internally and not exposed externally, so no one else will have a strong reference to the
    120       * values. Weak keys on the other hand can be useful in some scenarios.
    121       *
    122       * @since 7.0
    123       */
    124      @Beta
    125      public static <E> ConcurrentHashMultiset<E> create(
    126          GenericMapMaker<? super E, ? super Number> mapMaker) {
    127        return new ConcurrentHashMultiset<E>(mapMaker.<E, AtomicInteger>makeMap());
    128      }
    129
    130      /**
    131       * Creates an instance using {@code countMap} to store elements and their counts.
    132       *
    133       * <p>This instance will assume ownership of {@code countMap}, and other code
    134       * should not maintain references to the map or modify it in any way.
    135       *
    136       * @param countMap backing map for storing the elements in the multiset and
    137       *     their counts. It must be empty.
    138       * @throws IllegalArgumentException if {@code countMap} is not empty
    139       */
    140      @VisibleForTesting ConcurrentHashMultiset(ConcurrentMap<E, AtomicInteger> countMap) {
    141        checkArgument(countMap.isEmpty());
    142        this.countMap = countMap;
    143      }
    144
    145      // Query Operations
    146
    147      /**
    148       * Returns the number of occurrences of {@code element} in this multiset.
    149       *
    150       * @param element the element to look for
    151       * @return the nonnegative number of occurrences of the element
    152       */
    153      @Override public int count(@Nullable Object element) {
    154        AtomicInteger existingCounter = safeGet(element);
    155        return (existingCounter == null) ? 0 : existingCounter.get();
    156      }
    157
    158      /**
    159       * Depending on the type of the underlying map, map.get may throw NullPointerException or
    160       * ClassCastException, if the object is null or of the wrong type. We usually just want to treat
    161       * those cases as if the element isn't in the map, by catching the exceptions and returning null.
    162       */
    163      private AtomicInteger safeGet(Object element) {
    164        try {
    165          return countMap.get(element);
    166        } catch (NullPointerException e) {
    167          return null;
    168        } catch (ClassCastException e) {
    169          return null;
    170        }
    171      }
    172
    173      /**
    174       * {@inheritDoc}
    175       *
    176       * <p>If the data in the multiset is modified by any other threads during this method,
    177       * it is undefined which (if any) of these modifications will be reflected in the result.
    178       */
    179      @Override public int size() {
    180        long sum = 0L;
    181        for (AtomicInteger value : countMap.values()) {
    182          sum += value.get();
    183        }
    184        return Ints.saturatedCast(sum);
    185      }
    186
    187      /*
    188       * Note: the superclass toArray() methods assume that size() gives a correct
    189       * answer, which ours does not.
    190       */
    191
    192      @Override public Object[] toArray() {
    193        return snapshot().toArray();
    194      }
    195
    196      @Override public <T> T[] toArray(T[] array) {
    197        return snapshot().toArray(array);
    198      }
    199
    200      /*
    201       * We'd love to use 'new ArrayList(this)' or 'list.addAll(this)', but
    202       * either of these would recurse back to us again!
    203       */
    204      private List<E> snapshot() {
    205        List<E> list = Lists.newArrayListWithExpectedSize(size());
    206        for (Multiset.Entry<E> entry : entrySet()) {
    207          E element = entry.getElement();
    208          for (int i = entry.getCount(); i > 0; i--) {
    209            list.add(element);
    210          }
    211        }
    212        return list;
    213      }
    214
    215      // Modification Operations
    216
    217      /**
    218       * Adds a number of occurrences of the specified element to this multiset.
    219       *
    220       * @param element the element to add
    221       * @param occurrences the number of occurrences to add
    222       * @return the previous count of the element before the operation; possibly zero
    223       * @throws IllegalArgumentException if {@code occurrences} is negative, or if
    224       *     the resulting amount would exceed {@link Integer#MAX_VALUE}
    225       */
    226      @Override public int add(E element, int occurrences) {
    227        checkNotNull(element);
    228        if (occurrences == 0) {
    229          return count(element);
    230        }
    231        checkArgument(occurrences > 0, "Invalid occurrences: %s", occurrences);
    232
    233        while (true) {
    234          AtomicInteger existingCounter = safeGet(element);
    235          if (existingCounter == null) {
    236            existingCounter = countMap.putIfAbsent(element, new AtomicInteger(occurrences));
    237            if (existingCounter == null) {
    238              return 0;
    239            }
    240            // existingCounter != null: fall through to operate against the existing AtomicInteger
    241          }
    242
    243          while (true) {
    244            int oldValue = existingCounter.get();
    245            if (oldValue != 0) {
    246              try {
    247                int newValue = IntMath.checkedAdd(oldValue, occurrences);
    248                if (existingCounter.compareAndSet(oldValue, newValue)) {
    249                  // newValue can't == 0, so no need to check & remove
    250                  return oldValue;
    251                }
    252              } catch (ArithmeticException overflow) {
    253                throw new IllegalArgumentException("Overflow adding " + occurrences
    254                    + " occurrences to a count of " + oldValue);
    255              }
    256            } else {
    257              // In the case of a concurrent remove, we might observe a zero value, which means another
    258              // thread is about to remove (element, existingCounter) from the map. Rather than wait,
    259              // we can just do that work here.
    260              AtomicInteger newCounter = new AtomicInteger(occurrences);
    261              if ((countMap.putIfAbsent(element, newCounter) == null)
    262                  || countMap.replace(element, existingCounter, newCounter)) {
    263                return 0;
    264              }
    265              break;
    266            }
    267          }
    268
    269          // If we're still here, there was a race, so just try again.
    270        }
    271      }
    272
    273      /**
    274       * Removes a number of occurrences of the specified element from this multiset. If the multiset
    275       * contains fewer than this number of occurrences to begin with, all occurrences will be removed.
    276       *
    277       * @param element the element whose occurrences should be removed
    278       * @param occurrences the number of occurrences of the element to remove
    279       * @return the count of the element before the operation; possibly zero
    280       * @throws IllegalArgumentException if {@code occurrences} is negative
    281       */
    282      /*
    283       * TODO(cpovirk): remove and removeExactly currently accept null inputs only
    284       * if occurrences == 0. This satisfies both NullPointerTester and
    285       * CollectionRemoveTester.testRemove_nullAllowed, but it's not clear that it's
    286       * a good policy, especially because, in order for the test to pass, the
    287       * parameter must be misleadingly annotated as @Nullable. I suspect that
    288       * we'll want to remove @Nullable, add an eager checkNotNull, and loosen up
    289       * testRemove_nullAllowed.
    290       */
    291      @Override public int remove(@Nullable Object element, int occurrences) {
    292        if (occurrences == 0) {
    293          return count(element);
    294        }
    295        checkArgument(occurrences > 0, "Invalid occurrences: %s", occurrences);
    296
    297        AtomicInteger existingCounter = safeGet(element);
    298        if (existingCounter == null) {
    299          return 0;
    300        }
    301        while (true) {
    302          int oldValue = existingCounter.get();
    303          if (oldValue != 0) {
    304            int newValue = Math.max(0, oldValue - occurrences);
    305            if (existingCounter.compareAndSet(oldValue, newValue)) {
    306              if (newValue == 0) {
    307                // Just CASed to 0; remove the entry to clean up the map. If the removal fails,
    308                // another thread has already replaced it with a new counter, which is fine.
    309                countMap.remove(element, existingCounter);
    310              }
    311              return oldValue;
    312            }
    313          } else {
    314            return 0;
    315          }
    316        }
    317      }
    318
    319      /**
    320       * Removes exactly the specified number of occurrences of {@code element}, or makes no
    321       * change if this is not possible.
    322       *
    323       * <p>This method, in contrast to {@link #remove(Object, int)}, has no effect when the
    324       * element count is smaller than {@code occurrences}.
    325       *
    326       * @param element the element to remove
    327       * @param occurrences the number of occurrences of {@code element} to remove
    328       * @return {@code true} if the removal was possible (including if {@code occurrences} is zero)
    329       */
    330      public boolean removeExactly(@Nullable Object element, int occurrences) {
    331        if (occurrences == 0) {
    332          return true;
    333        }
    334        checkArgument(occurrences > 0, "Invalid occurrences: %s", occurrences);
    335
    336        AtomicInteger existingCounter = safeGet(element);
    337        if (existingCounter == null) {
    338          return false;
    339        }
    340        while (true) {
    341          int oldValue = existingCounter.get();
    342          if (oldValue < occurrences) {
    343            return false;
    344          }
    345          int newValue = oldValue - occurrences;
    346          if (existingCounter.compareAndSet(oldValue, newValue)) {
    347            if (newValue == 0) {
    348              // Just CASed to 0; remove the entry to clean up the map. If the removal fails,
    349              // another thread has already replaced it with a new counter, which is fine.
    350              countMap.remove(element, existingCounter);
    351            }
    352            return true;
    353          }
    354        }
    355      }
    356
    357      /**
    358       * Adds or removes occurrences of {@code element} such that the {@link #count} of the
    359       * element becomes {@code count}.
    360       *
    361       * @return the count of {@code element} in the multiset before this call
    362       * @throws IllegalArgumentException if {@code count} is negative
    363       */
    364      @Override public int setCount(E element, int count) {
    365        checkNotNull(element);
    366        checkNonnegative(count, "count");
    367        while (true) {
    368          AtomicInteger existingCounter = safeGet(element);
    369          if (existingCounter == null) {
    370            if (count == 0) {
    371              return 0;
    372            } else {
    373              existingCounter = countMap.putIfAbsent(element, new AtomicInteger(count));
    374              if (existingCounter == null) {
    375                return 0;
    376              }
    377              // existingCounter != null: fall through
    378            }
    379          }
    380
    381          while (true) {
    382            int oldValue = existingCounter.get();
    383            if (oldValue == 0) {
    384              if (count == 0) {
    385                return 0;
    386              } else {
    387                AtomicInteger newCounter = new AtomicInteger(count);
    388                if ((countMap.putIfAbsent(element, newCounter) == null)
    389                    || countMap.replace(element, existingCounter, newCounter)) {
    390                  return 0;
    391                }
    392              }
    393              break;
    394            } else {
    395              if (existingCounter.compareAndSet(oldValue, count)) {
    396                if (count == 0) {
    397                  // Just CASed to 0; remove the entry to clean up the map. If the removal fails,
    398                  // another thread has already replaced it with a new counter, which is fine.
    399                  countMap.remove(element, existingCounter);
    400                }
    401                return oldValue;
    402              }
    403            }
    404          }
    405        }
    406      }
    407
    408      /**
    409       * Sets the number of occurrences of {@code element} to {@code newCount}, but only if
    410       * the count is currently {@code expectedOldCount}. If {@code element} does not appear
    411       * in the multiset exactly {@code expectedOldCount} times, no changes will be made.
    412       *
    413       * @return {@code true} if the change was successful. This usually indicates
    414       *     that the multiset has been modified, but not always: in the case that
    415       *     {@code expectedOldCount == newCount}, the method will return {@code true} if
    416       *     the condition was met.
    417       * @throws IllegalArgumentException if {@code expectedOldCount} or {@code newCount} is negative
    418       */
    419      @Override public boolean setCount(E element, int expectedOldCount, int newCount) {
    420        checkNotNull(element);
    421        checkNonnegative(expectedOldCount, "oldCount");
    422        checkNonnegative(newCount, "newCount");
    423
    424        AtomicInteger existingCounter = safeGet(element);
    425        if (existingCounter == null) {
    426          if (expectedOldCount != 0) {
    427            return false;
    428          } else if (newCount == 0) {
    429            return true;
    430          } else {
    431            // if our write lost the race, it must have lost to a nonzero value, so we can stop
    432            return countMap.putIfAbsent(element, new AtomicInteger(newCount)) == null;
    433          }
    434        }
    435        int oldValue = existingCounter.get();
    436        if (oldValue == expectedOldCount) {
    437          if (oldValue == 0) {
    438            if (newCount == 0) {
    439              // Just observed a 0; try to remove the entry to clean up the map
    440              countMap.remove(element, existingCounter);
    441              return true;
    442            } else {
    443              AtomicInteger newCounter = new AtomicInteger(newCount);
    444              return (countMap.putIfAbsent(element, newCounter) == null)
    445                  || countMap.replace(element, existingCounter, newCounter);
    446            }
    447          } else {
    448            if (existingCounter.compareAndSet(oldValue, newCount)) {
    449              if (newCount == 0) {
    450                // Just CASed to 0; remove the entry to clean up the map. If the removal fails,
    451                // another thread has already replaced it with a new counter, which is fine.
    452                countMap.remove(element, existingCounter);
    453              }
    454              return true;
    455            }
    456          }
    457        }
    458        return false;
    459      }
    460
    461      // Views
    462
    463      @Override Set<E> createElementSet() {
    464        final Set<E> delegate = countMap.keySet();
    465        return new ForwardingSet<E>() {
    466          @Override protected Set<E> delegate() {
    467            return delegate;
    468          }
    469          @Override public boolean remove(Object object) {
    470            try {
    471              return delegate.remove(object);
    472            } catch (NullPointerException e) {
    473              return false;
    474            } catch (ClassCastException e) {
    475              return false;
    476            }
    477          }
    478          @Override public boolean removeAll(Collection<?> c) {
    479            return standardRemoveAll(c);
    480          }
    481        };
    482      }
    483
    484      private transient EntrySet entrySet;
    485
    486      @Override public Set<Multiset.Entry<E>> entrySet() {
    487        EntrySet result = entrySet;
    488        if (result == null) {
    489          entrySet = result = new EntrySet();
    490        }
    491        return result;
    492      }
    493
    494      @Override int distinctElements() {
    495        return countMap.size();
    496      }
    497
    498      @Override public boolean isEmpty() {
    499        return countMap.isEmpty();
    500      }
    501
    502      @Override Iterator<Entry<E>> entryIterator() {
    503        // AbstractIterator makes this fairly clean, but it doesn't support remove(). To support
    504        // remove(), we create an AbstractIterator, and then use ForwardingIterator to delegate to it.
    505        final Iterator<Entry<E>> readOnlyIterator =
    506            new AbstractIterator<Entry<E>>() {
    507              private Iterator<Map.Entry<E, AtomicInteger>> mapEntries = countMap.entrySet().iterator();
    508
    509              @Override protected Entry<E> computeNext() {
    510                while (true) {
    511                  if (!mapEntries.hasNext()) {
    512                    return endOfData();
    513                  }
    514                  Map.Entry<E, AtomicInteger> mapEntry = mapEntries.next();
    515                  int count = mapEntry.getValue().get();
    516                  if (count != 0) {
    517                    return Multisets.immutableEntry(mapEntry.getKey(), count);
    518                  }
    519                }
    520              }
    521            };
    522
    523        return new ForwardingIterator<Entry<E>>() {
    524          private Entry<E> last;
    525
    526          @Override protected Iterator<Entry<E>> delegate() {
    527            return readOnlyIterator;
    528          }
    529
    530          @Override public Entry<E> next() {
    531            last = super.next();
    532            return last;
    533          }
    534
    535          @Override public void remove() {
    536            checkState(last != null);
    537            ConcurrentHashMultiset.this.setCount(last.getElement(), 0);
    538            last = null;
    539          }
    540        };
    541      }
    542
    543      @Override public void clear() {
    544        countMap.clear();
    545      }
    546
    547      private class EntrySet extends AbstractMultiset<E>.EntrySet {
    548        @Override ConcurrentHashMultiset<E> multiset() {
    549          return ConcurrentHashMultiset.this;
    550        }
    551
    552        /*
    553         * Note: the superclass toArray() methods assume that size() gives a correct
    554         * answer, which ours does not.
    555         */
    556
    557        @Override public Object[] toArray() {
    558          return snapshot().toArray();
    559        }
    560
    561        @Override public <T> T[] toArray(T[] array) {
    562          return snapshot().toArray(array);
    563        }
    564
    565        private List<Multiset.Entry<E>> snapshot() {
    566          List<Multiset.Entry<E>> list = Lists.newArrayListWithExpectedSize(size());
    567          // Not Iterables.addAll(list, this), because that'll forward right back here.
    568          Iterators.addAll(list, iterator());
    569          return list;
    570        }
    571
    572        @Override public boolean remove(Object object) {
    573          if (object instanceof Multiset.Entry) {
    574            Multiset.Entry<?> entry = (Multiset.Entry<?>) object;
    575            Object element = entry.getElement();
    576            int entryCount = entry.getCount();
    577            if (entryCount != 0) {
    578              // Safe as long as we never add a new entry, which we won't.
    579              @SuppressWarnings("unchecked")
    580              Multiset<Object> multiset = (Multiset) multiset();
    581              return multiset.setCount(element, entryCount, 0);
    582            }
    583          }
    584          return false;
    585        }
    586      }
    587
    588      /**
    589       * @serialData the ConcurrentMap of elements and their counts.
    590       */
    591      private void writeObject(ObjectOutputStream stream) throws IOException {
    592        stream.defaultWriteObject();
    593        stream.writeObject(countMap);
    594      }
    595
    596      private void readObject(ObjectInputStream stream) throws IOException, ClassNotFoundException {
    597        stream.defaultReadObject();
    598        @SuppressWarnings("unchecked") // reading data stored by writeObject
    599        ConcurrentMap<E, Integer> deserializedCountMap =
    600            (ConcurrentMap<E, Integer>) stream.readObject();
    601        FieldSettersHolder.COUNT_MAP_FIELD_SETTER.set(this, deserializedCountMap);
    602      }
    603
    604      private static final long serialVersionUID = 1;
    605    }