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      
154      @Override
155      public int count(@Nullable Object element) {
156        AtomicInteger existingCounter = safeGet(element);
157        return (existingCounter == null) ? 0 : existingCounter.get();
158      }
159    
160      /**
161       * Depending on the type of the underlying map, map.get may throw NullPointerException or
162       * ClassCastException, if the object is null or of the wrong type. We usually just want to treat
163       * those cases as if the element isn't in the map, by catching the exceptions and returning null.
164       */
165      private AtomicInteger safeGet(Object element) {
166        try {
167          return countMap.get(element);
168        } catch (NullPointerException e) {
169          return null;
170        } catch (ClassCastException e) {
171          return null;
172        }
173      }
174    
175      /**
176       * {@inheritDoc}
177       *
178       * <p>If the data in the multiset is modified by any other threads during this method,
179       * it is undefined which (if any) of these modifications will be reflected in the result.
180       */
181      
182      @Override
183      public int size() {
184        long sum = 0L;
185        for (AtomicInteger value : countMap.values()) {
186          sum += value.get();
187        }
188        return Ints.saturatedCast(sum);
189      }
190    
191      /*
192       * Note: the superclass toArray() methods assume that size() gives a correct
193       * answer, which ours does not.
194       */
195    
196      
197      @Override
198      public Object[] toArray() {
199        return snapshot().toArray();
200      }
201    
202      
203      @Override
204      public <T> T[] toArray(T[] array) {
205        return snapshot().toArray(array);
206      }
207    
208      /*
209       * We'd love to use 'new ArrayList(this)' or 'list.addAll(this)', but
210       * either of these would recurse back to us again!
211       */
212      private List<E> snapshot() {
213        List<E> list = Lists.newArrayListWithExpectedSize(size());
214        for (Multiset.Entry<E> entry : entrySet()) {
215          E element = entry.getElement();
216          for (int i = entry.getCount(); i > 0; i--) {
217            list.add(element);
218          }
219        }
220        return list;
221      }
222    
223      // Modification Operations
224    
225      /**
226       * Adds a number of occurrences of the specified element to this multiset.
227       *
228       * @param element the element to add
229       * @param occurrences the number of occurrences to add
230       * @return the previous count of the element before the operation; possibly zero
231       * @throws IllegalArgumentException if {@code occurrences} is negative, or if
232       *     the resulting amount would exceed {@link Integer#MAX_VALUE}
233       */
234      @Override
235      public int add(E element, int occurrences) {
236        checkNotNull(element);
237        if (occurrences == 0) {
238          return count(element);
239        }
240        checkArgument(occurrences > 0, "Invalid occurrences: %s", occurrences);
241    
242        while (true) {
243          AtomicInteger existingCounter = safeGet(element);
244          if (existingCounter == null) {
245            existingCounter = countMap.putIfAbsent(element, new AtomicInteger(occurrences));
246            if (existingCounter == null) {
247              return 0;
248            }
249            // existingCounter != null: fall through to operate against the existing AtomicInteger
250          }
251    
252          while (true) {
253            int oldValue = existingCounter.get();
254            if (oldValue != 0) {
255              try {
256                int newValue = IntMath.checkedAdd(oldValue, occurrences);
257                if (existingCounter.compareAndSet(oldValue, newValue)) {
258                  // newValue can't == 0, so no need to check & remove
259                  return oldValue;
260                }
261              } catch (ArithmeticException overflow) {
262                throw new IllegalArgumentException("Overflow adding " + occurrences
263                    + " occurrences to a count of " + oldValue);
264              }
265            } else {
266              // In the case of a concurrent remove, we might observe a zero value, which means another
267              // thread is about to remove (element, existingCounter) from the map. Rather than wait,
268              // we can just do that work here.
269              AtomicInteger newCounter = new AtomicInteger(occurrences);
270              if ((countMap.putIfAbsent(element, newCounter) == null)
271                  || countMap.replace(element, existingCounter, newCounter)) {
272                return 0;
273              }
274              break;
275            }
276          }
277    
278          // If we're still here, there was a race, so just try again.
279        }
280      }
281    
282      /**
283       * Removes a number of occurrences of the specified element from this multiset. If the multiset
284       * contains fewer than this number of occurrences to begin with, all occurrences will be removed.
285       *
286       * @param element the element whose occurrences should be removed
287       * @param occurrences the number of occurrences of the element to remove
288       * @return the count of the element before the operation; possibly zero
289       * @throws IllegalArgumentException if {@code occurrences} is negative
290       */
291      /*
292       * TODO(cpovirk): remove and removeExactly currently accept null inputs only
293       * if occurrences == 0. This satisfies both NullPointerTester and
294       * CollectionRemoveTester.testRemove_nullAllowed, but it's not clear that it's
295       * a good policy, especially because, in order for the test to pass, the
296       * parameter must be misleadingly annotated as @Nullable. I suspect that
297       * we'll want to remove @Nullable, add an eager checkNotNull, and loosen up
298       * testRemove_nullAllowed.
299       */
300      @Override
301      public int remove(@Nullable Object element, int occurrences) {
302        if (occurrences == 0) {
303          return count(element);
304        }
305        checkArgument(occurrences > 0, "Invalid occurrences: %s", occurrences);
306    
307        AtomicInteger existingCounter = safeGet(element);
308        if (existingCounter == null) {
309          return 0;
310        }
311        while (true) {
312          int oldValue = existingCounter.get();
313          if (oldValue != 0) {
314            int newValue = Math.max(0, oldValue - occurrences);
315            if (existingCounter.compareAndSet(oldValue, newValue)) {
316              if (newValue == 0) {
317                // Just CASed to 0; remove the entry to clean up the map. If the removal fails,
318                // another thread has already replaced it with a new counter, which is fine.
319                countMap.remove(element, existingCounter);
320              }
321              return oldValue;
322            }
323          } else {
324            return 0;
325          }
326        }
327      }
328    
329      /**
330       * Removes exactly the specified number of occurrences of {@code element}, or makes no
331       * change if this is not possible.
332       *
333       * <p>This method, in contrast to {@link #remove(Object, int)}, has no effect when the
334       * element count is smaller than {@code occurrences}.
335       *
336       * @param element the element to remove
337       * @param occurrences the number of occurrences of {@code element} to remove
338       * @return {@code true} if the removal was possible (including if {@code occurrences} is zero)
339       */
340      public boolean removeExactly(@Nullable Object element, int occurrences) {
341        if (occurrences == 0) {
342          return true;
343        }
344        checkArgument(occurrences > 0, "Invalid occurrences: %s", occurrences);
345    
346        AtomicInteger existingCounter = safeGet(element);
347        if (existingCounter == null) {
348          return false;
349        }
350        while (true) {
351          int oldValue = existingCounter.get();
352          if (oldValue < occurrences) {
353            return false;
354          }
355          int newValue = oldValue - occurrences;
356          if (existingCounter.compareAndSet(oldValue, newValue)) {
357            if (newValue == 0) {
358              // Just CASed to 0; remove the entry to clean up the map. If the removal fails,
359              // another thread has already replaced it with a new counter, which is fine.
360              countMap.remove(element, existingCounter);
361            }
362            return true;
363          }
364        }
365      }
366    
367      /**
368       * Adds or removes occurrences of {@code element} such that the {@link #count} of the
369       * element becomes {@code count}.
370       *
371       * @return the count of {@code element} in the multiset before this call
372       * @throws IllegalArgumentException if {@code count} is negative
373       */
374      @Override
375      public int setCount(E element, int count) {
376        checkNotNull(element);
377        checkNonnegative(count, "count");
378        while (true) {
379          AtomicInteger existingCounter = safeGet(element);
380          if (existingCounter == null) {
381            if (count == 0) {
382              return 0;
383            } else {
384              existingCounter = countMap.putIfAbsent(element, new AtomicInteger(count));
385              if (existingCounter == null) {
386                return 0;
387              }
388              // existingCounter != null: fall through
389            }
390          }
391    
392          while (true) {
393            int oldValue = existingCounter.get();
394            if (oldValue == 0) {
395              if (count == 0) {
396                return 0;
397              } else {
398                AtomicInteger newCounter = new AtomicInteger(count);
399                if ((countMap.putIfAbsent(element, newCounter) == null)
400                    || countMap.replace(element, existingCounter, newCounter)) {
401                  return 0;
402                }
403              }
404              break;
405            } else {
406              if (existingCounter.compareAndSet(oldValue, count)) {
407                if (count == 0) {
408                  // Just CASed to 0; remove the entry to clean up the map. If the removal fails,
409                  // another thread has already replaced it with a new counter, which is fine.
410                  countMap.remove(element, existingCounter);
411                }
412                return oldValue;
413              }
414            }
415          }
416        }
417      }
418    
419      /**
420       * Sets the number of occurrences of {@code element} to {@code newCount}, but only if
421       * the count is currently {@code expectedOldCount}. If {@code element} does not appear
422       * in the multiset exactly {@code expectedOldCount} times, no changes will be made.
423       *
424       * @return {@code true} if the change was successful. This usually indicates
425       *     that the multiset has been modified, but not always: in the case that
426       *     {@code expectedOldCount == newCount}, the method will return {@code true} if
427       *     the condition was met.
428       * @throws IllegalArgumentException if {@code expectedOldCount} or {@code newCount} is negative
429       */
430      @Override
431      public boolean setCount(E element, int expectedOldCount, int newCount) {
432        checkNotNull(element);
433        checkNonnegative(expectedOldCount, "oldCount");
434        checkNonnegative(newCount, "newCount");
435    
436        AtomicInteger existingCounter = safeGet(element);
437        if (existingCounter == null) {
438          if (expectedOldCount != 0) {
439            return false;
440          } else if (newCount == 0) {
441            return true;
442          } else {
443            // if our write lost the race, it must have lost to a nonzero value, so we can stop
444            return countMap.putIfAbsent(element, new AtomicInteger(newCount)) == null;
445          }
446        }
447        int oldValue = existingCounter.get();
448        if (oldValue == expectedOldCount) {
449          if (oldValue == 0) {
450            if (newCount == 0) {
451              // Just observed a 0; try to remove the entry to clean up the map
452              countMap.remove(element, existingCounter);
453              return true;
454            } else {
455              AtomicInteger newCounter = new AtomicInteger(newCount);
456              return (countMap.putIfAbsent(element, newCounter) == null)
457                  || countMap.replace(element, existingCounter, newCounter);
458            }
459          } else {
460            if (existingCounter.compareAndSet(oldValue, newCount)) {
461              if (newCount == 0) {
462                // Just CASed to 0; remove the entry to clean up the map. If the removal fails,
463                // another thread has already replaced it with a new counter, which is fine.
464                countMap.remove(element, existingCounter);
465              }
466              return true;
467            }
468          }
469        }
470        return false;
471      }
472    
473      // Views
474    
475      
476      @Override
477      Set<E> createElementSet() {
478        final Set<E> delegate = countMap.keySet();
479        return new ForwardingSet<E>() {
480          
481          @Override
482          protected Set<E> delegate() {
483            return delegate;
484          }
485          
486          @Override
487          public boolean remove(Object object) {
488            try {
489              return delegate.remove(object);
490            } catch (NullPointerException e) {
491              return false;
492            } catch (ClassCastException e) {
493              return false;
494            }
495          }
496          
497          @Override
498          public boolean removeAll(Collection<?> c) {
499            return standardRemoveAll(c);
500          }
501        };
502      }
503    
504      private transient EntrySet entrySet;
505    
506      
507      @Override
508      public Set<Multiset.Entry<E>> entrySet() {
509        EntrySet result = entrySet;
510        if (result == null) {
511          entrySet = result = new EntrySet();
512        }
513        return result;
514      }
515    
516      
517      @Override
518      int distinctElements() {
519        return countMap.size();
520      }
521    
522      
523      @Override
524      public boolean isEmpty() {
525        return countMap.isEmpty();
526      }
527    
528      
529      @Override
530      Iterator<Entry<E>> entryIterator() {
531        // AbstractIterator makes this fairly clean, but it doesn't support remove(). To support
532        // remove(), we create an AbstractIterator, and then use ForwardingIterator to delegate to it.
533        final Iterator<Entry<E>> readOnlyIterator =
534            new AbstractIterator<Entry<E>>() {
535              private Iterator<Map.Entry<E, AtomicInteger>> mapEntries = countMap.entrySet().iterator();
536    
537              
538              @Override
539              protected Entry<E> computeNext() {
540                while (true) {
541                  if (!mapEntries.hasNext()) {
542                    return endOfData();
543                  }
544                  Map.Entry<E, AtomicInteger> mapEntry = mapEntries.next();
545                  int count = mapEntry.getValue().get();
546                  if (count != 0) {
547                    return Multisets.immutableEntry(mapEntry.getKey(), count);
548                  }
549                }
550              }
551            };
552    
553        return new ForwardingIterator<Entry<E>>() {
554          private Entry<E> last;
555    
556          
557          @Override
558          protected Iterator<Entry<E>> delegate() {
559            return readOnlyIterator;
560          }
561    
562          
563          @Override
564          public Entry<E> next() {
565            last = super.next();
566            return last;
567          }
568    
569          
570          @Override
571          public void remove() {
572            checkState(last != null);
573            ConcurrentHashMultiset.this.setCount(last.getElement(), 0);
574            last = null;
575          }
576        };
577      }
578    
579      
580      @Override
581      public void clear() {
582        countMap.clear();
583      }
584    
585      private class EntrySet extends AbstractMultiset<E>.EntrySet {
586        
587        @Override
588        ConcurrentHashMultiset<E> multiset() {
589          return ConcurrentHashMultiset.this;
590        }
591    
592        /*
593         * Note: the superclass toArray() methods assume that size() gives a correct
594         * answer, which ours does not.
595         */
596    
597        
598        @Override
599        public Object[] toArray() {
600          return snapshot().toArray();
601        }
602    
603        
604        @Override
605        public <T> T[] toArray(T[] array) {
606          return snapshot().toArray(array);
607        }
608    
609        private List<Multiset.Entry<E>> snapshot() {
610          List<Multiset.Entry<E>> list = Lists.newArrayListWithExpectedSize(size());
611          // Not Iterables.addAll(list, this), because that'll forward right back here.
612          Iterators.addAll(list, iterator());
613          return list;
614        }
615    
616        
617        @Override
618        public boolean remove(Object object) {
619          if (object instanceof Multiset.Entry) {
620            Multiset.Entry<?> entry = (Multiset.Entry<?>) object;
621            Object element = entry.getElement();
622            int entryCount = entry.getCount();
623            if (entryCount != 0) {
624              // Safe as long as we never add a new entry, which we won't.
625              @SuppressWarnings("unchecked")
626              Multiset<Object> multiset = (Multiset) multiset();
627              return multiset.setCount(element, entryCount, 0);
628            }
629          }
630          return false;
631        }
632      }
633    
634      /**
635       * @serialData the ConcurrentMap of elements and their counts.
636       */
637      private void writeObject(ObjectOutputStream stream) throws IOException {
638        stream.defaultWriteObject();
639        stream.writeObject(countMap);
640      }
641    
642      private void readObject(ObjectInputStream stream) throws IOException, ClassNotFoundException {
643        stream.defaultReadObject();
644        @SuppressWarnings("unchecked") // reading data stored by writeObject
645        ConcurrentMap<E, Integer> deserializedCountMap =
646            (ConcurrentMap<E, Integer>) stream.readObject();
647        FieldSettersHolder.COUNT_MAP_FIELD_SETTER.set(this, deserializedCountMap);
648      }
649    
650      private static final long serialVersionUID = 1;
651    }