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