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
17 package com.google.common.collect;
18
19 import static com.google.common.base.Preconditions.checkArgument;
20 import static com.google.common.base.Preconditions.checkNotNull;
21 import static com.google.common.base.Preconditions.checkState;
22 import static com.google.common.collect.CollectPreconditions.checkNonnegative;
23
24 import com.google.common.annotations.Beta;
25 import com.google.common.annotations.GwtIncompatible;
26 import com.google.common.annotations.VisibleForTesting;
27 import com.google.common.collect.Serialization.FieldSetter;
28 import com.google.common.math.IntMath;
29 import com.google.common.primitives.Ints;
30 import com.google.errorprone.annotations.CanIgnoreReturnValue;
31 import com.google.j2objc.annotations.WeakOuter;
32 import java.io.IOException;
33 import java.io.ObjectInputStream;
34 import java.io.ObjectOutputStream;
35 import java.io.Serializable;
36 import java.util.Collection;
37 import java.util.Iterator;
38 import java.util.List;
39 import java.util.Map;
40 import java.util.Set;
41 import java.util.concurrent.ConcurrentHashMap;
42 import java.util.concurrent.ConcurrentMap;
43 import java.util.concurrent.atomic.AtomicInteger;
44 import javax.annotation.CheckForNull;
45 import org.checkerframework.checker.nullness.qual.Nullable;
46
47 /**
48 * A multiset that supports concurrent modifications and that provides atomic versions of most
49 * {@code Multiset} operations (exceptions where noted). Null elements are not supported.
50 *
51 * <p>See the Guava User Guide article on <a href=
52 * "https://github.com/google/guava/wiki/NewCollectionTypesExplained#multiset"> {@code
53 * Multiset}</a>.
54 *
55 * @author Cliff L. Biffle
56 * @author mike nonemacher
57 * @since 2.0
58 */
59 @GwtIncompatible
60 @ElementTypesAreNonnullByDefault
61 public final class ConcurrentHashMultiset<E> extends AbstractMultiset<E> implements Serializable {
62
63 /*
64 * The ConcurrentHashMultiset's atomic operations are implemented primarily in terms of
65 * AtomicInteger's atomic operations, with some help from ConcurrentMap's atomic operations on
66 * creation and removal (including automatic removal of zeroes). If the modification of an
67 * AtomicInteger results in zero, we compareAndSet the value to zero; if that succeeds, we remove
68 * the entry from the Map. If another operation sees a zero in the map, it knows that the entry is
69 * about to be removed, so this operation may remove it (often by replacing it with a new
70 * AtomicInteger).
71 */
72
73 /** The number of occurrences of each element. */
74 private final transient ConcurrentMap<E, AtomicInteger> countMap;
75
76 // This constant allows the deserialization code to set a final field. This holder class
77 // makes sure it is not initialized unless an instance is deserialized.
78 private static class FieldSettersHolder {
79 static final FieldSetter<ConcurrentHashMultiset> COUNT_MAP_FIELD_SETTER =
80 Serialization.getFieldSetter(ConcurrentHashMultiset.class, "countMap");
81 }
82
83 /**
84 * Creates a new, empty {@code ConcurrentHashMultiset} using the default initial capacity, load
85 * factor, and concurrency settings.
86 */
87 public static <E> ConcurrentHashMultiset<E> create() {
88 // TODO(schmoe): provide a way to use this class with other (possibly arbitrary)
89 // ConcurrentMap implementors. One possibility is to extract most of this class into
90 // an AbstractConcurrentMapMultiset.
91 return new ConcurrentHashMultiset<E>(new ConcurrentHashMap<E, AtomicInteger>());
92 }
93
94 /**
95 * Creates a new {@code ConcurrentHashMultiset} containing the specified elements, using the
96 * default initial capacity, load factor, and concurrency settings.
97 *
98 * <p>This implementation is highly efficient when {@code elements} is itself a {@link Multiset}.
99 *
100 * @param elements the elements that the multiset should contain
101 */
102 public static <E> ConcurrentHashMultiset<E> create(Iterable<? extends E> elements) {
103 ConcurrentHashMultiset<E> multiset = ConcurrentHashMultiset.create();
104 Iterables.addAll(multiset, elements);
105 return multiset;
106 }
107
108 /**
109 * Creates a new, empty {@code ConcurrentHashMultiset} using {@code countMap} as the internal
110 * backing map.
111 *
112 * <p>This instance will assume ownership of {@code countMap}, and other code should not maintain
113 * references to the map or modify it in any way.
114 *
115 * <p>The returned multiset is serializable if the input map is.
116 *
117 * @param countMap backing map for storing the elements in the multiset and their counts. It must
118 * be empty.
119 * @throws IllegalArgumentException if {@code countMap} is not empty
120 * @since 20.0
121 */
122 @Beta
123 public static <E> ConcurrentHashMultiset<E> create(ConcurrentMap<E, AtomicInteger> countMap) {
124 return new ConcurrentHashMultiset<E>(countMap);
125 }
126
127 @VisibleForTesting
128 ConcurrentHashMultiset(ConcurrentMap<E, AtomicInteger> countMap) {
129 checkArgument(countMap.isEmpty(), "the backing map (%s) must be empty", countMap);
130 this.countMap = countMap;
131 }
132
133 // Query Operations
134
135 /**
136 * Returns the number of occurrences of {@code element} in this multiset.
137 *
138 * @param element the element to look for
139 * @return the nonnegative number of occurrences of the element
140 */
141 @Override
142 public int count(@CheckForNull Object element) {
143 AtomicInteger existingCounter = Maps.safeGet(countMap, element);
144 return (existingCounter == null) ? 0 : existingCounter.get();
145 }
146
147 /**
148 * {@inheritDoc}
149 *
150 * <p>If the data in the multiset is modified by any other threads during this method, it is
151 * undefined which (if any) of these modifications will be reflected in the result.
152 */
153 @Override
154 public int size() {
155 long sum = 0L;
156 for (AtomicInteger value : countMap.values()) {
157 sum += value.get();
158 }
159 return Ints.saturatedCast(sum);
160 }
161
162 /*
163 * Note: the superclass toArray() methods assume that size() gives a correct
164 * answer, which ours does not.
165 */
166
167 @Override
168 public Object[] toArray() {
169 return snapshot().toArray();
170 }
171
172 @Override
173 /*
174 * Our checker says "found: T[]; required: T[]." That sounds bogus. I discuss a possible reason
175 * for this error in https://github.com/jspecify/checker-framework/issues/10.
176 */
177 @SuppressWarnings("nullness")
178 public <T extends @Nullable Object> T[] toArray(T[] array) {
179 return snapshot().toArray(array);
180 }
181
182 /*
183 * We'd love to use 'new ArrayList(this)' or 'list.addAll(this)', but
184 * either of these would recurse back to us again!
185 */
186 private List<E> snapshot() {
187 List<E> list = Lists.newArrayListWithExpectedSize(size());
188 for (Multiset.Entry<E> entry : entrySet()) {
189 E element = entry.getElement();
190 for (int i = entry.getCount(); i > 0; i--) {
191 list.add(element);
192 }
193 }
194 return list;
195 }
196
197 // Modification Operations
198
199 /**
200 * Adds a number of occurrences of the specified element to this multiset.
201 *
202 * @param element the element to add
203 * @param occurrences the number of occurrences to add
204 * @return the previous count of the element before the operation; possibly zero
205 * @throws IllegalArgumentException if {@code occurrences} is negative, or if the resulting amount
206 * would exceed {@link Integer#MAX_VALUE}
207 */
208 @CanIgnoreReturnValue
209 @Override
210 public int add(E element, int occurrences) {
211 checkNotNull(element);
212 if (occurrences == 0) {
213 return count(element);
214 }
215 CollectPreconditions.checkPositive(occurrences, "occurences");
216
217 while (true) {
218 AtomicInteger existingCounter = Maps.safeGet(countMap, element);
219 if (existingCounter == null) {
220 existingCounter = countMap.putIfAbsent(element, new AtomicInteger(occurrences));
221 if (existingCounter == null) {
222 return 0;
223 }
224 // existingCounter != null: fall through to operate against the existing AtomicInteger
225 }
226
227 while (true) {
228 int oldValue = existingCounter.get();
229 if (oldValue != 0) {
230 try {
231 int newValue = IntMath.checkedAdd(oldValue, occurrences);
232 if (existingCounter.compareAndSet(oldValue, newValue)) {
233 // newValue can't == 0, so no need to check & remove
234 return oldValue;
235 }
236 } catch (ArithmeticException overflow) {
237 throw new IllegalArgumentException(
238 "Overflow adding " + occurrences + " occurrences to a count of " + oldValue);
239 }
240 } else {
241 // In the case of a concurrent remove, we might observe a zero value, which means another
242 // thread is about to remove (element, existingCounter) from the map. Rather than wait,
243 // we can just do that work here.
244 AtomicInteger newCounter = new AtomicInteger(occurrences);
245 if ((countMap.putIfAbsent(element, newCounter) == null)
246 || countMap.replace(element, existingCounter, newCounter)) {
247 return 0;
248 }
249 break;
250 }
251 }
252
253 // If we're still here, there was a race, so just try again.
254 }
255 }
256
257 /**
258 * Removes a number of occurrences of the specified element from this multiset. If the multiset
259 * contains fewer than this number of occurrences to begin with, all occurrences will be removed.
260 *
261 * @param element the element whose occurrences should be removed
262 * @param occurrences the number of occurrences of the element to remove
263 * @return the count of the element before the operation; possibly zero
264 * @throws IllegalArgumentException if {@code occurrences} is negative
265 */
266 /*
267 * TODO(cpovirk): remove and removeExactly currently accept null inputs only
268 * if occurrences == 0. This satisfies both NullPointerTester and
269 * CollectionRemoveTester.testRemove_nullAllowed, but it's not clear that it's
270 * a good policy, especially because, in order for the test to pass, the
271 * parameter must be misleadingly annotated as @Nullable. I suspect that
272 * we'll want to remove @Nullable, add an eager checkNotNull, and loosen up
273 * testRemove_nullAllowed.
274 */
275 @CanIgnoreReturnValue
276 @Override
277 public int remove(@CheckForNull Object element, int occurrences) {
278 if (occurrences == 0) {
279 return count(element);
280 }
281 CollectPreconditions.checkPositive(occurrences, "occurences");
282
283 AtomicInteger existingCounter = Maps.safeGet(countMap, element);
284 if (existingCounter == null) {
285 return 0;
286 }
287 while (true) {
288 int oldValue = existingCounter.get();
289 if (oldValue != 0) {
290 int newValue = Math.max(0, oldValue - occurrences);
291 if (existingCounter.compareAndSet(oldValue, newValue)) {
292 if (newValue == 0) {
293 // Just CASed to 0; remove the entry to clean up the map. If the removal fails,
294 // another thread has already replaced it with a new counter, which is fine.
295 countMap.remove(element, existingCounter);
296 }
297 return oldValue;
298 }
299 } else {
300 return 0;
301 }
302 }
303 }
304
305 /**
306 * Removes exactly the specified number of occurrences of {@code element}, or makes no change if
307 * this is not possible.
308 *
309 * <p>This method, in contrast to {@link #remove(Object, int)}, has no effect when the element
310 * count is smaller than {@code occurrences}.
311 *
312 * @param element the element to remove
313 * @param occurrences the number of occurrences of {@code element} to remove
314 * @return {@code true} if the removal was possible (including if {@code occurrences} is zero)
315 * @throws IllegalArgumentException if {@code occurrences} is negative
316 */
317 @CanIgnoreReturnValue
318 public boolean removeExactly(@CheckForNull Object element, int occurrences) {
319 if (occurrences == 0) {
320 return true;
321 }
322 CollectPreconditions.checkPositive(occurrences, "occurences");
323
324 AtomicInteger existingCounter = Maps.safeGet(countMap, element);
325 if (existingCounter == null) {
326 return false;
327 }
328 while (true) {
329 int oldValue = existingCounter.get();
330 if (oldValue < occurrences) {
331 return false;
332 }
333 int newValue = oldValue - occurrences;
334 if (existingCounter.compareAndSet(oldValue, newValue)) {
335 if (newValue == 0) {
336 // Just CASed to 0; remove the entry to clean up the map. If the removal fails,
337 // another thread has already replaced it with a new counter, which is fine.
338 countMap.remove(element, existingCounter);
339 }
340 return true;
341 }
342 }
343 }
344
345 /**
346 * Adds or removes occurrences of {@code element} such that the {@link #count} of the element
347 * becomes {@code count}.
348 *
349 * @return the count of {@code element} in the multiset before this call
350 * @throws IllegalArgumentException if {@code count} is negative
351 */
352 @CanIgnoreReturnValue
353 @Override
354 public int setCount(E element, int count) {
355 checkNotNull(element);
356 checkNonnegative(count, "count");
357 while (true) {
358 AtomicInteger existingCounter = Maps.safeGet(countMap, element);
359 if (existingCounter == null) {
360 if (count == 0) {
361 return 0;
362 } else {
363 existingCounter = countMap.putIfAbsent(element, new AtomicInteger(count));
364 if (existingCounter == null) {
365 return 0;
366 }
367 // existingCounter != null: fall through
368 }
369 }
370
371 while (true) {
372 int oldValue = existingCounter.get();
373 if (oldValue == 0) {
374 if (count == 0) {
375 return 0;
376 } else {
377 AtomicInteger newCounter = new AtomicInteger(count);
378 if ((countMap.putIfAbsent(element, newCounter) == null)
379 || countMap.replace(element, existingCounter, newCounter)) {
380 return 0;
381 }
382 }
383 break;
384 } else {
385 if (existingCounter.compareAndSet(oldValue, count)) {
386 if (count == 0) {
387 // Just CASed to 0; remove the entry to clean up the map. If the removal fails,
388 // another thread has already replaced it with a new counter, which is fine.
389 countMap.remove(element, existingCounter);
390 }
391 return oldValue;
392 }
393 }
394 }
395 }
396 }
397
398 /**
399 * Sets the number of occurrences of {@code element} to {@code newCount}, but only if the count is
400 * currently {@code expectedOldCount}. If {@code element} does not appear in the multiset exactly
401 * {@code expectedOldCount} times, no changes will be made.
402 *
403 * @return {@code true} if the change was successful. This usually indicates that the multiset has
404 * been modified, but not always: in the case that {@code expectedOldCount == newCount}, the
405 * method will return {@code true} if the condition was met.
406 * @throws IllegalArgumentException if {@code expectedOldCount} or {@code newCount} is negative
407 */
408 @CanIgnoreReturnValue
409 @Override
410 public boolean setCount(E element, int expectedOldCount, int newCount) {
411 checkNotNull(element);
412 checkNonnegative(expectedOldCount, "oldCount");
413 checkNonnegative(newCount, "newCount");
414
415 AtomicInteger existingCounter = Maps.safeGet(countMap, element);
416 if (existingCounter == null) {
417 if (expectedOldCount != 0) {
418 return false;
419 } else if (newCount == 0) {
420 return true;
421 } else {
422 // if our write lost the race, it must have lost to a nonzero value, so we can stop
423 return countMap.putIfAbsent(element, new AtomicInteger(newCount)) == null;
424 }
425 }
426 int oldValue = existingCounter.get();
427 if (oldValue == expectedOldCount) {
428 if (oldValue == 0) {
429 if (newCount == 0) {
430 // Just observed a 0; try to remove the entry to clean up the map
431 countMap.remove(element, existingCounter);
432 return true;
433 } else {
434 AtomicInteger newCounter = new AtomicInteger(newCount);
435 return (countMap.putIfAbsent(element, newCounter) == null)
436 || countMap.replace(element, existingCounter, newCounter);
437 }
438 } else {
439 if (existingCounter.compareAndSet(oldValue, newCount)) {
440 if (newCount == 0) {
441 // Just CASed to 0; remove the entry to clean up the map. If the removal fails,
442 // another thread has already replaced it with a new counter, which is fine.
443 countMap.remove(element, existingCounter);
444 }
445 return true;
446 }
447 }
448 }
449 return false;
450 }
451
452 // Views
453
454 @Override
455 Set<E> createElementSet() {
456 final Set<E> delegate = countMap.keySet();
457 return new ForwardingSet<E>() {
458 @Override
459 protected Set<E> delegate() {
460 return delegate;
461 }
462
463 @Override
464 public boolean contains(@CheckForNull Object object) {
465 return object != null && Collections2.safeContains(delegate, object);
466 }
467
468 @Override
469 public boolean containsAll(Collection<?> collection) {
470 return standardContainsAll(collection);
471 }
472
473 @Override
474 public boolean remove(@CheckForNull Object object) {
475 return object != null && Collections2.safeRemove(delegate, object);
476 }
477
478 @Override
479 public boolean removeAll(Collection<?> c) {
480 return standardRemoveAll(c);
481 }
482 };
483 }
484
485 @Override
486 Iterator<E> elementIterator() {
487 throw new AssertionError("should never be called");
488 }
489
490 /** @deprecated Internal method, use {@link #entrySet()}. */
491 @Deprecated
492 @Override
493 public Set<Multiset.Entry<E>> createEntrySet() {
494 return new EntrySet();
495 }
496
497 @Override
498 int distinctElements() {
499 return countMap.size();
500 }
501
502 @Override
503 public boolean isEmpty() {
504 return countMap.isEmpty();
505 }
506
507 @Override
508 Iterator<Entry<E>> entryIterator() {
509 // AbstractIterator makes this fairly clean, but it doesn't support remove(). To support
510 // remove(), we create an AbstractIterator, and then use ForwardingIterator to delegate to it.
511 final Iterator<Entry<E>> readOnlyIterator =
512 new AbstractIterator<Entry<E>>() {
513 private final Iterator<Map.Entry<E, AtomicInteger>> mapEntries =
514 countMap.entrySet().iterator();
515
516 @Override
517 protected Entry<E> computeNext() {
518 while (true) {
519 if (!mapEntries.hasNext()) {
520 return endOfData();
521 }
522 Map.Entry<E, AtomicInteger> mapEntry = mapEntries.next();
523 int count = mapEntry.getValue().get();
524 if (count != 0) {
525 return Multisets.immutableEntry(mapEntry.getKey(), count);
526 }
527 }
528 }
529 };
530
531 return new ForwardingIterator<Entry<E>>() {
532 @CheckForNull private Entry<E> last;
533
534 @Override
535 protected Iterator<Entry<E>> delegate() {
536 return readOnlyIterator;
537 }
538
539 @Override
540 public Entry<E> next() {
541 last = super.next();
542 return last;
543 }
544
545 @Override
546 public void remove() {
547 checkState(last != null, "no calls to next() since the last call to remove()");
548 ConcurrentHashMultiset.this.setCount(last.getElement(), 0);
549 last = null;
550 }
551 };
552 }
553
554 @Override
555 public Iterator<E> iterator() {
556 return Multisets.iteratorImpl(this);
557 }
558
559 @Override
560 public void clear() {
561 countMap.clear();
562 }
563
564 @WeakOuter
565 private class EntrySet extends AbstractMultiset<E>.EntrySet {
566 @Override
567 ConcurrentHashMultiset<E> multiset() {
568 return ConcurrentHashMultiset.this;
569 }
570
571 /*
572 * Note: the superclass toArray() methods assume that size() gives a correct
573 * answer, which ours does not.
574 */
575
576 @Override
577 public Object[] toArray() {
578 return snapshot().toArray();
579 }
580
581 @Override
582 /*
583 * Our checker says "found: T[]; required: T[]." That sounds bogus. I discuss a possible reason
584 * for this error in https://github.com/jspecify/checker-framework/issues/10.
585 */
586 @SuppressWarnings("nullness")
587 public <T extends @Nullable Object> T[] toArray(T[] array) {
588 return snapshot().toArray(array);
589 }
590
591 private List<Multiset.Entry<E>> snapshot() {
592 List<Multiset.Entry<E>> list = Lists.newArrayListWithExpectedSize(size());
593 // Not Iterables.addAll(list, this), because that'll forward right back here.
594 Iterators.addAll(list, iterator());
595 return list;
596 }
597 }
598
599 /** @serialData the ConcurrentMap of elements and their counts. */
600 private void writeObject(ObjectOutputStream stream) throws IOException {
601 stream.defaultWriteObject();
602 stream.writeObject(countMap);
603 }
604
605 private void readObject(ObjectInputStream stream) throws IOException, ClassNotFoundException {
606 stream.defaultReadObject();
607 @SuppressWarnings("unchecked") // reading data stored by writeObject
608 ConcurrentMap<E, Integer> deserializedCountMap =
609 (ConcurrentMap<E, Integer>) stream.readObject();
610 FieldSettersHolder.COUNT_MAP_FIELD_SETTER.set(this, deserializedCountMap);
611 }
612
613 private static final long serialVersionUID = 1;
614 }