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    }