Ejemplo n.º 1
0
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()
Ejemplo n.º 2
0
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