def on_connected(self, headers, body): """ Once the connection is established, and 'heart-beat' is found in the headers, we calculate the real heartbeat numbers (based on what the server sent and what was specified by the client) - if the heartbeats are not 0, we start up the heartbeat loop accordingly. :param dict headers: headers in the connection message :param body: the message body """ if "heart-beat" in headers: self.heartbeats = utils.calculate_heartbeats( headers["heart-beat"].replace(' ', '').split(','), self.heartbeats) logging.debug("Heartbeats calculated %s", str(self.heartbeats)) if self.heartbeats != (0, 0): self.send_sleep = self.heartbeats[0] / 1000 # by default, receive gets an additional grace of 50% # set a different heart-beat-receive-scale when creating the connection to override that self.receive_sleep = (self.heartbeats[1] / 1000) * self.heart_beat_receive_scale self.loop_sleep = max(1, int(min(self.send_sleep, self.receive_sleep) / 2.0)) logging.debug("Set receive_sleep to %s, send_sleep to %s, loop sleep to %s", self.receive_sleep, self.send_sleep, self.loop_sleep) # Give grace of receiving the first heartbeat self.received_heartbeat = monotonic() + self.receive_sleep self.running = True if self.heartbeat_thread is None: self.heartbeat_thread = utils.default_create_thread( self.__heartbeat_loop) self.heartbeat_thread.name = "StompHeartbeat%s" % \ getattr(self.heartbeat_thread, "name", "Thread")
def transmit(self, frame): """ Convert a frame object to a frame string and transmit to the server. :param Frame frame: the Frame object to transmit """ with self.__listeners_change_condition: listeners = sorted(self.listeners.items()) for (_, listener) in listeners: if not listener: continue try: listener.on_send(frame) except AttributeError: continue if frame.cmd == CMD_DISCONNECT and HDR_RECEIPT in frame.headers: self.__disconnect_receipt = frame.headers[HDR_RECEIPT] lines = convert_frame(frame) packed_frame = pack(lines) if logging.isEnabledFor(logging.DEBUG): logging.debug("Sending frame: %s", clean_lines(lines)) else: logging.info("Sending frame: %r", frame.cmd or "heartbeat") self.send(packed_frame)
def notify(self, frame_type, headers=None, body=None): """ Utility function for notifying listeners of incoming and outgoing messages :param str frame_type: the type of message :param dict headers: the map of headers associated with the message :param body: the content of the message """ if frame_type == 'receipt': # logic for wait-on-receipt notification receipt = headers['receipt-id'] receipt_value = self.__receipts.get(receipt) with self.__send_wait_condition: self.set_receipt(receipt, None) self.__send_wait_condition.notify() if receipt_value == CMD_DISCONNECT: self.set_connected(False) # received a stomp 1.1+ disconnect receipt if receipt == self.__disconnect_receipt: self.disconnect_socket() self.__disconnect_receipt = None elif frame_type == 'connected': self.set_connected(True) elif frame_type == 'disconnected': self.__notified_on_disconnect = True self.set_connected(False) with self.__listeners_change_condition: listeners = sorted(self.listeners.items()) for (_, listener) in listeners: if not listener: continue notify_func = getattr(listener, 'on_%s' % frame_type, None) if not notify_func: logging.debug("listener %s has no method on_%s", listener, frame_type) continue if frame_type in ('heartbeat', 'disconnected'): notify_func() continue if frame_type == 'connecting': notify_func(self.current_host_and_port) continue if frame_type == 'error' and not self.connected: with self.__connect_wait_condition: self.connection_error = True self.__connect_wait_condition.notify() rtn = notify_func(headers, body) if rtn: (headers, body) = rtn return (headers, body)
def start(self): logging.debug('Starting stomp server') self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) self.s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) self.s.bind((self.host, self.port)) self.s.listen(1) self.running = True thread = threading.Thread(None, self.run) thread.daemon = True thread.start() self.stopped = False logging.debug('Stomp server started')
def on_error(self, headers, body): """ Increment the error count. See :py:meth:`ConnectionListener.on_error` :param dict headers: headers in the message :param body: the message content """ if logging.isEnabledFor(logging.DEBUG): logging.debug("received an error %s [%s]", body, headers) else: logging.info("received an error %s", body) self.errors += 1
def receive(self): """ :rtype: bytes """ try: return self.socket.recv(self.__recv_bytes) except socket.error: _, e, _ = sys.exc_info() if get_errno(e) in (errno.EAGAIN, errno.EINTR): logging.debug("socket read interrupted, restarting") raise exception.InterruptedException() if self.is_connected(): raise
def notify(self, frame_type, frame=None): """ Utility function for notifying listeners of incoming and outgoing messages :param str frame_type: the type of message :param dict headers: the map of headers associated with the message :param body: the content of the message """ if frame_type == "receipt": # logic for wait-on-receipt notification receipt = frame.headers["receipt-id"] receipt_value = self.__receipts.get(receipt) with self.__send_wait_condition: self.set_receipt(receipt, None) self.__send_wait_condition.notify() if receipt_value == CMD_DISCONNECT: self.set_connected(False) # received a stomp 1.1+ disconnect receipt if receipt == self.__disconnect_receipt: self.disconnect_socket() self.__disconnect_receipt = None elif frame_type == "connected": self.set_connected(True) elif frame_type == "disconnected": self.__notified_on_disconnect = True self.set_connected(False) with self.__listeners_change_condition: listeners = sorted(self.listeners.items()) for (_, listener) in listeners: notify_func = getattr(listener, "on_%s" % frame_type, None) if not notify_func: logging.debug("listener %s has no method on_%s", listener, frame_type) continue if frame_type in ("heartbeat", "disconnected"): notify_func() continue if frame_type == "connecting": notify_func(self.current_host_and_port) continue if frame_type == "error" and not self.connected: with self.__connect_wait_condition: self.connection_error = True self.__connect_wait_condition.notify() notify_func(frame)
def stop(self): logging.debug('Stopping test server') if self.conn: try: self.conn.shutdown(socket.SHUT_WR) except Exception: pass self.conn.close() if self.s: self.s.close() self.running = False self.conn = None self.s = None self.stopped = True logging.debug('Connection stopped')
def process_frame(self, f, frame_str): """ :param Frame f: Frame object :param bytes frame_str: raw frame content """ frame_type = f.cmd.lower() if frame_type in ["connected", "message", "receipt", "error", "heartbeat"]: if frame_type == "message": self.notify("before_message", f) if logging.isEnabledFor(logging.DEBUG): logging.debug("Received frame: %r, headers=%r, body=%r", f.cmd, f.headers, f.body) else: logging.info("Received frame: %r, len(body)=%r", f.cmd, length(f.body)) self.notify(frame_type, f) else: logging.warning("Unknown response frame type: '%s' (frame length was %d)", frame_type, length(frame_str))
def run(self): self.conn, _ = self.s.accept() while self.running: try: _ = self.conn.recv(1024) frame = self.get_next_frame() if self.conn is None: break if frame is not None: self.conn.send(encode(frame)) except Exception: _, e, _ = sys.exc_info() logging.debug(e) break try: self.conn.close() except: pass self.stopped = True logging.debug('Run loop completed')
def run(self): self.conn, _ = self.s.accept() while self.running: try: _ = self.conn.recv(1024) frame = self.get_next_frame() if self.conn is None: break if frame is not None: logging.info("Stompserver sending frame %s", frame) self.conn.send(encode(frame)) except Exception: _, e, _ = sys.exc_info() logging.debug(e) break time.sleep(0.1) try: self.conn.close() except: pass self.stopped = True logging.info("Run loop completed")
def __heartbeat_loop(self): """ Main loop for sending (and monitoring received) heartbeats. """ logging.info('Starting heartbeat loop') now = monotonic() # Setup the initial due time for the outbound heartbeat if self.send_sleep != 0: self.next_outbound_heartbeat = now + self.send_sleep while self.running: now = monotonic() next_events = [] if self.next_outbound_heartbeat is not None: next_events.append(self.next_outbound_heartbeat - now) if self.receive_sleep != 0: t = self.received_heartbeat + self.receive_sleep - now if t > 0: next_events.append(t) sleep_time = min(next_events) if sleep_time > 0: terminate = self.heartbeat_terminate_event.wait(sleep_time) if terminate: break now = monotonic() if not self.transport.is_connected(): time.sleep(self.send_sleep) continue if self.send_sleep != 0 and now > self.next_outbound_heartbeat: logging.debug("Sending a heartbeat message at %s", now) try: self.transport.transmit(utils.Frame(None, {}, None)) except exception.NotConnectedException: logging.debug("Lost connection, unable to send heartbeat") except Exception: _, e, _ = sys.exc_info() logging.debug("Unable to send heartbeat, due to: %s", e) if self.receive_sleep != 0: diff_receive = now - self.received_heartbeat if diff_receive > self.receive_sleep: # heartbeat timeout logging.warning( "Heartbeat timeout: diff_receive=%s, time=%s, lastrec=%s", diff_receive, now, self.received_heartbeat) self.transport.set_connected(False) self.transport.disconnect_socket() self.transport.stop() for listener in self.transport.listeners.values(): listener.on_heartbeat_timeout() self.heartbeat_thread = None self.heartbeat_terminate_event.clear() logging.info('Heartbeat loop ended')
def on_heartbeat_timeout(self): """ Increment the heartbeat timeout. See :py:meth:`ConnectionListener.on_heartbeat_timeout` """ logging.debug("received heartbeat timeout") self.heartbeat_timeouts += 1
def attempt_connection(self): """ Try connecting to the (host, port) tuples specified at construction time. """ self.connection_error = False sleep_exp = 1 connect_count = 0 while self.running and self.socket is None and ( connect_count < self.__reconnect_attempts_max or self.__reconnect_attempts_max == -1): for host_and_port in self.__host_and_ports: try: logging.info("Attempting connection to host %s, port %s", host_and_port[0], host_and_port[1]) self.socket = socket.create_connection( host_and_port, self.__timeout) self.__enable_keepalive() need_ssl = self.__need_ssl(host_and_port) if need_ssl: # wrap socket ssl_params = self.get_ssl(host_and_port) if ssl_params['ca_certs']: cert_validation = ssl.CERT_REQUIRED else: cert_validation = ssl.CERT_NONE try: tls_context = ssl.create_default_context( cafile=ssl_params['ca_certs']) except AttributeError: tls_context = None if tls_context: # Wrap the socket for TLS certfile = ssl_params['cert_file'] keyfile = ssl_params['key_file'] password = ssl_params.get('password') if certfile and not keyfile: keyfile = certfile if certfile: tls_context.load_cert_chain( certfile, keyfile, password) if cert_validation is None or cert_validation == ssl.CERT_NONE: tls_context.check_hostname = False tls_context.verify_mode = cert_validation self.socket = tls_context.wrap_socket( self.socket, server_hostname=host_and_port[0]) else: # Old-style wrap_socket where we don't have a modern SSLContext (so no SNI) self.socket = ssl.wrap_socket( self.socket, keyfile=ssl_params['key_file'], certfile=ssl_params['cert_file'], cert_reqs=cert_validation, ca_certs=ssl_params['ca_certs'], ssl_version=ssl_params['ssl_version']) self.socket.settimeout(self.__timeout) if self.blocking is not None: self.socket.setblocking(self.blocking) # # Validate server cert # if need_ssl and ssl_params['cert_validator']: cert = self.socket.getpeercert() (ok, errmsg) = ssl_params['cert_validator']( cert, host_and_port[0]) if not ok: raise SSLError( "Server certificate validation failed: %s", errmsg) self.current_host_and_port = host_and_port logging.info("Established connection to host %s, port %s", host_and_port[0], host_and_port[1]) break except socket.error: self.socket = None connect_count += 1 logging.warning("Could not connect to host %s, port %s", host_and_port[0], host_and_port[1], exc_info=1) if self.socket is None: sleep_duration = (min(self.__reconnect_sleep_max, ( (self.__reconnect_sleep_initial / (1.0 + self.__reconnect_sleep_increase)) * math.pow(1.0 + self.__reconnect_sleep_increase, sleep_exp) )) * (1.0 + random.random() * self.__reconnect_sleep_jitter)) sleep_end = monotonic() + sleep_duration logging.debug( "Sleeping for %.1f seconds before attempting reconnect", sleep_duration) while self.running and monotonic() < sleep_end: time.sleep(0.2) if sleep_duration < self.__reconnect_sleep_max: sleep_exp += 1 if not self.socket: raise exception.ConnectFailedException()
def __read(self): """ Read the next frame(s) from the socket. :return: list of frames read :rtype: list(bytes) """ fastbuf = BytesIO() while self.running: try: try: c = self.receive() except exception.InterruptedException: logging.debug("socket read interrupted, restarting") continue except Exception: logging.debug("socket read error", exc_info=True) c = b'' if c is None or len(c) == 0: logging.debug("nothing received, raising CCE") raise exception.ConnectionClosedException() if c == b'\x0a' and not self.__recvbuf and not fastbuf.tell(): # # EOL to an empty receive buffer: treat as heartbeat. # Note that this may misdetect an optional EOL at end of frame as heartbeat in case the # previous receive() got a complete frame whose NUL at end of frame happened to be the # last byte of that read. But that should be harmless in practice. # fastbuf.close() return [c] fastbuf.write(c) if b'\x00' in c: # # Possible end of frame # break self.__recvbuf += fastbuf.getvalue() fastbuf.close() result = [] if self.__recvbuf and self.running: while True: pos = self.__recvbuf.find(b'\x00') if pos >= 0: frame = self.__recvbuf[0:pos] preamble_end_match = PREAMBLE_END_RE.search(frame) if preamble_end_match: preamble_end = preamble_end_match.start() content_length_match = BaseTransport.__content_length_re.search( frame[0:preamble_end]) if content_length_match: content_length = int( content_length_match.group('value')) content_offset = preamble_end_match.end() frame_size = content_offset + content_length if frame_size > len(frame): # # Frame contains NUL bytes, need to read more # if frame_size < len(self.__recvbuf): pos = frame_size frame = self.__recvbuf[0:pos] else: # # Haven't read enough data yet, exit loop and wait for more to arrive # break result.append(frame) pos += 1 # # Ignore optional EOLs at end of frame # while self.__recvbuf[pos:pos + 1] == b'\x0a': pos += 1 self.__recvbuf = self.__recvbuf[pos:] else: break return result