Commit 3a9205cc authored by Robin Obůrka's avatar Robin Obůrka Committed by Robin Obůrka

msgloop: Provide object oriented alternative for sn_main()

parent 81a309d5
...@@ -2,5 +2,5 @@ from sn.exceptions import * ...@@ -2,5 +2,5 @@ from sn.exceptions import *
from sn.messages import * from sn.messages import *
from sn.argparser import * from sn.argparser import *
from sn.network import * from sn.network import *
from sn.msgloop import sn_main from sn.msgloop import SNPipelineBox, SNGeneratorBox, SNTerminationBox
import sn.logging import sn.logging
...@@ -13,9 +13,6 @@ from .messages import encode_msg, parse_msg ...@@ -13,9 +13,6 @@ from .messages import encode_msg, parse_msg
from .exceptions import * from .exceptions import *
logger = logging.getLogger("sn_main")
class LoopHardFail(Exception): class LoopHardFail(Exception):
pass pass
...@@ -24,179 +21,200 @@ class LoopFail(Exception): ...@@ -24,179 +21,200 @@ class LoopFail(Exception):
pass pass
def sn_main(box_name, process, setup=None, teardown=None, before_first_request=None, argparser=None): class SNBox():
sn_ctx = SN(zmq.Context.instance(), argparser or get_arg_parser()) def __init__(self, box_name, argparser=None):
# Important provided values into box
self.name = box_name
self.logger = logging.getLogger(box_name)
# Internal context values
self.zmq_ctx = zmq.Context.instance()
self.sn_ctx = SN(self.zmq_ctx, argparser or get_arg_parser())
self.args = self.sn_ctx.args
# User data
self.context = None
# Error management
self.loop_continue = True
self.errors_in_row = 0
# Core methods - Will be implemented in non-abstract boxes
def check_configuration(self):
raise NotImplementedError("check_configuration")
def get_processed_message(self):
raise NotImplementedError("get_processed_message")
def process_result(self, result):
raise NotImplementedError("process_result")
# Public API for boxes - will be optionally implemented in final boxes
def setup(self):
return {}
def teardown(self):
pass
context = None def before_first_request(self):
try: pass
user_data = get_user_data(setup, sn_ctx.args)
context = build_context(box_name, sn_ctx, user_data) def process(self, msg_type, payload):
check_configuration(context, process) raise NotImplementedError("process")
logger.info("SN main starting loop for %s box", box_name) # Provided functionality - should be final implementation
register_signals(context) def run(self):
self.check_configuration()
try:
self.context = self.get_user_data()
_sn_main_loop(context, before_first_request, process) self.logger.info("SN main starting loop for %s box", self.name)
self.register_signals()
except LoopHardFail as e: self.before_first_request()
logger.error("Hard Fail of box: %s", context.name) self.run_loop()
logger.exception(e)
# Finally will be called, because sys.exit() raises exception that will be uncaught.
sys.exit(1)
except KeyboardInterrupt: except LoopHardFail as e:
pass self.logger.error("Hard Fail of box: %s", self.name)
self.logger.exception(e)
# Finally will be called, because sys.exit() raises exception that will be uncaught.
sys.exit(1)
finally: except KeyboardInterrupt:
if context: pass
# Is possible that context wasn't built yet (e.g. error in setup callback)
if teardown:
teardown(context)
teardown_context(context)
finally:
self.teardown_box()
self.teardown()
def get_user_data(setup, args): def get_user_data(self):
if setup: user_data = self.setup()
number_of_arguments = len((inspect.signature(setup).parameters))
if number_of_arguments == 1:
user_data = setup(args)
else:
# Let the function fail in implicit way in case that setup callback
# has more than 1 argument
user_data = setup()
if isinstance(user_data, dict): if isinstance(user_data, dict):
return user_data return SimpleNamespace(**user_data)
else: else:
raise SetupError("Setup function didn't return a dictionary") raise SetupError("Setup function didn't return a dictionary")
return {} def register_signals(self):
def signal_handler(signum, frame):
self.logger.info("Signal %s received", signum)
self.loop_continue = False
for sig in [ signal.SIGHUP, signal.SIGTERM, signal.SIGQUIT, signal.SIGABRT ]:
signal.signal(sig, signal_handler)
def build_context(box_name, sn_ctx, user_data): def teardown_box(self):
socket_recv = get_socket(sn_ctx, "in") self.zmq_ctx.destroy()
socket_send = get_socket(sn_ctx, "out")
ctx = { def run_loop(self):
"name": box_name, while self.loop_continue:
"logger": logging.getLogger(box_name), try:
"loop_continue": True, result = self.get_processed_message()
"errors_in_row": 0, self.process_result(result)
"sn_ctx": sn_ctx, self.errors_in_row = 0
"zmq_ctx": sn_ctx.context,
"args": sn_ctx.args,
"socket_recv": socket_recv,
"socket_send": socket_send,
}
for k, v in user_data.items(): except StopIteration:
if k in ctx: self.logger.info("Box %s raised StopIteration", self.name)
raise SetupError("Used reserved word in user_data: %s", k) break
else:
ctx[k] = v
return SimpleNamespace(**ctx) except (SetupError, NotImplementedError) as e:
raise e
except Exception as e:
self.logger.error("Uncaught exception from loop: %s", type(e).__name__)
self.logger.exception(e)
def get_socket(context, sock_name): self.errors_in_row += 1
socket = None if self.errors_in_row > 10:
try: raise LoopHardFail("Many errors in row.")
socket = context.get_socket(sock_name)
except UndefinedSocketError as e: # Helper methods
pass def get_socket(self, sock_name):
socket = None
try:
socket = self.sn_ctx.get_socket(sock_name)
return socket except UndefinedSocketError as e:
pass
return socket
def check_configuration(context, process):
if not context.socket_recv and not context.socket_send:
raise SetupError("Neither input nor output socket provided")
if not context.socket_recv and not inspect.isgeneratorfunction(process):
raise SetupError("Generator is expected for output-only box")
class SNPipelineBox(SNBox):
def __init__(self, box_name, argparser=None):
super().__init__(box_name, argparser)
self.socket_recv = self.get_socket("in")
self.socket_send = self.get_socket("out")
def teardown_context(context): def check_configuration(self):
if context.socket_recv: if not self.socket_recv and not self.socket_send:
context.socket_recv.close() raise SetupError("Neither input nor output socket provided")
if context.socket_send:
context.socket_send.close()
context.zmq_ctx.destroy()
def get_processed_message(self):
msg = self.socket_recv.recv_multipart()
msg_type, payload = parse_msg(msg)
def _sn_main_loop(context, before_first_request, process): return self.process(msg_type, payload)
if inspect.isgeneratorfunction(process):
get_processed_message = processed_message_from_generator(context, process)
else:
get_processed_message = processed_message_from_function(context, process)
if before_first_request: def process_result(self, result):
result = before_first_request(context) if not result:
if result: # The box hasn't any reasonable answer
process_result(context.socket_send, result) return
while context.loop_continue:
try: try:
result = get_processed_message() msg_type, payload = result
process_result(context.socket_send, result) msg_out = encode_msg(msg_type, payload)
context.errors_in_row = 0 self.socket_send.send_multipart(msg_out)
except StopIteration:
context.logger.info("Box %s raised StopIteration", context.name)
break
except SetupError as e: except (ValueError, InvalidMsgError):
raise e raise LoopFail("Generated broken output message. Possibly bug in box.")
except Exception as e:
logger.error("Uncaught exception from loop: %s", type(e).__name__)
logger.exception(e)
context.errors_in_row += 1 class SNGeneratorBox(SNBox):
if context.errors_in_row > 10: def __init__(self, box_name, argparser=None):
raise LoopHardFail("Many errors in row.") super().__init__(box_name, argparser)
self.socket_send = self.get_socket("out")
# Ensure about process() method before try to get iterator
self.check_configuration()
def processed_message_from_generator(context, process): self.process_iterator = self.process()
iterator = process(context)
def get_from_generator():
return next(iterator)
return get_from_generator def check_configuration(self):
if not self.socket_send:
raise SetupError("Output socket wasn't provided")
if not inspect.isgeneratorfunction(self.process):
raise SetupError("Generator is expected for output-only box")
def get_processed_message(self):
return next(self.process_iterator)
def processed_message_from_function(context, process): def process_result(self, result):
def get_from_function(): if not result:
msg = context.socket_recv.recv_multipart() # The box hasn't any reasonable answer
msg_type, payload = parse_msg(msg) return
return process(context, msg_type, payload)
return get_from_function
try:
msg_type, payload = result
msg_out = encode_msg(msg_type, payload)
self.socket_send.send_multipart(msg_out)
def process_result(socket_send, result): except (ValueError, InvalidMsgError):
if not result: raise LoopFail("Generated broken output message. Possibly bug in box.")
# The box is output-only or it hasn't any reasonable answer
return
if not socket_send:
raise SetupError("Box generated output but there is any output socket. Bad configuration?")
try: class SNTerminationBox(SNBox):
msg_type, payload = result def __init__(self, box_name, argparser=None):
msg_out = encode_msg(msg_type, payload) super().__init__(box_name, argparser)
socket_send.send_multipart(msg_out) self.socket_recv = self.get_socket("in")
except (ValueError, InvalidMsgError): def check_configuration(self):
raise LoopFail("Generated broken output message. Possibly bug in box.") if not self.socket_recv:
raise SetupError("Input socket wasn't provided")
def get_processed_message(self):
msg = self.socket_recv.recv_multipart()
msg_type, payload = parse_msg(msg)
def register_signals(context): return self.process(msg_type, payload)
def signal_handler(signum, frame):
context.logger.info("Signal %s received", signum)
context.loop_continue = False
for sig in [ signal.SIGHUP, signal.SIGTERM, signal.SIGQUIT, signal.SIGABRT ]: def process_result(self, result):
signal.signal(sig, signal_handler) if result:
raise LoopFail("Input-only box generated output message. Possibly bug in box.")
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment