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()
class ZTcpConnector(object):
DISCONNECTED = 0
CONNECTING = 1
CONNECTED = 2
MAX_RETRY_DELAY = 30
INIT_RETRY_DELAY = 1
def __init__(self, event_loop, server_addr):
self._event_loop = event_loop
self._server_addr = server_addr
self._connect = False
self._state = self.DISCONNECTED
self._channel = None
self._sock = None
self._new_connection_callback = None
self._connect_error_callback = None
self._retry_delay = self.INIT_RETRY_DELAY
def set_new_connection_callback(self, cb):
self._new_connection_callback = cb
def set_connect_error_callback(self, cb):
self._connect_error_callback = cb
def start(self):
self._connect = True
self._event_loop.run_in_loop(self._start_in_loop)
def restart(self):
self._event_loop.assert_in_loop_thread()
self._set_state(self.DISCONNECTED)
self._retry_delay = self.INIT_RETRY_DELAY
self._connect = True
self._start_in_loop()
def stop(self):
self._connect = False
self._event_loop.run_in_loop(self._stop_in_loop)
def server_addr(self):
return self._server_addr
def _set_state(self, s):
self._state = s
def _start_in_loop(self):
self._event_loop.assert_in_loop_thread()
assert self._state == self.DISCONNECTED
if self._connect:
self._do_connect()
def _stop_in_loop(self):
self._event_loop.assert_in_loop_thread()
if self._state == self.CONNECTING:
self._set_state(self.DISCONNECTED)
self._remove_and_reset_channel()
self._retry()
def _do_connect(self):
self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
self._sock.connect(self._server_addr)
self._connecting()
except socket.error as e:
err = e.errno
if err in (EINPROGRESS, EINTR, EISCONN):
self._connecting()
elif err in (EAGAIN, EADDRINUSE, EADDRNOTAVAIL, ECONNREFUSED, ENETUNREACH):
self._retry()
elif err in (EACCES, EPERM, EAFNOSUPPORT, EALREADY, EBADF, EFAULT, ENOTSOCK):
self._sock.close()
if self._connect_error_callback:
self._connect_error_callback(err)
else:
logging.error(
"ZTcpConnector._connect connect to %s with error code %d" % (str(self._server_addr), err)
)
else:
self._sock.close()
if self._connect_error_callback:
self._connect_error_callback(self._server_addr, err)
else:
logging.error(
"ZTcpConnector._connect connect to %s with unexpected error code %d"
% (str(self._server_addr), err)
)
def _connecting(self):
self._set_state(self.CONNECTING)
assert self._channel is None
self._channel = ZChannel(self._event_loop, self._sock.fileno())
self._channel.set_write_callback(self._handle_write)
self._channel.set_error_callback(self._handle_error)
self._channel.enable_writing()
def _handle_write(self):
logging.info("ZTcpConnector._handle_write %d" % (self._state))
if self._state == self.CONNECTING:
self._remove_and_reset_channel()
err = self._sock.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
if err != 0:
logging.warn("ZTcpConnector._handle_write SO_ERROR = %d" % err)
self._retry()
# FIXME check self connect
else:
self._set_state(self.CONNECTED)
if self._connect:
self._new_connection_callback(self._sock)
else:
self._sock.close()
else:
assert self._state == self.DISCONNECTED
def _handle_error(self):
assert self._state == self.CONNECTING
self._remove_and_reset_channel()
err = self._sock.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
logging.warn("ZTcpConnector._handle_error SO_ERROR = %d" % err)
self.retry()
def _retry(self):
import time
self._sock.close()
self._set_state(self.DISCONNECTED)
if self._connect:
logging.info(
"ZTcpConnector._retry - retring connecting to %r in %d" % (self._server_addr, self._retry_delay)
)
# FIXME use timer
# add in the timer queue
time.sleep(self._retry_delay)
self._retry_delay *= 2
if self._retry_delay > self.MAX_RETRY_DELAY:
self._retry_delay = self.MAX_RETRY_DELAY
def _remove_and_reset_channel(self):
self._channel.disable_all()
self._channel.unregister()
# Can't reset _channel here, because we are inside ZChannel::handle_event
self._event_loop.queue_in_loop(self._reset_channel)
def _reset_channel(self):
self._channel = None