<s> from __future__ import division <EOL> """<STR_LIT>""" <EOL> __license__ = """<STR_LIT>""" <EOL> __all__ = [ "<STR_LIT>" ] <EOL> from collections import deque <EOL> import logging <EOL> import platform <EOL> import sys <EOL> import threading <EOL> import traceback <EOL> import weakref <EOL> from . common import CompressionType <EOL> from . exceptions import ( <EOL> ERROR_CODES , <EOL> KafkaException , <EOL> InvalidMessageSize , <EOL> MessageSizeTooLarge , <EOL> NoBrokersAvailableError , <EOL> NotLeaderForPartition , <EOL> ProducerQueueFullError , <EOL> ProducerStoppedException , <EOL> SocketDisconnectedError , <EOL> ) <EOL> from . partitioners import random_partitioner <EOL> from . protocol import Message , ProduceRequest <EOL> from . utils . compat import iteritems , itervalues , Empty <EOL> log = logging . getLogger ( __name__ ) <EOL> class Producer ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , <EOL> cluster , <EOL> topic , <EOL> partitioner = random_partitioner , <EOL> compression = CompressionType . NONE , <EOL> max_retries = <NUM_LIT:3> , <EOL> retry_backoff_ms = <NUM_LIT:100> , <EOL> required_acks = <NUM_LIT:1> , <EOL> ack_timeout_ms = <NUM_LIT:10> * <NUM_LIT:1000> , <EOL> max_queued_messages = <NUM_LIT> , <EOL> min_queued_messages = <NUM_LIT> , <EOL> linger_ms = <NUM_LIT:5> * <NUM_LIT:1000> , <EOL> block_on_queue_full = True , <EOL> max_request_size = <NUM_LIT> , <EOL> sync = False , <EOL> delivery_reports = False , <EOL> auto_start = True ) : <EOL> """<STR_LIT>""" <EOL> self . _cluster = cluster <EOL> self . _topic = topic <EOL> self . _partitioner = partitioner <EOL> self . _compression = compression <EOL> if self . _compression == CompressionType . SNAPPY and platform . python_implementation == "<STR_LIT>" : <EOL> log . warning ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> self . _max_retries = max_retries <EOL> self . _retry_backoff_ms = retry_backoff_ms <EOL> self . _required_acks = required_acks <EOL> self . _ack_timeout_ms = ack_timeout_ms <EOL> self . _max_queued_messages = max_queued_messages <EOL> self . _min_queued_messages = max ( <NUM_LIT:1> , min_queued_messages if not sync else <NUM_LIT:1> ) <EOL> self . _linger_ms = linger_ms <EOL> self . _block_on_queue_full = block_on_queue_full <EOL> self . _max_request_size = max_request_size <EOL> self . _synchronous = sync <EOL> self . _worker_exception = None <EOL> self . _worker_trace_logged = False <EOL> self . _owned_brokers = None <EOL> self . _delivery_reports = ( _DeliveryReportQueue ( self . _cluster . handler ) <EOL> if delivery_reports or self . _synchronous <EOL> else _DeliveryReportNone ( ) ) <EOL> self . _auto_start = auto_start <EOL> self . _running = False <EOL> self . _update_lock = self . _cluster . handler . Lock ( ) <EOL> if self . _auto_start : <EOL> self . start ( ) <EOL> def __del__ ( self ) : <EOL> log . debug ( "<STR_LIT>" . format ( self ) ) <EOL> self . stop ( ) <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> def __repr__ ( self ) : <EOL> return "<STR_LIT>" . format ( <EOL> module = self . __class__ . __module__ , <EOL> name = self . __class__ . __name__ , <EOL> id_ = hex ( id ( self ) ) <EOL> ) <EOL> def __enter__ ( self ) : <EOL> """<STR_LIT>""" <EOL> return self <EOL> def __exit__ ( self , exc_type , exc_value , traceback ) : <EOL> """<STR_LIT>""" <EOL> self . stop ( ) <EOL> def start ( self ) : <EOL> """<STR_LIT>""" <EOL> if not self . _running : <EOL> self . _setup_owned_brokers ( ) <EOL> self . _running = True <EOL> self . _raise_worker_exceptions ( ) <EOL> def _update ( self ) : <EOL> """<STR_LIT>""" <EOL> with self . _update_lock : <EOL> self . _cluster . update ( ) <EOL> queued_messages = self . _setup_owned_brokers ( ) <EOL> if len ( queued_messages ) : <EOL> log . debug ( "<STR_LIT>" , <EOL> len ( queued_messages ) ) <EOL> for message in queued_messages : <EOL> self . _produce ( message ) <EOL> def _setup_owned_brokers ( self ) : <EOL> """<STR_LIT>""" <EOL> queued_messages = [ ] <EOL> if self . _owned_brokers is not None : <EOL> brokers = list ( self . _owned_brokers . keys ( ) ) <EOL> for broker in brokers : <EOL> owned_broker = self . _owned_brokers . pop ( broker ) <EOL> owned_broker . stop ( ) <EOL> while True : <EOL> batch = owned_broker . flush ( self . _linger_ms , self . _max_request_size ) <EOL> if not batch : <EOL> break <EOL> queued_messages . extend ( batch ) <EOL> self . _owned_brokers = { } <EOL> for partition in self . _topic . partitions . values ( ) : <EOL> if partition . leader . id not in self . _owned_brokers : <EOL> self . _owned_brokers [ partition . leader . id ] = OwnedBroker ( <EOL> self , partition . leader ) <EOL> return queued_messages <EOL> def stop ( self ) : <EOL> """<STR_LIT>""" <EOL> def get_queue_readers ( ) : <EOL> if not self . _owned_brokers : <EOL> return [ ] <EOL> return [ owned_broker . _queue_reader_worker <EOL> for owned_broker in self . _owned_brokers . values ( ) <EOL> if owned_broker . running ] <EOL> def stop_owned_brokers ( ) : <EOL> self . _wait_all ( ) <EOL> if self . _owned_brokers is not None : <EOL> for owned_broker in self . _owned_brokers . values ( ) : <EOL> owned_broker . stop ( ) <EOL> while self . _running : <EOL> queue_readers = get_queue_readers ( ) <EOL> stop_owned_brokers ( ) <EOL> if len ( queue_readers ) == <NUM_LIT:0> : <EOL> self . _running = False <EOL> else : <EOL> for queue_reader in queue_readers : <EOL> queue_reader . join ( ) <EOL> def produce ( self , message , partition_key = None ) : <EOL> """<STR_LIT>""" <EOL> if not ( isinstance ( partition_key , bytes ) or partition_key is None ) : <EOL> raise TypeError ( "<STR_LIT>" <EOL> "<STR_LIT>" , type ( partition_key ) ) <EOL> if not ( isinstance ( message , bytes ) or message is None ) : <EOL> raise TypeError ( "<STR_LIT>" <EOL> "<STR_LIT>" , type ( message ) ) <EOL> if not self . _running : <EOL> raise ProducerStoppedException ( ) <EOL> partitions = list ( self . _topic . partitions . values ( ) ) <EOL> partition_id = self . _partitioner ( partitions , partition_key ) . id <EOL> msg = Message ( value = message , <EOL> partition_key = partition_key , <EOL> partition_id = partition_id , <EOL> delivery_report_q = self . _delivery_reports . queue ) <EOL> self . _produce ( msg ) <EOL> if self . _synchronous : <EOL> while True : <EOL> self . _cluster . handler . sleep ( ) <EOL> try : <EOL> reported_msg , exc = self . get_delivery_report ( timeout = <NUM_LIT:1> ) <EOL> break <EOL> except Empty : <EOL> continue <EOL> assert reported_msg is msg <EOL> if exc is not None : <EOL> raise exc <EOL> self . _raise_worker_exceptions ( ) <EOL> def get_delivery_report ( self , block = True , timeout = None ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> return self . _delivery_reports . queue . get ( block , timeout ) <EOL> except AttributeError : <EOL> raise KafkaException ( "<STR_LIT>" ) <EOL> def _produce ( self , message ) : <EOL> """<STR_LIT>""" <EOL> success = False <EOL> while not success : <EOL> leader_id = self . _topic . partitions [ message . partition_id ] . leader . id <EOL> if leader_id in self . _owned_brokers : <EOL> self . _owned_brokers [ leader_id ] . enqueue ( message ) <EOL> success = True <EOL> else : <EOL> success = False <EOL> def _send_request ( self , message_batch , owned_broker ) : <EOL> """<STR_LIT>""" <EOL> req = ProduceRequest ( <EOL> compression_type = self . _compression , <EOL> required_acks = self . _required_acks , <EOL> timeout = self . _ack_timeout_ms <EOL> ) <EOL> req . delivered = <NUM_LIT:0> <EOL> for msg in message_batch : <EOL> req . add_message ( msg , self . _topic . name , msg . partition_id ) <EOL> log . debug ( "<STR_LIT>" , <EOL> len ( message_batch ) , owned_broker . broker . id ) <EOL> def _get_partition_msgs ( partition_id , req ) : <EOL> """<STR_LIT>""" <EOL> return ( <EOL> mset <EOL> for topic , partitions in iteritems ( req . msets ) <EOL> for p_id , mset in iteritems ( partitions ) <EOL> if p_id == partition_id <EOL> ) <EOL> def mark_as_delivered ( message_batch ) : <EOL> owned_broker . increment_messages_pending ( - <NUM_LIT:1> * len ( message_batch ) ) <EOL> req . delivered += len ( message_batch ) <EOL> for msg in message_batch : <EOL> self . _delivery_reports . put ( msg ) <EOL> try : <EOL> response = owned_broker . broker . produce_messages ( req ) <EOL> if self . _required_acks == <NUM_LIT:0> : <EOL> mark_as_delivered ( message_batch ) <EOL> return <EOL> to_retry = [ ] <EOL> for topic , partitions in iteritems ( response . topics ) : <EOL> for partition , presponse in iteritems ( partitions ) : <EOL> if presponse . err == <NUM_LIT:0> : <EOL> mark_as_delivered ( req . msets [ topic ] [ partition ] . messages ) <EOL> continue <EOL> if presponse . err == NotLeaderForPartition . ERROR_CODE : <EOL> self . _update ( ) <EOL> info = "<STR_LIT>" . format ( <EOL> topic , <EOL> partition , <EOL> owned_broker . broker . host , <EOL> owned_broker . broker . port , <EOL> presponse . err ) <EOL> log . warning ( info ) <EOL> exc = ERROR_CODES [ presponse . err ] ( info ) <EOL> to_retry . extend ( <EOL> ( mset , exc ) <EOL> for mset in _get_partition_msgs ( partition , req ) ) <EOL> except SocketDisconnectedError as exc : <EOL> log . warning ( '<STR_LIT>' , <EOL> owned_broker . broker . host , <EOL> owned_broker . broker . port ) <EOL> try : <EOL> self . _update ( ) <EOL> except NoBrokersAvailableError : <EOL> log . warning ( "<STR_LIT>" ) <EOL> to_retry = [ <EOL> ( mset , exc ) <EOL> for topic , partitions in iteritems ( req . msets ) <EOL> for p_id , mset in iteritems ( partitions ) <EOL> ] <EOL> log . debug ( "<STR_LIT>" . format ( <EOL> req . delivered , len ( message_batch ) , owned_broker . broker . id ) ) <EOL> if to_retry : <EOL> self . _cluster . handler . sleep ( self . _retry_backoff_ms / <NUM_LIT:1000> ) <EOL> owned_broker . increment_messages_pending ( - <NUM_LIT:1> * len ( to_retry ) ) <EOL> for mset , exc in to_retry : <EOL> non_recoverable = type ( exc ) in ( InvalidMessageSize , <EOL> MessageSizeTooLarge ) <EOL> for msg in mset . messages : <EOL> if ( non_recoverable or msg . produce_attempt >= self . _max_retries ) : <EOL> self . _delivery_reports . put ( msg , exc ) <EOL> log . error ( "<STR_LIT>" % exc ) <EOL> else : <EOL> msg . produce_attempt += <NUM_LIT:1> <EOL> self . _produce ( msg ) <EOL> def _wait_all ( self ) : <EOL> """<STR_LIT>""" <EOL> log . info ( "<STR_LIT>" ) <EOL> while any ( q . message_is_pending ( ) for q in itervalues ( self . _owned_brokers ) ) : <EOL> self . _cluster . handler . sleep ( <NUM_LIT> ) <EOL> self . _raise_worker_exceptions ( ) <EOL> class OwnedBroker ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , producer , broker , auto_start = True ) : <EOL> self . producer = weakref . proxy ( producer ) <EOL> self . broker = broker <EOL> self . lock = self . producer . _cluster . handler . RLock ( ) <EOL> self . flush_ready = self . producer . _cluster . handler . Event ( ) <EOL> self . slot_available = self . producer . _cluster . handler . Event ( ) <EOL> self . queue = deque ( ) <EOL> self . messages_pending = <NUM_LIT:0> <EOL> self . running = True <EOL> self . _auto_start = auto_start <EOL> if self . _auto_start : <EOL> self . start ( ) <EOL> def start ( self ) : <EOL> def queue_reader ( ) : <EOL> while self . running : <EOL> try : <EOL> batch = self . flush ( self . producer . _linger_ms , self . producer . _max_request_size ) <EOL> if batch : <EOL> self . producer . _send_request ( batch , self ) <EOL> except Exception : <EOL> self . producer . _worker_exception = sys . exc_info ( ) <EOL> break <EOL> log . info ( "<STR_LIT>" , self . broker . host , <EOL> self . broker . port ) <EOL> log . info ( "<STR_LIT>" , self . broker . id ) <EOL> name = "<STR_LIT>" . format ( self . broker . id ) <EOL> self . _queue_reader_worker = self . producer . _cluster . handler . spawn ( <EOL> queue_reader , name = name ) <EOL> def stop ( self ) : <EOL> self . running = False <EOL> def increment_messages_pending ( self , amnt ) : <EOL> with self . lock : <EOL> self . messages_pending += amnt <EOL> self . messages_pending = max ( <NUM_LIT:0> , self . messages_pending ) <EOL> def message_is_pending ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . messages_pending > <NUM_LIT:0> <EOL> def enqueue ( self , message ) : <EOL> """<STR_LIT>""" <EOL> self . _wait_for_slot_available ( ) <EOL> with self . lock : <EOL> self . queue . appendleft ( message ) <EOL> self . increment_messages_pending ( <NUM_LIT:1> ) <EOL> if len ( self . queue ) >= self . producer . _min_queued_messages : <EOL> if not self . flush_ready . is_set ( ) : <EOL> self . flush_ready . set ( ) <EOL> def flush ( self , linger_ms , max_request_size , release_pending = False ) : <EOL> """<STR_LIT>""" <EOL> self . _wait_for_flush_ready ( linger_ms ) <EOL> with self . lock : <EOL> batch = [ ] <EOL> batch_size_in_bytes = <NUM_LIT:0> <EOL> while len ( self . queue ) > <NUM_LIT:0> : <EOL> peeked_message = self . queue [ - <NUM_LIT:1> ] <EOL> if peeked_message and peeked_message . value is not None : <EOL> if len ( peeked_message ) > max_request_size : <EOL> exc = MessageSizeTooLarge ( <EOL> "<STR_LIT>" . format ( max_request_size ) <EOL> ) <EOL> log . warning ( exc ) <EOL> message = self . queue . pop ( ) <EOL> if peeked_message . delivery_report_q is not None : <EOL> peeked_message . delivery_report_q . put ( ( message , exc ) ) <EOL> self . increment_messages_pending ( - <NUM_LIT:1> ) <EOL> continue <EOL> elif batch_size_in_bytes + len ( peeked_message ) > max_request_size : <EOL> log . debug ( "<STR_LIT>" ) <EOL> self . flush_ready . set ( ) <EOL> break <EOL> message = self . queue . pop ( ) <EOL> batch_size_in_bytes += len ( message ) <EOL> batch . append ( message ) <EOL> if release_pending : <EOL> self . increment_messages_pending ( - <NUM_LIT:1> * len ( batch ) ) <EOL> if not self . slot_available . is_set ( ) : <EOL> self . slot_available . set ( ) <EOL> return batch <EOL> def _wait_for_flush_ready ( self , linger_ms ) : <EOL> """<STR_LIT>""" <EOL> if len ( self . queue ) < self . producer . _min_queued_messages : <EOL> with self . lock : <EOL> if len ( self . queue ) < self . producer . _min_queued_messages : <EOL> self . flush_ready . clear ( ) <EOL> if linger_ms > <NUM_LIT:0> : <EOL> self . flush_ready . wait ( ( linger_ms / <NUM_LIT:1000> ) ) <EOL> def _wait_for_slot_available ( self ) : <EOL> """<STR_LIT>""" <EOL> if len ( self . queue ) >= self . producer . _max_queued_messages : <EOL> with self . lock : <EOL> if len ( self . queue ) >= self . producer . _max_queued_messages : <EOL> self . slot_available . clear ( ) <EOL> if self . producer . _block_on_queue_full : <EOL> self . slot_available . wait ( ) <EOL> else : <EOL> raise ProducerQueueFullError ( "<STR_LIT>" , <EOL> self . broker . id ) <EOL> class _DeliveryReportQueue ( threading . local ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , handler ) : <EOL> self . queue = handler . Queue ( ) <EOL> @ staticmethod <EOL> def put ( msg , exc = None ) : <EOL> msg . delivery_report_q . put ( ( msg , exc ) ) <EOL> class _DeliveryReportNone ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self ) : <EOL> self . queue = None <EOL> @ staticmethod <EOL> def put ( msg , exc = None ) : <EOL> return </s>
