From 8cb27f105276d9588e278ee459abaf638602f4c1 Mon Sep 17 00:00:00 2001
From: Ross Nicoll <ross.nicoll@r3.com>
Date: Thu, 2 Feb 2017 16:10:10 +0000
Subject: [PATCH] Switch to using CompositeKey inside network map cache

Switch to using CompositeKey inside network map cache, so that look up of node info from keys
can be done in a more sensible manner.
---
 .../node/services/network/InMemoryNetworkMapCache.kt   | 10 +++++-----
 .../net/corda/testing/node/MockNetworkMapCache.kt      |  8 ++++----
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/node/src/main/kotlin/net/corda/node/services/network/InMemoryNetworkMapCache.kt b/node/src/main/kotlin/net/corda/node/services/network/InMemoryNetworkMapCache.kt
index 892194c24f..ae30fe3cfd 100644
--- a/node/src/main/kotlin/net/corda/node/services/network/InMemoryNetworkMapCache.kt
+++ b/node/src/main/kotlin/net/corda/node/services/network/InMemoryNetworkMapCache.kt
@@ -54,10 +54,10 @@ open class InMemoryNetworkMapCache : SingletonSerializeAsToken(), NetworkMapCach
     override val mapServiceRegistered: ListenableFuture<Unit> get() = _registrationFuture
 
     private var registeredForPush = false
-    protected var registeredNodes: MutableMap<Party, NodeInfo> = Collections.synchronizedMap(HashMap<Party, NodeInfo>())
+    protected var registeredNodes: MutableMap<CompositeKey, NodeInfo> = Collections.synchronizedMap(HashMap())
 
     override fun getPartyInfo(party: Party): PartyInfo? {
-        val node = registeredNodes[party]
+        val node = registeredNodes[party.owningKey]
         if (node != null) {
             return PartyInfo.Node(node)
         }
@@ -71,7 +71,7 @@ open class InMemoryNetworkMapCache : SingletonSerializeAsToken(), NetworkMapCach
         return null
     }
 
-    override fun getNodeByLegalIdentityKey(compositeKey: CompositeKey): NodeInfo? = registeredNodes[Party("", compositeKey)]
+    override fun getNodeByLegalIdentityKey(compositeKey: CompositeKey): NodeInfo? = registeredNodes[compositeKey]
 
     override fun track(): Pair<List<NodeInfo>, Observable<MapChange>> {
         synchronized(_changed) {
@@ -113,7 +113,7 @@ open class InMemoryNetworkMapCache : SingletonSerializeAsToken(), NetworkMapCach
 
     override fun addNode(node: NodeInfo) {
         synchronized(_changed) {
-            val previousNode = registeredNodes.put(node.legalIdentity, node)
+            val previousNode = registeredNodes.put(node.legalIdentity.owningKey, node)
             if (previousNode == null) {
                 changePublisher.onNext(MapChange.Added(node))
             } else if (previousNode != node) {
@@ -124,7 +124,7 @@ open class InMemoryNetworkMapCache : SingletonSerializeAsToken(), NetworkMapCach
 
     override fun removeNode(node: NodeInfo) {
         synchronized(_changed) {
-            registeredNodes.remove(node.legalIdentity)
+            registeredNodes.remove(node.legalIdentity.owningKey)
             changePublisher.onNext(MapChange.Removed(node))
         }
     }
diff --git a/test-utils/src/main/kotlin/net/corda/testing/node/MockNetworkMapCache.kt b/test-utils/src/main/kotlin/net/corda/testing/node/MockNetworkMapCache.kt
index a961f264a1..c200fc3cf8 100644
--- a/test-utils/src/main/kotlin/net/corda/testing/node/MockNetworkMapCache.kt
+++ b/test-utils/src/main/kotlin/net/corda/testing/node/MockNetworkMapCache.kt
@@ -21,8 +21,8 @@ class MockNetworkMapCache() : InMemoryNetworkMapCache() {
     init {
         val mockNodeA = NodeInfo(MockAddress("bankC:8080"), Party("Bank C", DummyPublicKey("Bank C")))
         val mockNodeB = NodeInfo(MockAddress("bankD:8080"), Party("Bank D", DummyPublicKey("Bank D")))
-        registeredNodes[mockNodeA.legalIdentity] = mockNodeA
-        registeredNodes[mockNodeB.legalIdentity] = mockNodeB
+        registeredNodes[mockNodeA.legalIdentity.owningKey] = mockNodeA
+        registeredNodes[mockNodeB.legalIdentity.owningKey] = mockNodeB
         runWithoutMapService()
     }
 
@@ -32,7 +32,7 @@ class MockNetworkMapCache() : InMemoryNetworkMapCache() {
      */
     @VisibleForTesting
     fun addRegistration(node: NodeInfo) {
-        registeredNodes[node.legalIdentity] = node
+        registeredNodes[node.legalIdentity.owningKey] = node
     }
 
     /**
@@ -41,6 +41,6 @@ class MockNetworkMapCache() : InMemoryNetworkMapCache() {
      */
     @VisibleForTesting
     fun deleteRegistration(legalIdentity: Party) : Boolean {
-        return registeredNodes.remove(legalIdentity) != null
+        return registeredNodes.remove(legalIdentity.owningKey) != null
     }
 }