Merge pull request #26006 from WillemKauf/manual-backport-25932-v24.3.x-5

[v24.3.x] `storage`: decompress batch during `batch_timequery()` (Manual backport)

Author: WillemKauf

src/v/cloud_storage/remote_partition.cc

@@ -1280,16 +1280,17 @@ remote_partition::timequery(storage::timequery_config cfg) {
     auto translating_reader = co_await make_reader(config);
 
     // Read one batch from the reader to learn the offset
-    model::record_batch_reader::storage_t data
-      = co_await model::consume_reader_to_memory(
-        std::move(translating_reader.reader), model::no_timeout);
+    auto batches = co_await model::consume_reader_to_memory(
+      std::move(translating_reader.reader), model::no_timeout);
 
-    auto& batches = std::get<model::record_batch_reader::data_t>(data);
     vlog(_ctxlog.debug, "timequery: {} batches", batches.size());
 
     if (batches.size()) {
-        co_return storage::batch_timequery(
-          *(batches.begin()), cfg.min_offset, cfg.time, cfg.max_offset);
+        co_return co_await storage::batch_timequery(
+          std::move(*(batches.begin())),
+          cfg.min_offset,
+          cfg.time,
+          cfg.max_offset);
     } else {
         co_return std::nullopt;
     }

src/v/kafka/server/tests/list_offsets_test.cc

@@ -222,10 +222,8 @@ FIXTURE_TEST(list_offsets_by_time, redpanda_thread_fixture) {
     // Arbitrary synthetic timestamp for start of produce
     auto base_timestamp = 100000;
 
-    for (long i = 0; i < batch_count; ++i) {
-        // Mixture of compressed and uncompressed, they have distinct offset
-        // lookup behavior when searching by timequery, which will be
-        // validated
+    for (size_t i = 0; i < batch_count; ++i) {
+        // Mixture of compressed and uncompressed batches.
         bool compressed = i % 3 == 0;
         batches.push_back(make_random_batch(model::test::record_batch_spec{
           // after queries below.
@@ -279,30 +277,17 @@ FIXTURE_TEST(list_offsets_by_time, redpanda_thread_fixture) {
           }},
         });
 
-        const auto& batch = batches[i];
         auto resp_midbatch
           = client.dispatch(std::move(req2), kafka::api_version(1)).get();
         BOOST_REQUIRE_EQUAL(resp_midbatch.data.topics.size(), 1);
         BOOST_REQUIRE_EQUAL(resp_midbatch.data.topics[0].partitions.size(), 1);
-        if (batch.compressed()) {
-            // Compressed batch: result will point to start of batch, slightly
-            // earlier than the query timestamp
-            BOOST_CHECK(
-              resp_midbatch.data.topics[0].partitions[0].timestamp
-              == model::timestamp(base_timestamp + i * record_count));
-            BOOST_CHECK(
-              resp_midbatch.data.topics[0].partitions[0].offset
-              == model::offset(i * record_count));
-        } else {
-            // Uncompressed batch: result should have seeked to correct record
-            BOOST_CHECK(
-              resp_midbatch.data.topics[0].partitions[0].timestamp
-              == model::timestamp(
-                base_timestamp + i * record_count + record_offset));
-            BOOST_CHECK(
-              resp_midbatch.data.topics[0].partitions[0].offset
-              == model::offset(i * record_count + record_offset));
-        }
+        BOOST_CHECK(
+          resp_midbatch.data.topics[0].partitions[0].timestamp
+          == model::timestamp(
+            base_timestamp + i * record_count + record_offset));
+        BOOST_CHECK(
+          resp_midbatch.data.topics[0].partitions[0].offset
+          == model::offset(i * record_count + record_offset));
     }
 
     client.stop().then([&client] { client.shutdown(); }).get();

src/v/model/record.h

@@ -25,6 +25,7 @@
 #include "serde/rw/iobuf.h"
 #include "serde/rw/rw.h"
 
+#include <seastar/core/loop.hh>
 #include <seastar/core/smp.hh>
 #include <seastar/util/optimized_optional.hh>
 
@@ -871,7 +872,18 @@ class record_batch
     ss::future<> for_each_record_async(Func f) const {
         auto it = record_batch_iterator::create(*this);
         while (it.has_next()) {
-            co_await ss::futurize_invoke(f, it.next());
+            if constexpr (std::is_same_v<
+                            ss::futurize_t<
+                              std::invoke_result_t<Func, model::record>>,
+                            ss::future<ss::stop_iteration>>) {
+                ss::stop_iteration s = co_await ss::futurize_invoke(
+                  f, it.next());
+                if (s == ss::stop_iteration::yes) {
+                    co_return;
+                }
+            } else {
+                co_await ss::futurize_invoke(f, it.next());
+            }
         }
     }
 

src/v/storage/disk_log_impl.cc

@@ -2572,23 +2572,20 @@ ss::future<std::optional<timequery_result>>
 disk_log_impl::timequery(timequery_config cfg) {
     vassert(!_closed, "timequery on closed log - {}", *this);
     if (_segs.empty()) {
-        return ss::make_ready_future<std::optional<timequery_result>>();
-    }
-    return make_reader(cfg).then([cfg](model::record_batch_reader reader) {
-        return model::consume_reader_to_memory(
-                 std::move(reader), model::no_timeout)
-          .then([cfg](model::record_batch_reader::storage_t st) {
-              using ret_t = std::optional<timequery_result>;
-              auto& batches = std::get<model::record_batch_reader::data_t>(st);
-              if (
-                !batches.empty()
-                && batches.front().header().max_timestamp >= cfg.time) {
-                  return ret_t(batch_timequery(
-                    batches.front(), cfg.min_offset, cfg.time, cfg.max_offset));
-              }
-              return ret_t();
-          });
-    });
+        co_return std::nullopt;
+    }
+
+    auto reader = co_await make_reader(cfg);
+    auto batches = co_await model::consume_reader_to_memory(
+      std::move(reader), model::no_timeout);
+
+    if (
+      !batches.empty() && batches.front().header().max_timestamp >= cfg.time) {
+        co_return co_await batch_timequery(
+          std::move(batches.front()), cfg.min_offset, cfg.time, cfg.max_offset);
+    }
+
+    co_return std::nullopt;
 }
 
 ss::future<> disk_log_impl::remove_segment_permanently(

src/v/storage/log_reader.cc

@@ -18,6 +18,7 @@
 #include "storage/logger.h"
 #include "storage/offset_translator_state.h"
 #include "storage/parser_errc.h"
+#include "storage/parser_utils.h"
 #include "storage/segment_set.h"
 #include "storage/types.h"
 
@@ -581,8 +582,8 @@ bool log_reader::is_done() {
            || is_finished_offset(_lease->range, _config.start_offset);
 }
 
-timequery_result batch_timequery(
-  const model::record_batch& b,
+ss::future<timequery_result> batch_timequery(
+  model::record_batch b,
   model::offset min_offset,
   model::timestamp t,
   model::offset max_offset) {
@@ -595,23 +596,25 @@ timequery_result batch_timequery(
     // records in the batch have different timestamps.
     model::offset result_o = b.base_offset();
     model::timestamp result_t = b.header().first_timestamp;
-    if (!b.compressed()) {
-        b.for_each_record(
-          [&result_o, &result_t, &b, query_interval, t](
-            const model::record& r) -> ss::stop_iteration {
-              auto record_o = model::offset{r.offset_delta()} + b.base_offset();
-              auto record_t = model::timestamp(
-                b.header().first_timestamp() + r.timestamp_delta());
-              if (record_t >= t && query_interval.contains(record_o)) {
-                  result_o = record_o;
-                  result_t = record_t;
-                  return ss::stop_iteration::yes;
-              } else {
-                  return ss::stop_iteration::no;
-              }
-          });
-    }
-    return {result_o, result_t};
+    auto batch = co_await internal::decompress_batch(std::move(b));
+    co_await batch.for_each_record_async(
+      [&result_o, &result_t, &batch, query_interval, t](
+        const model::record& r) -> ss::future<ss::stop_iteration> {
+          auto record_o = model::offset{r.offset_delta()} + batch.base_offset();
+          auto record_t = model::timestamp(
+            batch.header().first_timestamp() + r.timestamp_delta());
+          if (record_t >= t && query_interval.contains(record_o)) {
+              result_o = record_o;
+              result_t = record_t;
+              return ss::make_ready_future<ss::stop_iteration>(
+                ss::stop_iteration::yes);
+          } else {
+              return ss::make_ready_future<ss::stop_iteration>(
+                ss::stop_iteration::no);
+          }
+      });
+
+    co_return timequery_result{result_o, result_t};
 }
 
 } // namespace storage

src/v/storage/log_reader.h

@@ -297,8 +297,8 @@ class log_reader final : public model::record_batch_reader::impl {
  * \param t The timestamp to search for
  * \param max_offset The maximum offset to consider
  */
-timequery_result batch_timequery(
-  const model::record_batch& b,
+ss::future<timequery_result> batch_timequery(
+  model::record_batch b,
   model::offset min_offset,
   model::timestamp t,
   model::offset max_offset);