From 5565b3e80db960511f5b7de2a6831b13d58a9767 Mon Sep 17 00:00:00 2001
From: Michele Sollecito <michele.sollecito@gmail.com>
Date: Tue, 1 May 2018 19:32:29 +0700
Subject: [PATCH] [CORDA-1411]: Prevent MappedSchema caching from leaking
 memory. (#3042)

---
 .../net/corda/core/schemas/PersistentTypes.kt | 20 +++++++++++++++++++
 docs/source/changelog.rst                     |  1 +
 .../internal/persistence/CordaPersistence.kt  |  3 ++-
 .../persistence/HibernateConfiguration.kt     |  7 +++----
 4 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/core/src/main/kotlin/net/corda/core/schemas/PersistentTypes.kt b/core/src/main/kotlin/net/corda/core/schemas/PersistentTypes.kt
index fbf9b3ccf9..799ed787c7 100644
--- a/core/src/main/kotlin/net/corda/core/schemas/PersistentTypes.kt
+++ b/core/src/main/kotlin/net/corda/core/schemas/PersistentTypes.kt
@@ -42,6 +42,26 @@ open class MappedSchema(schemaFamily: Class<*>,
                         val mappedTypes: Iterable<Class<*>>) {
     val name: String = schemaFamily.name
     override fun toString(): String = "${this.javaClass.simpleName}(name=$name, version=$version)"
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (javaClass != other?.javaClass) return false
+
+        other as MappedSchema
+
+        if (version != other.version) return false
+        if (mappedTypes != other.mappedTypes) return false
+        if (name != other.name) return false
+
+        return true
+    }
+
+    override fun hashCode(): Int {
+        var result = version
+        result = 31 * result + mappedTypes.hashCode()
+        result = 31 * result + name.hashCode()
+        return result
+    }
 }
 //DOCEND MappedSchema
 
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index 53dadbeead..bb7eb2e01d 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -6,6 +6,7 @@ release, see :doc:`upgrade-notes`.
 
 Unreleased
 ==========
+* Avoided a memory leak deriving from incorrect MappedSchema caching strategy.
 
 * Added program line argument ``on-unknown-config-keys`` to allow specifying behaviour on unknown node configuration property keys.
   Values are: [FAIL, WARN, IGNORE], default to FAIL if unspecified.
diff --git a/node-api/src/main/kotlin/net/corda/nodeapi/internal/persistence/CordaPersistence.kt b/node-api/src/main/kotlin/net/corda/nodeapi/internal/persistence/CordaPersistence.kt
index 4884c21832..0c8c5724ed 100644
--- a/node-api/src/main/kotlin/net/corda/nodeapi/internal/persistence/CordaPersistence.kt
+++ b/node-api/src/main/kotlin/net/corda/nodeapi/internal/persistence/CordaPersistence.kt
@@ -25,7 +25,8 @@ data class DatabaseConfig(
         val initialiseSchema: Boolean = true,
         val serverNameTablePrefix: String = "",
         val transactionIsolationLevel: TransactionIsolationLevel = TransactionIsolationLevel.REPEATABLE_READ,
-        val exportHibernateJMXStatistics: Boolean = false
+        val exportHibernateJMXStatistics: Boolean = false,
+        val mappedSchemaCacheSize: Long = 100
 )
 
 // This class forms part of the node config and so any changes to it must be handled with care
diff --git a/node-api/src/main/kotlin/net/corda/nodeapi/internal/persistence/HibernateConfiguration.kt b/node-api/src/main/kotlin/net/corda/nodeapi/internal/persistence/HibernateConfiguration.kt
index ea3c11d252..acc995d3d1 100644
--- a/node-api/src/main/kotlin/net/corda/nodeapi/internal/persistence/HibernateConfiguration.kt
+++ b/node-api/src/main/kotlin/net/corda/nodeapi/internal/persistence/HibernateConfiguration.kt
@@ -1,5 +1,6 @@
 package net.corda.nodeapi.internal.persistence
 
+import com.github.benmanes.caffeine.cache.Caffeine
 import net.corda.core.internal.castIfPossible
 import net.corda.core.schemas.MappedSchema
 import net.corda.core.utilities.contextLogger
@@ -21,7 +22,6 @@ import org.hibernate.type.descriptor.sql.BlobTypeDescriptor
 import org.hibernate.type.descriptor.sql.VarbinaryTypeDescriptor
 import java.lang.management.ManagementFactory
 import java.sql.Connection
-import java.util.concurrent.ConcurrentHashMap
 import javax.management.ObjectName
 import javax.persistence.AttributeConverter
 
@@ -35,8 +35,7 @@ class HibernateConfiguration(
         private val logger = contextLogger()
     }
 
-    // TODO: make this a guava cache or similar to limit ability for this to grow forever.
-    private val sessionFactories = ConcurrentHashMap<Set<MappedSchema>, SessionFactory>()
+    private val sessionFactories = Caffeine.newBuilder().maximumSize(databaseConfig.mappedSchemaCacheSize).build<Set<MappedSchema>, SessionFactory>()
 
     val sessionFactoryForRegisteredSchemas = schemas.let {
         logger.info("Init HibernateConfiguration for schemas: $it")
@@ -44,7 +43,7 @@ class HibernateConfiguration(
     }
 
     /** @param key must be immutable, not just read-only. */
-    fun sessionFactoryForSchemas(key: Set<MappedSchema>) = sessionFactories.computeIfAbsent(key, { makeSessionFactoryForSchemas(key) })
+    fun sessionFactoryForSchemas(key: Set<MappedSchema>): SessionFactory = sessionFactories.get(key, ::makeSessionFactoryForSchemas)!!
 
     private fun makeSessionFactoryForSchemas(schemas: Set<MappedSchema>): SessionFactory {
         logger.info("Creating session factory for schemas: $schemas")