k/d/rpc/client: Add retry_with_leader_mitigation to consume

Signed-off-by: Oren Leiman <oren.leiman@redpanda.com>

Author: oleiman

src/v/kafka/data/rpc/client.cc

@@ -583,43 +583,56 @@ client::do_get_single_partition_offsets_once(model::topic_partition tp) {
 }
 
 ss::future<result<consume_reply, cluster::errc>> client::consume(
+  model::topic_partition tp,
+  kafka::offset start_offset,
+  kafka::offset max_offset,
+  size_t min_bytes,
+  size_t max_bytes,
+  model::timeout_clock::duration timeout) {
+    co_return co_await retry_with_leader_mitigation(
+      [this, tp, start_offset, max_offset, min_bytes, max_bytes, timeout]() {
+          return do_consume_once(
+            tp, start_offset, max_offset, min_bytes, max_bytes, timeout);
+      });
+}
+
+ss::future<result<consume_reply, cluster::errc>> client::do_consume_once(
   model::topic_partition tp,
   kafka::offset start_offset,
   kafka::offset max_offset,
   size_t min_bytes,
   size_t max_bytes,
   model::timeout_clock::duration timeout) {
     using ret_t = result<consume_reply, cluster::errc>;
+    model::topic_namespace_view tp_ns(model::kafka_namespace, tp.topic);
 
-    // Check if topic exists first
-    auto topic_cfg = _metadata_cache->find_topic_cfg(
-      model::topic_namespace_view(model::kafka_namespace, tp.topic));
+    auto topic_cfg = _metadata_cache->find_topic_cfg(tp_ns);
     if (!topic_cfg) {
         consume_reply reply;
         reply.tp = tp;
         reply.err = cluster::errc::topic_not_exists;
         co_return ret_t(std::move(reply));
     }
 
-    // Find the leader for this partition
-    auto ktp = model::ktp(tp.topic, tp.partition);
-    auto leader = _leaders->get_leader_node(
-      model::topic_namespace_view(model::kafka_namespace, tp.topic),
-      tp.partition);
-
+    auto leader = _leaders->get_leader_node(tp_ns, tp.partition);
     if (!leader) {
-        consume_reply reply;
-        reply.tp = tp;
-        reply.err = cluster::errc::not_leader;
-        co_return ret_t(std::move(reply));
+        co_return ret_t(cluster::errc::not_leader);
     }
 
     consume_request req(
       tp, start_offset, max_offset, min_bytes, max_bytes, timeout);
 
+    auto is_retriable = [](cluster::errc ec) -> bool {
+        return ec == cluster::errc::not_leader || ec == cluster::errc::timeout
+               || ec == cluster::errc::partition_operation_failed;
+    };
+
     // If leader is local, call local service
     if (*leader == _self) {
         auto reply = co_await _local_service->local().consume(std::move(req));
+        if (is_retriable(reply.err)) {
+            co_return ret_t(reply.err);
+        }
         co_return ret_t(std::move(reply));
     }
 
@@ -643,12 +656,13 @@ ss::future<result<consume_reply, cluster::errc>> client::consume(
           });
 
     if (result.has_error()) {
-        consume_reply reply;
-        reply.tp = tp;
-        reply.err = map_errc(result.assume_error());
-        co_return ret_t(std::move(reply));
+        co_return ret_t(map_errc(result.assume_error()));
     }
 
-    co_return ret_t(std::move(result.value()));
+    auto reply = std::move(result).value();
+    if (is_retriable(reply.err)) {
+        co_return ret_t(reply.err);
+    }
+    co_return ret_t(std::move(reply));
 }
 } // namespace kafka::data::rpc

src/v/kafka/data/rpc/client.h

@@ -121,6 +121,14 @@ class client {
     get_remote_partition_offsets(
       model::node_id, chunked_vector<topic_partitions> topics);
 
+    ss::future<result<consume_reply, cluster::errc>> do_consume_once(
+      model::topic_partition,
+      kafka::offset start_offset,
+      kafka::offset max_offset,
+      size_t min_bytes,
+      size_t max_bytes,
+      model::timeout_clock::duration timeout);
+
     template<typename Func>
     std::invoke_result_t<Func> retry(Func&&);
 

src/v/kafka/data/rpc/test/kafka_data_rpc_test.cc

@@ -455,6 +455,24 @@ TEST_P(KafkaDataRpcTest, GetSinglePartitionOffsetsRetries) {
     EXPECT_EQ(res.value().high_watermark, kafka::offset(0));
     EXPECT_EQ(res.value().last_stable_offset, kafka::offset(-1));
 }
+
+TEST_P(KafkaDataRpcTest, ConsumeRetries) {
+    auto ntp = make_ntp("retry_consume");
+    create_topic(model::topic_namespace(ntp.ns, ntp.tp.topic));
+
+    auto batches = record_batches::make();
+    cluster::errc ec = produce(ntp, batches);
+    ASSERT_EQ(ec, cluster::errc::success);
+
+    // Inject timeouts on the first 2 consume attempts; the basic retry
+    // policy should back off and succeed on the 3rd attempt.
+    set_errors_to_inject(2);
+    auto r = consume(ntp.tp, kafka::offset(0), kafka::offset::max());
+    ASSERT_TRUE(r.has_value());
+    EXPECT_EQ(r.value().err, cluster::errc::success);
+    EXPECT_EQ(r.value().batches.size(), batches.size());
+}
+
 INSTANTIATE_TEST_SUITE_P(
   WorksLocallyAndRemotely,
   KafkaDataRpcTest,