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
