feat kafka: reuse consumer after user exception

commit_hash:aa170d722e6d81556d581b37da4e1afc33297b2c

Author: jdillinger

kafka/functional_tests/balanced_consumer_groups/static_config.yaml

@@ -29,7 +29,8 @@ components_manager:
               - test-topic
             auto_offset_reset: smallest
             max_batch_size: 1
-            restart_after_failure_delay: 1s
+            poll_timeout: 200ms
+            restart_after_failure_delay: 100ms
             security_protocol: PLAINTEXT
 
         kafka-consumer-second: *kafka_consumer
@@ -40,7 +41,8 @@ components_manager:
               - test-cooperative-topic
             auto_offset_reset: smallest
             max_batch_size: 1
-            restart_after_failure_delay: 1s
+            poll_timeout: 200ms
+            restart_after_failure_delay: 100ms
             security_protocol: PLAINTEXT
             rd_kafka_custom_options:
                 partition.assignment.strategy: cooperative-sticky

kafka/functional_tests/balanced_consumer_groups/tests/test_faulty_consumer.py

@@ -36,6 +36,10 @@ def second_consumer_received(_data):
     def second_consumer_subscribed(_data):
         logging.info('Second consumer subscribed')
 
+    @testpoint('tp_kafka-consumer-second_revoked')
+    def second_consumer_revoked(_data):
+        logging.info('Second consumer revoked')
+
     @testpoint('tp_kafka-consumer-second_polled')
     def second_consumer_polled(_data):
         logging.info('Second consumer polled')
@@ -66,7 +70,11 @@ def second_consumer_failed(_data):
         CONSUMERS[1],
     )
     assert len(first_consumer_messages) == 1 and len(second_consumer_messages) == 1
+    first_consumer_messages = first_consumer_messages[0]['partition']
     second_consumer_partition = second_consumer_messages[0]['partition']
+    logging.info(
+        f'First consumer was subscribed to {first_consumer_messages} partition',
+    )
     logging.info(
         f'Second consumer was subscribed to {second_consumer_partition} partition',
     )
@@ -77,12 +85,17 @@ def second_consumer_failed(_data):
         MESSAGE_TO_FAIL,
         second_consumer_partition,
     )
-    await second_consumer_polled.wait_call()
-    await second_consumer_failed.wait_call()
+    for _ in range(0, 2):
+        await second_consumer_polled.wait_call()
+        await second_consumer_failed.wait_call()
+
+    assert not second_consumer_subscribed.has_calls
+    assert not second_consumer_revoked.has_calls
 
     await stop_consumers(service_client, [CONSUMERS[1]])
 
     await first_consumer_revoked.wait_call()
+    await second_consumer_revoked.wait_call()
     await first_consumer_subscribed.wait_call()
     await first_consumer_subscribed.wait_call()
 
@@ -143,10 +156,12 @@ def second_consumer_failed(_data):
         CONSUMERS[1],
     )
     assert len(second_consumer_messages) == 2
+    assert {'key-1', 'key-2'} == set(parse_message_keys(second_consumer_messages))
 
     await kafka_producer.send(TOPIC, 'key-3', MESSAGE_TO_FAIL, 0)
-    await second_consumer_polled.wait_call()
-    await second_consumer_failed.wait_call()
+    for _ in range(0, 2):
+        await second_consumer_polled.wait_call()
+        await second_consumer_failed.wait_call()
 
     await make_non_faulty(service_client, CONSUMERS[1])
 

kafka/include/userver/kafka/impl/consumer.hpp

@@ -10,6 +10,7 @@
 #include <userver/kafka/impl/consumer_params.hpp>
 #include <userver/kafka/impl/holders.hpp>
 #include <userver/kafka/impl/stats.hpp>
+#include <userver/utils/span.hpp>
 #include <userver/utils/statistics/writer.hpp>
 #include <userver/utils/zstring_view.hpp>
 
@@ -85,6 +86,11 @@ class Consumer final {
         std::chrono::milliseconds timeout
     ) const;
 
+    /// @brief Seeks multiple topic partitions to the given offsets in one call.
+    /// @param params Topic, partition_id and offset for each seek.
+    /// @param timeout Timeout for the operation.
+    void MultiSeek(utils::span<const SeekParams> params, std::chrono::milliseconds timeout) const;
+
     /// @brief Seeks the \b partition_id for the specified \b topic to the begginning offset .
     /// @see ConsumerScope::SeekToBeginning for better commitment process
     void SeekToBeginning(utils::zstring_view topic, std::uint32_t partition_id, std::chrono::milliseconds timeout)

kafka/include/userver/kafka/impl/consumer_params.hpp

@@ -3,12 +3,20 @@
 #include <chrono>
 
 #include <userver/logging/level.hpp>
+#include <userver/utils/zstring_view.hpp>
 #include <userver/yaml_config/fwd.hpp>
 
 USERVER_NAMESPACE_BEGIN
 
 namespace kafka::impl {
 
+/// @brief Topic, partition and offset for a single seek (e.g. in MultiSeek).
+struct SeekParams {
+    utils::zstring_view topic;
+    std::uint32_t partition_id;
+    std::uint64_t offset;
+};
+
 /// @brief Specifies the logging format for the message key.
 enum class MessageKeyLogFormat {
     kPlainText,  ///< Log message key as is

kafka/src/kafka/impl/consumer.cpp

@@ -1,6 +1,8 @@
 #include <userver/kafka/impl/consumer.hpp>
 
+#include <map>
 #include <string_view>
+#include <vector>
 
 #include <userver/engine/sleep.hpp>
 #include <userver/formats/json/value_builder.hpp>
@@ -129,8 +131,34 @@ void Consumer::RunConsuming(ConsumerScope::Callback callback) {
             consumer_->AccountMessageBatchProcessingSucceeded(polled_messages);
             TESTPOINT(fmt::format("tp_{}", name_), {});
         } catch (const std::exception& e) {
+            LOG_ERROR("Messages processing failed in consumer into client callback: {}", e.what());
             consumer_->AccountMessageBatchProcessingFailed(polled_messages);
-            throw;
+
+            // Seek back to the start of the failed batch so the same messages
+            // are redelivered on the next poll (no consumer recreation).
+            // Seek with numeric offset is local-only in librdkafka: it sets
+            // next_fetch_position; no broker request until the next Fetch (poll).
+            using Key = std::pair<utils::zstring_view, std::uint32_t>;
+            std::map<Key, std::uint64_t> min_offset_by_partition;
+            for (const auto& msg : polled_messages) {
+                Key key{msg.GetTopic(), static_cast<std::uint32_t>(msg.GetPartition())};
+                const auto offset = static_cast<std::uint64_t>(msg.GetOffset());
+                auto it = min_offset_by_partition.find(key);
+                if (it == min_offset_by_partition.end()) {
+                    min_offset_by_partition[key] = offset;
+                } else {
+                    it->second = std::min(it->second, offset);
+                }
+            }
+            std::vector<SeekParams> seek_params;
+            seek_params.reserve(min_offset_by_partition.size());
+            for (const auto& [key, offset] : min_offset_by_partition) {
+                seek_params.push_back({key.first, key.second, offset});
+            }
+            constexpr auto kSeekAfterFailureTimeout = std::chrono::seconds(20);
+            MultiSeek(seek_params, kSeekAfterFailureTimeout);
+            LOG_WARNING("Seeked back to reprocess failed batch");
+            CallErrorTestpoint(fmt::format("tp_error_{}", name_), e.what());
         }
     }
 }
@@ -219,6 +247,16 @@ void Consumer::Seek(
     }).Get();
 }
 
+void Consumer::MultiSeek(utils::span<const SeekParams> params, std::chrono::milliseconds timeout) const {
+    UINVARIANT(processing_.load(), "Message processing is not currently started");
+
+    utils::Async(consumer_blocking_task_processor_, "consumer_multi_seek", [this, params, timeout] {
+        ExtendCurrentSpan();
+
+        consumer_->MultiSeek(params, timeout);
+    }).Get();
+}
+
 void Consumer::SeekToBeginning(utils::zstring_view topic, std::uint32_t partition_id, std::chrono::milliseconds timeout)
     const {
     UINVARIANT(processing_.load(), "Message processing is not currently started");

kafka/src/kafka/impl/consumer_impl.cpp

@@ -14,6 +14,7 @@
 #include <userver/tracing/span.hpp>
 #include <userver/utils/encoding/hex.hpp>
 #include <userver/utils/fast_scope_guard.hpp>
+#include <userver/utils/numeric_cast.hpp>
 #include <userver/utils/span.hpp>
 
 #include <kafka/impl/error_buffer.hpp>
@@ -310,6 +311,7 @@ void ConsumerImpl::LogCallback(const char* facility, const char* message, int lo
 void ConsumerImpl::RebalanceCallback(rd_kafka_resp_err_t err, const rd_kafka_topic_partition_list_t* partitions) {
     tracing::Span span{"rebalance_callback"};
     span.AddTag("kafka_callback", "rebalance_callback");
+    span.AddTag("kafka_consumer", name_);
 
     LOG(execution_params_.operation_log_level,
         "Consumer group rebalanced ('{}' protocol)",
@@ -729,6 +731,20 @@ void ConsumerImpl::Seek(
     SeekToOffset(topic, partition_id, static_cast<std::int64_t>(offset), timeout);
 }
 
+void ConsumerImpl::MultiSeek(utils::span<const SeekParams> params, std::chrono::milliseconds timeout) const {
+    if (params.empty()) {
+        return;
+    }
+    for (const auto& p : params) {
+        if (p.offset > static_cast<std::uint64_t>(std::numeric_limits<std::int64_t>::max())) {
+            throw SeekInvalidArgumentException(
+                fmt::format("Offset value have to be <= std::int64_t::max(). offset: {}", p.offset)
+            );
+        }
+    }
+    SeekToOffsets(params, timeout);
+}
+
 void ConsumerImpl::SeekToOffset(
     utils::zstring_view topic,
     std::uint32_t partition_id,
@@ -768,13 +784,7 @@ void ConsumerImpl::SeekToOffset(
     const auto* err =
         rd_kafka_seek_partitions(consumer_.GetHandle(), topic_partitions_list.GetHandle(), ToRdKafkaTimeout(deadline));
     if (err == nullptr) {
-        LOG_INFO(
-            "Seeked to offset: {}"
-            " for partition: {} topic: {} successfully",
-            offset,
-            partition_id,
-            topic
-        );
+        LOG_INFO("Seeked to offset: {} for partition: {} topic: {} successfully", offset, partition_id, topic);
         return;
     }
 
@@ -786,6 +796,50 @@ void ConsumerImpl::SeekToOffset(
     ));
 }
 
+void ConsumerImpl::SeekToOffsets(utils::span<const SeekParams> params, std::chrono::milliseconds timeout) const {
+    if (params.empty()) {
+        return;
+    }
+    if (timeout.count() <= 0) {
+        throw SeekInvalidArgumentException(fmt::format("Timeout value have to be > 0. value: {}ms", timeout.count()));
+    }
+
+    const auto deadline = engine::Deadline::FromDuration(timeout);
+
+    // Seek is local in librdkafka (next_fetch_position). No poll needed; next Fetch uses new position.
+    TopicPartitionsListHolder topic_partitions_list{
+        rd_kafka_topic_partition_list_new(utils::numeric_cast<int>(params.size()))
+    };
+    for (const auto& p : params) {
+        rd_kafka_topic_partition_t* part = rd_kafka_topic_partition_list_add(
+            topic_partitions_list.GetHandle(),
+            p.topic.c_str(),
+            static_cast<std::int32_t>(p.partition_id)
+        );
+        part->offset = static_cast<std::int64_t>(p.offset);
+    }
+
+    PrintTopicPartitionsList(topic_partitions_list.GetHandle(), [](const rd_kafka_topic_partition_t& partition) {
+        return fmt::format(
+            "Partition {} for topic '{}' seeking to offset: {}",
+            partition.partition,
+            partition.topic,
+            partition.offset
+        );
+    });
+
+    const auto* err =
+        rd_kafka_seek_partitions(consumer_.GetHandle(), topic_partitions_list.GetHandle(), ToRdKafkaTimeout(deadline));
+    if (err == nullptr) {
+        LOG_INFO("MultiSeek: seeked {} partition(s) successfully", params.size());
+        return;
+    }
+
+    throw SeekException(
+        fmt::format("Failed to multi-seek {} partition(s). err: {}", params.size(), rd_kafka_error_string(err))
+    );
+}
+
 void ConsumerImpl::SeekToEnd(utils::zstring_view topic, std::uint32_t partition_id, std::chrono::milliseconds timeout)
     const {
     SeekToOffset(topic, partition_id, RD_KAFKA_OFFSET_END, timeout);

kafka/src/kafka/impl/consumer_impl.hpp

@@ -12,6 +12,7 @@
 #include <userver/kafka/message.hpp>
 #include <userver/kafka/offset_range.hpp>
 #include <userver/kafka/rebalance_types.hpp>
+#include <userver/utils/span.hpp>
 #include <userver/utils/zstring_view.hpp>
 
 #include <kafka/impl/holders_aliases.hpp>
@@ -76,6 +77,11 @@ class ConsumerImpl final {
     /// @brief Seeks the partition ID for the specified \b topic to the end offset .
     void SeekToEnd(utils::zstring_view topic, std::uint32_t partition_id, std::chrono::milliseconds timeout) const;
 
+    /// @brief Seeks multiple topic partitions to the given offsets in one call.
+    /// @param params Topic, partition_id and offset for each seek.
+    /// @param timeout Timeout for the operation.
+    void MultiSeek(utils::span<const SeekParams> params, std::chrono::milliseconds timeout) const;
+
     /// @brief Effectively calls `PollMessage` until `deadline` is reached
     /// and no more than `max_batch_size` messages polled.
     MessageBatch PollBatch(std::size_t max_batch_size, engine::Deadline deadline);
@@ -146,6 +152,9 @@ class ConsumerImpl final {
         std::chrono::milliseconds timeout
     ) const;
 
+    /// @brief Seeks multiple topic partitions to the given offsets (internal, expects validated inputs).
+    void SeekToOffsets(utils::span<const SeekParams> params, std::chrono::milliseconds timeout) const;
+
     /// @brief Callback which is called after succeeded/failed commit.
     /// Currently, used for logging purposes.
     void OffsetCommitCallback(rd_kafka_resp_err_t err, const rd_kafka_topic_partition_list_s* committed_offsets) const;