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