<s> from __future__ import absolute_import <EOL> from collections import namedtuple <EOL> import logging <EOL> from multiprocessing import Process , Manager as MPManager <EOL> import time <EOL> import warnings <EOL> from six . moves import queue <EOL> from . . common import KafkaError <EOL> from . base import ( <EOL> Consumer , <EOL> AUTO_COMMIT_MSG_COUNT , AUTO_COMMIT_INTERVAL , <EOL> NO_MESSAGES_WAIT_TIME_SECONDS , <EOL> FULL_QUEUE_WAIT_TIME_SECONDS , <EOL> MAX_BACKOFF_SECONDS , <EOL> ) <EOL> from . simple import SimpleConsumer <EOL> log = logging . getLogger ( __name__ ) <EOL> Events = namedtuple ( "<STR_LIT>" , [ "<STR_LIT:start>" , "<STR_LIT>" , "<STR_LIT>" ] ) <EOL> def _mp_consume ( client , group , topic , message_queue , size , events , ** consumer_options ) : <EOL> """<STR_LIT>""" <EOL> interval = <NUM_LIT:1> <EOL> while not events . exit . is_set ( ) : <EOL> try : <EOL> client . reinit ( ) <EOL> consumer = SimpleConsumer ( client , group , topic , <EOL> auto_commit = False , <EOL> auto_commit_every_n = None , <EOL> auto_commit_every_t = None , <EOL> ** consumer_options ) <EOL> consumer . provide_partition_info ( ) <EOL> while True : <EOL> events . start . wait ( ) <EOL> if events . exit . is_set ( ) : <EOL> break <EOL> count = <NUM_LIT:0> <EOL> message = consumer . get_message ( ) <EOL> if message : <EOL> while True : <EOL> try : <EOL> message_queue . put ( message , timeout = FULL_QUEUE_WAIT_TIME_SECONDS ) <EOL> break <EOL> except queue . Full : <EOL> if events . exit . is_set ( ) : break <EOL> count += <NUM_LIT:1> <EOL> if count == size . value : <EOL> events . pause . wait ( ) <EOL> else : <EOL> time . sleep ( NO_MESSAGES_WAIT_TIME_SECONDS ) <EOL> consumer . stop ( ) <EOL> except KafkaError as e : <EOL> log . error ( "<STR_LIT>" % ( e , interval ) ) <EOL> time . sleep ( interval ) <EOL> interval = interval * <NUM_LIT:2> if interval * <NUM_LIT:2> < MAX_BACKOFF_SECONDS else MAX_BACKOFF_SECONDS <EOL> class MultiProcessConsumer ( Consumer ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , client , group , topic , <EOL> partitions = None , <EOL> auto_commit = True , <EOL> auto_commit_every_n = AUTO_COMMIT_MSG_COUNT , <EOL> auto_commit_every_t = AUTO_COMMIT_INTERVAL , <EOL> num_procs = <NUM_LIT:1> , <EOL> partitions_per_proc = <NUM_LIT:0> , <EOL> ** simple_consumer_options ) : <EOL> warnings . warn ( '<STR_LIT>' <EOL> '<STR_LIT>' , <EOL> DeprecationWarning ) <EOL> super ( MultiProcessConsumer , self ) . __init__ ( <EOL> client , group , topic , <EOL> partitions = partitions , <EOL> auto_commit = auto_commit , <EOL> auto_commit_every_n = auto_commit_every_n , <EOL> auto_commit_every_t = auto_commit_every_t ) <EOL> manager = MPManager ( ) <EOL> self . queue = manager . Queue ( <NUM_LIT> ) <EOL> self . events = Events ( <EOL> start = manager . Event ( ) , <EOL> exit = manager . Event ( ) , <EOL> pause = manager . Event ( ) ) <EOL> self . size = manager . Value ( '<STR_LIT:i>' , <NUM_LIT:0> ) <EOL> partitions = list ( self . offsets . copy ( ) . keys ( ) ) <EOL> if partitions_per_proc : <EOL> num_procs = len ( partitions ) / partitions_per_proc <EOL> if num_procs * partitions_per_proc < len ( partitions ) : <EOL> num_procs += <NUM_LIT:1> <EOL> chunks = [ partitions [ proc : : num_procs ] for proc in range ( num_procs ) ] <EOL> self . procs = [ ] <EOL> for chunk in chunks : <EOL> options = { '<STR_LIT>' : list ( chunk ) } <EOL> if simple_consumer_options : <EOL> simple_consumer_options . pop ( '<STR_LIT>' , None ) <EOL> options . update ( simple_consumer_options ) <EOL> args = ( client . copy ( ) , self . group , self . topic , self . queue , <EOL> self . size , self . events ) <EOL> proc = Process ( target = _mp_consume , args = args , kwargs = options ) <EOL> proc . daemon = True <EOL> proc . start ( ) <EOL> self . procs . append ( proc ) <EOL> def __repr__ ( self ) : <EOL> return '<STR_LIT>' % ( self . group , self . topic , len ( self . procs ) ) <EOL> def stop ( self ) : <EOL> self . events . exit . set ( ) <EOL> self . events . pause . set ( ) <EOL> self . events . start . set ( ) <EOL> for proc in self . procs : <EOL> proc . join ( ) <EOL> proc . terminate ( ) <EOL> super ( MultiProcessConsumer , self ) . stop ( ) <EOL> def __iter__ ( self ) : <EOL> """<STR_LIT>""" <EOL> self . size . value = <NUM_LIT:0> <EOL> self . events . pause . set ( ) <EOL> while True : <EOL> self . events . start . set ( ) <EOL> try : <EOL> partition , message = self . queue . get ( block = True , timeout = <NUM_LIT:1> ) <EOL> except queue . Empty : <EOL> break <EOL> self . offsets [ partition ] = message . offset + <NUM_LIT:1> <EOL> self . events . start . clear ( ) <EOL> self . count_since_commit += <NUM_LIT:1> <EOL> self . _auto_commit ( ) <EOL> yield message <EOL> self . events . start . clear ( ) <EOL> def get_messages ( self , count = <NUM_LIT:1> , block = True , timeout = <NUM_LIT:10> ) : <EOL> """<STR_LIT>""" <EOL> messages = [ ] <EOL> self . size . value = count <EOL> self . events . pause . clear ( ) <EOL> if timeout is not None : <EOL> max_time = time . time ( ) + timeout <EOL> new_offsets = { } <EOL> while count > <NUM_LIT:0> and ( timeout is None or timeout > <NUM_LIT:0> ) : <EOL> if self . queue . empty ( ) : <EOL> self . events . start . set ( ) <EOL> block_next_call = block is True or block > len ( messages ) <EOL> try : <EOL> partition , message = self . queue . get ( block_next_call , <EOL> timeout ) <EOL> except queue . Empty : <EOL> break <EOL> _msg = ( partition , message ) if self . partition_info else message <EOL> messages . append ( _msg ) <EOL> new_offsets [ partition ] = message . offset + <NUM_LIT:1> <EOL> count -= <NUM_LIT:1> <EOL> if timeout is not None : <EOL> timeout = max_time - time . time ( ) <EOL> self . size . value = <NUM_LIT:0> <EOL> self . events . start . clear ( ) <EOL> self . events . pause . set ( ) <EOL> self . offsets . update ( new_offsets ) <EOL> self . count_since_commit += len ( messages ) <EOL> self . _auto_commit ( ) <EOL> return messages </s>
