Merge pull request #4319 from axel22/fix-si-7943

Fix SI-7943 -- make `TrieMap.getOrElseUpdate` atomic.

3 files changed, 61 insertions, 0 deletions

diff --git a/src/library/scala/collection/concurrent/TrieMap.scala b/src/library/scala/collection/concurrent/TrieMap.scala
index fccc1d81b9..bcfea7a463 100644
--- a/src/library/scala/collection/concurrent/TrieMap.scala
+++ b/src/library/scala/collection/concurrent/TrieMap.scala
@@ -873,6 +873,44 @@ extends scala.collection.concurrent.Map[K, V]
     insertifhc(k, hc, v, INode.KEY_ABSENT)
   }
 
+  // TODO once computeIfAbsent is added to concurrent.Map,
+  // move the comment there and tweak the 'at most once' part
+  /** If the specified key is not already in the map, computes its value using
+   *  the given thunk `op` and enters it into the map.
+   *
+   *  Since concurrent maps cannot contain `null` for keys or values,
+   *  a `NullPointerException` is thrown if the thunk `op`
+   *  returns `null`.
+   *
+   *  If the specified mapping function throws an exception,
+   *  that exception is rethrown.
+   *  
+   *  Note: This method will invoke op at most once.
+   *  However, `op` may be invoked without the result being added to the map if
+   *  a concurrent process is also trying to add a value corresponding to the
+   *  same key `k`.
+   *
+   *  @param k      the key to modify
+   *  @param op     the expression that computes the value
+   *  @return       the newly added value
+   */
+  override def getOrElseUpdate(k: K, op: =>V): V = {
+    val oldv = lookup(k)
+    if (oldv != null) oldv.asInstanceOf[V]
+    else {
+      val v = op
+      if (v == null) {
+        throw new NullPointerException("Concurrent TrieMap values cannot be null.")
+      } else {
+        val hc = computeHash(k)
+        insertifhc(k, hc, v, INode.KEY_ABSENT) match {
+          case Some(oldv) => oldv
+          case None => v
+        }
+      }
+    }
+  }
+
   def remove(k: K, v: V): Boolean = {
     val hc = computeHash(k)
     removehc(k, v, hc).nonEmpty
diff --git a/src/library/scala/collection/mutable/MapLike.scala b/src/library/scala/collection/mutable/MapLike.scala
index af28df1b88..44af886cf5 100644
--- a/src/library/scala/collection/mutable/MapLike.scala
+++ b/src/library/scala/collection/mutable/MapLike.scala
@@ -178,6 +178,10 @@ trait MapLike[A, B, +This <: MapLike[A, B, This] with Map[A, B]]
    *
    *  Otherwise, computes value from given expression `op`, stores with key
    *  in map and returns that value.
+   *
+   *  Concurrent map implementations may evaluate the expression `op`
+   *  multiple times, or may evaluate `op` without inserting the result.
+   *  
    *  @param  key the key to test
    *  @param  op  the computation yielding the value to associate with `key`, if
    *              `key` is previously unbound.
diff --git a/test/files/scalacheck/Ctrie.scala b/test/files/scalacheck/Ctrie.scala
index 714f1c3b09..eef9d06f37 100644
--- a/test/files/scalacheck/Ctrie.scala
+++ b/test/files/scalacheck/Ctrie.scala
@@ -186,6 +186,25 @@ object Test extends Properties("concurrent.TrieMap") {
     })
   }
 
+  property("concurrent getOrElseUpdate") = forAll(threadCounts, sizes) {
+    (p, sz) =>
+    val totalInserts = new java.util.concurrent.atomic.AtomicInteger
+    val ct = new TrieMap[Wrap, String]
+
+    val results = inParallel(p) {
+      idx =>
+      (0 until sz) foreach {
+        i =>
+        val v = ct.getOrElseUpdate(Wrap(i), idx + ":" + i)
+        if (v == idx + ":" + i) totalInserts.incrementAndGet()
+      }
+    }
+
+    (totalInserts.get == sz) && ((0 until sz) forall {
+      case i => ct(Wrap(i)).split(":")(1).toInt == i
+    })
+  }
+
 }