<s> from __future__ import division <EOL> """<STR_LIT>""" <EOL> __license__ = """<STR_LIT>""" <EOL> __all__ = [ "<STR_LIT>" ] <EOL> import itertools <EOL> import logging <EOL> import socket <EOL> import sys <EOL> import time <EOL> import traceback <EOL> from uuid import uuid4 <EOL> import weakref <EOL> from kazoo . client import KazooClient <EOL> from kazoo . handlers . gevent import SequentialGeventHandler <EOL> from kazoo . exceptions import NoNodeException , NodeExistsError <EOL> from kazoo . recipe . watchers import ChildrenWatch <EOL> from . common import OffsetType <EOL> from . exceptions import KafkaException , PartitionOwnedError , ConsumerStoppedException <EOL> from . handlers import GEventHandler <EOL> from . simpleconsumer import SimpleConsumer <EOL> from . utils . compat import range , get_bytes , itervalues , iteritems , get_string <EOL> try : <EOL> from . import rdkafka <EOL> except ImportError : <EOL> rdkafka = False <EOL> log = logging . getLogger ( __name__ ) <EOL> def _catch_thread_exception ( fn ) : <EOL> """<STR_LIT>""" <EOL> def wrapped ( self , * args , ** kwargs ) : <EOL> try : <EOL> ret = fn ( self , * args , ** kwargs ) <EOL> except Exception : <EOL> self . _worker_exception = sys . exc_info ( ) <EOL> else : <EOL> return ret <EOL> return wrapped <EOL> class BalancedConsumer ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , <EOL> topic , <EOL> cluster , <EOL> consumer_group , <EOL> fetch_message_max_bytes = <NUM_LIT> * <NUM_LIT> , <EOL> num_consumer_fetchers = <NUM_LIT:1> , <EOL> auto_commit_enable = False , <EOL> auto_commit_interval_ms = <NUM_LIT> * <NUM_LIT:1000> , <EOL> queued_max_messages = <NUM_LIT> , <EOL> fetch_min_bytes = <NUM_LIT:1> , <EOL> fetch_wait_max_ms = <NUM_LIT:100> , <EOL> offsets_channel_backoff_ms = <NUM_LIT:1000> , <EOL> offsets_commit_max_retries = <NUM_LIT:5> , <EOL> auto_offset_reset = OffsetType . EARLIEST , <EOL> consumer_timeout_ms = - <NUM_LIT:1> , <EOL> rebalance_max_retries = <NUM_LIT:5> , <EOL> rebalance_backoff_ms = <NUM_LIT:2> * <NUM_LIT:1000> , <EOL> zookeeper_connection_timeout_ms = <NUM_LIT:6> * <NUM_LIT:1000> , <EOL> zookeeper_connect = '<STR_LIT>' , <EOL> zookeeper = None , <EOL> auto_start = True , <EOL> reset_offset_on_start = False , <EOL> post_rebalance_callback = None , <EOL> use_rdkafka = False , <EOL> compacted_topic = False ) : <EOL> """<STR_LIT>""" <EOL> self . _cluster = cluster <EOL> if not isinstance ( consumer_group , bytes ) : <EOL> raise TypeError ( "<STR_LIT>" ) <EOL> self . _consumer_group = consumer_group <EOL> self . _topic = topic <EOL> self . _auto_commit_enable = auto_commit_enable <EOL> self . _auto_commit_interval_ms = auto_commit_interval_ms <EOL> self . _fetch_message_max_bytes = fetch_message_max_bytes <EOL> self . _fetch_min_bytes = fetch_min_bytes <EOL> self . _rebalance_max_retries = rebalance_max_retries <EOL> self . _num_consumer_fetchers = num_consumer_fetchers <EOL> self . _queued_max_messages = queued_max_messages <EOL> self . _fetch_wait_max_ms = fetch_wait_max_ms <EOL> self . _rebalance_backoff_ms = rebalance_backoff_ms <EOL> self . _consumer_timeout_ms = consumer_timeout_ms <EOL> self . _offsets_channel_backoff_ms = offsets_channel_backoff_ms <EOL> self . _offsets_commit_max_retries = offsets_commit_max_retries <EOL> self . _auto_offset_reset = auto_offset_reset <EOL> self . _zookeeper_connect = zookeeper_connect <EOL> self . _zookeeper_connection_timeout_ms = zookeeper_connection_timeout_ms <EOL> self . _reset_offset_on_start = reset_offset_on_start <EOL> self . _post_rebalance_callback = post_rebalance_callback <EOL> self . _generation_id = - <NUM_LIT:1> <EOL> self . _running = False <EOL> self . _worker_exception = None <EOL> self . _worker_trace_logged = False <EOL> self . _is_compacted_topic = compacted_topic <EOL> if not rdkafka and use_rdkafka : <EOL> raise ImportError ( "<STR_LIT>" ) <EOL> if isinstance ( self . _cluster . handler , GEventHandler ) and use_rdkafka : <EOL> raise ImportError ( "<STR_LIT>" ) <EOL> self . _use_rdkafka = rdkafka and use_rdkafka <EOL> self . _rebalancing_lock = cluster . handler . Lock ( ) <EOL> self . _consumer = None <EOL> self . _consumer_id = get_bytes ( "<STR_LIT>" . format ( <EOL> hostname = socket . gethostname ( ) , <EOL> uuid = uuid4 ( ) <EOL> ) ) <EOL> self . _setting_watches = True <EOL> self . _topic_path = '<STR_LIT>' . format ( <EOL> group = self . _consumer_group , <EOL> topic = self . _topic . name ) <EOL> self . _consumer_id_path = '<STR_LIT>' . format ( <EOL> group = self . _consumer_group ) <EOL> self . _zookeeper = None <EOL> self . _owns_zookeeper = zookeeper is None <EOL> if zookeeper is not None : <EOL> self . _zookeeper = zookeeper <EOL> if auto_start is True : <EOL> self . start ( ) <EOL> def __del__ ( self ) : <EOL> log . debug ( "<STR_LIT>" . format ( self ) ) <EOL> if self . _running : <EOL> self . stop ( ) <EOL> def __repr__ ( self ) : <EOL> return "<STR_LIT>" . format ( <EOL> module = self . __class__ . __module__ , <EOL> name = self . __class__ . __name__ , <EOL> id_ = hex ( id ( self ) ) , <EOL> group = self . _consumer_group <EOL> ) <EOL> def _raise_worker_exceptions ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . _worker_exception is not None : <EOL> _ , ex , tb = self . _worker_exception <EOL> if not self . _worker_trace_logged : <EOL> self . _worker_trace_logged = True <EOL> log . error ( "<STR_LIT>" , <EOL> "<STR_LIT>" . join ( traceback . format_tb ( tb ) ) ) <EOL> raise ex <EOL> @ property <EOL> def topic ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _topic <EOL> @ property <EOL> def partitions ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _consumer . partitions if self . _consumer else dict ( ) <EOL> @ property <EOL> def _partitions ( self ) : <EOL> """<STR_LIT>""" <EOL> return set ( <EOL> [ ] if self . partitions is None else itervalues ( self . partitions ) ) <EOL> @ property <EOL> def held_offsets ( self ) : <EOL> """<STR_LIT>""" <EOL> if not self . _consumer : <EOL> return None <EOL> return self . _consumer . held_offsets <EOL> def start ( self ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> if self . _zookeeper is None : <EOL> self . _setup_zookeeper ( self . _zookeeper_connect , <EOL> self . _zookeeper_connection_timeout_ms ) <EOL> self . _zookeeper . ensure_path ( self . _topic_path ) <EOL> self . _add_self ( ) <EOL> self . _running = True <EOL> self . _set_watches ( ) <EOL> self . _rebalance ( ) <EOL> except Exception : <EOL> log . exception ( "<STR_LIT>" ) <EOL> self . stop ( ) <EOL> def stop ( self ) : <EOL> """<STR_LIT>""" <EOL> log . debug ( "<STR_LIT>" . format ( self ) ) <EOL> with self . _rebalancing_lock : <EOL> self . _running = False <EOL> if self . _consumer is not None : <EOL> self . _consumer . stop ( ) <EOL> if self . _owns_zookeeper : <EOL> self . _zookeeper . stop ( ) <EOL> else : <EOL> self . _remove_partitions ( self . _get_held_partitions ( ) ) <EOL> try : <EOL> self . _zookeeper . delete ( self . _path_self ) <EOL> except NoNodeException : <EOL> pass <EOL> def _setup_zookeeper ( self , zookeeper_connect , timeout ) : <EOL> """<STR_LIT>""" <EOL> kazoo_kwargs = { '<STR_LIT>' : timeout / <NUM_LIT:1000> } <EOL> if isinstance ( self . _cluster . handler , GEventHandler ) : <EOL> kazoo_kwargs [ '<STR_LIT>' ] = SequentialGeventHandler ( ) <EOL> self . _zookeeper = KazooClient ( zookeeper_connect , ** kazoo_kwargs ) <EOL> self . _zookeeper . start ( ) <EOL> def _setup_internal_consumer ( self , partitions = None , start = True ) : <EOL> """<STR_LIT>""" <EOL> if partitions is None : <EOL> partitions = [ ] <EOL> if partitions != self . _partitions : <EOL> cns = self . _get_internal_consumer ( partitions = list ( partitions ) , start = start ) <EOL> if self . _post_rebalance_callback is not None : <EOL> old_offsets = ( self . _consumer . held_offsets <EOL> if self . _consumer else dict ( ) ) <EOL> new_offsets = cns . held_offsets <EOL> try : <EOL> reset_offsets = self . _post_rebalance_callback ( <EOL> self , old_offsets , new_offsets ) <EOL> except Exception : <EOL> log . exception ( "<STR_LIT>" ) <EOL> self . _worker_exception = sys . exc_info ( ) <EOL> return False <EOL> if reset_offsets : <EOL> cns . reset_offsets ( partition_offsets = [ <EOL> ( cns . partitions [ id_ ] , offset ) for <EOL> ( id_ , offset ) in iteritems ( reset_offsets ) ] ) <EOL> self . _consumer = cns <EOL> return True <EOL> def _get_internal_consumer ( self , partitions = None , start = True ) : <EOL> """<STR_LIT>""" <EOL> if partitions is None : <EOL> partitions = [ ] <EOL> reset_offset_on_start = self . _reset_offset_on_start <EOL> if self . _consumer is not None : <EOL> self . _consumer . stop ( ) <EOL> reset_offset_on_start = False <EOL> Cls = ( rdkafka . RdKafkaSimpleConsumer <EOL> if self . _use_rdkafka else SimpleConsumer ) <EOL> return Cls ( <EOL> self . _topic , <EOL> self . _cluster , <EOL> consumer_group = self . _consumer_group , <EOL> partitions = partitions , <EOL> auto_commit_enable = self . _auto_commit_enable , <EOL> auto_commit_interval_ms = self . _auto_commit_interval_ms , <EOL> fetch_message_max_bytes = self . _fetch_message_max_bytes , <EOL> fetch_min_bytes = self . _fetch_min_bytes , <EOL> num_consumer_fetchers = self . _num_consumer_fetchers , <EOL> queued_max_messages = self . _queued_max_messages , <EOL> fetch_wait_max_ms = self . _fetch_wait_max_ms , <EOL> consumer_timeout_ms = self . _consumer_timeout_ms , <EOL> offsets_channel_backoff_ms = self . _offsets_channel_backoff_ms , <EOL> offsets_commit_max_retries = self . _offsets_commit_max_retries , <EOL> auto_offset_reset = self . _auto_offset_reset , <EOL> reset_offset_on_start = reset_offset_on_start , <EOL> auto_start = start , <EOL> compacted_topic = self . _is_compacted_topic , <EOL> generation_id = self . _generation_id , <EOL> consumer_id = self . _consumer_id <EOL> ) <EOL> def _decide_partitions ( self , participants , consumer_id = None ) : <EOL> """<STR_LIT>""" <EOL> p_to_str = lambda p : '<STR_LIT:->' . join ( [ str ( p . topic . name ) , str ( p . leader . id ) , str ( p . id ) ] ) <EOL> all_parts = self . _topic . partitions . values ( ) <EOL> all_parts = sorted ( all_parts , key = p_to_str ) <EOL> participants = sorted ( participants ) <EOL> idx = participants . index ( consumer_id or self . _consumer_id ) <EOL> parts_per_consumer = len ( all_parts ) // len ( participants ) <EOL> remainder_ppc = len ( all_parts ) % len ( participants ) <EOL> start = parts_per_consumer * idx + min ( idx , remainder_ppc ) <EOL> num_parts = parts_per_consumer + ( <NUM_LIT:0> if ( idx + <NUM_LIT:1> > remainder_ppc ) else <NUM_LIT:1> ) <EOL> new_partitions = itertools . islice ( all_parts , start , start + num_parts ) <EOL> new_partitions = set ( new_partitions ) <EOL> log . info ( '<STR_LIT>' , <EOL> self . _consumer_id , len ( participants ) , len ( all_parts ) , <EOL> len ( new_partitions ) ) <EOL> log . debug ( '<STR_LIT>' , [ p_to_str ( p ) for p in new_partitions ] ) <EOL> return new_partitions <EOL> def _get_participants ( self ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> consumer_ids = self . _zookeeper . get_children ( self . _consumer_id_path ) <EOL> except NoNodeException : <EOL> log . debug ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> return [ ] <EOL> participants = [ ] <EOL> for id_ in consumer_ids : <EOL> try : <EOL> topic , stat = self . _zookeeper . get ( "<STR_LIT>" % ( self . _consumer_id_path , id_ ) ) <EOL> if topic == self . _topic . name : <EOL> participants . append ( get_bytes ( id_ ) ) <EOL> except NoNodeException : <EOL> pass <EOL> participants = sorted ( participants ) <EOL> return participants <EOL> def _build_watch_callback ( self , fn , proxy ) : <EOL> """<STR_LIT>""" <EOL> def _callback ( children ) : <EOL> try : <EOL> proxy . __repr__ ( ) <EOL> except ReferenceError : <EOL> return False <EOL> return fn ( proxy , children ) <EOL> return _callback <EOL> def _set_watches ( self ) : <EOL> """<STR_LIT>""" <EOL> proxy = weakref . proxy ( self ) <EOL> _brokers_changed = self . _build_watch_callback ( BalancedConsumer . _brokers_changed , proxy ) <EOL> _topics_changed = self . _build_watch_callback ( BalancedConsumer . _topics_changed , proxy ) <EOL> _consumers_changed = self . _build_watch_callback ( BalancedConsumer . _consumers_changed , proxy ) <EOL> self . _setting_watches = True <EOL> broker_path = '<STR_LIT>' <EOL> try : <EOL> self . _broker_watcher = ChildrenWatch ( <EOL> self . _zookeeper , broker_path , <EOL> _brokers_changed <EOL> ) <EOL> except NoNodeException : <EOL> raise Exception ( <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> % broker_path ) <EOL> self . _topics_watcher = ChildrenWatch ( <EOL> self . _zookeeper , <EOL> '<STR_LIT>' , <EOL> _topics_changed <EOL> ) <EOL> self . _consumer_watcher = ChildrenWatch ( <EOL> self . _zookeeper , self . _consumer_id_path , <EOL> _consumers_changed <EOL> ) <EOL> self . _setting_watches = False <EOL> def _add_self ( self ) : <EOL> """<STR_LIT>""" <EOL> participants = self . _get_participants ( ) <EOL> if len ( self . _topic . partitions ) <= len ( participants ) : <EOL> raise KafkaException ( "<STR_LIT>" ) <EOL> self . _zookeeper . create ( <EOL> self . _path_self , self . _topic . name , ephemeral = True , makepath = True ) <EOL> @ property <EOL> def _path_self ( self ) : <EOL> """<STR_LIT>""" <EOL> return '<STR_LIT>' . format ( <EOL> path = self . _consumer_id_path , <EOL> id_ = get_string ( self . _consumer_id ) <EOL> ) <EOL> def _update_member_assignment ( self ) : <EOL> """<STR_LIT>""" <EOL> for i in range ( self . _rebalance_max_retries ) : <EOL> try : <EOL> participants = self . _get_participants ( ) <EOL> if self . _consumer_id not in participants : <EOL> self . _add_self ( ) <EOL> participants . append ( self . _consumer_id ) <EOL> new_partitions = self . _decide_partitions ( participants ) <EOL> if not new_partitions : <EOL> log . warning ( "<STR_LIT>" , <EOL> self . _consumer_id ) <EOL> current_zk_parts = self . _get_held_partitions ( ) <EOL> self . _remove_partitions ( current_zk_parts - new_partitions ) <EOL> self . _add_partitions ( new_partitions - current_zk_parts ) <EOL> if self . _setup_internal_consumer ( new_partitions ) : <EOL> log . info ( '<STR_LIT>' ) <EOL> break <EOL> except PartitionOwnedError as ex : <EOL> if i == self . _rebalance_max_retries - <NUM_LIT:1> : <EOL> log . warning ( '<STR_LIT>' , <EOL> ex . partition , i ) <EOL> raise <EOL> log . info ( '<STR_LIT>' , ex . partition ) <EOL> self . _cluster . handler . sleep ( i * ( self . _rebalance_backoff_ms / <NUM_LIT:1000> ) ) <EOL> def _rebalance ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . _consumer is not None : <EOL> self . commit_offsets ( ) <EOL> with self . _rebalancing_lock : <EOL> if not self . _running : <EOL> raise ConsumerStoppedException <EOL> log . info ( '<STR_LIT>' % ( <EOL> self . _consumer_id , self . _topic . name ) ) <EOL> self . _update_member_assignment ( ) <EOL> def _path_from_partition ( self , p ) : <EOL> """<STR_LIT>""" <EOL> return "<STR_LIT>" % ( self . _topic_path , p . leader . id , p . id ) <EOL> def _remove_partitions ( self , partitions ) : <EOL> """<STR_LIT>""" <EOL> for p in partitions : <EOL> self . _zookeeper . delete ( self . _path_from_partition ( p ) ) <EOL> def _add_partitions ( self , partitions ) : <EOL> """<STR_LIT>""" <EOL> for p in partitions : <EOL> try : <EOL> self . _zookeeper . create ( <EOL> self . _path_from_partition ( p ) , <EOL> value = get_bytes ( self . _consumer_id ) , <EOL> ephemeral = True <EOL> ) <EOL> except NodeExistsError : <EOL> raise PartitionOwnedError ( p ) <EOL> def _get_held_partitions ( self ) : <EOL> """<STR_LIT>""" <EOL> zk_partition_ids = set ( ) <EOL> all_partitions = self . _zookeeper . get_children ( self . _topic_path ) <EOL> for partition_slug in all_partitions : <EOL> try : <EOL> owner_id , stat = self . _zookeeper . get ( <EOL> '<STR_LIT>' . format ( <EOL> path = self . _topic_path , slug = partition_slug ) ) <EOL> if owner_id == get_bytes ( self . _consumer_id ) : <EOL> zk_partition_ids . add ( int ( partition_slug . split ( '<STR_LIT:->' ) [ <NUM_LIT:1> ] ) ) <EOL> except NoNodeException : <EOL> pass <EOL> return set ( self . _topic . partitions [ _id ] for _id in zk_partition_ids ) <EOL> @ _catch_thread_exception <EOL> def _brokers_changed ( self , brokers ) : <EOL> if not self . _running : <EOL> return False <EOL> if self . _setting_watches : <EOL> return <EOL> log . debug ( "<STR_LIT>" . format ( <EOL> self . _consumer_id ) ) <EOL> self . _rebalance ( ) <EOL> @ _catch_thread_exception <EOL> def _consumers_changed ( self , consumers ) : <EOL> if not self . _running : <EOL> return False <EOL> if self . _setting_watches : <EOL> return <EOL> log . debug ( "<STR_LIT>" . format ( <EOL> self . _consumer_id ) ) <EOL> self . _rebalance ( ) <EOL> @ _catch_thread_exception <EOL> def _topics_changed ( self , topics ) : <EOL> if not self . _running : <EOL> return False <EOL> if self . _setting_watches : <EOL> return <EOL> log . debug ( "<STR_LIT>" . format ( <EOL> self . _consumer_id ) ) <EOL> self . _rebalance ( ) <EOL> def reset_offsets ( self , partition_offsets = None ) : <EOL> """<STR_LIT>""" <EOL> self . _raise_worker_exceptions ( ) <EOL> if not self . _consumer : <EOL> raise ConsumerStoppedException ( "<STR_LIT>" ) <EOL> self . _consumer . reset_offsets ( partition_offsets = partition_offsets ) <EOL> def consume ( self , block = True ) : <EOL> """<STR_LIT>""" <EOL> def consumer_timed_out ( ) : <EOL> """<STR_LIT>""" <EOL> if self . _consumer_timeout_ms == - <NUM_LIT:1> : <EOL> return False <EOL> disp = ( time . time ( ) - self . _last_message_time ) * <NUM_LIT> <EOL> return disp > self . _consumer_timeout_ms <EOL> message = None <EOL> self . _last_message_time = time . time ( ) <EOL> while message is None and not consumer_timed_out ( ) : <EOL> self . _raise_worker_exceptions ( ) <EOL> try : <EOL> message = self . _consumer . consume ( block = block ) <EOL> except ( ConsumerStoppedException , AttributeError ) : <EOL> if not self . _running : <EOL> raise ConsumerStoppedException <EOL> continue <EOL> if message : <EOL> self . _last_message_time = time . time ( ) <EOL> if not block : <EOL> return message <EOL> return message <EOL> def __iter__ ( self ) : <EOL> """<STR_LIT>""" <EOL> while True : <EOL> message = self . consume ( block = True ) <EOL> if not message : <EOL> raise StopIteration <EOL> yield message <EOL> def commit_offsets ( self ) : <EOL> """<STR_LIT>""" <EOL> self . _raise_worker_exceptions ( ) <EOL> if not self . _consumer : <EOL> raise KafkaException ( "<STR_LIT>" ) <EOL> return self . _consumer . commit_offsets ( ) </s>
