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 org.checkerframework.checker.nullness.compatqual.NullableDecl;
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"> {@code
052 * 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 initial capacity, load
083   * 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 the
094   * 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(@NullableDecl 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, it is
149   * 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 the resulting amount
199   *     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 @NullableDecl. I suspect that
265   * we'll want to remove @NullableDecl, add an eager checkNotNull, and loosen up
266   * testRemove_nullAllowed.
267   */
268  @CanIgnoreReturnValue
269  @Override
270  public int remove(@NullableDecl 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 change if
300   * this is not possible.
301   *
302   * <p>This method, in contrast to {@link #remove(Object, int)}, has no effect when the element
303   * 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(@NullableDecl 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 element
340   * 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 the count is
393   * currently {@code expectedOldCount}. If {@code element} does not appear in the multiset exactly
394   * {@code expectedOldCount} times, no changes will be made.
395   *
396   * @return {@code true} if the change was successful. This usually indicates that the multiset has
397   *     been modified, but not always: in the case that {@code expectedOldCount == newCount}, the
398   *     method will return {@code true} if the condition was met.
399   * @throws IllegalArgumentException if {@code expectedOldCount} or {@code newCount} is negative
400   */
401  @CanIgnoreReturnValue
402  @Override
403  public boolean setCount(E element, int expectedOldCount, int newCount) {
404    checkNotNull(element);
405    checkNonnegative(expectedOldCount, "oldCount");
406    checkNonnegative(newCount, "newCount");
407
408    AtomicInteger existingCounter = Maps.safeGet(countMap, element);
409    if (existingCounter == null) {
410      if (expectedOldCount != 0) {
411        return false;
412      } else if (newCount == 0) {
413        return true;
414      } else {
415        // if our write lost the race, it must have lost to a nonzero value, so we can stop
416        return countMap.putIfAbsent(element, new AtomicInteger(newCount)) == null;
417      }
418    }
419    int oldValue = existingCounter.get();
420    if (oldValue == expectedOldCount) {
421      if (oldValue == 0) {
422        if (newCount == 0) {
423          // Just observed a 0; try to remove the entry to clean up the map
424          countMap.remove(element, existingCounter);
425          return true;
426        } else {
427          AtomicInteger newCounter = new AtomicInteger(newCount);
428          return (countMap.putIfAbsent(element, newCounter) == null)
429              || countMap.replace(element, existingCounter, newCounter);
430        }
431      } else {
432        if (existingCounter.compareAndSet(oldValue, newCount)) {
433          if (newCount == 0) {
434            // Just CASed to 0; remove the entry to clean up the map. If the removal fails,
435            // another thread has already replaced it with a new counter, which is fine.
436            countMap.remove(element, existingCounter);
437          }
438          return true;
439        }
440      }
441    }
442    return false;
443  }
444
445  // Views
446
447  @Override
448  Set<E> createElementSet() {
449    final Set<E> delegate = countMap.keySet();
450    return new ForwardingSet<E>() {
451      @Override
452      protected Set<E> delegate() {
453        return delegate;
454      }
455
456      @Override
457      public boolean contains(@NullableDecl Object object) {
458        return object != null && Collections2.safeContains(delegate, object);
459      }
460
461      @Override
462      public boolean containsAll(Collection<?> collection) {
463        return standardContainsAll(collection);
464      }
465
466      @Override
467      public boolean remove(Object object) {
468        return object != null && Collections2.safeRemove(delegate, object);
469      }
470
471      @Override
472      public boolean removeAll(Collection<?> c) {
473        return standardRemoveAll(c);
474      }
475    };
476  }
477
478  /** @deprecated Internal method, use {@link #entrySet()}. */
479  @Deprecated
480  @Override
481  public Set<Multiset.Entry<E>> createEntrySet() {
482    return new EntrySet();
483  }
484
485  @Override
486  int distinctElements() {
487    return countMap.size();
488  }
489
490  @Override
491  public boolean isEmpty() {
492    return countMap.isEmpty();
493  }
494
495  @Override
496  Iterator<Entry<E>> entryIterator() {
497    // AbstractIterator makes this fairly clean, but it doesn't support remove(). To support
498    // remove(), we create an AbstractIterator, and then use ForwardingIterator to delegate to it.
499    final Iterator<Entry<E>> readOnlyIterator =
500        new AbstractIterator<Entry<E>>() {
501          private final Iterator<Map.Entry<E, AtomicInteger>> mapEntries =
502              countMap.entrySet().iterator();
503
504          @Override
505          protected Entry<E> computeNext() {
506            while (true) {
507              if (!mapEntries.hasNext()) {
508                return endOfData();
509              }
510              Map.Entry<E, AtomicInteger> mapEntry = mapEntries.next();
511              int count = mapEntry.getValue().get();
512              if (count != 0) {
513                return Multisets.immutableEntry(mapEntry.getKey(), count);
514              }
515            }
516          }
517        };
518
519    return new ForwardingIterator<Entry<E>>() {
520      private Entry<E> last;
521
522      @Override
523      protected Iterator<Entry<E>> delegate() {
524        return readOnlyIterator;
525      }
526
527      @Override
528      public Entry<E> next() {
529        last = super.next();
530        return last;
531      }
532
533      @Override
534      public void remove() {
535        checkRemove(last != null);
536        ConcurrentHashMultiset.this.setCount(last.getElement(), 0);
537        last = null;
538      }
539    };
540  }
541
542  @Override
543  public void clear() {
544    countMap.clear();
545  }
546
547  @WeakOuter
548  private class EntrySet extends AbstractMultiset<E>.EntrySet {
549    @Override
550    ConcurrentHashMultiset<E> multiset() {
551      return ConcurrentHashMultiset.this;
552    }
553
554    /*
555     * Note: the superclass toArray() methods assume that size() gives a correct
556     * answer, which ours does not.
557     */
558
559    @Override
560    public Object[] toArray() {
561      return snapshot().toArray();
562    }
563
564    @Override
565    public <T> T[] toArray(T[] array) {
566      return snapshot().toArray(array);
567    }
568
569    private List<Multiset.Entry<E>> snapshot() {
570      List<Multiset.Entry<E>> list = Lists.newArrayListWithExpectedSize(size());
571      // Not Iterables.addAll(list, this), because that'll forward right back here.
572      Iterators.addAll(list, iterator());
573      return list;
574    }
575  }
576
577  /** @serialData the ConcurrentMap of elements and their counts. */
578  private void writeObject(ObjectOutputStream stream) throws IOException {
579    stream.defaultWriteObject();
580    stream.writeObject(countMap);
581  }
582
583  private void readObject(ObjectInputStream stream) throws IOException, ClassNotFoundException {
584    stream.defaultReadObject();
585    @SuppressWarnings("unchecked") // reading data stored by writeObject
586    ConcurrentMap<E, Integer> deserializedCountMap =
587        (ConcurrentMap<E, Integer>) stream.readObject();
588    FieldSettersHolder.COUNT_MAP_FIELD_SETTER.set(this, deserializedCountMap);
589  }
590
591  private static final long serialVersionUID = 1;
592}