ENT-2848 More contract attachment caching to avoid database queries slowing the node down (#4433)

core/src/main/kotlin/net/corda/core/node/services/AttachmentStorage.kt

@@ -78,12 +78,21 @@ interface AttachmentStorage {
 
     /**
      * Find the Attachment Id of the contract attachment with the highest version for a given contract class name
-     * from trusted upload sources.
+     * from trusted upload sources.  If both a signed and unsigned attachment exist, prefer the signed one.
      *
      * @param contractClassName The fully qualified name of the contract class.
      * @param minContractVersion The minimum contract version that should be returned.
      * @return the [AttachmentId] of the contract, or null if none meet the criteria.
      */
     fun getContractAttachmentWithHighestContractVersion(contractClassName: String, minContractVersion: Int): AttachmentId?
+
+    /**
+     * Find the Attachment Ids of the contract attachments for a given contract class name
+     * from trusted upload sources.
+     *
+     * @param contractClassName The fully qualified name of the contract class.
+     * @return the [AttachmentId]s of the contract attachments, or an empty set if none meet the criteria.
+     */
+    fun getContractAttachments(contractClassName: String): Set<AttachmentId>
 }
 

core/src/main/kotlin/net/corda/core/transactions/TransactionBuilder.kt

@@ -256,8 +256,7 @@ open class TransactionBuilder @JvmOverloads constructor(
         val outputHashConstraints = outputStates?.filter { it.constraint is HashAttachmentConstraint } ?: emptyList()
         val outputSignatureConstraints = outputStates?.filter { it.constraint is SignatureAttachmentConstraint } ?: emptyList()
         if (inputsHashConstraints.isNotEmpty() && (outputHashConstraints.isNotEmpty() || outputSignatureConstraints.isNotEmpty())) {
-            val attachmentQueryCriteria = AttachmentQueryCriteria.AttachmentsQueryCriteria(contractClassNamesCondition = Builder.equal(listOf(contractClassName)))
+            val attachmentIds = services.attachments.getContractAttachments(contractClassName)
-            val attachmentIds = services.attachments.queryAttachments(attachmentQueryCriteria)
             // only switchover if we have both signed and unsigned attachments for the given contract class name
             if (attachmentIds.isNotEmpty() && attachmentIds.size == 2)  {
                 val attachmentsToUse = attachmentIds.map {

node/src/main/kotlin/net/corda/node/services/persistence/NodeAttachmentService.kt

@@ -10,6 +10,7 @@ import net.corda.core.CordaRuntimeException
 import net.corda.core.contracts.Attachment
 import net.corda.core.contracts.ContractAttachment
 import net.corda.core.contracts.ContractClassName
+import net.corda.core.contracts.Version
 import net.corda.core.crypto.SecureHash
 import net.corda.core.crypto.sha256
 import net.corda.core.internal.*
@@ -54,7 +55,6 @@ class NodeAttachmentService(
         cacheFactory: NamedCacheFactory,
         private val database: CordaPersistence
 ) : AttachmentStorageInternal, SingletonSerializeAsToken() {
-
     // This is to break the circular dependency.
     lateinit var servicesForResolution: ServicesForResolution
 
@@ -84,7 +84,7 @@ class NodeAttachmentService(
     }
 
     @Entity
-    @Table(name = "${NODE_DATABASE_PREFIX}attachments", indexes = [Index(name = "att_id_idx", columnList = "att_id")])
+    @Table(name = "${NODE_DATABASE_PREFIX}attachments", indexes = [(Index(name = "att_id_idx", columnList = "att_id"))])
     class DBAttachment(
             @Id
             @Column(name = "att_id", nullable = false)
@@ -403,10 +403,29 @@ class NodeAttachmentService(
         }
     }
 
-    private val contractsCache = InfrequentlyMutatedCache<String, NavigableMap<Int, AttachmentId>>("NodeAttachmentService_contractAttachmentVersions", cacheFactory)
+    // Holds onto a signed and/or unsigned attachment (at least one or the other).
+    private data class AttachmentIds(val signed: AttachmentId?, val unsigned: AttachmentId?) {
+        init {
+            // One of them at least must exist.
+            check(signed != null || unsigned != null)
+        }
 
-    override fun getContractAttachmentWithHighestContractVersion(contractClassName: String, minContractVersion: Int): AttachmentId? {
+        fun toList(): List<AttachmentId> =
-        val versions: NavigableMap<Int, AttachmentId> = contractsCache.get(contractClassName) { name ->
+                if(signed != null) {
+                    if(unsigned != null) {
+                        listOf(signed, unsigned)
+                    } else listOf(signed)
+                } else listOf(unsigned!!)
+    }
+
+    /**
+     * This caches contract attachment versions by contract class name.  For each version, we support one signed and one unsigned attachment, since that is allowed.
+     *
+     * It is correctly invalidated as new attachments are uploaded.
+     */
+    private val contractsCache = InfrequentlyMutatedCache<ContractClassName, NavigableMap<Version, AttachmentIds>>("NodeAttachmentService_contractAttachmentVersions", cacheFactory)
+
+    private fun getContractAttachmentVersions(contractClassName: String): NavigableMap<Version, AttachmentIds> = contractsCache.get(contractClassName) { name ->
         val attachmentQueryCriteria = AttachmentQueryCriteria.AttachmentsQueryCriteria(contractClassNamesCondition = Builder.equal(listOf(name)),
                 versionCondition = Builder.greaterThanOrEqual(0), uploaderCondition = Builder.`in`(TRUSTED_UPLOADERS))
         val attachmentSort = AttachmentSort(listOf(AttachmentSort.AttachmentSortColumn(AttachmentSort.AttachmentSortAttribute.VERSION, Sort.Direction.DESC)))
@@ -426,9 +445,26 @@ class NodeAttachmentService(
             val query = session.createQuery(criteriaQuery)
 
             // execution
-                TreeMap(query.resultList.map { it.version to AttachmentId.parse(it.attId) }.toMap())
+            TreeMap(query.resultList.groupBy { it.version }.map { makeAttachmentIds(it) }.toMap())
         }
     }
-        return versions.tailMap(minContractVersion, true).lastEntry()?.value
+
+    private fun makeAttachmentIds(it: Map.Entry<Int, List<DBAttachment>>): Pair<Version, AttachmentIds> {
+        check(it.value.size <= 2)
+        val signed = it.value.filter { it.signers?.isNotEmpty() ?: false }.map { AttachmentId.parse(it.attId) }.singleOrNull()
+        val unsigned = it.value.filter { it.signers?.isEmpty() ?: true }.map { AttachmentId.parse(it.attId) }.singleOrNull()
+        return it.key to AttachmentIds(signed, unsigned)
     }
+
+    override fun getContractAttachmentWithHighestContractVersion(contractClassName: String, minContractVersion: Int): AttachmentId? {
+        val versions: NavigableMap<Version, AttachmentIds> = getContractAttachmentVersions(contractClassName)
+        val newestAttachmentIds = versions.tailMap(minContractVersion, true).lastEntry()?.value
+        return newestAttachmentIds?.toList()?.first()
+    }
+
+    override fun getContractAttachments(contractClassName: String): Set<AttachmentId> {
+        val versions: NavigableMap<Version, AttachmentIds> = getContractAttachmentVersions(contractClassName)
+        return versions.values.flatMap { it.toList() }.toSet()
+    }
+
 }

testing/test-utils/src/main/kotlin/net/corda/testing/services/MockAttachmentStorage.kt

@@ -29,7 +29,6 @@ import java.util.jar.JarInputStream
  * A mock implementation of [AttachmentStorage] for use within tests
  */
 class MockAttachmentStorage : AttachmentStorage, SingletonSerializeAsToken() {
-
     private data class ContractAttachmentMetadata(val name: ContractClassName, val version: Int, val isSigned: Boolean)
 
     private val _files = HashMap<SecureHash, Pair<Attachment, ByteArray>>()
@@ -123,4 +122,9 @@ class MockAttachmentStorage : AttachmentStorage, SingletonSerializeAsToken() {
         val attachmentSort = AttachmentSort(listOf(AttachmentSort.AttachmentSortColumn(AttachmentSort.AttachmentSortAttribute.VERSION, Sort.Direction.DESC)))
         return queryAttachments(attachmentQueryCriteria, attachmentSort).firstOrNull()
     }
+
+    override fun getContractAttachments(contractClassName: String): Set<AttachmentId> {
+        val attachmentQueryCriteria = AttachmentQueryCriteria.AttachmentsQueryCriteria(contractClassNamesCondition = Builder.equal(listOf(contractClassName)))
+        return queryAttachments(attachmentQueryCriteria).toSet()
+    }
 }