From a35ae91b677b36793e95a559e3c1add72779ce2a Mon Sep 17 00:00:00 2001
From: Anders Ingemann <anders@ingemann.de>
Date: Mon, 24 Nov 2014 18:54:31 +0100
Subject: [PATCH] Simplify? Hopefully...

---
 bootstrapvz/base/log.py                    | 29 ++++++++++++++++++++++
 bootstrapvz/base/remote/__init__.py        | 19 ++++++++++----
 bootstrapvz/base/remote/callback.py        |  9 ++-----
 bootstrapvz/base/remote/log.py             | 29 ----------------------
 bootstrapvz/base/remote/server.py          |  2 +-
 bootstrapvz/base/remote/ssh_rpc_manager.py |  6 -----
 6 files changed, 46 insertions(+), 48 deletions(-)
 delete mode 100644 bootstrapvz/base/remote/log.py

diff --git a/bootstrapvz/base/log.py b/bootstrapvz/base/log.py
index b6096e3..b406763 100644
--- a/bootstrapvz/base/log.py
+++ b/bootstrapvz/base/log.py
@@ -88,3 +88,32 @@ class FileFormatter(logging.Formatter):
 	"""
 	def format(self, record):
 		return super(FileFormatter, self).format(record)
+
+
+class LogForwarder(logging.Handler):
+
+	def __init__(self, level=logging.NOTSET):
+		self.server = None
+		super(LogForwarder, self).__init__(level)
+
+	def set_server(self, server):
+		self.server = server
+
+	def emit(self, record):
+		if self.server is not None:
+			if record.exc_info is not None:
+				import traceback
+				exc_type, exc_value, exc_traceback = record.exc_info
+				record.exc_info = traceback.print_exception(exc_type, exc_value, exc_traceback)
+			# TODO: Use serpent instead
+			import pickle
+			self.server.handle(pickle.dumps(record))
+
+
+class LogServer(object):
+
+	def handle(self, pickled_record):
+		import pickle
+		record = pickle.loads(pickled_record)
+		log = logging.getLogger()
+		log.handle(record)
diff --git a/bootstrapvz/base/remote/__init__.py b/bootstrapvz/base/remote/__init__.py
index 7367efb..b7ba576 100644
--- a/bootstrapvz/base/remote/__init__.py
+++ b/bootstrapvz/base/remote/__init__.py
@@ -32,12 +32,21 @@ def main():
 	try:
 		manager.start()
 		server = manager.rpc_server
+		from callback import CallbackServer
+		callback_server = CallbackServer(manager.local_callback_port)
+		from bootstrapvz.base.log import LogServer
+		log_server = LogServer()
+		try:
+			callback_server.start(log_server)
+			server.set_log_server(log_server)
 
-		# Everything has been set up, begin the bootstrapping process
-		server.run(manifest,
-		           debug=opts['--debug'],
-		           pause_on_error=False,
-		           dry_run=opts['--dry-run'])
+			# Everything has been set up, begin the bootstrapping process
+			server.run(manifest,
+			           debug=opts['--debug'],
+			           pause_on_error=False,
+			           dry_run=opts['--dry-run'])
+		finally:
+			callback_server.stop()
 	finally:
 		manager.stop()
 
diff --git a/bootstrapvz/base/remote/callback.py b/bootstrapvz/base/remote/callback.py
index 9b77684..82d9060 100644
--- a/bootstrapvz/base/remote/callback.py
+++ b/bootstrapvz/base/remote/callback.py
@@ -6,14 +6,11 @@ class CallbackServer(object):
 		self.listen_port = listen_port
 		self.stop_serving = False
 
-		from log import LogServer
-		self.log_server = LogServer()
-
-	def start(self, rpc_server):
+	def start(self, log_server):
 		import Pyro4
 		Pyro4.config.COMMTIMEOUT = 0.5
 		daemon = Pyro4.Daemon('localhost', port=self.listen_port, unixsocket=None)
-		daemon.register(self.log_server)
+		daemon.register(log_server)
 
 		def serve():
 			daemon.requestLoop(loopCondition=lambda: not self.stop_serving)
@@ -21,8 +18,6 @@ class CallbackServer(object):
 		self.thread = Thread(target=serve)
 		self.thread.start()
 
-		rpc_server.set_log_server(self.log_server)
-
 	def stop(self):
 		self.stop_serving = True
 		if hasattr(self, 'thread'):
diff --git a/bootstrapvz/base/remote/log.py b/bootstrapvz/base/remote/log.py
deleted file mode 100644
index 443b3b4..0000000
--- a/bootstrapvz/base/remote/log.py
+++ /dev/null
@@ -1,29 +0,0 @@
-import logging
-import pickle
-
-
-class LogForwarder(logging.Handler):
-
-	def __init__(self, level=logging.NOTSET):
-		self.server = None
-		super(LogForwarder, self).__init__(level)
-
-	def set_server(self, server):
-		self.server = server
-
-	def emit(self, record):
-		if self.server is not None:
-			if record.exc_info is not None:
-				import traceback
-				exc_type, exc_value, exc_traceback = record.exc_info
-				record.exc_info = traceback.print_exception(exc_type, exc_value, exc_traceback)
-			self.server.handle(pickle.dumps(record))
-
-
-class LogServer(object):
-
-	def handle(self, pickled_record):
-		import logging
-		log = logging.getLogger()
-		record = pickle.loads(pickled_record)
-		log.handle(record)
diff --git a/bootstrapvz/base/remote/server.py b/bootstrapvz/base/remote/server.py
index 76f2832..2a834db 100644
--- a/bootstrapvz/base/remote/server.py
+++ b/bootstrapvz/base/remote/server.py
@@ -8,7 +8,7 @@ def main():
 
 def setup_logging():
 	import logging
-	from log import LogForwarder
+	from bootstrapvz.base.log import LogForwarder
 	log_forwarder = LogForwarder()
 	root = logging.getLogger()
 	root.addHandler(log_forwarder)
diff --git a/bootstrapvz/base/remote/ssh_rpc_manager.py b/bootstrapvz/base/remote/ssh_rpc_manager.py
index f359655..7143916 100644
--- a/bootstrapvz/base/remote/ssh_rpc_manager.py
+++ b/bootstrapvz/base/remote/ssh_rpc_manager.py
@@ -55,12 +55,6 @@ class SSHRPCManager(object):
 			self.process.terminate()
 			raise e
 
-		from callback import CallbackServer
-		self.callback_server = CallbackServer(self.local_callback_port)
-		self.callback_server.start(self.rpc_server)
-
 	def stop(self):
 		self.rpc_server.stop()
-		if hasattr(self, 'callback_server'):
-			self.callback_server.stop()
 		self.process.terminate()