Simplify? Hopefully...

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)

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'])
+		finally:
+			callback_server.stop()
 	finally:
 		manager.stop()
 

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'):

bootstrapvz/base/remote/log.py

@@ -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)

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)

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()