k/d/rpc/client: Introduce get_single_partition_offsets

Reliable version that tries to mitigate not_leader errors

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

Author: oleiman

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

@@ -524,6 +524,64 @@ client::get_remote_partition_offsets(
     co_return ret_t(std::move(result.value().partition_offsets));
 }
 
+ss::future<result<partition_offsets, cluster::errc>>
+client::get_single_partition_offsets(model::topic_partition tp) {
+    co_return co_await retry_with_leader_mitigation(
+      [this, tp]() { return do_get_single_partition_offsets_once(tp); });
+}
+
+ss::future<result<partition_offsets, cluster::errc>>
+client::do_get_single_partition_offsets_once(model::topic_partition tp) {
+    using ret_t = result<partition_offsets, cluster::errc>;
+    model::topic_namespace_view tp_ns(model::kafka_namespace, tp.topic);
+
+    auto topic_cfg = _metadata_cache->find_topic_cfg(tp_ns);
+    if (!topic_cfg) {
+        co_return ret_t(cluster::errc::topic_not_exists);
+    }
+
+    auto leader = _leaders->get_leader_node(tp_ns, tp.partition);
+    if (!leader) {
+        co_return ret_t(cluster::errc::not_leader);
+    }
+
+    // Build the single-partition request in the same wire format the
+    // existing get_offsets RPC expects.
+    chunked_vector<topic_partitions> topics;
+    topic_partitions tps;
+    tps.topic = tp.topic;
+    tps.partitions.push_back(tp.partition);
+    topics.push_back(std::move(tps));
+
+    partition_offsets_map offsets_map;
+    if (*leader == _self) {
+        offsets_map = co_await _local_service->local().get_offsets(
+          std::move(topics));
+    } else {
+        auto remote_result = co_await get_remote_partition_offsets(
+          *leader, std::move(topics));
+        if (remote_result.has_error()) {
+            co_return ret_t(remote_result.error());
+        }
+        offsets_map = std::move(remote_result.value());
+    }
+
+    // Extract the single result.
+    auto topic_it = offsets_map.find(tp.topic);
+    if (topic_it == offsets_map.end()) {
+        co_return ret_t(cluster::errc::not_leader);
+    }
+    auto part_it = topic_it->second.find(tp.partition);
+    if (part_it == topic_it->second.end()) {
+        co_return ret_t(cluster::errc::not_leader);
+    }
+    auto& por = part_it->second;
+    if (por.err != cluster::errc::success) {
+        co_return ret_t(por.err);
+    }
+    co_return ret_t(por.offsets);
+}
+
 ss::future<result<consume_reply, cluster::errc>> client::consume(
   model::topic_partition tp,
   kafka::offset start_offset,

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

@@ -94,6 +94,12 @@ class client {
     ss::future<result<partition_offsets_map, cluster::errc>>
       get_partition_offsets(chunked_vector<topic_partitions>);
 
+    /// Get offsets for a single partition. Retries with leadership
+    /// mitigation — suitable for latency-sensitive callers like schema
+    /// registry that always target a single known partition.
+    ss::future<result<partition_offsets, cluster::errc>>
+      get_single_partition_offsets(model::topic_partition);
+
     ss::future<result<consume_reply, cluster::errc>> consume(
       model::topic_partition,
       kafka::offset start_offset,
@@ -108,6 +114,9 @@ class client {
     ss::future<produce_reply>
       do_remote_produce(model::node_id, produce_request);
 
+    ss::future<result<partition_offsets, cluster::errc>>
+      do_get_single_partition_offsets_once(model::topic_partition);
+
     ss::future<result<partition_offsets_map, cluster::errc>>
     get_remote_partition_offsets(
       model::node_id, chunked_vector<topic_partitions> topics);

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

@@ -174,6 +174,14 @@ class KafkaDataRpcTest : public ::testing::TestWithParam<test_parameters> {
           .offsets.high_watermark;
     }
 
+    result<partition_offsets, cluster::errc>
+    get_single_partition_offsets(model::topic_partition tp) {
+        return _kd->client()
+          .local()
+          .get_single_partition_offsets(std::move(tp))
+          .get();
+    }
+
     result<consume_reply, cluster::errc> consume(
       model::topic_partition tp,
       kafka::offset start_offset,
@@ -405,6 +413,48 @@ TEST_P(KafkaDataRpcTest, ProduceWithLeaderMitigationRetries) {
     ASSERT_TRUE(r.last_offset.has_value());
     EXPECT_EQ(get_hwm(ntp), kafka::offset(model::next_offset(*r.last_offset)));
 }
+
+TEST_P(KafkaDataRpcTest, ClientCanGetSinglePartitionOffsets) {
+    auto ntp = make_ntp("single_offsets");
+    create_topic(model::topic_namespace(ntp.ns, ntp.tp.topic));
+
+    auto res = get_single_partition_offsets(ntp.tp);
+    ASSERT_TRUE(res.has_value());
+    // No records produced — HWM is 0 (next of -1), LSO is -1
+    EXPECT_EQ(res.value().high_watermark, kafka::offset(0));
+    EXPECT_EQ(res.value().last_stable_offset, kafka::offset(-1));
+
+    // After producing a record, offsets should advance.
+    auto batch = model::test::make_random_batch({.count = 1, .records = 1});
+    auto pr = produce_with_leader_mitigation(ntp, std::move(batch));
+    ASSERT_EQ(pr.ec, cluster::errc::success);
+
+    auto res2 = get_single_partition_offsets(ntp.tp);
+    ASSERT_TRUE(res2.has_value());
+    EXPECT_EQ(
+      res2.value().high_watermark,
+      kafka::offset(model::next_offset(*pr.last_offset)));
+}
+
+TEST_P(KafkaDataRpcTest, GetSinglePartitionOffsetsReturnsTopicNotExists) {
+    auto ntp = make_ntp("does-not-exist");
+    auto res = get_single_partition_offsets(ntp.tp);
+    ASSERT_TRUE(res.has_error());
+    EXPECT_EQ(res.error(), cluster::errc::topic_not_exists);
+}
+
+TEST_P(KafkaDataRpcTest, GetSinglePartitionOffsetsRetries) {
+    auto ntp = make_ntp("retry_offsets");
+    create_topic(model::topic_namespace(ntp.ns, ntp.tp.topic));
+
+    // Inject timeouts on the first 2 attempts; the basic retry policy
+    // should back off and succeed on the 3rd attempt.
+    set_errors_to_inject(2);
+    auto res = get_single_partition_offsets(ntp.tp);
+    ASSERT_TRUE(res.has_value());
+    EXPECT_EQ(res.value().high_watermark, kafka::offset(0));
+    EXPECT_EQ(res.value().last_stable_offset, kafka::offset(-1));
+}
 INSTANTIATE_TEST_SUITE_P(
   WorksLocallyAndRemotely,
   KafkaDataRpcTest,