From e9c2af661ddccd9cb6334594cb1f0f6a6acadad7 Mon Sep 17 00:00:00 2001
From: Chris Rankin <chris.rankin@r3.com>
Date: Wed, 19 Feb 2020 15:26:01 +0000
Subject: [PATCH] NOTICK: Tidy up driver shutdown so that resources are always
 released on error.

---
 .../testing/node/internal/DriverDSLImpl.kt    | 52 +++++++++----------
 1 file changed, 26 insertions(+), 26 deletions(-)

diff --git a/testing/node-driver/src/main/kotlin/net/corda/testing/node/internal/DriverDSLImpl.kt b/testing/node-driver/src/main/kotlin/net/corda/testing/node/internal/DriverDSLImpl.kt
index 48c602d359..b7e5c9123f 100644
--- a/testing/node-driver/src/main/kotlin/net/corda/testing/node/internal/DriverDSLImpl.kt
+++ b/testing/node-driver/src/main/kotlin/net/corda/testing/node/internal/DriverDSLImpl.kt
@@ -1129,18 +1129,18 @@ fun <DI : DriverDSL, D : InternalDriverDSL, A> genericDriver(
         coerce: (D) -> DI,
         dsl: DI.() -> A
 ): A {
-    val serializationEnv = setDriverSerialization(driverDsl.cordappsClassLoader)
-    val shutdownHook = addShutdownHook(driverDsl::shutdown)
-    try {
-        driverDsl.start()
-        return dsl(coerce(driverDsl))
-    } catch (exception: Throwable) {
-        DriverDSLImpl.log.error("Driver shutting down because of exception", exception)
-        throw exception
-    } finally {
-        driverDsl.shutdown()
-        shutdownHook.cancel()
-        serializationEnv?.close()
+    setDriverSerialization(driverDsl.cordappsClassLoader).use { _ ->
+        val shutdownHook = addShutdownHook(driverDsl::shutdown)
+        try {
+            driverDsl.start()
+            return dsl(coerce(driverDsl))
+        } catch (exception: Throwable) {
+            DriverDSLImpl.log.error("Driver shutting down because of exception", exception)
+            throw exception
+        } finally {
+            driverDsl.shutdown()
+            shutdownHook.cancel()
+        }
     }
 }
 
@@ -1157,8 +1157,8 @@ fun <DI : DriverDSL, D : InternalDriverDSL, A> genericDriver(
         driverDslWrapper: (DriverDSLImpl) -> D,
         coerce: (D) -> DI, dsl: DI.() -> A
 ): A {
-    val serializationEnv = setDriverSerialization()
-    val driverDsl = driverDslWrapper(
+    setDriverSerialization().use { _ ->
+        val driverDsl = driverDslWrapper(
             DriverDSLImpl(
                 portAllocation = defaultParameters.portAllocation,
                 debugPortAllocation = defaultParameters.debugPortAllocation,
@@ -1180,18 +1180,18 @@ fun <DI : DriverDSL, D : InternalDriverDSL, A> genericDriver(
                 djvmCordaSource = defaultParameters.djvmCordaSource,
                 environmentVariables = defaultParameters.environmentVariables
             )
-    )
-    val shutdownHook = addShutdownHook(driverDsl::shutdown)
-    try {
-        driverDsl.start()
-        return dsl(coerce(driverDsl))
-    } catch (exception: Throwable) {
-        DriverDSLImpl.log.error("Driver shutting down because of exception", exception)
-        throw exception
-    } finally {
-        driverDsl.shutdown()
-        shutdownHook.cancel()
-        serializationEnv?.close()
+        )
+        val shutdownHook = addShutdownHook(driverDsl::shutdown)
+        try {
+            driverDsl.start()
+            return dsl(coerce(driverDsl))
+        } catch (exception: Throwable) {
+            DriverDSLImpl.log.error("Driver shutting down because of exception", exception)
+            throw exception
+        } finally {
+            driverDsl.shutdown()
+            shutdownHook.cancel()
+        }
     }
 }