View Javadoc
1   /*
2    * Copyright (C) 2007 The Guava Authors
3    *
4    * Licensed under the Apache License, Version 2.0 (the "License");
5    * you may not use this file except in compliance with the License.
6    * You may obtain a copy of the License at
7    *
8    * http://www.apache.org/licenses/LICENSE-2.0
9    *
10   * Unless required by applicable law or agreed to in writing, software
11   * distributed under the License is distributed on an "AS IS" BASIS,
12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13   * See the License for the specific language governing permissions and
14   * limitations under the License.
15   */
16  
17  package com.google.common.collect;
18  
19  import static com.google.common.base.Preconditions.checkArgument;
20  import static com.google.common.base.Preconditions.checkNotNull;
21  import static com.google.common.collect.CollectPreconditions.checkNonnegative;
22  import static com.google.common.collect.CollectPreconditions.checkRemove;
23  
24  import com.google.common.annotations.Beta;
25  import com.google.common.annotations.GwtIncompatible;
26  import com.google.common.annotations.VisibleForTesting;
27  import com.google.common.collect.Serialization.FieldSetter;
28  import com.google.common.math.IntMath;
29  import com.google.common.primitives.Ints;
30  import com.google.errorprone.annotations.CanIgnoreReturnValue;
31  import com.google.j2objc.annotations.WeakOuter;
32  import java.io.IOException;
33  import java.io.ObjectInputStream;
34  import java.io.ObjectOutputStream;
35  import java.io.Serializable;
36  import java.util.Collection;
37  import java.util.Iterator;
38  import java.util.List;
39  import java.util.Map;
40  import java.util.Set;
41  import java.util.concurrent.ConcurrentHashMap;
42  import java.util.concurrent.ConcurrentMap;
43  import java.util.concurrent.atomic.AtomicInteger;
44  import javax.annotation.Nullable;
45  
46  /**
47   * A multiset that supports concurrent modifications and that provides atomic versions of most
48   * {@code Multiset} operations (exceptions where noted). Null elements are not supported.
49   *
50   * <p>See the Guava User Guide article on <a href=
51   * "https://github.com/google/guava/wiki/NewCollectionTypesExplained#multiset">
52   * {@code Multiset}</a>.
53   *
54   * @author Cliff L. Biffle
55   * @author mike nonemacher
56   * @since 2.0
57   */
58  @GwtIncompatible
59  public final class ConcurrentHashMultiset<E> extends AbstractMultiset<E> implements Serializable {
60  
61    /*
62     * The ConcurrentHashMultiset's atomic operations are implemented primarily in terms of
63     * AtomicInteger's atomic operations, with some help from ConcurrentMap's atomic operations on
64     * creation and removal (including automatic removal of zeroes). If the modification of an
65     * AtomicInteger results in zero, we compareAndSet the value to zero; if that succeeds, we remove
66     * the entry from the Map. If another operation sees a zero in the map, it knows that the entry is
67     * about to be removed, so this operation may remove it (often by replacing it with a new
68     * AtomicInteger).
69     */
70  
71    /** The number of occurrences of each element. */
72    private final transient ConcurrentMap<E, AtomicInteger> countMap;
73  
74    // This constant allows the deserialization code to set a final field. This holder class
75    // makes sure it is not initialized unless an instance is deserialized.
76    private static class FieldSettersHolder {
77      static final FieldSetter<ConcurrentHashMultiset> COUNT_MAP_FIELD_SETTER =
78          Serialization.getFieldSetter(ConcurrentHashMultiset.class, "countMap");
79    }
80  
81    /**
82     * Creates a new, empty {@code ConcurrentHashMultiset} using the default
83     * initial capacity, load factor, and concurrency settings.
84     */
85    public static <E> ConcurrentHashMultiset<E> create() {
86      // TODO(schmoe): provide a way to use this class with other (possibly arbitrary)
87      // ConcurrentMap implementors. One possibility is to extract most of this class into
88      // an AbstractConcurrentMapMultiset.
89      return new ConcurrentHashMultiset<E>(new ConcurrentHashMap<E, AtomicInteger>());
90    }
91  
92    /**
93     * Creates a new {@code ConcurrentHashMultiset} containing the specified elements, using
94     * the default initial capacity, load factor, and concurrency settings.
95     *
96     * <p>This implementation is highly efficient when {@code elements} is itself a {@link Multiset}.
97     *
98     * @param elements the elements that the multiset should contain
99     */
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 }