path: root/kafka/client_async.py

import heapq
import itertools
import logging
import random
import select
import sys
import time

import six

import kafka.common as Errors # TODO: make Errors a separate class

from .cluster import ClusterMetadata
from .conn import BrokerConnection, ConnectionStates, collect_hosts
from .future import Future
from .protocol.metadata import MetadataRequest
from .protocol.produce import ProduceRequest

log = logging.getLogger(__name__)


class KafkaClient(object):
    """
    A network client for asynchronous request/response network i/o.
    This is an internal class used to implement the
    user-facing producer and consumer clients.

    This class is not thread-safe!
    """
    _bootstrap_servers = 'localhost'
    _client_id = 'kafka-python-0.10.0'
    _reconnect_backoff_ms = 50
    _retry_backoff_ms = 100
    _send_buffer_bytes = 131072
    _receive_buffer_bytes = 32768
    _request_timeout_ms = 40000
    _max_in_flight_requests_per_connection=5

    def __init__(self, **kwargs):
        for config in (
            'client_id', 'max_in_flight_requests_per_connection',
            'reconnect_backoff_ms', 'retry_backoff_ms',
            'send_buffer_bytes', 'receive_buffer_bytes',
            'request_timeout_ms', 'bootstrap_servers'
        ):
            if config in kwargs:
                setattr(self, '_' + config, kwargs.pop(config))

        self.cluster = ClusterMetadata(**kwargs)
        self._topics = set() # empty set will fetch all topic metadata
        self._metadata_refresh_in_progress = False
        self._conns = {}
        self._connecting = set()
        self._delayed_tasks = DelayedTaskQueue()
        self._last_bootstrap = 0
        self._bootstrap_fails = 0
        self._bootstrap(collect_hosts(self._bootstrap_servers))

    def _bootstrap(self, hosts):
        # Exponential backoff if bootstrap fails
        backoff_ms = self._reconnect_backoff_ms * 2 ** self._bootstrap_fails
        next_at = self._last_bootstrap + backoff_ms / 1000.0
        now = time.time()
        if next_at > now:
            log.debug("Sleeping %0.4f before bootstrapping again", next_at - now)
            time.sleep(next_at - now)
        self._last_bootstrap = time.time()

        metadata_request = MetadataRequest([])
        for host, port in hosts:
            log.debug("Attempting to bootstrap via node at %s:%s", host, port)
            bootstrap = BrokerConnection(
                host, port,
                client_id=self._client_id,
                receive_buffer_bytes=self._receive_buffer_bytes,
                send_buffer_bytes=self._send_buffer_bytes,
                request_timeout_ms=self._request_timeout_ms,
                max_in_flight_requests_per_connection=self._max_in_flight_requests_per_connection,
                reconnect_backoff_ms=self._reconnect_backoff_ms
            )
            bootstrap.connect()
            while bootstrap.state is ConnectionStates.CONNECTING:
                bootstrap.connect()
            if bootstrap.state is not ConnectionStates.CONNECTED:
                bootstrap.close()
                continue
            future = bootstrap.send(metadata_request)
            while not future.is_done:
                bootstrap.recv()
            if future.failed():
                bootstrap.close()
                continue
            self.cluster.update_metadata(future.value)

            # A cluster with no topics can return no broker metadata
            # in that case, we should keep the bootstrap connection
            if not len(self.cluster.brokers()):
                self._conns['bootstrap'] = bootstrap
            self._bootstrap_fails = 0
            break
        # No bootstrap found...
        else:
            log.error('Unable to bootstrap from %s', hosts)
            # Max exponential backoff is 2^12, x4000 (50ms -> 200s)
            self._bootstrap_fails = min(self._bootstrap_fails + 1, 12)

    def _can_connect(self, node_id):
        if node_id not in self._conns:
            if self.cluster.broker_metadata(node_id):
                return True
            return False
        conn = self._conns[node_id]
        return conn.state is ConnectionStates.DISCONNECTED and not conn.blacked_out()

    def _initiate_connect(self, node_id):
        """Initiate a connection to the given node"""
        broker = self.cluster.broker_metadata(node_id)
        if not broker:
            raise Errors.IllegalArgumentError('Broker %s not found in current cluster metadata', node_id)

        if node_id not in self._conns:
            log.debug("Initiating connection to node %s at %s:%s",
                      node_id, broker.host, broker.port)
            self._conns[node_id] = BrokerConnection(
                broker.host, broker.port,
                client_id=self._client_id,
                receive_buffer_bytes=self._receive_buffer_bytes,
                send_buffer_bytes=self._send_buffer_bytes,
                request_timeout_ms=self._request_timeout_ms,
                max_in_flight_requests_per_connection=self._max_in_flight_requests_per_connection,
                reconnect_backoff_ms=self._reconnect_backoff_ms
            )
        return self._finish_connect(node_id)

    def _finish_connect(self, node_id):
        if node_id not in self._conns:
            raise Errors.IllegalArgumentError('Node %s not found in connections', node_id)
        state = self._conns[node_id].connect()
        if state is ConnectionStates.CONNECTING:
            self._connecting.add(node_id)
        elif node_id in self._connecting:
            log.debug("Node %s connection state is %s", node_id, state)
            self._connecting.remove(node_id)
        return state

    def ready(self, node_id):
        """
        Begin connecting to the given node, return true if we are already
        connected and ready to send to that node.

        @param node_id The id of the node to check
        @return True if we are ready to send to the given node
        """
        if self.is_ready(node_id):
            return True

        if self._can_connect(node_id):
            # if we are interested in sending to a node
            # and we don't have a connection to it, initiate one
            self._initiate_connect(node_id)

        if node_id in self._connecting:
            self._finish_connect(node_id)

        return self.is_ready(node_id)

    def close(self, node_id=None):
        """Closes the connection to a particular node (if there is one).

        @param node_id The id of the node
        """
        if node_id is None:
            for conn in self._conns.values():
                conn.close()
        elif node_id in self._conns:
            self._conns[node_id].close()
        else:
            log.warning("Node %s not found in current connection list; skipping", node_id)
            return

    def connection_delay(self, node_id):
        """
        Returns the number of milliseconds to wait, based on the connection
        state, before attempting to send data. When disconnected, this respects
        the reconnect backoff time. When connecting or connected, this handles
        slow/stalled connections.

        @param node_id The id of the node to check
        @return The number of milliseconds to wait.
        """
        if node_id not in self._conns:
            return 0

        conn = self._conns[node_id]
        time_waited_ms = time.time() - (conn.last_attempt or 0)
        if conn.state is ConnectionStates.DISCONNECTED:
            return max(self._reconnect_backoff_ms - time_waited_ms, 0)
        else:
            return sys.maxint

    def connection_failed(self, node_id):
        """
        Check if the connection of the node has failed, based on the connection
        state. Such connection failures are usually transient and can be resumed
        in the next ready(node) call, but there are cases where transient
        failures need to be caught and re-acted upon.

        @param node_id the id of the node to check
        @return true iff the connection has failed and the node is disconnected
        """
        if node_id not in self._conns:
            return False
        return self._conns[node_id].state is ConnectionStates.DISCONNECTED

    def is_ready(self, node_id):
        """
        Check if the node with the given id is ready to send more requests.

        @param node_id The id of the node
        @return true if the node is ready
        """
        # if we need to update our metadata now declare all requests unready to
        # make metadata requests first priority
        if not self._metadata_refresh_in_progress and not self.cluster.ttl() == 0:
            if self._can_send_request(node_id):
                return True
        return False

    def _can_send_request(self, node_id):
        if node_id not in self._conns:
            return False
        conn = self._conns[node_id]
        return conn.connected() and conn.can_send_more()

    def send(self, node_id, request):
        """
        Send the given request. Requests can only be sent out to ready nodes.

        @param node destination node
        @param request The request
        @param now The current timestamp
        """
        if not self._can_send_request(node_id):
            raise Errors.IllegalStateError("Attempt to send a request to node %s which is not ready." % node_id)

        # Every request gets a response, except one special case:
        expect_response = True
        if isinstance(request, ProduceRequest) and request.required_acks == 0:
            expect_response = False

        return self._conns[node_id].send(request, expect_response=expect_response)

    def poll(self, timeout_ms=None, future=None):
        """Do actual reads and writes to sockets.

        @param timeout_ms The maximum amount of time to wait (in ms) for
                          responses if there are none available immediately.
                          Must be non-negative. The actual timeout will be the
                          minimum of timeout, request timeout and metadata
                          timeout. If unspecified, default to request_timeout_ms
        @param future Optionally block until the provided future completes.
        @return The list of responses received.
        """
        if timeout_ms is None:
            timeout_ms = self._request_timeout_ms

        responses = []

        # Loop for futures, break after first loop if None
        while True:

            # Attempt to complete pending connections
            for node_id in list(self._connecting):
                self._finish_connect(node_id)

            # Send a metadata request if needed
            metadata_timeout = self._maybe_refresh_metadata()

            # Send scheduled tasks
            for task in self._delayed_tasks.pop_ready():
                try:
                    task()
                except Exception as e:
                    log.error("Task %s failed: %s", task, e)

            timeout = min(timeout_ms, metadata_timeout, self._request_timeout_ms)
            timeout /= 1000.0

            responses.extend(self._poll(timeout))
            if not future or future.is_done:
                break

        return responses

    def _poll(self, timeout):
        # select on reads across all connected sockets, blocking up to timeout
        sockets = [conn._sock for conn in six.itervalues(self._conns)
                   if (conn.state is ConnectionStates.CONNECTED and
                       conn.in_flight_requests)]
        if sockets:
            select.select(sockets, [], [], timeout)

        responses = []
        # list, not iterator, because inline callbacks may add to self._conns
        for conn in list(self._conns.values()):
            if conn.state is ConnectionStates.CONNECTING:
                conn.connect()

            if conn.in_flight_requests:
                response = conn.recv() # This will run callbacks / errbacks
                if response:
                    responses.append(response)
        return responses

    def in_flight_request_count(self, node_id=None):
        """Get the number of in-flight requests"""
        if node_id is not None:
            if node_id not in self._conns:
                return 0
            return len(self._conns[node_id].in_flight_requests)
        else:
            return sum([len(conn.in_flight_requests) for conn in self._conns.values()])

    def least_loaded_node(self):
        """
        Choose the node with the fewest outstanding requests which is at least
        eligible for connection. This method will prefer a node with an
        existing connection, but will potentially choose a node for which we
        don't yet have a connection if all existing connections are in use.
        This method will never choose a node for which there is no existing
        connection and from which we have disconnected within the reconnect
        backoff period.

        @return The node_id with the fewest in-flight requests.
        """
        nodes = list(self._conns.keys())
        random.shuffle(nodes)
        inflight = sys.maxint
        found = None
        for node_id in nodes:
            conn = self._conns[node_id]
            curr_inflight = len(conn.in_flight_requests)
            if curr_inflight == 0 and conn.connected():
                # if we find an established connection with no in-flight requests we can stop right away
                return node_id
            elif not conn.blacked_out() and curr_inflight < inflight:
                # otherwise if this is the best we have found so far, record that
                inflight = curr_inflight
                found = node_id

        if found is not None:
            return found

        # if we found no connected node, return a disconnected one
        log.debug("No connected nodes found. Trying disconnected nodes.")
        for node_id in nodes:
            if not self._conns[node_id].is_blacked_out():
                return node_id

        # if still no luck, look for a node not in self._conns yet
        log.debug("No luck. Trying all broker metadata")
        for broker in self.cluster.brokers():
            if broker.nodeId not in self._conns:
                return broker.nodeId

        # Last option: try to bootstrap again
        log.error('No nodes found in metadata -- retrying bootstrap')
        self._bootstrap(collect_hosts(self._bootstrap_servers))
        return None

    def set_topics(self, topics):
        """
        Set specific topics to track for metadata

        Returns a future that will complete after metadata request/response
        """
        if set(topics).difference(self._topics):
            future = self.cluster.request_update()
        else:
            future = Future().success(set(topics))
        self._topics = set(topics)
        return future

    # request metadata update on disconnect and timedout
    def _maybe_refresh_metadata(self):
        """Send a metadata request if needed"""
        ttl = self.cluster.ttl()
        if ttl > 0:
            return ttl

        if self._metadata_refresh_in_progress:
            return sys.maxint

        node_id = self.least_loaded_node()

        if self._can_send_request(node_id):
            request = MetadataRequest(list(self._topics))
            log.debug("Sending metadata request %s to node %s", request, node_id)
            future = self.send(node_id, request)
            future.add_callback(self.cluster.update_metadata)
            future.add_errback(self.cluster.failed_update)

            self._metadata_refresh_in_progress = True
            def refresh_done(val_or_error):
                self._metadata_refresh_in_progress = False
            future.add_callback(refresh_done)
            future.add_errback(refresh_done)

        elif self._can_connect(node_id):
            log.debug("Initializing connection to node %s for metadata request", node_id)
            self._initiate_connect(node_id)

        return 0

    def schedule(self, task, at):
        """
        Schedule a new task to be executed at the given time.

        This is "best-effort" scheduling and should only be used for coarse
        synchronization. A task cannot be scheduled for multiple times
        simultaneously; any previously scheduled instance of the same task
        will be cancelled.

        @param task The task to be scheduled -- function or implement __call__
        @param at Epoch seconds when it should run (see time.time())
        @returns Future
        """
        return self._delayed_tasks.add(task, at)

    def unschedule(self, task):
        """
        Unschedule a task. This will remove all instances of the task from the task queue.
        This is a no-op if the task is not scheduled.

        @param task The task to be unscheduled.
        """
        self._delayed_tasks.remove(task)


class DelayedTaskQueue(object):
    # see https://docs.python.org/2/library/heapq.html
    def __init__(self):
        self._tasks = [] # list of entries arranged in a heap
        self._task_map = {} # mapping of tasks to entries
        self._counter = itertools.count() # unique sequence count

    def add(self, task, at):
        """Add a task to run at a later time

        task: anything
        at: seconds from epoch to schedule task (see time.time())
        """
        if task in self._task_map:
            self.remove(task)
        count = next(self._counter)
        future = Future()
        entry = [at, count, (task, future)]
        self._task_map[task] = entry
        heapq.heappush(self._tasks, entry)
        return future

    def remove(self, task):
        """Remove a previously scheduled task

        Raises KeyError if task is not found
        """
        entry = self._task_map.pop(task)
        task, future = entry[-1]
        future.failure(Errors.Cancelled)
        entry[-1] = 'REMOVED'

    def _drop_removed(self):
        while self._tasks and self._tasks[0][-1] is 'REMOVED':
            at, count, task = heapq.heappop(self._tasks)

    def _pop_next(self):
        self._drop_removed()
        if not self._tasks:
            raise KeyError('pop from an empty DelayedTaskQueue')
        _, _, maybe_task = heapq.heappop(self._tasks)
        if maybe_task is 'REMOVED':
            raise ValueError('popped a removed tasks from queue - bug')
        else:
            task, future = maybe_task
        del self._task_map[task]
        return task

    def next_at(self):
        """Number of seconds until next task is ready"""
        self._drop_removed()
        if not self._tasks:
            return sys.maxint
        else:
            return max(self._tasks[0][0] - time.time(), 0)

    def pop_ready(self):
        """Pop and return a list of all ready (task, future) tuples"""
        self._drop_removed()
        ready_tasks = []
        while self._tasks and self._tasks[0][0] < time.time():
            ready_tasks.append(self._pop_next())
        return ready_tasks