Пример #1
0
    def upgrade_to_snappy(self):
        assert SnappySocket, 'snappy requires the python-snappy package'

        # in order to upgrade to Snappy we need to use whatever IOStream
        # is currently in place (normal or SSL)...
        #
        # first read any compressed bytes the existing IOStream might have
        # already buffered and use that to bootstrap the SnappySocket, then
        # monkey patch the existing IOStream by replacing its socket
        # with a wrapper that will automagically handle compression.
        existing_data = self.stream._consume(self.stream._read_buffer_size)
        self.socket = SnappySocket(self.socket)
        self.socket.bootstrap(existing_data)
        self.stream.socket = self.socket
Пример #2
0
class AsyncConn(EventedMixin):
    """
    Low level object representing a TCP connection to nsqd.

    When a message on this connection is requeued and the requeue delay has not been specified,
    it calculates the delay automatically by an increasing multiple of ``requeue_delay``.

    Generates the following events that can be listened to with :meth:`nsq.AsyncConn.on`:

     * ``connect``
     * ``close``
     * ``error``
     * ``identify``
     * ``identify_response``
     * ``heartbeat``
     * ``ready``
     * ``message``
     * ``response``
     * ``backoff``
     * ``resume``

    :param host: the host to connect to

    :param port: the post to connect to

    :param timeout: the timeout for read/write operations (in seconds)

    :param heartbeat_interval: the amount of time in seconds to negotiate with the connected
        producers to send heartbeats (requires nsqd 0.2.19+)

    :param requeue_delay: the base multiple used when calculating requeue delay
        (multiplied by # of attempts)

    :param tls_v1: enable TLS v1 encryption (requires nsqd 0.2.22+)

    :param tls_options: dictionary of options to pass to `ssl.wrap_socket()
        <http://docs.python.org/2/library/ssl.html#ssl.wrap_socket>`_ as ``**kwargs``

    :param snappy: enable Snappy stream compression (requires nsqd 0.2.23+)

    :param output_buffer_size: size of the buffer (in bytes) used by nsqd for buffering writes
        to this connection

    :param output_buffer_timeout: timeout (in ms) used by nsqd before flushing buffered writes
        (set to 0 to disable).  **Warning**: configuring clients with an extremely low (``< 25ms``)
        ``output_buffer_timeout`` has a significant effect on ``nsqd`` CPU usage (particularly
        with ``> 50`` clients connected).

    :param sample_rate: take only a sample of the messages being sent to the client. Not setting
        this or setting it to 0 will ensure you get all the messages destined for the client.
        Sample rate can be greater than 0 or less than 100 and the client will receive that
        percentage of the message traffic. (requires nsqd 0.2.25+)

     :param user_agent: a string identifying the agent for this client in the spirit of
        HTTP (default: "<client_library_name>/<version>") (requires nsqd 0.2.25+)

    """
    def __init__(self, host, port, timeout=1.0, heartbeat_interval=30, requeue_delay=90,
                 tls_v1=False, tls_options=None, snappy=False, user_agent=None,
                 output_buffer_size=16 * 1024, output_buffer_timeout=250, sample_rate=0):
        assert isinstance(host, (str, unicode))
        assert isinstance(port, int)
        assert isinstance(timeout, float)
        assert isinstance(tls_options, (dict, None.__class__))
        assert isinstance(heartbeat_interval, int) and heartbeat_interval >= 1
        assert isinstance(requeue_delay, int) and requeue_delay >= 0
        assert isinstance(output_buffer_size, int) and output_buffer_size >= 0
        assert isinstance(output_buffer_timeout, int) and output_buffer_timeout >= 0
        assert isinstance(sample_rate, int) and sample_rate >= 0 and sample_rate < 100
        assert tls_v1 and ssl or not tls_v1, \
            'tls_v1 requires Python 2.6+ or Python 2.5 w/ pip install ssl'

        self.state = 'INIT'
        self.host = host
        self.port = port
        self.timeout = timeout
        self.last_recv_timestamp = time.time()
        self.last_msg_timestamp = time.time()
        self.in_flight = 0
        self.rdy = 0
        self.rdy_timeout = None
        # for backwards compatibility when interacting with older nsqd
        # (pre 0.2.20), default this to their hard-coded max
        self.max_rdy_count = 2500
        self.tls_v1 = tls_v1
        self.tls_options = tls_options
        self.snappy = snappy
        self.hostname = socket.gethostname()
        self.short_hostname = self.hostname.split('.')[0]
        self.heartbeat_interval = heartbeat_interval
        self.requeue_delay = requeue_delay

        self.output_buffer_size = output_buffer_size
        self.output_buffer_timeout = output_buffer_timeout
        self.sample_rate = sample_rate
        self.user_agent = user_agent

        if self.user_agent is None:
            self.user_agent = 'pynsq/%s' % __version__

        super(AsyncConn, self).__init__()

    @property
    def id(self):
        return str(self)

    def __str__(self):
        return self.host + ':' + str(self.port)

    def connect(self):
        if self.state not in ['INIT', 'DISCONNECTED']:
            return

        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.settimeout(self.timeout)
        self.socket.setblocking(0)

        self.stream = tornado.iostream.IOStream(self.socket)
        self.stream.set_close_callback(self._socket_close)

        self.state = 'CONNECTING'
        self.on('connect', self._on_connect)
        self.on('data', self._on_data)

        self.stream.connect((self.host, self.port), self._connect_callback)

    def _connect_callback(self):
        self.state = 'CONNECTED'
        self.stream.write(nsq.MAGIC_V2)
        self._start_read()
        self.trigger('connect', conn=self)

    def _start_read(self):
        self.stream.read_bytes(4, self._read_size)

    def _socket_close(self):
        self.state = 'DISCONNECTED'
        self.trigger('close', conn=self)

    def close(self):
        self.stream.close()

    def _read_size(self, data):
        try:
            size = struct.unpack('>l', data)[0]
            self.stream.read_bytes(size, self._read_body)
        except Exception:
            self.close()
            self.trigger('error', nsq.IntegrityError('failed to unpack size'))

    def _read_body(self, data):
        try:
            self.trigger('data', conn=self, data=data)
        except Exception:
            logging.exception('uncaught exception in data event')
        tornado.ioloop.IOLoop.instance().add_callback(self._start_read)

    def send(self, data):
        self.stream.write(data)

    def upgrade_to_tls(self, options=None):
        assert ssl, 'tls_v1 requires Python 2.6+ or Python 2.5 w/ pip install ssl'

        # in order to upgrade to TLS we need to *replace* the IOStream...
        #
        # first remove the event handler for the currently open socket
        # so that when we add the socket to the new SSLIOStream below,
        # it can re-add the appropriate event handlers.
        tornado.ioloop.IOLoop.instance().remove_handler(self.socket.fileno())

        opts = {
            'cert_reqs': ssl.CERT_REQUIRED,
            'ca_certs': tornado.simple_httpclient._DEFAULT_CA_CERTS
        }
        opts.update(options or {})
        self.socket = ssl.wrap_socket(self.socket, ssl_version=ssl.PROTOCOL_TLSv1,
                                      do_handshake_on_connect=False, **opts)

        self.stream = tornado.iostream.SSLIOStream(self.socket)
        self.stream.set_close_callback(self._socket_close)

        # now that the IOStream has been swapped we can kickstart
        # the SSL handshake
        self.stream._do_ssl_handshake()

    def upgrade_to_snappy(self):
        assert SnappySocket, 'snappy requires the python-snappy package'

        # in order to upgrade to Snappy we need to use whatever IOStream
        # is currently in place (normal or SSL)...
        #
        # first read any compressed bytes the existing IOStream might have
        # already buffered and use that to bootstrap the SnappySocket, then
        # monkey patch the existing IOStream by replacing its socket
        # with a wrapper that will automagically handle compression.
        existing_data = self.stream._consume(self.stream._read_buffer_size)
        self.socket = SnappySocket(self.socket)
        self.socket.bootstrap(existing_data)
        self.stream.socket = self.socket

    def send_rdy(self, value):
        try:
            self.send(nsq.ready(value))
        except Exception, e:
            self.close()
            self.trigger('error', conn=self,
                         error=nsq.SendError('failed to send RDY %d' % value, e))
            return False
        self.last_rdy = value
        self.rdy = value
        return True