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 }