Using additional params for MP consumer child process

kafka/consumer/multiprocess.py

@@ -2,6 +2,8 @@
 
 import logging
 import time
+
+from collections import namedtuple
 from multiprocessing import Process, Queue as MPQueue, Event, Value
 
 try:
@@ -15,10 +17,11 @@
 )
 from .simple import Consumer, SimpleConsumer
 
-log = logging.getLogger("kafka")
+Events = namedtuple("Events", ["start", "pause", "exit"])
 
+log = logging.getLogger("kafka")
 
-def _mp_consume(client, group, topic, chunk, queue, start, exit, pause, size):
+def _mp_consume(client, group, topic, queue, size, events, **consumer_options):
  """
  A child process worker which consumes messages based on the
  notifications given by the controller process
@@ -34,20 +37,20 @@ def _mp_consume(client, group, topic, chunk, queue, start, exit, pause, size):
  # We will start consumers without auto-commit. Auto-commit will be
  # done by the master controller process.
  consumer = SimpleConsumer(client, group, topic,
- partitions=chunk,
  auto_commit=False,
  auto_commit_every_n=None,
- auto_commit_every_t=None)
+ auto_commit_every_t=None,
+ **consumer_options)
 
  # Ensure that the consumer provides the partition information
  consumer.provide_partition_info()
 
  while True:
  # Wait till the controller indicates us to start consumption
- start.wait()
+ events.start.wait()
 
  # If we are asked to quit, do so
- if exit.is_set():
+ if events.exit.is_set():
  break
 
  # Consume messages and add them to the queue. If the controller
@@ -65,7 +68,7 @@ def _mp_consume(client, group, topic, chunk, queue, start, exit, pause, size):
  # loop consuming all available messages before the controller
  # can reset the 'start' event
  if count == size.value:
- pause.wait()
+ events.pause.wait()
 
  else:
  # In case we did not receive any message, give up the CPU for
@@ -105,7 +108,8 @@ class MultiProcessConsumer(Consumer):
  def __init__(self, client, group, topic, auto_commit=True,
  auto_commit_every_n=AUTO_COMMIT_MSG_COUNT,
  auto_commit_every_t=AUTO_COMMIT_INTERVAL,
- num_procs=1, partitions_per_proc=0):
+ num_procs=1, partitions_per_proc=0,
+ simple_consumer_options=None):
 
  # Initiate the base consumer class
  super(MultiProcessConsumer, self).__init__(
@@ -118,9 +122,10 @@ def __init__(self, client, group, topic, auto_commit=True,
  # Variables for managing and controlling the data flow from
  # consumer child process to master
  self.queue = MPQueue(1024) # Child consumers dump messages into this
- self.start = Event() # Indicates the consumers to start fetch
- self.exit = Event() # Requests the consumers to shutdown
- self.pause = Event() # Requests the consumers to pause fetch
+ self.events = Events(
+ start = Event(), # Indicates the consumers to start fetch
+ exit = Event(), # Requests the consumers to shutdown
+ pause = Event()) # Requests the consumers to pause fetch
  self.size = Value('i', 0) # Indicator of number of messages to fetch
 
  # dict.keys() returns a view in py3 + it's not a thread-safe operation
@@ -143,12 +148,13 @@ def __init__(self, client, group, topic, auto_commit=True,
 
  self.procs = []
  for chunk in chunks:
- args = (client.copy(),
- group, topic, chunk,
- self.queue, self.start, self.exit,
- self.pause, self.size)
+ options = {'partitions': list(chunk)}
+ if simple_consumer_options:
+ options.update(simple_consumer_options)
 
- proc = Process(target=_mp_consume, args=args)
+ args = (client.copy(), group, topic, self.queue,
+ self.size, self.events)
+ proc = Process(target=_mp_consume, args=args, kwargs=options)
  proc.daemon = True
  proc.start()
  self.procs.append(proc)
@@ -159,9 +165,9 @@ def __repr__(self):
 
  def stop(self):
  # Set exit and start off all waiting consumers
- self.exit.set()
- self.pause.set()
- self.start.set()
+ self.events.exit.set()
+ self.events.pause.set()
+ self.events.start.set()
 
  for proc in self.procs:
  proc.join()
@@ -176,10 +182,10 @@ def __iter__(self):
  # Trigger the consumer procs to start off.
  # We will iterate till there are no more messages available
  self.size.value = 0
- self.pause.set()
+ self.events.pause.set()
 
  while True:
- self.start.set()
+ self.events.start.set()
  try:
  # We will block for a small while so that the consumers get
  # a chance to run and put some messages in the queue
@@ -191,12 +197,12 @@ def __iter__(self):
 
  # Count, check and commit messages if necessary
  self.offsets[partition] = message.offset + 1
- self.start.clear()
+ self.events.start.clear()
  self.count_since_commit += 1
  self._auto_commit()
  yield message
 
- self.start.clear()
+ self.events.start.clear()
 
  def get_messages(self, count=1, block=True, timeout=10):
  """
@@ -216,7 +222,7 @@ def get_messages(self, count=1, block=True, timeout=10):
  # necessary, but these will not be committed to kafka. Also, the extra
  # messages can be provided in subsequent runs
  self.size.value = count
- self.pause.clear()
+ self.events.pause.clear()
 
  if timeout is not None:
  max_time = time.time() + timeout
@@ -228,7 +234,7 @@ def get_messages(self, count=1, block=True, timeout=10):
  # go into overdrive and keep consuming thousands of
  # messages when the user might need only a few
  if self.queue.empty():
- self.start.set()
+ self.events.start.set()
 
  try:
  partition, message = self.queue.get(block, timeout)
@@ -242,8 +248,8 @@ def get_messages(self, count=1, block=True, timeout=10):
  timeout = max_time - time.time()
 
  self.size.value = 0
- self.start.clear()
- self.pause.set()
+ self.events.start.clear()
+ self.events.pause.set()
 
  # Update and commit offsets if necessary
  self.offsets.update(new_offsets)

test/test_consumer_integration.py

@@ -63,6 +63,8 @@ def consumer(self, **kwargs):
 
  if consumer_class == SimpleConsumer:
  kwargs.setdefault('iter_timeout', 0)
+ elif consumer_class == MultiProcessConsumer:
+ kwargs.setdefault('simple_consumer_options', {'iter_timeout': 0})
 
  return consumer_class(self.client, group, topic, **kwargs)
 
@@ -243,7 +245,8 @@ def test_multi_proc_pending(self):
  self.send_messages(0, range(0, 10))
  self.send_messages(1, range(10, 20))
 
- consumer = MultiProcessConsumer(self.client, "group1", self.topic, auto_commit=False)
+ consumer = MultiProcessConsumer(self.client, "group1", self.topic, auto_commit=False,
+ simple_consumer_options={'iter_timeout': 0})
 
  self.assertEqual(consumer.pending(), 20)
  self.assertEqual(consumer.pending(partitions=[0]), 10)