class ZTcpAcceptor(object):
def __init__(self, event_loop, listen_addr):
self._event_loop = event_loop
self._accept_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self._accept_socket.setblocking(False)
self._accept_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
self._accept_socket.bind(listen_addr)
self._accept_channel = ZChannel(event_loop, self._accept_socket.fileno())
self._accept_channel.set_read_callback(self.handle_read)
self._new_connection_callback = None
self._listening = False
def listen(self):
self._event_loop.assert_in_loop_thread()
self._listening = True
self._accept_socket.listen(socket.SOMAXCONN)
self._accept_channel.enable_reading()
def listening(self):
return self._listening
def handle_read(self, receive_time):
self._event_loop.assert_in_loop_thread()
try:
sock, peer_addr = self._accept_socket.accept()
sock.setblocking(False)
if self._new_connection_callback:
self._new_connection_callback(sock, peer_addr)
else:
conn.close()
except socket.error as e:
logging.error('ZTcpAcceptor.handle_read (%s)' %str(e))
def set_new_connection_callback(self, cb):
self._new_connection_callback = cb
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 ZEventLoop(object):
def __init__(self):
self._poller = default_poller(self)
self._tid = current_thread().ident
self._looping = False
self._quit = False
self._doing_pending_jobs = False
self._pending_jobs = []
self._mutex = Lock()
self._handling_event = False
# FIXME Hack for Windows
if system() != 'Windows':
(rfd, wfd) = pipe()
self._wakeup_rfd = rfd
self._wakeup_wfd = wfd
self._wakeup_channel = ZChannel(self, self._wakeup_rfd)
self._wakeup_channel.set_read_callback(self._handle_read)
self._wakeup_channel.set_close_callback(self._handle_close)
self._wakeup_channel.enable_reading()
def loop(self):
assert not self._looping
self.assert_in_loop_thread()
self._looping = True
active_channels = []
while not self._quit:
del active_channels[:]
try:
poll_return_time = self._poller.poll(active_channels, 1)
self._handling_event = True
for ch in active_channels:
ch.handle_event(poll_return_time)
self._handling_event = False
self._do_pending_jobs()
except Exception as e:
exc_type, exc_value, exc_traceback = exc_info()
logging.error('ZEventLoop.loop encouting error %s'
%(''.join(format_exception(exc_type, exc_value, exc_traceback))))
def quit(self):
self._quit = True
if not self.is_in_loop_thread():
self._wakeup()
def run_in_loop(self, cb):
if self.is_in_loop_thread():
cb()
else:
self.queue_in_loop(cb)
def queue_in_loop(self, cb):
with self._mutex:
self._pending_jobs.append(cb)
if not self.is_in_loop_thread() or self._doing_pending_jobs:
self._wakeup()
def update_channel(self, channel):
assert channel.owner_loop() == self
self.assert_in_loop_thread()
self._poller.update(channel)
def unregister_channel(self, channel):
assert channel.owner_loop() == self
self.assert_in_loop_thread()
assert channel.has_no_events()
self._poller.unregister(channel)
def assert_in_loop_thread(self):
assert self.is_in_loop_thread()
def is_in_loop_thread(self):
return self._tid == current_thread().ident
def _wakeup(self):
if system() != 'Windows':
cnt = write(self._wakeup_wfd, '1')
if cnt != 1:
logging.error('ZEventLoop._wakeup() failed to wakeup')
def _handle_read(self, receive_time):
data = read(self._wakeup_rfd, 1)
if int(data) != 1:
logging.error('ZEventLoop._handle_read() failed to read')
def _handle_close(self):
def unregister_wakeup_channel():
logging.error('ZTcpAcceptor._handle_close unregister wakeup channel')
self.unregister_channel(self._wakeup_channel)
self.run_in_loop(unregister_wakeup_channel)
def _do_pending_jobs(self):
self._doing_pending_jobs = True
jobs = []
with self._mutex:
jobs.extend(self._pending_jobs)
del self._pending_jobs[:]
for job in jobs:
job()
self._doing_pending_jobs = False