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.checkState;
021    import static com.google.common.collect.Multisets.checkNonnegative;
022    
023    import com.google.common.annotations.Beta;
024    import com.google.common.annotations.VisibleForTesting;
025    import com.google.common.collect.Serialization.FieldSetter;
026    import com.google.common.primitives.Ints;
027    
028    import java.io.IOException;
029    import java.io.ObjectInputStream;
030    import java.io.ObjectOutputStream;
031    import java.io.Serializable;
032    import java.util.Iterator;
033    import java.util.List;
034    import java.util.Map;
035    import java.util.Set;
036    import java.util.concurrent.ConcurrentHashMap;
037    import java.util.concurrent.ConcurrentMap;
038    import java.util.concurrent.atomic.AtomicInteger;
039    
040    import javax.annotation.Nullable;
041    
042    /**
043     * A multiset that supports concurrent modifications and that provides atomic
044     * versions of most {@code Multiset} operations (exceptions where noted). Null
045     * elements are not supported.
046     *
047     * @author Cliff L. Biffle
048     * @author mike nonemacher
049     * @since 2.0 (imported from Google Collections Library)
050     */
051    public final class ConcurrentHashMultiset<E> extends AbstractMultiset<E> implements Serializable {
052    
053      /*
054       * The ConcurrentHashMultiset's atomic operations are implemented primarily in terms of
055       * AtomicInteger's atomic operations, with some help from ConcurrentMap's atomic operations on
056       * creation and removal (including automatic removal of zeroes). If the modification of an
057       * AtomicInteger results in zero, we compareAndSet the value to zero; if that succeeds, we remove
058       * the entry from the Map. If another operation sees a zero in the map, it knows that the entry is
059       * about to be removed, so this operation may remove it (often by replacing it with a new
060       * AtomicInteger).
061       */
062    
063      /** The number of occurrences of each element. */
064      private final transient ConcurrentMap<E, AtomicInteger> countMap;
065    
066      // This constant allows the deserialization code to set a final field. This
067      // holder class makes sure it is not initialized unless an instance is
068      // deserialized.
069      private static class FieldSettersHolder {
070        static final FieldSetter<ConcurrentHashMultiset> COUNT_MAP_FIELD_SETTER =
071            Serialization.getFieldSetter(ConcurrentHashMultiset.class, "countMap");
072      }
073    
074      /**
075       * Creates a new, empty {@code ConcurrentHashMultiset} using the default
076       * initial capacity, load factor, and concurrency settings.
077       */
078      public static <E> ConcurrentHashMultiset<E> create() {
079        // TODO(schmoe): provide a way to use this class with other (possibly arbitrary)
080        // ConcurrentMap implementors. One possibility is to extract most of this class into
081        // an AbstractConcurrentMapMultiset.
082        return new ConcurrentHashMultiset<E>(new ConcurrentHashMap<E, AtomicInteger>());
083      }
084    
085      /**
086       * Creates a new {@code ConcurrentHashMultiset} containing the specified
087       * elements, using the default initial capacity, load factor, and concurrency
088       * settings.
089       * 
090       * <p>This implementation is highly efficient when {@code elements} is itself
091       * a {@link Multiset}.
092       *
093       * @param elements the elements that the multiset should contain
094       */
095      public static <E> ConcurrentHashMultiset<E> create(Iterable<? extends E> elements) {
096        ConcurrentHashMultiset<E> multiset = ConcurrentHashMultiset.create();
097        Iterables.addAll(multiset, elements);
098        return multiset;
099      }
100    
101      /**
102       * Creates a new, empty {@code ConcurrentHashMultiset} using {@code mapMaker}
103       * to construct the internal backing map.
104       *
105       * <p>If this {@link MapMaker} is configured to use entry eviction of any kind, this eviction
106       * applies to all occurrences of a given element as a single unit. However, most updates to the
107       * multiset do not count as map updates at all, since we're usually just mutating the value
108       * stored in the map, so {@link MapMaker#expireAfterAccess} makes sense (evict the entry that
109       * was queried or updated longest ago), but {@link MapMaker#expireAfterWrite} doesn't, because
110       * the eviction time is measured from when we saw the first occurrence of the object.
111       *
112       * <p>The returned multiset is serializable but any serialization caveats
113       * given in {@code MapMaker} apply.
114       *
115       * <p>Finally, soft/weak values can be used but are not very useful: the values are created
116       * internally and not exposed externally, so no one else will have a strong reference to the
117       * values. Weak keys on the other hand can be useful in some scenarios.
118       * 
119       * @since 7.0
120       */
121      @Beta
122      public static <E> ConcurrentHashMultiset<E> create(
123          GenericMapMaker<? super E, ? super Number> mapMaker) {
124        return new ConcurrentHashMultiset<E>(mapMaker.<E, AtomicInteger>makeMap());
125      }
126    
127      /**
128       * Creates an instance using {@code countMap} to store elements and their
129       * counts.
130       *
131       * <p>This instance will assume ownership of {@code countMap}, and other code
132       * should not maintain references to the map or modify it in any way.
133       *
134       * @param countMap backing map for storing the elements in the multiset and
135       *     their counts. It must be empty.
136       * @throws IllegalArgumentException if {@code countMap} is not empty
137       */
138      @VisibleForTesting ConcurrentHashMultiset(ConcurrentMap<E, AtomicInteger> countMap) {
139        checkArgument(countMap.isEmpty());
140        this.countMap = countMap;
141      }
142    
143      // Query Operations
144    
145      /**
146       * Returns the number of occurrences of {@code element} in this multiset.
147       *
148       * @param element the element to look for
149       * @return the nonnegative number of occurrences of the element
150       */
151      @Override public int count(@Nullable Object element) {
152        AtomicInteger existingCounter = safeGet(element);
153        return (existingCounter == null) ? 0 : existingCounter.get();
154      }
155    
156      /**
157       * Depending on the type of the underlying map, map.get may throw NullPointerException or
158       * ClassCastException, if the object is null or of the wrong type. We usually just want to treat
159       * those cases as if the element isn't in the map, by catching the exceptions and returning null.
160       */
161      private AtomicInteger safeGet(Object element) {
162        try {
163          return countMap.get(element);
164        } catch (NullPointerException e) {
165          return null;
166        } catch (ClassCastException e) {
167          return null;
168        }
169      }
170    
171      /**
172       * {@inheritDoc}
173       *
174       * <p>If the data in the multiset is modified by any other threads during this
175       * method, it is undefined which (if any) of these modifications will be
176       * reflected in the result.
177       */
178      @Override public int size() {
179        long sum = 0L;
180        for (AtomicInteger value : countMap.values()) {
181          sum += value.get();
182        }
183        return Ints.saturatedCast(sum);
184      }
185    
186      /*
187       * Note: the superclass toArray() methods assume that size() gives a correct
188       * answer, which ours does not.
189       */
190    
191      @Override public Object[] toArray() {
192        return snapshot().toArray();
193      }
194    
195      @Override public <T> T[] toArray(T[] array) {
196        return snapshot().toArray(array);
197      }
198    
199      /*
200       * We'd love to use 'new ArrayList(this)' or 'list.addAll(this)', but
201       * either of these would recurse back to us again!
202       */
203      private List<E> snapshot() {
204        List<E> list = Lists.newArrayListWithExpectedSize(size());
205        for (Multiset.Entry<E> entry : entrySet()) {
206          E element = entry.getElement();
207          for (int i = entry.getCount(); i > 0; i--) {
208            list.add(element);
209          }
210        }
211        return list;
212      }
213    
214      // Modification Operations
215    
216      /**
217       * Adds a number of occurrences of the specified element to this multiset.
218       *
219       * @param element the element to add
220       * @param occurrences the number of occurrences to add
221       * @return the previous count of the element before the operation; possibly
222       *     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        if (occurrences == 0) {
228          return count(element);
229        }
230        checkArgument(occurrences > 0, "Invalid occurrences: %s", occurrences);
231    
232        while (true) {
233          AtomicInteger existingCounter = safeGet(element);
234          if (existingCounter == null) {
235            existingCounter = countMap.putIfAbsent(element, new AtomicInteger(occurrences));
236            if (existingCounter == null) {
237              return 0;
238            }
239            // existingCounter != null: fall through to operate against the existing AtomicInteger
240          }
241    
242          while (true) {
243            int oldValue = existingCounter.get();
244            if (oldValue != 0) {
245              checkArgument(occurrences <= Integer.MAX_VALUE - oldValue,
246                  "Overflow adding %s occurrences to a count of %s",
247                  occurrences, oldValue);
248              int newValue = oldValue + occurrences;
249              if (existingCounter.compareAndSet(oldValue, newValue)) {
250                // newValue can't == 0, so no need to check & remove
251                return oldValue;
252              }
253            } else {
254              // In the case of a concurrent remove, we might observe a zero value, which means another
255              // thread is about to remove (element, existingCounter) from the map. Rather than wait,
256              // we can just do that work here.
257              AtomicInteger newCounter = new AtomicInteger(occurrences);
258              if ((countMap.putIfAbsent(element, newCounter) == null)
259                  || countMap.replace(element, existingCounter, newCounter)) {
260                return 0;
261              }
262              break;
263            }
264          }
265    
266          // If we're still here, there was a race, so just try again.
267        }
268      }
269    
270      /**
271       * Removes a number of occurrences of the specified element from this
272       * multiset. If the multiset contains fewer than this number of occurrences to
273       * begin with, all occurrences will be removed.
274       *
275       * @param element the element whose occurrences should be removed
276       * @param occurrences the number of occurrences of the element to remove
277       * @return the count of the element before the operation; possibly zero
278       * @throws IllegalArgumentException if {@code occurrences} is negative
279       */
280      @Override public int remove(@Nullable Object element, int occurrences) {
281        if (occurrences == 0) {
282          return count(element);
283        }
284        checkArgument(occurrences > 0, "Invalid occurrences: %s", occurrences);
285    
286        AtomicInteger existingCounter = safeGet(element);
287        if (existingCounter == null) {
288          return 0;
289        }
290        while (true) {
291          int oldValue = existingCounter.get();
292          if (oldValue != 0) {
293            int newValue = Math.max(0, oldValue - occurrences);
294            if (existingCounter.compareAndSet(oldValue, newValue)) {
295              if (newValue == 0) {
296                // Just CASed to 0; remove the entry to clean up the map. If the removal fails,
297                // another thread has already replaced it with a new counter, which is fine.
298                countMap.remove(element, existingCounter);
299              }
300              return oldValue;
301            }
302          } else {
303            return 0;
304          }
305        }
306      }
307    
308      /**
309       * Removes exactly the specified number of occurrences of {@code element}, or
310       * makes no change if this is not possible.
311       *
312       * <p>This method, in contrast to {@link #remove(Object, int)}, has no effect
313       * when the element count is smaller than {@code occurrences}.
314       *
315       * @param element the element to remove
316       * @param occurrences the number of occurrences of {@code element} to remove
317       * @return {@code true} if the removal was possible (including if {@code
318       *     occurrences} is zero)
319       */
320      public boolean removeExactly(@Nullable Object element, int occurrences) {
321        if (occurrences == 0) {
322          return true;
323        }
324        checkArgument(occurrences > 0, "Invalid occurrences: %s", occurrences);
325    
326        AtomicInteger existingCounter = safeGet(element);
327        if (existingCounter == null) {
328          return false;
329        }
330        while (true) {
331          int oldValue = existingCounter.get();
332          if (oldValue < occurrences) {
333            return false;
334          }
335          int newValue = oldValue - occurrences;
336          if (existingCounter.compareAndSet(oldValue, newValue)) {
337            if (newValue == 0) {
338              // Just CASed to 0; remove the entry to clean up the map. If the removal fails,
339              // another thread has already replaced it with a new counter, which is fine.
340              countMap.remove(element, existingCounter);
341            }
342            return true;
343          }
344        }
345      }
346    
347      /**
348       * Adds or removes occurrences of {@code element} such that the {@link #count}
349       * of the element becomes {@code count}.
350       *
351       * @return the count of {@code element} in the multiset before this call
352       * @throws IllegalArgumentException if {@code count} is negative
353       */
354      @Override public int setCount(E element, int count) {
355        checkNonnegative(count, "count");
356        while (true) {
357          AtomicInteger existingCounter = safeGet(element);
358          if (existingCounter == null) {
359            if (count == 0) {
360              return 0;
361            } else {
362              existingCounter = countMap.putIfAbsent(element, new AtomicInteger(count));
363              if (existingCounter == null) {
364                return 0;
365              }
366              // existingCounter != null: fall through
367            }
368          }
369    
370          while (true) {
371            int oldValue = existingCounter.get();
372            if (oldValue == 0) {
373              if (count == 0) {
374                return 0;
375              } else {
376                AtomicInteger newCounter = new AtomicInteger(count);
377                if ((countMap.putIfAbsent(element, newCounter) == null)
378                    || countMap.replace(element, existingCounter, newCounter)) {
379                  return 0;
380                }
381              }
382              break;
383            } else {
384              if (existingCounter.compareAndSet(oldValue, count)) {
385                if (count == 0) {
386                  // Just CASed to 0; remove the entry to clean up the map. If the removal fails,
387                  // another thread has already replaced it with a new counter, which is fine.
388                  countMap.remove(element, existingCounter);
389                }
390                return oldValue;
391              }
392            }
393          }
394        }
395      }
396    
397      /**
398       * Sets the number of occurrences of {@code element} to {@code newCount}, but
399       * only if the count is currently {@code expectedOldCount}. If {@code element} does
400       * not appear in the multiset exactly {@code expectedOldCount} times, no changes will
401       * be made.
402       *
403       * @return {@code true} if the change was successful. This usually indicates
404       *     that the multiset has been modified, but not always: in the case that
405       *     {@code expectedOldCount == newCount}, the method will return {@code true} if
406       *     the condition was met.
407       * @throws IllegalArgumentException if {@code expectedOldCount} or {@code newCount} is negative
408       */
409      @Override public boolean setCount(E element, int expectedOldCount, int newCount) {
410        checkNonnegative(expectedOldCount, "oldCount");
411        checkNonnegative(newCount, "newCount");
412    
413        AtomicInteger existingCounter = safeGet(element);
414        if (existingCounter == null) {
415          if (expectedOldCount != 0) {
416            return false;
417          } else if (newCount == 0) {
418            return true;
419          } else {
420            // if our write lost the race, it must have lost to a nonzero value, so we can stop
421            return countMap.putIfAbsent(element, new AtomicInteger(newCount)) == null;
422          }
423        }
424        int oldValue = existingCounter.get();
425        if (oldValue == expectedOldCount) {
426          if (oldValue == 0) {
427            if (newCount == 0) {
428              // Just observed a 0; try to remove the entry to clean up the map
429              countMap.remove(element, existingCounter);
430              return true;
431            } else {
432              AtomicInteger newCounter = new AtomicInteger(newCount);
433              return (countMap.putIfAbsent(element, newCounter) == null)
434                  || countMap.replace(element, existingCounter, newCounter);
435            }
436          } else {
437            if (existingCounter.compareAndSet(oldValue, newCount)) {
438              if (newCount == 0) {
439                // Just CASed to 0; remove the entry to clean up the map. If the removal fails,
440                // another thread has already replaced it with a new counter, which is fine.
441                countMap.remove(element, existingCounter);
442              }
443              return true;
444            }
445          }
446        }
447        return false;
448      }
449    
450      // Views
451    
452      @Override Set<E> createElementSet() {
453        final Set<E> delegate = countMap.keySet();
454        return new ForwardingSet<E>() {
455          @Override protected Set<E> delegate() {
456            return delegate;
457          }
458          @Override public boolean remove(Object object) {
459            try {
460              return delegate.remove(object);
461            } catch (NullPointerException e) {
462              return false;
463            } catch (ClassCastException e) {
464              return false;
465            }
466          }
467        };
468      }
469    
470      private transient EntrySet entrySet;
471    
472      @Override public Set<Multiset.Entry<E>> entrySet() {
473        EntrySet result = entrySet;
474        if (result == null) {
475          entrySet = result = new EntrySet();
476        }
477        return result;
478      }
479    
480      @Override int distinctElements() {
481        return countMap.size();
482      }
483    
484      @Override public boolean isEmpty() {
485        return countMap.isEmpty();
486      }
487    
488      @Override Iterator<Entry<E>> entryIterator() {
489        // AbstractIterator makes this fairly clean, but it doesn't support remove(). To support
490        // remove(), we create an AbstractIterator, and then use ForwardingIterator to delegate to it.
491        final Iterator<Entry<E>> readOnlyIterator =
492            new AbstractIterator<Entry<E>>() {
493              private Iterator<Map.Entry<E, AtomicInteger>> mapEntries = countMap.entrySet().iterator();
494    
495              @Override protected Entry<E> computeNext() {
496                while (true) {
497                  if (!mapEntries.hasNext()) {
498                    return endOfData();
499                  }
500                  Map.Entry<E, AtomicInteger> mapEntry = mapEntries.next();
501                  int count = mapEntry.getValue().get();
502                  if (count != 0) {
503                    return Multisets.immutableEntry(mapEntry.getKey(), count);
504                  }
505                }
506              }
507            };
508    
509        return new ForwardingIterator<Entry<E>>() {
510          private Entry<E> last;
511    
512          @Override protected Iterator<Entry<E>> delegate() {
513            return readOnlyIterator;
514          }
515    
516          @Override public Entry<E> next() {
517            last = super.next();
518            return last;
519          }
520    
521          @Override public void remove() {
522            checkState(last != null);
523            ConcurrentHashMultiset.this.setCount(last.getElement(), 0);
524            last = null;
525          }
526        };
527      }
528    
529      @Override public void clear() {
530        countMap.clear();
531      }
532    
533      private class EntrySet extends AbstractMultiset<E>.EntrySet {
534        @Override ConcurrentHashMultiset<E> multiset() {
535          return ConcurrentHashMultiset.this;
536        }
537    
538        /*
539         * Note: the superclass toArray() methods assume that size() gives a correct
540         * answer, which ours does not.
541         */
542    
543        @Override public Object[] toArray() {
544          return snapshot().toArray();
545        }
546    
547        @Override public <T> T[] toArray(T[] array) {
548          return snapshot().toArray(array);
549        }
550    
551        private List<Multiset.Entry<E>> snapshot() {
552          List<Multiset.Entry<E>> list = Lists.newArrayListWithExpectedSize(size());
553          // Not Iterables.addAll(list, this), because that'll forward right back here.
554          Iterators.addAll(list, iterator());
555          return list;
556        }
557    
558        @Override public boolean remove(Object object) {
559          if (object instanceof Multiset.Entry) {
560            Multiset.Entry<?> entry = (Multiset.Entry<?>) object;
561            Object element = entry.getElement();
562            int entryCount = entry.getCount();
563            if (entryCount != 0) {
564              // Safe as long as we never add a new entry, which we won't.
565              @SuppressWarnings("unchecked")
566              Multiset<Object> multiset = (Multiset) multiset();
567              return multiset.setCount(element, entryCount, 0);
568            }
569          }
570          return false;
571        }
572      }
573    
574      /**
575       * @serialData the ConcurrentMap of elements and their counts.
576       */
577      private void writeObject(ObjectOutputStream stream) throws IOException {
578        stream.defaultWriteObject();
579        stream.writeObject(countMap);
580      }
581    
582      private void readObject(ObjectInputStream stream) throws IOException, ClassNotFoundException {
583        stream.defaultReadObject();
584        @SuppressWarnings("unchecked") // reading data stored by writeObject
585        ConcurrentMap<E, Integer> deserializedCountMap =
586            (ConcurrentMap<E, Integer>) stream.readObject();
587        FieldSettersHolder.COUNT_MAP_FIELD_SETTER.set(this, deserializedCountMap);
588      }
589    
590      private static final long serialVersionUID = 1;
591    }
592