def receive_data(self):
"""Constantly receives data from the ZMQ socket and handles it accordingly.
"""
while True:
if self.thread_exit:
break
try:
message_bytes = self.comm_sock.recv()
except zmq.ZMQError:
continue
message = CommunicationPacket()
message.ParseFromString(message_bytes)
if message.type == CommunicationPacket.EXIT:
break
workflow = self.__get_workflow_by_execution_id(
message.workflow_execution_id)
if workflow:
if message.type == CommunicationPacket.PAUSE:
workflow.pause()
elif message.type == CommunicationPacket.ABORT:
workflow.abort()
return
def receive_data(self):
"""Constantly receives data from the ZMQ socket and handles it accordingly.
"""
while True:
if self.thread_exit:
break
try:
message_bytes = self.comm_sock.recv()
except zmq.ZMQError:
continue
message = CommunicationPacket()
message.ParseFromString(message_bytes)
if message.type == CommunicationPacket.EXIT:
break
workflow = self.__get_workflow_by_execution_uid(message.workflow_execution_uid)
if workflow:
if message.type == CommunicationPacket.PAUSE:
workflow.pause()
elif message.type == CommunicationPacket.RESUME:
workflow.resume()
elif message.type == CommunicationPacket.TRIGGER:
trigger_data = json.loads(message.data_in)
if len(message.arguments) > 0:
arguments = []
for arg in message.arguments:
arguments.append(Argument(**(MessageToDict(arg, preserving_proto_field_name=True))))
trigger_data["arguments"] = arguments
workflow.send_data_to_action(trigger_data)
return
def receive_communications(self):
"""Constantly receives data from the ZMQ socket and handles it accordingly"""
logger.info('Starting workflow communication receiver')
while not self.exit:
try:
message_bytes = self.comm_sock.recv()
except zmq.ZMQError:
continue
message = CommunicationPacket()
try:
message.ParseFromString(message_bytes)
except DecodeError:
logger.error(
'Worker communication handler could not decode communication packet'
)
else:
message_type = message.type
if message_type == CommunicationPacket.WORKFLOW:
logger.debug(
'Worker received workflow communication packet')
yield WorkerCommunicationMessageData(
WorkerCommunicationMessageType.workflow,
self._format_workflow_message_data(
message.workflow_control_message))
elif message_type == CommunicationPacket.CASE:
logger.debug('Workflow received case communication packet')
yield WorkerCommunicationMessageData(
WorkerCommunicationMessageType.case,
self._format_case_message_data(
message.case_control_message))
elif message_type == CommunicationPacket.EXIT:
logger.info('Worker received exit message')
break
raise StopIteration
def test_receive_exit(self):
receiver = self.get_receiver()
message = CommunicationPacket()
message.type = CommunicationPacket.EXIT
with patch.object(receiver.comm_sock, 'recv', return_value=message.SerializeToString()):
message_generator = receiver.receive_communications()
with self.assertRaises(StopIteration):
message = next(message_generator)
def send_exit_to_worker_comms(self):
"""Sends the exit message over the communication sockets, otherwise worker receiver threads will hang
"""
message = CommunicationPacket()
message.type = CommunicationPacket.EXIT
message_bytes = message.SerializeToString()
for worker in self.workers:
self.comm_socket.send_multipart([worker, message_bytes])
def test_receive_exit(self):
receiver = self.get_receiver()
message = CommunicationPacket()
message.type = CommunicationPacket.EXIT
with patch.object(receiver.comm_sock,
'recv',
return_value=message.SerializeToString()):
message_generator = receiver.receive_communications()
with self.assertRaises(StopIteration):
message = next(message_generator)
def check_receive_workflow_communication_message(self, proto_message_type, data_message_type):
receiver = self.get_receiver()
message = CommunicationPacket()
message.type = CommunicationPacket.WORKFLOW
message.workflow_control_message.type = proto_message_type
uid = str(uuid4())
message.workflow_control_message.workflow_execution_id = uid
expected = WorkerCommunicationMessageData(
WorkerCommunicationMessageType.workflow,
WorkflowCommunicationMessageData(data_message_type, uid))
self.check_receive_communication_message(receiver, message, expected)
def _create_case_update_message(case_id, message_type, subscriptions=None):
message = CommunicationPacket()
message.type = CommunicationPacket.CASE
message.case_control_message.id = case_id
message.case_control_message.type = message_type
subscriptions = subscriptions or []
for subscription in subscriptions:
sub = message.case_control_message.subscriptions.add()
sub.id = subscription.id
sub.events.extend(subscription.events)
return message
def check_receive_workflow_communication_message(self, proto_message_type,
data_message_type):
receiver = self.get_receiver()
message = CommunicationPacket()
message.type = CommunicationPacket.WORKFLOW
message.workflow_control_message.type = proto_message_type
uid = str(uuid4())
message.workflow_control_message.workflow_execution_id = uid
expected = WorkerCommunicationMessageData(
WorkerCommunicationMessageType.workflow,
WorkflowCommunicationMessageData(data_message_type, uid))
self.check_receive_communication_message(receiver, message, expected)
def resume_workflow(self, workflow_execution_uid):
"""Resumes a workflow that has previously been paused.
Args:
workflow_execution_uid (str): The execution UID of the workflow.
"""
logger.info('Resuming workflow {0}'.format(workflow_execution_uid))
if workflow_execution_uid in self.workflow_comms:
message = CommunicationPacket()
message.type = CommunicationPacket.RESUME
message.workflow_execution_uid = workflow_execution_uid
message_bytes = message.SerializeToString()
self.comm_socket.send_multipart(
[self.workflow_comms[workflow_execution_uid], message_bytes])
return True
else:
return False
def pause_workflow(self, workflow_execution_uid):
"""Pauses a workflow currently executing.
Args:
workflow_execution_uid (str): The execution UID of the workflow.
"""
logger.info('Pausing workflow {0}'.format(workflow_execution_uid))
if workflow_execution_uid in self.workflow_comms:
message = CommunicationPacket()
message.type = CommunicationPacket.PAUSE
message.workflow_execution_uid = workflow_execution_uid
message_bytes = message.SerializeToString()
self.comm_socket.send_multipart(
[self.workflow_comms[workflow_execution_uid], message_bytes])
return True
else:
return False
def to_received_message(message_bytes):
"""Constantly receives data from the ZMQ socket and handles it accordingly"""
message = CommunicationPacket()
try:
message.ParseFromString(message_bytes)
except DecodeError:
logger.error(
'Worker communication handler could not decode communication packet'
)
else:
message_type = message.type
if message_type == CommunicationPacket.WORKFLOW:
logger.debug('Worker received workflow communication packet')
return WorkerCommunicationMessageData(
WorkerCommunicationMessageType.workflow,
ProtobufWorkflowCommunicationConverter.
_format_workflow_message_data(
message.workflow_control_message))
elif message_type == CommunicationPacket.EXIT:
logger.info('Worker received exit message')
return None
def send_data_to_trigger(self, data_in, workflow_uids, arguments=None):
"""Sends the data_in to the workflows specified in workflow_uids.
Args:
data_in (dict): Data to be used to match against the triggers for an Action awaiting data.
workflow_uids (list[str]): A list of workflow execution UIDs to send this data to.
arguments (list[Argument]): An optional list of Arguments to update for an
Action awaiting data for a trigger. Defaults to None.
"""
data = dict()
data['data_in'] = data_in
message = CommunicationPacket()
message.type = CommunicationPacket.TRIGGER
message.data_in = json.dumps(data)
if arguments:
for argument in arguments:
arg = message.arguments.add()
arg.name = argument.name
for field in ('value', 'reference', 'selection'):
val = getattr(argument, field)
if val is not None:
if not isinstance(val, string_types):
try:
setattr(arg, field, json.dumps(val))
except ValueError:
setattr(arg, field, str(val))
else:
setattr(arg, field, val)
for uid in workflow_uids:
if uid in self.workflow_comms:
message.workflow_execution_uid = uid
message_bytes = message.SerializeToString()
self.comm_socket.send_multipart(
[self.workflow_comms[uid], message_bytes])
def create_worker_exit_message():
"""Creates a message to tell the Workers to exit
"""
message = CommunicationPacket()
message.type = CommunicationPacket.EXIT
return message.SerializeToString()
def _create_workflow_control_message(control_type, workflow_execution_id):
message = CommunicationPacket()
message.type = CommunicationPacket.WORKFLOW
message.workflow_control_message.type = control_type
message.workflow_control_message.workflow_execution_id = workflow_execution_id
return message.SerializeToString()
def create_worker_exit_message():
"""Creates a message to tell the Workers to exit
"""
message = CommunicationPacket()
message.type = CommunicationPacket.EXIT
return message.SerializeToString()
def send_exit_to_worker_comms(self):
"""Sends the exit message over the communication sockets, otherwise worker receiver threads will hang"""
message = CommunicationPacket()
message.type = CommunicationPacket.EXIT
self._send_message(message)
def _create_workflow_control_message(control_type, workflow_execution_id):
message = CommunicationPacket()
message.type = CommunicationPacket.WORKFLOW
message.workflow_control_message.type = control_type
message.workflow_control_message.workflow_execution_id = workflow_execution_id
return message.SerializeToString()