ENT-2378 ENT-2377 Minimise database access during transaction creation and recording for public key handling (#3812)

* ENT-2378 ENT-2377 minimise database access by making more use of caches for filterMyKeys, and don't make a redundant database check for existence of new random keys.

* ENT-2378 Unit test for new method in PersistentIdentityService.

node/src/main/kotlin/net/corda/node/internal/AbstractNode.kt

@@ -314,6 +314,7 @@ abstract class AbstractNode<S>(val configuration: NodeConfiguration,
         }
 
         val (nodeInfo, signedNodeInfo) = nodeInfoAndSigned
+        identityService.ourNames = nodeInfo.legalIdentities.map { it.name }.toSet()
         services.start(nodeInfo, netParams)
         networkMapUpdater.start(trustRoot, signedNetParams.raw.hash, signedNodeInfo.raw.hash)
         startMessagingService(rpcOps, nodeInfo, myNotaryIdentity, netParams)
@@ -748,7 +749,7 @@ abstract class AbstractNode<S>(val configuration: NodeConfiguration,
         }
     }
 
-    protected open fun makeKeyManagementService(identityService: IdentityService): KeyManagementServiceInternal {
+    protected open fun makeKeyManagementService(identityService: PersistentIdentityService): KeyManagementServiceInternal {
         // Place the long term identity key in the KMS. Eventually, this is likely going to be separated again because
         // the KMS is meant for derived temporary keys used in transactions, and we're not supposed to sign things with
         // the identity key. But the infrastructure to make that easy isn't here yet.

node/src/main/kotlin/net/corda/node/services/api/IdentityServiceInternal.kt

@@ -21,10 +21,21 @@ interface IdentityServiceInternal : IdentityService {
     }
 
     /** This method exists so it can be mocked with doNothing, rather than having to make up a possibly invalid return value. */
-    fun justVerifyAndRegisterIdentity(identity: PartyAndCertificate) {
+    fun justVerifyAndRegisterIdentity(identity: PartyAndCertificate, isNewRandomIdentity: Boolean = false) {
-        verifyAndRegisterIdentity(identity)
+        verifyAndRegisterIdentity(identity, isNewRandomIdentity)
     }
 
+    /**
+     * Verify and then store an identity.
+     *
+     * @param identity a party and the certificate path linking them to the network trust root.
+     * @param isNewRandomIdentity true if the identity will not have been registered before (e.g. because it is randomly generated by ourselves).
+     * @return the issuing entity, if known.
+     * @throws IllegalArgumentException if the certificate path is invalid.
+     */
+    @Throws(CertificateExpiredException::class, CertificateNotYetValidException::class, InvalidAlgorithmParameterException::class)
+    fun verifyAndRegisterIdentity(identity: PartyAndCertificate, isNewRandomIdentity: Boolean): PartyAndCertificate?
+
     fun partiesFromName(query: String, exactMatch: Boolean, x500name: CordaX500Name, results: LinkedHashSet<Party>, party: Party) {
         val components = listOfNotNull(x500name.commonName, x500name.organisationUnit, x500name.organisation, x500name.locality, x500name.state, x500name.country)
         components.forEach { component ->
@@ -48,9 +59,10 @@ interface IdentityServiceInternal : IdentityService {
      *
      * @param trustAnchor The trust anchor that will verify the identity's validity
      * @param identity The identity to verify
+     * @param isNewRandomIdentity true if the identity will not have been registered before (e.g. because it is randomly generated by ourselves).
      */
     @Throws(CertificateExpiredException::class, CertificateNotYetValidException::class, InvalidAlgorithmParameterException::class)
-    fun verifyAndRegisterIdentity(trustAnchor: TrustAnchor, identity: PartyAndCertificate): PartyAndCertificate? {
+    fun verifyAndRegisterIdentity(trustAnchor: TrustAnchor, identity: PartyAndCertificate, isNewRandomIdentity: Boolean = false): PartyAndCertificate? {
         // Validate the chain first, before we do anything clever with it
         val identityCertChain = identity.certPath.x509Certificates
         try {
@@ -71,8 +83,8 @@ interface IdentityServiceInternal : IdentityService {
             val firstPath = X509Utilities.buildCertPath(identityCertChain.slice(idx until identityCertChain.size))
             verifyAndRegisterIdentity(trustAnchor, PartyAndCertificate(firstPath))
         }
-        return registerIdentity(identity)
+        return registerIdentity(identity, isNewRandomIdentity)
     }
 
-    fun registerIdentity(identity: PartyAndCertificate): PartyAndCertificate?
+    fun registerIdentity(identity: PartyAndCertificate, isNewRandomIdentity: Boolean = false): PartyAndCertificate?
 }

node/src/main/kotlin/net/corda/node/services/identity/InMemoryIdentityService.kt

@@ -42,7 +42,10 @@ class InMemoryIdentityService(identities: List<PartyAndCertificate> = emptyList(
     @Throws(CertificateExpiredException::class, CertificateNotYetValidException::class, InvalidAlgorithmParameterException::class)
     override fun verifyAndRegisterIdentity(identity: PartyAndCertificate): PartyAndCertificate? = verifyAndRegisterIdentity(trustAnchor, identity)
 
-    override fun registerIdentity(identity: PartyAndCertificate): PartyAndCertificate? {
+    @Throws(CertificateExpiredException::class, CertificateNotYetValidException::class, InvalidAlgorithmParameterException::class)
+    override fun verifyAndRegisterIdentity(identity: PartyAndCertificate, isNewRandomIdentity: Boolean): PartyAndCertificate? = verifyAndRegisterIdentity(trustAnchor, identity)
+
+    override fun registerIdentity(identity: PartyAndCertificate, isNewRandomIdentity: Boolean): PartyAndCertificate? {
         val identityCertChain = identity.certPath.x509Certificates
         log.trace { "Registering identity $identity" }
         keyToParties[identity.owningKey] = identity

node/src/main/kotlin/net/corda/node/services/identity/PersistentIdentityService.kt

@@ -122,16 +122,26 @@ class PersistentIdentityService : SingletonSerializeAsToken(), IdentityServiceIn
 
     @Throws(CertificateExpiredException::class, CertificateNotYetValidException::class, InvalidAlgorithmParameterException::class)
     override fun verifyAndRegisterIdentity(identity: PartyAndCertificate): PartyAndCertificate? {
+        return verifyAndRegisterIdentity(identity, false)
+    }
+
+    @Throws(CertificateExpiredException::class, CertificateNotYetValidException::class, InvalidAlgorithmParameterException::class)
+    override fun verifyAndRegisterIdentity(identity: PartyAndCertificate, isNewRandomIdentity: Boolean): PartyAndCertificate? {
         return database.transaction {
-            verifyAndRegisterIdentity(trustAnchor, identity)
+            verifyAndRegisterIdentity(trustAnchor, identity, isNewRandomIdentity)
         }
     }
 
-    override fun registerIdentity(identity: PartyAndCertificate): PartyAndCertificate? {
+    override fun registerIdentity(identity: PartyAndCertificate, isNewRandomIdentity: Boolean): PartyAndCertificate? {
         val identityCertChain = identity.certPath.x509Certificates
         log.debug { "Registering identity $identity" }
         val key = mapToKey(identity)
-        keyToParties.addWithDuplicatesAllowed(key, identity)
+        if (isNewRandomIdentity) {
+            // Because this is supposed to be new and random, there's no way we have it in the database already, so skip the pessimistic check.
+            keyToParties.set(key, identity)
+        } else {
+            keyToParties.addWithDuplicatesAllowed(key, identity)
+        }
         // Always keep the first party we registered, as that's the well known identity
         principalToParties.addWithDuplicatesAllowed(identity.name, key, false)
         val parentId = mapToKey(identityCertChain[1].publicKey)
@@ -168,4 +178,15 @@ class PersistentIdentityService : SingletonSerializeAsToken(), IdentityServiceIn
 
     @Throws(UnknownAnonymousPartyException::class)
     override fun assertOwnership(party: Party, anonymousParty: AnonymousParty) = database.transaction { super.assertOwnership(party, anonymousParty) }
+
+    lateinit var ourNames: Set<CordaX500Name>
+
+    // Allows us to cheaply eliminate keys we know belong to others by using the cache contents without triggering loading.
+    fun stripCachedPeerKeys(keys: Iterable<PublicKey>): Iterable<PublicKey> {
+        return keys.filter {
+            val party = keyToParties.getIfCached(mapToKey(it))?.party?.name
+            party == null || party in ourNames
+        }
+    }
+
 }

node/src/main/kotlin/net/corda/node/services/keys/KMSUtils.kt

@@ -46,7 +46,7 @@ fun freshCertificate(identityService: IdentityService,
     val ourCertPath = X509Utilities.buildCertPath(ourCertificate, issuer.certPath.x509Certificates)
     val anonymisedIdentity = PartyAndCertificate(ourCertPath)
     if (identityService is IdentityServiceInternal) {
-        identityService.justVerifyAndRegisterIdentity(anonymisedIdentity)
+        identityService.justVerifyAndRegisterIdentity(anonymisedIdentity, true)
     } else {
         identityService.verifyAndRegisterIdentity(anonymisedIdentity)
     }

node/src/main/kotlin/net/corda/node/services/keys/PersistentKeyManagementService.kt

@@ -2,9 +2,9 @@ package net.corda.node.services.keys
 
 import net.corda.core.crypto.*
 import net.corda.core.identity.PartyAndCertificate
-import net.corda.core.node.services.IdentityService
 import net.corda.core.serialization.SingletonSerializeAsToken
 import net.corda.core.utilities.MAX_HASH_HEX_SIZE
+import net.corda.node.services.identity.PersistentIdentityService
 import net.corda.node.utilities.AppendOnlyPersistentMap
 import net.corda.nodeapi.internal.persistence.CordaPersistence
 import net.corda.nodeapi.internal.persistence.NODE_DATABASE_PREFIX
@@ -25,9 +25,8 @@ import javax.persistence.Lob
  *
  * This class needs database transactions to be in-flight during method calls and init.
  */
-class PersistentKeyManagementService(val identityService: IdentityService,
+class PersistentKeyManagementService(val identityService: PersistentIdentityService,
                                      private val database: CordaPersistence) : SingletonSerializeAsToken(), KeyManagementServiceInternal {
-
     @Entity
     @javax.persistence.Table(name = "${NODE_DATABASE_PREFIX}our_key_pairs")
     class PersistentKey(
@@ -70,7 +69,7 @@ class PersistentKeyManagementService(val identityService: IdentityService,
     override val keys: Set<PublicKey> get() = database.transaction { keysMap.allPersisted().map { it.first }.toSet() }
 
     override fun filterMyKeys(candidateKeys: Iterable<PublicKey>): Iterable<PublicKey> = database.transaction {
-        candidateKeys.filter { keysMap[it] != null }
+        identityService.stripCachedPeerKeys(candidateKeys).filter { keysMap[it] != null } // TODO: bulk cache access.
     }
 
     override fun freshKey(): PublicKey {

node/src/main/kotlin/net/corda/node/utilities/AppendOnlyPersistentMap.kt

@@ -136,6 +136,14 @@ abstract class AppendOnlyPersistentMapBase<K, V, E, out EK>(
 
     operator fun contains(key: K) = get(key) != null
 
+    /**
+     * Allow checking the cache content without falling back to database if there's a miss.
+     *
+     * @param key The cache key
+     * @return The value in the cache, or null if not present.
+     */
+    fun getIfCached(key: K): V? = cache.getIfPresent(key!!)?.value
+
     /**
      * Removes all of the mappings from this map and underlying storage. The map will be empty after this call returns.
      * WARNING!! The method is not thread safe.

node/src/test/kotlin/net/corda/node/services/identity/PersistentIdentityServiceTests.kt

@@ -54,6 +54,7 @@ class PersistentIdentityServiceTests {
                 identityService::wellKnownPartyFromAnonymous
         )
         identityService.database = database
+        identityService.ourNames = setOf(ALICE_NAME)
         identityService.start(DEV_ROOT_CA.certificate)
     }
 
@@ -93,6 +94,25 @@ class PersistentIdentityServiceTests {
         assertNull(identityService.wellKnownPartyFromX500Name(ALICE.name))
     }
 
+    @Test
+    fun `stripping others when none registered does not strip`() {
+        assertEquals(identityService.stripCachedPeerKeys(listOf(BOB_PUBKEY)).first(), BOB_PUBKEY)
+    }
+
+    @Test
+    fun `stripping others when only us registered does not strip`() {
+        identityService.verifyAndRegisterIdentity(ALICE_IDENTITY)
+        assertEquals(identityService.stripCachedPeerKeys(listOf(BOB_PUBKEY)).first(), BOB_PUBKEY)
+    }
+
+    @Test
+    fun `stripping others when us and others registered does not strip us`() {
+        identityService.verifyAndRegisterIdentity(ALICE_IDENTITY)
+        identityService.verifyAndRegisterIdentity(BOB_IDENTITY)
+        val stripped = identityService.stripCachedPeerKeys(listOf(ALICE_PUBKEY, BOB_PUBKEY))
+        assertEquals(stripped.single(), ALICE_PUBKEY)
+    }
+
     @Test
     fun `get identity by substring match`() {
         identityService.verifyAndRegisterIdentity(ALICE_IDENTITY)

node/src/test/kotlin/net/corda/node/services/persistence/HibernateConfigurationTest.kt

@@ -33,8 +33,8 @@ import net.corda.node.internal.configureDatabase
 import net.corda.node.services.api.IdentityServiceInternal
 import net.corda.node.services.api.WritableTransactionStorage
 import net.corda.node.services.schema.ContractStateAndRef
-import net.corda.node.services.schema.PersistentStateService
 import net.corda.node.services.schema.NodeSchemaService
+import net.corda.node.services.schema.PersistentStateService
 import net.corda.node.services.vault.NodeVaultService
 import net.corda.node.services.vault.VaultSchemaV1
 import net.corda.nodeapi.internal.persistence.CordaPersistence
@@ -119,7 +119,7 @@ class HibernateConfigurationTest {
 
             // `consumeCash` expects we can self-notarise transactions
             services = object : MockServices(cordappPackages, BOB_NAME, rigorousMock<IdentityServiceInternal>().also {
-                doNothing().whenever(it).justVerifyAndRegisterIdentity(argThat { name == BOB_NAME })
+                doNothing().whenever(it).justVerifyAndRegisterIdentity(argThat { name == BOB_NAME }, any())
             }, generateKeyPair(), dummyNotary.keyPair) {
                 override val vaultService = NodeVaultService(Clock.systemUTC(), keyManagementService, servicesForResolution, database, schemaService).apply { start() }
                 override fun recordTransactions(statesToRecord: StatesToRecord, txs: Iterable<SignedTransaction>) {

node/src/test/kotlin/net/corda/node/services/vault/NodeVaultServiceTest.kt

@@ -1,6 +1,7 @@
 package net.corda.node.services.vault
 
 import co.paralleluniverse.fibers.Suspendable
+import com.nhaarman.mockito_kotlin.any
 import com.nhaarman.mockito_kotlin.argThat
 import com.nhaarman.mockito_kotlin.doNothing
 import com.nhaarman.mockito_kotlin.whenever
@@ -600,7 +601,7 @@ class NodeVaultServiceTest {
         assertEquals(services.identityService.partyFromKey(identity.owningKey), identity.party)
         val anonymousIdentity = services.keyManagementService.freshKeyAndCert(identity, false)
         val thirdPartyServices = MockServices(emptyList(), MEGA_CORP.name, rigorousMock<IdentityServiceInternal>().also {
-            doNothing().whenever(it).justVerifyAndRegisterIdentity(argThat { name == MEGA_CORP.name })
+            doNothing().whenever(it).justVerifyAndRegisterIdentity(argThat { name == MEGA_CORP.name }, any())
         })
         val thirdPartyIdentity = thirdPartyServices.keyManagementService.freshKeyAndCert(thirdPartyServices.myInfo.singleIdentityAndCert(), false)
         val amount = Amount(1000, Issued(BOC.ref(1), GBP))

testing/node-driver/src/main/kotlin/net/corda/testing/node/internal/InternalMockNetwork.kt

@@ -22,7 +22,6 @@ import net.corda.core.messaging.SingleMessageRecipient
 import net.corda.core.node.NetworkParameters
 import net.corda.core.node.NodeInfo
 import net.corda.core.node.NotaryInfo
-import net.corda.core.node.services.IdentityService
 import net.corda.core.serialization.SerializationWhitelist
 import net.corda.core.utilities.NetworkHostAndPort
 import net.corda.core.utilities.contextLogger
@@ -37,6 +36,7 @@ import net.corda.node.services.api.FlowStarter
 import net.corda.node.services.api.ServiceHubInternal
 import net.corda.node.services.api.StartedNodeServices
 import net.corda.node.services.config.*
+import net.corda.node.services.identity.PersistentIdentityService
 import net.corda.node.services.keys.E2ETestKeyManagementService
 import net.corda.node.services.keys.KeyManagementServiceInternal
 import net.corda.node.services.messaging.Message
@@ -386,7 +386,7 @@ open class InternalMockNetwork(defaultParameters: MockNetworkParameters = MockNe
             (network as MockNodeMessagingService).spy = spy
         }
 
-        override fun makeKeyManagementService(identityService: IdentityService): KeyManagementServiceInternal {
+        override fun makeKeyManagementService(identityService: PersistentIdentityService): KeyManagementServiceInternal {
             return E2ETestKeyManagementService(identityService)
         }