From b7fb87ed832f751f5729fad3fd227990899cf9f7 Mon Sep 17 00:00:00 2001
From: "rick.parker" <rick.parker@r3cev.com>
Date: Fri, 10 Nov 2017 11:40:12 +0000
Subject: [PATCH] Easy all in one launch of UI + SSH tunnels.

---
 .../main/kotlin/net/corda/jmeter/Launcher.kt  | 16 ++++++++++-
 .../src/main/kotlin/net/corda/jmeter/Ssh.kt   | 27 +++++++++++++------
 2 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/tools/jmeter/src/main/kotlin/net/corda/jmeter/Launcher.kt b/tools/jmeter/src/main/kotlin/net/corda/jmeter/Launcher.kt
index a108da4b89..827e76cef4 100644
--- a/tools/jmeter/src/main/kotlin/net/corda/jmeter/Launcher.kt
+++ b/tools/jmeter/src/main/kotlin/net/corda/jmeter/Launcher.kt
@@ -13,6 +13,7 @@ class Launcher {
         @JvmStatic
         fun main(args: Array<String>) {
             val logger = LoggerFactory.getLogger(this::class.java)
+            logger.info("Launcher called with ${args.toList()}")
             val jmeter = JMeter()
             val capsuleDir = System.getProperty("capsule.dir")
             if (capsuleDir != null) {
@@ -44,8 +45,21 @@ class Launcher {
                 }
                 jmeter.start(arrayOf("-s", "-p", (capsuleDirPath / "jmeter.properties").toString()) + extraArgs + args)
             } else {
-                jmeter.start(args)
+                jmeter.start(maybeOpenSshTunnels(args))
             }
         }
+
+        private fun maybeOpenSshTunnels(args: Array<String>): Array<String> {
+            var index = 0
+            for (arg in args) {
+                if (arg == "-Xssh") {
+                    // start ssh
+                    Ssh.main(args.copyOfRange(index + 1, args.size), false)
+                    return if (index == 0) emptyArray() else args.copyOfRange(0, index)
+                }
+                index++
+            }
+            return args
+        }
     }
 }
\ No newline at end of file
diff --git a/tools/jmeter/src/main/kotlin/net/corda/jmeter/Ssh.kt b/tools/jmeter/src/main/kotlin/net/corda/jmeter/Ssh.kt
index 74a0fa1036..47060966ca 100644
--- a/tools/jmeter/src/main/kotlin/net/corda/jmeter/Ssh.kt
+++ b/tools/jmeter/src/main/kotlin/net/corda/jmeter/Ssh.kt
@@ -4,6 +4,7 @@ import com.jcraft.jsch.JSch
 import com.jcraft.jsch.Session
 import net.corda.core.utilities.NetworkHostAndPort
 import net.corda.loadtest.setupJSchWithSshAgent
+import net.corda.nodeapi.internal.addShutdownHook
 import org.slf4j.LoggerFactory
 import java.io.BufferedReader
 import java.io.InputStreamReader
@@ -15,7 +16,8 @@ class Ssh {
         val log = LoggerFactory.getLogger(this::class.java)
 
         @JvmStatic
-        fun main(args: Array<String>) {
+        @JvmOverloads
+        fun main(args: Array<String>, wait: Boolean = true) {
             val userName = System.getProperty("user.name")
             val jsch = setupJSchWithSshAgent()
             val sessions = mutableListOf<Session>()
@@ -58,14 +60,23 @@ class Ssh {
                 // ssh ${remoteHostAndPort.host} -R 0.0.0.0:clientRmiLocalPort:localhost:clientRmiLocalPort -N
                 createInboundTunnel(session, NetworkHostAndPort("0.0.0.0", clientRmiLocalPort), NetworkHostAndPort("localhost", clientRmiLocalPort))
             }
-            val input = BufferedReader(InputStreamReader(System.`in`))
 
-            do {
-                log.info("Type 'quit' to exit cleanly.")
-            } while (input.readLine() != "quit")
-            sessions.forEach {
-                log.info("Closing tunnels for ${it.host}")
-                it.disconnect()
+            if (wait) {
+                val input = BufferedReader(InputStreamReader(System.`in`))
+                do {
+                    log.info("Type 'quit' to exit cleanly.")
+                } while (input.readLine() != "quit")
+                sessions.forEach {
+                    log.info("Closing tunnels for ${it.host}")
+                    it.disconnect()
+                }
+            } else {
+                addShutdownHook {
+                    sessions.forEach {
+                        log.info("Closing tunnels for ${it.host}")
+                        it.disconnect()
+                    }
+                }
             }
         }