def test_compacted_topic_consumption(self): # Compacted topics can have offsets skipped client = AIOKafkaClient( loop=self.loop, bootstrap_servers=[]) client.ready = mock.MagicMock() client.ready.side_effect = asyncio.coroutine(lambda a: True) client.force_metadata_update = mock.MagicMock() client.force_metadata_update.side_effect = asyncio.coroutine( lambda: False) client.send = mock.MagicMock() subscriptions = SubscriptionState(loop=self.loop) fetcher = Fetcher(client, subscriptions, loop=self.loop) tp = TopicPartition('test', 0) req = FetchRequest( -1, # replica_id 100, 100, [(tp.topic, [(tp.partition, 155, 100000)])]) builder = LegacyRecordBatchBuilder( magic=1, compression_type=0, batch_size=99999999) builder.append(160, value=b"12345", key=b"1", timestamp=None) builder.append(162, value=b"23456", key=b"2", timestamp=None) builder.append(167, value=b"34567", key=b"3", timestamp=None) batch = bytes(builder.build()) resp = FetchResponse( [('test', [( 0, 0, 3000, # partition, error_code, highwater_offset batch # Batch raw bytes )])]) subscriptions.assign_from_user({tp}) assignment = subscriptions.subscription.assignment tp_state = assignment.state_value(tp) client.send.side_effect = asyncio.coroutine(lambda n, r: resp) tp_state.seek(155) fetcher._in_flight.add(0) needs_wake_up = yield from fetcher._proc_fetch_request( assignment, 0, req) self.assertEqual(needs_wake_up, True) buf = fetcher._records[tp] # Test successful getone, the closest in batch offset=160 first = buf.getone() self.assertEqual(tp_state.position, 161) self.assertEqual( (first.value, first.key, first.offset), (b"12345", b"1", 160)) # Test successful getmany second, third = buf.getall() self.assertEqual(tp_state.position, 168) self.assertEqual( (second.value, second.key, second.offset), (b"23456", b"2", 162)) self.assertEqual( (third.value, third.key, third.offset), (b"34567", b"3", 167))
async def test_fetcher__update_fetch_positions(self): client = AIOKafkaClient(bootstrap_servers=[]) subscriptions = SubscriptionState() fetcher = Fetcher(client, subscriptions) self.add_cleanup(fetcher.close) # Disable background task fetcher._fetch_task.cancel() try: await fetcher._fetch_task except asyncio.CancelledError: pass fetcher._fetch_task = create_task(asyncio.sleep(1000000)) partition = TopicPartition('test', 0) offsets = {partition: OffsetAndTimestamp(12, -1)} async def _proc_offset_request(node_id, topic_data): return offsets fetcher._proc_offset_request = mock.Mock() fetcher._proc_offset_request.side_effect = _proc_offset_request def reset_assignment(): subscriptions.assign_from_user({partition}) assignment = subscriptions.subscription.assignment tp_state = assignment.state_value(partition) return assignment, tp_state assignment, tp_state = reset_assignment() self.assertIsNone(tp_state._position) # CASE: reset from committed # In basic case we will need to wait for committed update_task = create_task( fetcher._update_fetch_positions(assignment, 0, [partition]), ) await asyncio.sleep(0.1) self.assertFalse(update_task.done()) # Will continue only after committed is resolved tp_state.update_committed(OffsetAndMetadata(4, "")) needs_wakeup = await update_task self.assertFalse(needs_wakeup) self.assertEqual(tp_state._position, 4) self.assertEqual(fetcher._proc_offset_request.call_count, 0) # CASE: will not query committed if position already present await fetcher._update_fetch_positions(assignment, 0, [partition]) self.assertEqual(tp_state._position, 4) self.assertEqual(fetcher._proc_offset_request.call_count, 0) # CASE: awaiting_reset for the partition tp_state.await_reset(OffsetResetStrategy.LATEST) self.assertIsNone(tp_state._position) await fetcher._update_fetch_positions(assignment, 0, [partition]) self.assertEqual(tp_state._position, 12) self.assertEqual(fetcher._proc_offset_request.call_count, 1) # CASE: seeked while waiting for committed to be resolved assignment, tp_state = reset_assignment() update_task = create_task( fetcher._update_fetch_positions(assignment, 0, [partition]), ) await asyncio.sleep(0.1) self.assertFalse(update_task.done()) tp_state.seek(8) tp_state.update_committed(OffsetAndMetadata(4, "")) await update_task self.assertEqual(tp_state._position, 8) self.assertEqual(fetcher._proc_offset_request.call_count, 1) # CASE: awaiting_reset during waiting for committed assignment, tp_state = reset_assignment() update_task = create_task( fetcher._update_fetch_positions(assignment, 0, [partition]), ) await asyncio.sleep(0.1) self.assertFalse(update_task.done()) tp_state.await_reset(OffsetResetStrategy.LATEST) tp_state.update_committed(OffsetAndMetadata(4, "")) await update_task self.assertEqual(tp_state._position, 12) self.assertEqual(fetcher._proc_offset_request.call_count, 2) # CASE: reset using default strategy if committed offset undefined assignment, tp_state = reset_assignment() loop = get_running_loop() loop.call_later(0.01, tp_state.update_committed, OffsetAndMetadata(-1, "")) await fetcher._update_fetch_positions(assignment, 0, [partition]) self.assertEqual(tp_state._position, 12) self.assertEqual(fetcher._records, {}) # CASE: set error if _default_reset_strategy = OffsetResetStrategy.NONE assignment, tp_state = reset_assignment() loop.call_later(0.01, tp_state.update_committed, OffsetAndMetadata(-1, "")) fetcher._default_reset_strategy = OffsetResetStrategy.NONE needs_wakeup = await fetcher._update_fetch_positions( assignment, 0, [partition]) self.assertTrue(needs_wakeup) self.assertIsNone(tp_state._position) self.assertIsInstance(fetcher._records[partition], FetchError) fetcher._records.clear() # CASE: if _proc_offset_request errored, we will retry on another spin fetcher._proc_offset_request.side_effect = UnknownError() assignment, tp_state = reset_assignment() tp_state.await_reset(OffsetResetStrategy.LATEST) await fetcher._update_fetch_positions(assignment, 0, [partition]) self.assertIsNone(tp_state._position) self.assertTrue(tp_state.awaiting_reset) # CASE: reset 2 partitions separately, 1 will raise, 1 will get # committed fetcher._proc_offset_request.side_effect = _proc_offset_request partition2 = TopicPartition('test', 1) subscriptions.assign_from_user({partition, partition2}) assignment = subscriptions.subscription.assignment tp_state = assignment.state_value(partition) tp_state2 = assignment.state_value(partition2) tp_state.await_reset(OffsetResetStrategy.LATEST) loop.call_later(0.01, tp_state2.update_committed, OffsetAndMetadata(5, "")) await fetcher._update_fetch_positions(assignment, 0, [partition, partition2]) self.assertEqual(tp_state.position, 12) self.assertEqual(tp_state2.position, 5)
async def test_compacted_topic_consumption(self): # Compacted topics can have offsets skipped client = AIOKafkaClient(bootstrap_servers=[]) async def ready(conn): return True def force_metadata_update(): fut = create_future() fut.set_result(True) return fut client.ready = mock.MagicMock() client.ready.side_effect = ready client.force_metadata_update = mock.MagicMock() client.force_metadata_update.side_effect = force_metadata_update client.send = mock.MagicMock() subscriptions = SubscriptionState() fetcher = Fetcher(client, subscriptions) tp = TopicPartition('test', 0) req = FetchRequest( -1, # replica_id 100, 100, [(tp.topic, [(tp.partition, 155, 100000)])]) builder = LegacyRecordBatchBuilder(magic=1, compression_type=0, batch_size=99999999) builder.append(160, value=b"12345", key=b"1", timestamp=None) builder.append(162, value=b"23456", key=b"2", timestamp=None) builder.append(167, value=b"34567", key=b"3", timestamp=None) batch = bytes(builder.build()) resp = FetchResponse([( 'test', [( 0, 0, 3000, # partition, error_code, highwater_offset batch # Batch raw bytes )])]) async def send(node, ready): return resp subscriptions.assign_from_user({tp}) assignment = subscriptions.subscription.assignment tp_state = assignment.state_value(tp) client.send.side_effect = send tp_state.seek(155) fetcher._in_flight.add(0) needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req) self.assertEqual(needs_wake_up, True) buf = fetcher._records[tp] # Test successful getone, the closest in batch offset=160 first = buf.getone() self.assertEqual(tp_state.position, 161) self.assertEqual((first.value, first.key, first.offset), (b"12345", b"1", 160)) # Test successful getmany second, third = buf.getall() self.assertEqual(tp_state.position, 168) self.assertEqual((second.value, second.key, second.offset), (b"23456", b"2", 162)) self.assertEqual((third.value, third.key, third.offset), (b"34567", b"3", 167))
async def test_proc_fetch_request(self): client = AIOKafkaClient(bootstrap_servers=[]) subscriptions = SubscriptionState() fetcher = Fetcher(client, subscriptions, auto_offset_reset="latest") tp = TopicPartition('test', 0) tp_info = (tp.topic, [(tp.partition, 4, 100000)]) req = FetchRequest( -1, # replica_id 100, 100, [tp_info]) async def ready(conn): return True def force_metadata_update(): fut = create_future() fut.set_result(False) return fut client.ready = mock.MagicMock() client.ready.side_effect = ready client.force_metadata_update = mock.MagicMock() client.force_metadata_update.side_effect = force_metadata_update client.send = mock.MagicMock() builder = LegacyRecordBatchBuilder(magic=1, compression_type=0, batch_size=99999999) builder.append(offset=4, value=b"test msg", key=None, timestamp=None) raw_batch = bytes(builder.build()) fetch_response = FetchResponse([('test', [(0, 0, 9, raw_batch)])]) async def send(node, request): nonlocal fetch_response return fetch_response client.send.side_effect = send subscriptions.assign_from_user({tp}) assignment = subscriptions.subscription.assignment tp_state = assignment.state_value(tp) # The partition has no active position, so will ignore result needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req) self.assertEqual(needs_wake_up, False) self.assertEqual(fetcher._records, {}) # The partition's position does not match request's fetch offset subscriptions.seek(tp, 0) needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req) self.assertEqual(needs_wake_up, False) self.assertEqual(fetcher._records, {}) subscriptions.seek(tp, 4) needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req) self.assertEqual(needs_wake_up, True) buf = fetcher._records[tp] self.assertEqual(buf.getone().value, b"test msg") # If position changed after fetch request passed subscriptions.seek(tp, 4) needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req) subscriptions.seek(tp, 10) self.assertIsNone(buf.getone()) # If assignment is lost after fetch request passed subscriptions.seek(tp, 4) needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req) subscriptions.unsubscribe() self.assertIsNone(buf.getone()) subscriptions.assign_from_user({tp}) assignment = subscriptions.subscription.assignment tp_state = assignment.state_value(tp) # error -> no partition found (UnknownTopicOrPartitionError) subscriptions.seek(tp, 4) fetcher._records.clear() fetch_response = FetchResponse([('test', [(0, 3, 9, raw_batch)])]) cc = client.force_metadata_update.call_count needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req) self.assertEqual(needs_wake_up, False) self.assertEqual(client.force_metadata_update.call_count, cc + 1) # error -> topic auth failed (TopicAuthorizationFailedError) fetch_response = FetchResponse([('test', [(0, 29, 9, raw_batch)])]) needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req) self.assertEqual(needs_wake_up, True) with self.assertRaises(TopicAuthorizationFailedError): await fetcher.next_record([]) # error -> unknown fetch_response = FetchResponse([('test', [(0, -1, 9, raw_batch)])]) needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req) self.assertEqual(needs_wake_up, False) # error -> offset out of range with offset strategy fetch_response = FetchResponse([('test', [(0, 1, 9, raw_batch)])]) needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req) self.assertEqual(needs_wake_up, False) self.assertEqual(tp_state.has_valid_position, False) self.assertEqual(tp_state.awaiting_reset, True) self.assertEqual(tp_state.reset_strategy, OffsetResetStrategy.LATEST) # error -> offset out of range without offset strategy subscriptions.seek(tp, 4) fetcher._default_reset_strategy = OffsetResetStrategy.NONE needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req) self.assertEqual(needs_wake_up, True) with self.assertRaises(OffsetOutOfRangeError): await fetcher.next_record([]) await fetcher.close()
def test_fetcher__update_fetch_positions(self): client = AIOKafkaClient( loop=self.loop, bootstrap_servers=[]) subscriptions = SubscriptionState(loop=self.loop) fetcher = Fetcher(client, subscriptions, loop=self.loop) self.add_cleanup(fetcher.close) # Disable backgroud task fetcher._fetch_task.cancel() try: yield from fetcher._fetch_task except asyncio.CancelledError: pass fetcher._fetch_task = ensure_future( asyncio.sleep(1000000, loop=self.loop), loop=self.loop) partition = TopicPartition('test', 0) offsets = {partition: OffsetAndTimestamp(12, -1)} @asyncio.coroutine def _proc_offset_request(node_id, topic_data): return offsets fetcher._proc_offset_request = mock.Mock() fetcher._proc_offset_request.side_effect = _proc_offset_request def reset_assignment(): subscriptions.assign_from_user({partition}) assignment = subscriptions.subscription.assignment tp_state = assignment.state_value(partition) return assignment, tp_state assignment, tp_state = reset_assignment() self.assertIsNone(tp_state._position) # CASE: reset from committed # In basic case we will need to wait for committed update_task = ensure_future( fetcher._update_fetch_positions(assignment, 0, [partition]), loop=self.loop ) yield from asyncio.sleep(0.1, loop=self.loop) self.assertFalse(update_task.done()) # Will continue only after committed is resolved tp_state.update_committed(OffsetAndMetadata(4, "")) needs_wakeup = yield from update_task self.assertFalse(needs_wakeup) self.assertEqual(tp_state._position, 4) self.assertEqual(fetcher._proc_offset_request.call_count, 0) # CASE: will not query committed if position already present yield from fetcher._update_fetch_positions(assignment, 0, [partition]) self.assertEqual(tp_state._position, 4) self.assertEqual(fetcher._proc_offset_request.call_count, 0) # CASE: awaiting_reset for the partition tp_state.await_reset(OffsetResetStrategy.LATEST) self.assertIsNone(tp_state._position) yield from fetcher._update_fetch_positions(assignment, 0, [partition]) self.assertEqual(tp_state._position, 12) self.assertEqual(fetcher._proc_offset_request.call_count, 1) # CASE: seeked while waiting for committed to be resolved assignment, tp_state = reset_assignment() update_task = ensure_future( fetcher._update_fetch_positions(assignment, 0, [partition]), loop=self.loop ) yield from asyncio.sleep(0.1, loop=self.loop) self.assertFalse(update_task.done()) tp_state.seek(8) tp_state.update_committed(OffsetAndMetadata(4, "")) yield from update_task self.assertEqual(tp_state._position, 8) self.assertEqual(fetcher._proc_offset_request.call_count, 1) # CASE: awaiting_reset during waiting for committed assignment, tp_state = reset_assignment() update_task = ensure_future( fetcher._update_fetch_positions(assignment, 0, [partition]), loop=self.loop ) yield from asyncio.sleep(0.1, loop=self.loop) self.assertFalse(update_task.done()) tp_state.await_reset(OffsetResetStrategy.LATEST) tp_state.update_committed(OffsetAndMetadata(4, "")) yield from update_task self.assertEqual(tp_state._position, 12) self.assertEqual(fetcher._proc_offset_request.call_count, 2) # CASE: reset using default strategy if committed offset undefined assignment, tp_state = reset_assignment() self.loop.call_later( 0.01, tp_state.update_committed, OffsetAndMetadata(-1, "")) yield from fetcher._update_fetch_positions(assignment, 0, [partition]) self.assertEqual(tp_state._position, 12) self.assertEqual(fetcher._records, {}) # CASE: set error if _default_reset_strategy = OffsetResetStrategy.NONE assignment, tp_state = reset_assignment() self.loop.call_later( 0.01, tp_state.update_committed, OffsetAndMetadata(-1, "")) fetcher._default_reset_strategy = OffsetResetStrategy.NONE needs_wakeup = yield from fetcher._update_fetch_positions( assignment, 0, [partition]) self.assertTrue(needs_wakeup) self.assertIsNone(tp_state._position) self.assertIsInstance(fetcher._records[partition], FetchError) fetcher._records.clear() # CASE: if _proc_offset_request errored, we will retry on another spin fetcher._proc_offset_request.side_effect = UnknownError() assignment, tp_state = reset_assignment() tp_state.await_reset(OffsetResetStrategy.LATEST) yield from fetcher._update_fetch_positions(assignment, 0, [partition]) self.assertIsNone(tp_state._position) self.assertTrue(tp_state.awaiting_reset) # CASE: reset 2 partitions separately, 1 will raise, 1 will get # committed fetcher._proc_offset_request.side_effect = _proc_offset_request partition2 = TopicPartition('test', 1) subscriptions.assign_from_user({partition, partition2}) assignment = subscriptions.subscription.assignment tp_state = assignment.state_value(partition) tp_state2 = assignment.state_value(partition2) tp_state.await_reset(OffsetResetStrategy.LATEST) self.loop.call_later( 0.01, tp_state2.update_committed, OffsetAndMetadata(5, "")) yield from fetcher._update_fetch_positions( assignment, 0, [partition, partition2]) self.assertEqual(tp_state.position, 12) self.assertEqual(tp_state2.position, 5)
def test_proc_fetch_request(self): client = AIOKafkaClient( loop=self.loop, bootstrap_servers=[]) subscriptions = SubscriptionState(loop=self.loop) fetcher = Fetcher( client, subscriptions, auto_offset_reset="latest", loop=self.loop) tp = TopicPartition('test', 0) tp_info = (tp.topic, [(tp.partition, 4, 100000)]) req = FetchRequest( -1, # replica_id 100, 100, [tp_info]) client.ready = mock.MagicMock() client.ready.side_effect = asyncio.coroutine(lambda a: True) client.force_metadata_update = mock.MagicMock() client.force_metadata_update.side_effect = asyncio.coroutine( lambda: False) client.send = mock.MagicMock() builder = LegacyRecordBatchBuilder( magic=1, compression_type=0, batch_size=99999999) builder.append(offset=4, value=b"test msg", key=None, timestamp=None) raw_batch = bytes(builder.build()) client.send.side_effect = asyncio.coroutine( lambda n, r: FetchResponse( [('test', [(0, 0, 9, raw_batch)])])) subscriptions.assign_from_user({tp}) assignment = subscriptions.subscription.assignment tp_state = assignment.state_value(tp) # The partition has no active position, so will ignore result needs_wake_up = yield from fetcher._proc_fetch_request( assignment, 0, req) self.assertEqual(needs_wake_up, False) self.assertEqual(fetcher._records, {}) # The partition's position does not match request's fetch offset subscriptions.seek(tp, 0) needs_wake_up = yield from fetcher._proc_fetch_request( assignment, 0, req) self.assertEqual(needs_wake_up, False) self.assertEqual(fetcher._records, {}) subscriptions.seek(tp, 4) needs_wake_up = yield from fetcher._proc_fetch_request( assignment, 0, req) self.assertEqual(needs_wake_up, True) buf = fetcher._records[tp] self.assertEqual(buf.getone().value, b"test msg") # If position changed after fetch request passed subscriptions.seek(tp, 4) needs_wake_up = yield from fetcher._proc_fetch_request( assignment, 0, req) subscriptions.seek(tp, 10) self.assertIsNone(buf.getone()) # If assignment is lost after fetch request passed subscriptions.seek(tp, 4) needs_wake_up = yield from fetcher._proc_fetch_request( assignment, 0, req) subscriptions.unsubscribe() self.assertIsNone(buf.getone()) subscriptions.assign_from_user({tp}) assignment = subscriptions.subscription.assignment tp_state = assignment.state_value(tp) # error -> no partition found (UnknownTopicOrPartitionError) subscriptions.seek(tp, 4) fetcher._records.clear() client.send.side_effect = asyncio.coroutine( lambda n, r: FetchResponse( [('test', [(0, 3, 9, raw_batch)])])) cc = client.force_metadata_update.call_count needs_wake_up = yield from fetcher._proc_fetch_request( assignment, 0, req) self.assertEqual(needs_wake_up, False) self.assertEqual(client.force_metadata_update.call_count, cc + 1) # error -> topic auth failed (TopicAuthorizationFailedError) client.send.side_effect = asyncio.coroutine( lambda n, r: FetchResponse( [('test', [(0, 29, 9, raw_batch)])])) needs_wake_up = yield from fetcher._proc_fetch_request( assignment, 0, req) self.assertEqual(needs_wake_up, True) with self.assertRaises(TopicAuthorizationFailedError): yield from fetcher.next_record([]) # error -> unknown client.send.side_effect = asyncio.coroutine( lambda n, r: FetchResponse( [('test', [(0, -1, 9, raw_batch)])])) needs_wake_up = yield from fetcher._proc_fetch_request( assignment, 0, req) self.assertEqual(needs_wake_up, False) # error -> offset out of range with offset strategy client.send.side_effect = asyncio.coroutine( lambda n, r: FetchResponse( [('test', [(0, 1, 9, raw_batch)])])) needs_wake_up = yield from fetcher._proc_fetch_request( assignment, 0, req) self.assertEqual(needs_wake_up, False) self.assertEqual(tp_state.has_valid_position, False) self.assertEqual(tp_state.awaiting_reset, True) self.assertEqual(tp_state.reset_strategy, OffsetResetStrategy.LATEST) # error -> offset out of range without offset strategy subscriptions.seek(tp, 4) fetcher._default_reset_strategy = OffsetResetStrategy.NONE needs_wake_up = yield from fetcher._proc_fetch_request( assignment, 0, req) self.assertEqual(needs_wake_up, True) with self.assertRaises(OffsetOutOfRangeError): yield from fetcher.next_record([]) yield from fetcher.close()