class ZTcpConnection(object):
DISCONNECTED = 0
CONNECTING = 1
CONNECTED = 2
DISCONNECTING = 3
def __init__(self, event_loop, name, sock, peer_addr):
self._event_loop = event_loop
self._name = name
self._socket = sock
self._peer_addr = peer_addr
self._channel = ZChannel(event_loop, sock.fileno())
self._channel.set_read_callback(self.handle_read)
self._channel.set_write_callback(self.handle_write)
self._channel.set_close_callback(self.handle_close)
self._channel.set_error_callback(self.handle_error)
self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
self._high_watermark = 64 * 1024 * 1024
self._state = self.CONNECTING
self._connection_callback = None
self._message_callback = None
self._close_callback = None
self._write_complete_callback = None
self._high_watermark_callback = None
self._context = None
self._input_buffer = ZBuffer()
self._output_buffer = ZBuffer()
def name(self):
return self._name
def get_event_loop(self):
return self._event_loop
def set_connection_callback(self, cb):
self._connection_callback = cb
def set_message_callback(self, cb):
self._message_callback = cb
def set_close_callback(self, cb):
self._close_callback = cb
def set_write_complete_callback(self, cb):
self._write_complete_callback = cb
def set_high_watermark_callback(self, cb):
self._high_watermark_callback = cb
def _write_complete_callback_wrapper(self):
self._write_complete_callback(self)
def set_context(self, ct):
self._context = ct
def get_context(self):
return self._context
def send(self, data):
def send_in_loop_wrapper():
self.send_in_loop(data)
if self._state == self.CONNECTED:
if self._event_loop.is_in_loop_thread():
self.send_in_loop(data)
else:
self._event_loop.run_in_loop(send_in_loop_wrapper)
def send_in_loop(self, data):
self._event_loop.assert_in_loop_thread()
if self._state == self.DISCONNECTED:
logging.warn('disconnected connection, give up writing')
return
siz = len(data)
remaining = siz
has_fatal_error = False
try:
if not self._channel.is_writing() and self._output_buffer.readable_bytes() == 0:
nwritten = self._socket.send(data)
remaining = siz - nwritten
if remaining == 0 and self._write_complete_callback:
self._event_loop.queue_in_loop(self._write_complete_callback_wrapper)
except socket.error as e:
logging.error('ZTcpConnection.send_in_loop fail to write with error %s' %str(e))
if e.errno != errno.EWOULDBLOCK:
if e.errno == errno.EPIPE:
has_fatal_error = True
if not has_fatal_error and remaining > 0:
logging.info('ZTcpConnection.send_in_loop I am going to write more data')
old_len = self._output_buffer.readable_bytes()
if old_len + remaining >= self._high_watermark and \
old_len < self._high_watermark and \
self._high_watermark_callback:
siz = old_len + remaining
def high_watermark_callback_wrapper():
self._high_watermark_callback(self, siz)
self._event_loop.queue_in_loop(_high_watermark_callback_wrapper)
self._output_buffer.append(data[nwritten:])
if not self._channel.is_writting():
self._channel.enable_writing()
def handle_read(self, receive_time):
self._event_loop.assert_in_loop_thread()
# FIXME recv_into
try:
data = self._socket.recv(65536)
if data:
self._input_buffer.append(data)
self._message_callback(self, self._input_buffer, receive_time)
else:
self.handle_close()
except socket.error as e:
logging.error('ZTcpConnection.handle_read fail to read with error %s' %str(e))
self.handle_error()
def handle_write(self):
self._event_loop.assert_in_loop_thread()
if self._channel.is_writing():
try:
n = self._socket.send(self._ouput_buffer.tostring())
self._output_buffer.retrieve(n)
if self._output_buffer.readable_bytes() == 0:
self._channel.disable_writing()
if self._write_complete_callback:
self._event_loop.queue_in_loop(self._write_complete_callback_wrapper)
if self._state == self.DISCONNECTING:
self.shutdown_in_loop()
except socket.error as e:
logging.error('ZTcpConnection.handle_write fail to write with error %s' %str(e))
else:
logging.info('ZTcpConnection.handle_write connection is down, no more writing')
def _set_state(self, st):
self._state = st
def connected(self):
return self._state == self.CONNECTED
def handle_close(self):
self._event_loop.assert_in_loop_thread()
assert self._state == self.CONNECTED or self._state == self.DISCONNECTING
self._set_state(self.DISCONNECTED)
self._channel.disable_all()
self._connection_callback(self)
self._close_callback(self)
def handle_error(self):
logging.error('ZTcpConnection.handle_error error happened')
def shutdown_write(self):
if self._state == self.CONNECTED:
self._set_state(self.DISCONNECTING)
self._event_loop.run_in_loop(self.shutdown_write_in_loop)
def shutdown(self):
if self._state == self.CONNECTED:
self._set_state(self.DISCONNECTING)
self._event_loop.run_in_loop(self.shutdown_in_loop)
def shutdown_write_in_loop(self):
self._event_loop.assert_in_loop_thread()
# FIXME what if channel is writing ?
if not self._channel.is_writing():
self._socket.shutdown(socket.SHUT_WR)
def shutdown_in_loop(self):
self._event_loop.assert_in_loop_thread()
# FIXME what if channel is writing ?
if not self._channel.is_writing():
self._socket.shutdown(socket.SHUT_RDWR)
def set_tcp_no_delay(on):
if on:
self._socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
else:
self._socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 0)
def connection_established(self):
self._event_loop.assert_in_loop_thread()
assert self._state == self.CONNECTING
self._set_state(self.CONNECTED)
self._channel.tie(self)
self._channel.enable_reading()
self._connection_callback(self)
def connection_destroyed(self):
self._event_loop.assert_in_loop_thread()
if self._state == self.CONNECTED:
self._set_state(self.DISCONNECTED)
self._channel.disable_all()
self._connection_callback(self)
self._channel.unregister()
def send(self, tcp_conn, message_str):
buf = ZBuffer()
buf.append(message_str)
buf.prepend_int32(socket.htonl(len(message_str)))
tcp_conn.send(buf.retrieve_all_as_string())