def _check_retrans(self, request_id, event):
if request_id in self._retrans_map:
retrans = self._retrans_map[request_id]
if retrans.should_retrans(event):
dcm_events.register_callback(
self.incoming_parent_q_message,
args=[request_id, retrans.request_doc])
def send(self, doc):
with self._lock:
t = doc['type']
request_id = doc['request_id']
_g_logger.debug("Fake conn sending " + t)
self._check_retrans(request_id, t)
if t == message_types.MessageTypes.ACK:
# no reply required here
return
elif t == message_types.MessageTypes.NACK:
# no reply required here
return
elif t == message_types.MessageTypes.REPLY:
payload = doc['payload']
self._writer.write(json.dumps(payload) + '\n')
self._writer.flush()
if self._reply_ignore_count == 0:
# we must ACK the reply
reply_ack = {
"type": message_types.MessageTypes.ACK,
"request_id": doc["request_id"],
"message_id": doc["message_id"],
}
dcm_events.register_callback(
self.incoming_parent_q_message,
args=[doc["request_id"], reply_ack])
else:
self._reply_ignore_count -= 1
else:
raise test_exceptions.AgentTestException(
"type %s should never happen" % t)
def _sm_requested_reply_received(self, **kwargs):
"""
This is the standard case. After a request is made, and it receives an
acknowledgment, it later receives a reply. Here the reply will be
stored and the state machine will move on.
"""
message = kwargs['message']
if self._message_timer is not None:
msg = ("In the REQUESTED state the message ID should not be in the"
" list")
_g_logger.error(msg)
if self._reply_doc is not None:
msg = ("There should be exactly 1 reply received. Thus is the "
"reply_doc attribute is not None something we terribly "
"wrong.")
_g_logger.error(msg)
_g_logger.debug("The incoming reply is %s" % str(message))
self._reply_doc = message
if self._reply_callback is not None:
args = [message]
if self._reply_args:
args.extend(self._reply_args)
dcm_events.register_callback(
self._user_reply_callback,
args=self._reply_args, kwargs=self._reply_kwargs)
def _read_from_file(self):
buf = self._reader.readline().strip()
if not buf:
return
_g_logger.debug("read message " + buf)
ba = buf.split()
command = ba.pop(0)
arguments = ba
message_id = utils.new_message_id()
request_id = utils.new_message_id()
request_doc = {
'type': message_types.MessageTypes.REQUEST,
'request_id': request_id,
'message_id': message_id,
'payload': {'command': command, 'arguments': arguments}
}
# check for any retrans requests of this message
if len(self._retrans) > self._request_number:
rt = self._retrans[self._request_number]
self._retrans_map[request_id] = rt
rt.set_request_doc(request_doc)
self._request_number += 1
self._check_retrans(request_id, message_types.MessageTypes.REQUEST)
dcm_events.register_callback(
self.recv_obj.incoming_parent_q_message, args=[request_doc])
def run(self):
_g_logger.info("Job runner %s thread starting." % self.getName())
done = False
while not done:
try:
work = self._queue.get(True)
if work.quit:
done = True
continue
try:
_g_logger.debug("Running the long job %s:%s" %
(work.name, work.request_id))
job_reply = JobReply(work.job_id)
dcm_events.register_callback(
self._job_update_callback, args=[job_reply])
plugin = plugin_loader.load_python_module(
work.items_map["module_name"],
self._conf,
work.request_id,
work.items_map,
work.name,
work.arguments)
reply_obj = plugin.run()
job_reply.reply_doc = reply_obj.get_reply_doc()
except Exception as ex:
_g_logger.exception("An error occurred")
job_reply.error = str(ex)
job_reply.job_status = JobStatus.ERROR
else:
if job_reply.reply_doc is None:
job_reply.job_status = JobStatus.COMPLETE
elif job_reply.reply_doc["return_code"] == 0:
job_reply.job_status = JobStatus.COMPLETE
else:
job_reply.job_status = JobStatus.ERROR
job_reply.error = job_reply.reply_doc["message"]
finally:
job_reply.end_date = calendar.timegm(time.gmtime())
dcm_events.register_callback(
self._job_update_callback, args=[job_reply])
_g_logger.debug("Completed the long job %s:%s "
"STATUS=%s" % (work.name, work.request_id,
job_reply.job_status))
except queue.Empty:
_g_logger.exception("The queue was empty. This shouldn't "
"happen often")
except Exception as ex:
_g_logger.exception("Something went wrong processing the job")
finally:
self._current_job = None
_g_logger.info("Job runner %s thread ending." % self.getName())
def _sm_acked_cancel_received(self, **kwargs):
"""
A cancel is received from the remote end. We simply notify the user
of the request and allow the user to act upon it.
"""
dcm_events.register_callback(
self._cancel_callback,
args=self._cancel_callback_args,
kwargs=self._cancel_callback_kwargs)
def incoming_request(self, reply_obj):
payload = reply_obj.get_message_payload()
_g_logger.debug("Incoming request %s" % str(payload))
request_id = reply_obj.get_request_id()
_g_logger.info("Creating a request ID %s" % request_id)
items_map = plugin_loader.parse_plugin_doc(
self._conf, payload["command"])
dcm_logger.log_to_dcm_console_incoming_message(
job_name=payload["command"])
immediate = "immediate" in items_map
long_runner = "longer_runner" in items_map
if "longer_runner" in payload:
long_runner = bool(payload["longer_runner"])
# we ack first. This will write it to the persistent store before
# sending the message so the agent will have it for restarts
reply_obj.ack(None, None, None)
if long_runner:
try:
dj = self._long_runner.start_new_job(
self._conf,
request_id,
items_map,
payload["command"],
payload["arguments"])
except BaseException as ex:
reply_obj = plugin_base.PluginReply(
1,
error_message=urllib.parse.quote(str(ex).encode('utf-8')))
else:
payload_doc = dj.get_message_payload()
reply_obj = plugin_base.PluginReply(
0, reply_type="job_description", reply_object=payload_doc)
reply_doc = reply_obj.get_reply_doc()
wr = WorkReply(request_id, reply_doc)
dcm_events.register_callback(
self.work_complete_callback, args=[wr])
elif immediate:
items_map["long_runner"] = self._long_runner
reply_doc = _run_plugin(self._conf,
items_map,
request_id,
payload["command"],
payload["arguments"])
wr = WorkReply(request_id, reply_doc)
dcm_events.register_callback(
self.work_complete_callback, args=[wr])
else:
workload = WorkLoad(request_id, payload, items_map)
self.worker_q.put(workload)
_g_logger.debug(
"The request %s has been set to send an ACK" % request_id)
def _sm_successful_handshake(self):
"""
This is the standard case when a handshake is successfully processed
"""
_g_logger.debug("The handshake was successfully processed")
while not self.pre_hs_message_queue.empty:
incoming_data = self.pre_hs_message_queue.get()
dcm_events.register_callback(
self._receive_callback, args=[incoming_data])
self._backoff.activity()
def set_conn(self, conf, conn):
self._conn = conn
self._conf = conf
if conn is None:
return
for msg in self._unsent_msgs:
dcm_events.register_callback(
send_log_to_dcm_callback, kwargs={"conn": self._conn,
"message": msg})
self._unsent_msgs = []
def emit(self, record):
msg = self.format(record)
if self._conn is None:
self._unsent_msgs.append(msg)
else:
dcm_events.register_callback(
send_log_to_dcm_callback, kwargs={"conn": self._conn,
"token": "",
"message": msg,
"level": record.levelname})
def _sm_requesting_cancel_received(self, **kwargs):
"""
A cancel message flows over the wire after the request is received
but before it is acknowledged. Here we will tell the user about the
cancel. It is important that the cancel notification comes after
the message received notification.
"""
dcm_events.register_callback(
self._cancel_callback,
args=self._cancel_callback_args,
kwargs=self._cancel_callback_kwargs)
def _sm_connect(self):
try:
self._ws = _WebSocketClient(
self, self._server_url, self._receive_callback,
protocols=['dcm'], heartbeat_freq=self._heartbeat_freq,
ssl_options=self._ssl_options)
dcm_events.register_callback(
self._forming_connection_thread, in_thread=True)
except Exception as ex:
_g_logger.exception("Failed to connect to %s" % self._server_url)
self._throw_error(ex, notify=False)
self._cond.notify()
def _register_connect(self):
_g_logger.debug("Registering a connection to DCM")
if self._connect_timer is not None:
raise exceptions.AgentRuntimeException(
"There is already a connection registered")
self._connect_timer = dcm_events.register_callback(
self.event_connect_timeout,
delay=self._backoff.seconds_until_ready())
def run(self):
try:
_g_logger.info("Worker %s thread starting." % self.getName())
done = False
while not done:
try:
workload = self.worker_queue.get()
if workload is None:
continue
if workload.quit:
done = True
self.worker_queue.task_done()
continue
# setup message logging
with tracer.RequestTracer(workload.request_id):
reply_doc = _run_plugin(self._conf,
workload.items_map,
workload.request_id,
workload.payload["command"],
workload.payload["arguments"])
self.worker_queue.task_done()
_g_logger.debug(
"Adding the reply document to the reply "
"queue " + str(reply_doc))
work_reply = WorkReply(workload.request_id, reply_doc)
dcm_events.register_callback(
self._reply_callback, args=[work_reply])
_g_logger.info("Reply message sent for command " +
workload.payload["command"])
except queue.Empty:
pass
except:
_g_logger.exception(
"Something went wrong processing the queue")
raise
finally:
_g_logger.info("Worker %s thread ending." % self.getName())
def _sm_received_hs(self, incoming_data=None):
"""
The handshake has arrived
"""
try:
# if the handshake is rejected an exception will be thrown
hs = self._handshake_manager.incoming_document(incoming_data)
_g_logger.debug("We received a handshake with reply code %d"
% hs.reply_type)
if hs.reply_type != handshake.HandshakeIncomingReply.REPLY_CODE_SUCCESS:
_g_logger.warn("The handshake was rejected.")
if hs.reply_type == handshake.HandshakeIncomingReply.REPLY_CODE_FORCE_BACKOFF:
_g_logger.info("Backing off for %f seconds"
% float(hs.force_backoff))
self._backoff.force_backoff_time(hs.force_backoff)
self._ws.close()
ex = exceptions.AgentHandshakeException(hs)
self._throw_error(ex)
else:
dcm_events.register_callback(self.event_successful_handshake,
kwargs={"hs": hs})
self._cond.notify()
except Exception as ex:
self._throw_error(ex)
def _send_reply_ack(self):
# the ack reuses the message id that it is acknowledging
# this is here just to help track exchanges
self.ack_sender += 1
if self.ack_sender > 1:
pass
message_id = self._reply_doc['message_id']
reply_ack_doc = {'request_id': self._request_id,
'message_id': message_id,
'what': 'REPLY_ACK',
'ack_sender': self.ack_sender,
'type': types.MessageTypes.ACK}
self.send_doc(reply_ack_doc)
if self._completion_timer is not None:
self._completion_timer.cancel()
self._completion_timer = dcm_events.register_callback(
self.ack_sent_timeout,
kwargs={'timer': self._completion_timer},
delay=self._cleanup_timeout)
def _sm_open_incoming_message(self, incoming_data=None):
_g_logger.debug("New message received")
dcm_events.register_callback(
self._receive_callback, args=[incoming_data])
self._backoff.activity()
def _throw_error(self, exception, notify=True):
_g_logger.warning("throwing error %s" % str(exception))
dcm_events.register_callback(self.event_error,
kwargs={"exception": exception})
if notify:
self._cond.notify()
def send(self, doc):
dcm_events.register_callback(self._request.incoming_message,
args=[doc])
def send(self, doc):
dcm_events.register_callback(
self._request.incoming_message, args=[doc])
def send(self, conn):
_g_logger.info("Resending reply to %s" % self._send_doc["request_id"])
self._send_doc["entity"] = "timer"
conn.send(self._send_doc)
self._timer = dcm_events.register_callback(self._cb, args=[self], delay=self._timeout)
def _throw_error(self, exception, notify=True):
_g_logger.warning("throwing error %s" % str(exception))
dcm_events.register_callback(self.event_error,
kwargs={"exception": exception})
if notify:
self._cond.notify()
def _send_timeout(self):
self._timer = dcm_events.register_callback(self.timeout,
delay=self._timeout)
_g_logger.debug("Sending the alert message " + str(self.doc))
self._conn.send(self.doc)
def _sm_open_incoming_message(self, incoming_data=None):
_g_logger.debug("New message received")
dcm_events.register_callback(self._receive_callback,
args=[incoming_data])
self._backoff.activity()
def start(self):
if self._stopped:
return
self._timer = dcm_events.register_callback(
self.clean_sweep, delay=self._sweep_time)
def job_complete(self, job_id):
if self._conf.jobs_retain_job_time == 0:
return
dcm_events.register_callback(self._job_cleanup,
args=[job_id],
delay=self._conf.jobs_retain_job_time)
def _send_timeout(self):
self._timer = dcm_events.register_callback(
self.timeout, delay=self._timeout)
_g_logger.debug("Sending the alert message " + str(self.doc))
self._conn.send(self.doc)
def job_complete(self, job_id):
if self._conf.jobs_retain_job_time == 0:
return
dcm_events.register_callback(self._job_cleanup,
args=[job_id],
delay=self._conf.jobs_retain_job_time)
