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