bin/dblink in dblink-0.2 vs bin/dblink in dblink-0.3

@@ -108,17 +108,44 @@
 
 if CLI_OPTS[:verbose]
   puts "Config folder: #{pgb_runner.tmp_folder}"
 end
 
+pgb_runner_at_exit = Proc.new do
+  print "Stoping pgbouncer ..."
+  pgb_runner.stop_async
+  puts
+end
+
+at_exit do
+  pgb_runner_at_exit.call
+end
+
 print "Starting pgbouncer "
 pgb_runner.run_async
 pgb_pid = pgb_runner.wait_for_pid
 puts "- " + "OK".green + " (pid: #{pgb_pid})"
 
 print "Starting SSH tonnel "
 tonnel = TonnelRunner.new(pgb_runner.pgb_port, tmp_dir: pgb_runner.tmp_folder, verbose: CLI_OPTS[:verbose])
+
+tonnel_at_exit = Proc.new do
+  print "Stoping ssh ..."
+  tonnel.stop_async
+  puts
+
+  tonnel.runner_thread.kill
+
+  puts "Removing tmp dir ..."
+  FileUtils.rm_rf(pgb_runner.tmp_folder)
+end
+
+at_exit do
+  pgb_runner_at_exit.call
+  tonnel_at_exit.call
+end
+
 tonnel.run_async
 tonnel.wait_remote_port_allocated
 
 puts "- " + "OK".green + " (pid: #{tonnel.runner_pid})"
 puts "Remote port: #{tonnel.remote_port}"
@@ -166,20 +193,13 @@
   puts "    #{web_response['web_url']}"
   puts
 end
 
 at_exit do
-  print "Stoping pgbouncer ..."
-  pgb_runner.stop_async
-  puts
-  print "Stoping ssh ..."
-  tonnel.stop_async
-  puts
+  pgb_runner_at_exit.call
+  tonnel_at_exit.call
 
   tonnel.runner_thread.kill
-
-  puts "Removing tmp dir ..."
-  FileUtils.rm_rf(pgb_runner.tmp_folder)
 
   exit!
 end
 
 tonnel.runner_thread.value