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