Fix RMap in Exercise 6 of Chapter 6

- RMap.apply does not throw `NoSuchElementException'

- Fix memory leaks

Author: ssmylh

src/main/scala/org/learningconcurrency/exercises/ch6/Ex6.scala

@@ -18,32 +18,60 @@ import rx.lang.scala._
 object Ex6 extends App {
 
   class RMap[K, V] {
-    private[this] val map = scala.collection.mutable.Map[K, Subject[V]]()
-
-    def update(k: K, v: V): Unit = map.get(k) match {
-      case Some(s) => s.onNext(v)
-      case _ =>
-        val s = Subject[V]()
-        map(k) = s
-        s.onNext(v)
+    import scala.collection._
+    private[this] val allSubscribers = mutable.Map[K, (Subject[V], mutable.Set[Subscriber[V]])]()
+    private[this] val map = mutable.Map[K, V]()
+
+    def update(k: K, v: V): Unit = {
+      map(k) = v
+      allSubscribers.get(k) match {
+        case Some(s) => s._1.onNext(v)
+        case _ =>
+      }
     }
 
-    /* This method throws `NoSuchElementException` if the key does not exist in the map. */
-    def apply(k: K): Observable[V] = map.get(k).get
+    def apply(k: K): Observable[V] = Observable[V] { subscriber =>
+      val (subject, subscribers) =
+        allSubscribers.getOrElseUpdate(k, (Subject[V](), mutable.Set.empty[Subscriber[V]]))
+      subscribers += subscriber
+
+      val subscription = subject.subscribe(subscriber)
+
+      subscriber.add(Subscription {
+        subscription.unsubscribe()
+
+        subscribers -= subscriber
+        if (subscribers.isEmpty) {
+          allSubscribers -= k
+        }
+      })
+    }
+
+    /* return true if there is at least one subscriber which subscribes to the updates of the specific key. */
+    def hasSubscribers(k: K): Boolean = allSubscribers.get(k).isDefined
   }
 
   import scala.collection.mutable.ListBuffer
 
   val rmap = new RMap[String, Int]()
-  rmap("a") = 1
 
-  val o = rmap("a")
-  val buf = ListBuffer.empty[Int]
-  o.subscribe(buf += _)
+  val key = "a"
+  val o = rmap(key)
+  assert(rmap.hasSubscribers(key) == false)
+
+  val buf1 = ListBuffer.empty[Int]
+  val subscription1 = o.subscribe(buf1 += _)
+  val buf2 = ListBuffer.empty[Int]
+  val subscription2 = o.subscribe(buf2 += _)
 
-  rmap("a") = 2
-  rmap("a") = 3
+  rmap(key) = 1
+  rmap(key) = 2
+  assert(buf1 == ListBuffer(1, 2), buf1)
+  assert(buf2 == ListBuffer(1, 2), buf2)
 
-  assert(buf == ListBuffer(2, 3), buf)
+  subscription1.unsubscribe()
+  assert(rmap.hasSubscribers(key))
+  subscription2.unsubscribe()
+  assert(rmap.hasSubscribers(key) == false)
 
 }