Ensure that ServiceHub.WithEntityManager has a database transaction a…

…vailable (corda#5413) * tidy up withEntityManager * address rick review comment * api break
AkaSkyKnight · Aug 30, 2019 · cd0d5c7 · cd0d5c7
1 parent 136600c
commit cd0d5c7
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 10 deletions.
diff --git a/.ci/api-current.txt b/.ci/api-current.txt
@@ -3520,7 +3520,6 @@ public interface net.corda.core.node.ServiceHub extends net.corda.core.node.Serv
   @NotNull
   public abstract net.corda.core.contracts.StateAndRef<T> toStateAndRef(net.corda.core.contracts.StateRef)
   public abstract void withEntityManager(java.util.function.Consumer<javax.persistence.EntityManager>)
-  @NotNull
   public abstract T withEntityManager(kotlin.jvm.functions.Function1<? super javax.persistence.EntityManager, ? extends T>)
 ##
 @DoNotImplement
@@ -7816,7 +7815,6 @@ public class net.corda.testing.node.MockServices extends java.lang.Object implem
   @NotNull
   public net.corda.core.contracts.StateAndRef<T> toStateAndRef(net.corda.core.contracts.StateRef)
   public void withEntityManager(java.util.function.Consumer<javax.persistence.EntityManager>)
-  @NotNull
   public T withEntityManager(kotlin.jvm.functions.Function1<? super javax.persistence.EntityManager, ? extends T>)
   public static final net.corda.testing.node.MockServices$Companion Companion
 ##

diff --git a/core/src/main/kotlin/net/corda/core/node/ServiceHub.kt b/core/src/main/kotlin/net/corda/core/node/ServiceHub.kt
@@ -376,7 +376,7 @@ interface ServiceHub : ServicesForResolution {
      *
      * @param block a lambda function with access to an [EntityManager].
      */
-    fun <T : Any> withEntityManager(block: EntityManager.() -> T): T
+    fun <T : Any?> withEntityManager(block: EntityManager.() -> T): T
 
     /**
      * Exposes the Java Persistence API (JPA) to flows via a restricted [EntityManager]. This method can be used to

diff --git a/node/src/main/kotlin/net/corda/node/internal/AbstractNode.kt b/node/src/main/kotlin/net/corda/node/internal/AbstractNode.kt
@@ -409,7 +409,9 @@ abstract class AbstractNode<S>(val configuration: NodeConfiguration,
             // Shut down the SMM so no Fibers are scheduled.
             runOnStop += { smm.stop(acceptableLiveFiberCountOnStop()) }
             (smm as? StateMachineManagerInternal)?.let {
-                val flowMonitor = FlowMonitor({ smm.snapshot().filter { flow -> flow !in smm.flowHospital }.toSet() }, configuration.flowMonitorPeriodMillis, configuration.flowMonitorSuspensionLoggingThresholdMillis)
+                val flowMonitor = FlowMonitor({
+                    smm.snapshot().filter { flow -> flow !in smm.flowHospital }.toSet()
+                }, configuration.flowMonitorPeriodMillis, configuration.flowMonitorSuspensionLoggingThresholdMillis)
                 runOnStop += flowMonitor::stop
                 flowMonitor.start()
             }
@@ -596,7 +598,7 @@ abstract class AbstractNode<S>(val configuration: NodeConfiguration,
         // Needed because of bug CORDA-2653 - some Corda services can utilise third-party libraries that require access to
         // the Thread context class loader
 
-        val oldContextClassLoader : ClassLoader? = Thread.currentThread().contextClassLoader
+        val oldContextClassLoader: ClassLoader? = Thread.currentThread().contextClassLoader
         try {
             Thread.currentThread().contextClassLoader = cordappLoader.appClassLoader
 
@@ -1032,12 +1034,16 @@ abstract class AbstractNode<S>(val configuration: NodeConfiguration,
 
         override fun jdbcSession(): Connection = database.createSession()
 
-        override fun <T : Any> withEntityManager(block: EntityManager.() -> T): T {
-            return block(contextTransaction.restrictedEntityManager)
+        override fun <T : Any?> withEntityManager(block: EntityManager.() -> T): T {
+            return database.transaction {
+                block(restrictedEntityManager)
+            }
         }
 
         override fun withEntityManager(block: Consumer<EntityManager>) {
-            block.accept(contextTransaction.restrictedEntityManager)
+            withEntityManager {
+                block.accept(this)
+            }
         }
 
         // allows services to register handlers to be informed when the node stop method is called

diff --git a/testing/node-driver/src/main/kotlin/net/corda/testing/node/MockServices.kt b/testing/node-driver/src/main/kotlin/net/corda/testing/node/MockServices.kt
@@ -224,7 +224,7 @@ open class MockServices private constructor(
 
                 override fun jdbcSession(): Connection = persistence.createSession()
 
-                override fun <T : Any> withEntityManager(block: EntityManager.() -> T): T {
+                override fun <T : Any?> withEntityManager(block: EntityManager.() -> T): T {
                     return block(contextTransaction.restrictedEntityManager)
                 }
 
@@ -440,7 +440,7 @@ open class MockServices private constructor(
 
     override fun jdbcSession(): Connection = throw UnsupportedOperationException()
 
-    override fun <T : Any> withEntityManager(block: EntityManager.() -> T): T {
+    override fun <T : Any?> withEntityManager(block: EntityManager.() -> T): T {
         throw UnsupportedOperationException()
     }