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