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.base.Preconditions.checkState;
022import static com.google.common.collect.Multisets.checkNonnegative;
023
024import com.google.common.annotations.Beta;
025import com.google.common.annotations.VisibleForTesting;
026import com.google.common.collect.Serialization.FieldSetter;
027import com.google.common.math.IntMath;
028import com.google.common.primitives.Ints;
029
030import java.io.IOException;
031import java.io.ObjectInputStream;
032import java.io.ObjectOutputStream;
033import java.io.Serializable;
034import java.util.Collection;
035import java.util.Iterator;
036import java.util.List;
037import java.util.Map;
038import java.util.Set;
039import java.util.concurrent.ConcurrentHashMap;
040import java.util.concurrent.ConcurrentMap;
041import java.util.concurrent.atomic.AtomicInteger;
042
043import 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 */
057public 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}