Skip to content

Allow consumers to control start from beginning behavior #352

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 2 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions kafka/consumer/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,8 @@ class Consumer(object):
"""
def __init__(self, client, group, topic, partitions=None, auto_commit=True,
auto_commit_every_n=AUTO_COMMIT_MSG_COUNT,
auto_commit_every_t=AUTO_COMMIT_INTERVAL):
auto_commit_every_t=AUTO_COMMIT_INTERVAL,
start_from_beginning=True):

self.client = client
self.topic = topic
Expand All @@ -67,7 +68,7 @@ def __init__(self, client, group, topic, partitions=None, auto_commit=True,
self.commit)
self.commit_timer.start()

if auto_commit:
if auto_commit or not start_from_beginning:
self.fetch_last_known_offsets(partitions)
else:
for partition in partitions:
Expand Down
17 changes: 13 additions & 4 deletions kafka/consumer/multiprocess.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
log = logging.getLogger("kafka")


def _mp_consume(client, group, topic, chunk, queue, start, exit, pause, size):
def _mp_consume(client, group, topic, chunk, start_from_beginning, queue, start, exit, pause, size):
"""
A child process worker which consumes messages based on the
notifications given by the controller process
Expand All @@ -37,7 +37,8 @@ def _mp_consume(client, group, topic, chunk, queue, start, exit, pause, size):
partitions=chunk,
auto_commit=False,
auto_commit_every_n=None,
auto_commit_every_t=None)
auto_commit_every_t=None,
start_from_beginning=start_from_beginning)

# Ensure that the consumer provides the partition information
consumer.provide_partition_info()
Expand Down Expand Up @@ -105,15 +106,16 @@ 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, start_from_beginning=True):

# Initiate the base consumer class
super(MultiProcessConsumer, self).__init__(
client, group, topic,
partitions=None,
auto_commit=auto_commit,
auto_commit_every_n=auto_commit_every_n,
auto_commit_every_t=auto_commit_every_t)
auto_commit_every_t=auto_commit_every_t,
start_from_beginning=start_from_beginning)

# Variables for managing and controlling the data flow from
# consumer child process to master
Expand Down Expand Up @@ -141,10 +143,17 @@ def __init__(self, client, group, topic, auto_commit=True,
# The final set of chunks
chunks = [partitions[proc::num_procs] for proc in range(num_procs)]

# though the child workers will not auto commit, they should
# not start from offset = 0 if it is set.
child_start_from_beginning = start_from_beginning
if auto_commit:
child_start_from_beginning = False

self.procs = []
for chunk in chunks:
args = (client.copy(),
group, topic, chunk,
child_start_from_beginning,
self.queue, self.start, self.exit,
self.pause, self.size)

Expand Down
6 changes: 4 additions & 2 deletions kafka/consumer/simple.py
Original file line number Diff line number Diff line change
Expand Up @@ -114,13 +114,15 @@ def __init__(self, client, group, topic, auto_commit=True, partitions=None,
buffer_size=FETCH_BUFFER_SIZE_BYTES,
max_buffer_size=MAX_FETCH_BUFFER_SIZE_BYTES,
iter_timeout=None,
auto_offset_reset='largest'):
auto_offset_reset='largest',
start_from_beginning=True):
super(SimpleConsumer, self).__init__(
client, group, topic,
partitions=partitions,
auto_commit=auto_commit,
auto_commit_every_n=auto_commit_every_n,
auto_commit_every_t=auto_commit_every_t)
auto_commit_every_t=auto_commit_every_t,
start_from_beginning=start_from_beginning)

if max_buffer_size is not None and buffer_size > max_buffer_size:
raise ValueError("buffer_size (%d) is greater than "
Expand Down
35 changes: 35 additions & 0 deletions test/test_consumer_integration.py
Original file line number Diff line number Diff line change
Expand Up @@ -326,6 +326,41 @@ def test_offset_behavior__resuming_behavior(self):
consumer1.stop()
consumer2.stop()

@kafka_versions("0.8.1", "0.8.1.1", "0.8.2.0")
def test_multi_process_offset_behavior__resuming_behavior(self):
self.send_messages(0, range(0, 100))
self.send_messages(1, range(100, 200))

# Start a consumer
consumer1 = self.consumer(
consumer=MultiProcessConsumer,
auto_commit_every_t = None,
auto_commit_every_n = 20,
)

# Grab the first 195 messages
output_msgs1 = []
idx = 0
for message in consumer1:
if idx == 195:
break
idx +=1
output_msgs1.append(message.message.value)
self.assert_message_count(output_msgs1, 195)

# The total offset across both partitions should be at 180
consumer2 = self.consumer(
consumer=MultiProcessConsumer,
auto_commit_every_t = None,
auto_commit_every_n = 20,
)

# 181-200
self.assert_message_count([ message for message in consumer2 ], 20)

consumer1.stop()
consumer2.stop()

# TODO: Make this a unit test -- should not require integration
@kafka_versions("all")
def test_fetch_buffer_size(self):
Expand Down