Merge pull request #25498 from bharathv/v251x-oom-backport

[v25.1.x] [backport] dl/translation: fixes for OOM handling #25423

Author: bharathv

src/v/datalake/data_writer_interface.cc

@@ -15,25 +15,42 @@ namespace datalake {
 std::ostream& operator<<(std::ostream& os, const writer_error& ev) {
     switch (ev) {
     case writer_error::ok:
-        os << "Ok";
-        break;
+        return os << "Ok";
     case writer_error::parquet_conversion_error:
-        os << "Parquet Conversion Error";
-        break;
+        return os << "Parquet Conversion Error";
     case writer_error::file_io_error:
-        os << "File IO Error";
-        break;
+        return os << "File IO Error";
     case writer_error::no_data:
-        os << "No data";
-        break;
+        return os << "No data";
     case writer_error::flush_error:
-        os << "Flush failed";
-        break;
+        return os << "Flush failed";
+    case writer_error::oom_error:
+        return os << "Memory exhausted";
+    case writer_error::time_limit_exceeded:
+        return os << "Time limit exceeded";
+    case writer_error::shutting_down:
+        return os << "Shutting down";
+    case writer_error::unknown_error:
+        return os << "Unknown error";
     }
-    return os;
 }
 std::string data_writer_error_category::message(int ev) const {
     return fmt::to_string(static_cast<writer_error>(ev));
 }
 
+writer_error map_to_writer_error(reservation_error reservation_err) {
+    switch (reservation_err) {
+    case ok:
+        return writer_error::ok;
+    case shutting_down:
+        return writer_error::shutting_down;
+    case out_of_memory:
+        return writer_error::oom_error;
+    case time_quota_exceeded:
+        return writer_error::time_limit_exceeded;
+    case unknown:
+        return writer_error::unknown_error;
+    }
+}
+
 } // namespace datalake

src/v/datalake/data_writer_interface.h

@@ -26,6 +26,10 @@ enum class writer_error {
     file_io_error,
     no_data,
     flush_error,
+    oom_error,
+    time_limit_exceeded,
+    shutting_down,
+    unknown_error,
 };
 std::ostream& operator<<(std::ostream&, const writer_error&);
 
@@ -44,8 +48,17 @@ inline std::error_code make_error_code(writer_error e) noexcept {
     return {static_cast<int>(e), data_writer_error_category::error_category()};
 }
 
+enum reservation_error {
+    ok = 0,
+    shutting_down = 1,
+    out_of_memory = 2,
+    time_quota_exceeded = 3,
+    unknown = 4,
+};
+
+writer_error map_to_writer_error(reservation_error);
 /**
- * Interface to track memory used by the parquet writer. The reservations are
+ * Interface to track memory used by the parquet writers. The reservations are
  * held until the tracker object is alive or release is explicitly called.
  */
 class writer_mem_tracker {
@@ -59,14 +72,15 @@ class writer_mem_tracker {
     virtual ~writer_mem_tracker() = default;
 
     /**
-     * Notify the mem tracker of current memory usage. The writer may
-     * choose to compress/shrink memory upon which the tracker must be
-     * notified of the current usage. May not be called concurrently with
-     * other methods.
+     * Reserves passed input bytes.
      */
-    virtual ss::future<>
-    update_current_memory_usage(size_t current_bytes_usage, ss::abort_source&)
-      = 0;
+    virtual ss::future<reservation_error>
+    reserve_bytes(size_t bytes, ss::abort_source&) noexcept = 0;
+
+    /**
+     * Frees up passed input bytes.
+     */
+    virtual ss::future<> free_bytes(size_t bytes, ss::abort_source&) = 0;
 
     /**
      * Releases all the reservations. After this caller, the reserved bytes
@@ -173,7 +187,8 @@ class parquet_file_writer_factory {
 
     virtual ss::future<
       result<std::unique_ptr<parquet_file_writer>, writer_error>>
-    create_writer(const iceberg::struct_type& /* schema */) = 0;
+    create_writer(const iceberg::struct_type& /* schema */, ss::abort_source&)
+      = 0;
 };
 
 } // namespace datalake

src/v/datalake/local_parquet_file_writer.cc

@@ -10,6 +10,7 @@
 
 #include "datalake/local_parquet_file_writer.h"
 
+#include "base/units.h"
 #include "base/vlog.h"
 #include "datalake/logger.h"
 
@@ -167,17 +168,40 @@ local_parquet_file_writer_factory::local_parquet_file_writer_factory(
   local_path base_directory,
   ss::sstring file_name_prefix,
   ss::shared_ptr<parquet_ostream_factory> writer_factory,
-  std::unique_ptr<writer_mem_tracker> mem_tracker)
+  writer_mem_tracker& mem_tracker)
   : _base_directory(std::move(base_directory))
   , _file_name_prefix(std::move(file_name_prefix))
   , _writer_factory(std::move(writer_factory))
-  , _mem_tracker(std::move(mem_tracker)) {}
+  , _mem_tracker(mem_tracker) {}
 
 ss::future<result<std::unique_ptr<parquet_file_writer>, writer_error>>
 local_parquet_file_writer_factory::create_writer(
-  const iceberg::struct_type& schema) {
+  const iceberg::struct_type& schema, ss::abort_source& as) {
+    // There is a per writer cost associated which includes stuff like
+    // - local path string
+    // - associated partition key
+    // - schema
+    // - stats tracked about the writer
+    // - data structure overhead
+    //
+    // This limit is in place to avoid an explosion of writer instances,
+    // example partition_by(offset) which creates a writer per offset.
+    //
+    // Additionally one other contributor per writer is the buffer used
+    // in the output stream which defaults to 8_KiB, which is only released
+    // on output stream close().
+    //
+    // TODO: This is just a conservative estimate to prevent pathological cases
+    // of too many writers, needs empirical evaluation to determine the correct
+    // sizing.
+    static constexpr size_t WRITER_RESERVATION_OVERHEAD = 10_KiB;
+    auto reservation_err = co_await _mem_tracker.reserve_bytes(
+      WRITER_RESERVATION_OVERHEAD, as);
+    if (reservation_err != reservation_error::ok) {
+        co_return map_to_writer_error(reservation_err);
+    }
     auto writer = std::make_unique<local_parquet_file_writer>(
-      create_filename(), _writer_factory, *_mem_tracker);
+      create_filename(), _writer_factory, _mem_tracker);
 
     auto res = co_await writer->initialize(schema);
     if (res.has_error()) {

src/v/datalake/local_parquet_file_writer.h

@@ -65,18 +65,18 @@ class local_parquet_file_writer_factory : public parquet_file_writer_factory {
       local_path base_directory,
       ss::sstring file_name_prefix,
       ss::shared_ptr<parquet_ostream_factory>,
-      std::unique_ptr<writer_mem_tracker>);
+      writer_mem_tracker&);
 
     ss::future<result<std::unique_ptr<parquet_file_writer>, writer_error>>
-    create_writer(const iceberg::struct_type& schema) final;
+    create_writer(const iceberg::struct_type& schema, ss::abort_source&) final;
 
 private:
     local_path create_filename() const;
 
     local_path _base_directory;
     ss::sstring _file_name_prefix;
     ss::shared_ptr<parquet_ostream_factory> _writer_factory;
-    std::unique_ptr<writer_mem_tracker> _mem_tracker;
+    writer_mem_tracker& _mem_tracker;
 };
 
 } // namespace datalake

src/v/datalake/partitioning_writer.cc

@@ -57,7 +57,7 @@ ss::future<writer_error> partitioning_writer::add_data(
     }
     auto writer_iter = writers_.find(pk);
     if (writer_iter == writers_.end()) {
-        auto writer_res = co_await writer_factory_.create_writer(type_);
+        auto writer_res = co_await writer_factory_.create_writer(type_, as);
         if (writer_res.has_error()) {
             vlog(
               datalake_log.error,

src/v/datalake/record_multiplexer.cc

@@ -25,6 +25,29 @@
 
 #include <seastar/core/loop.hh>
 
+namespace {
+
+// Recoverable errors are the class of errors that donot leave the underlying
+// writers in a bad shape. Upon recoverable errors the translator may choose to
+// flush and continue as if nothing happened, so we preserve the state to
+// facilitate that.
+bool is_recoverable_error(datalake::writer_error err) {
+    switch (err) {
+    case datalake::writer_error::ok:
+    case datalake::writer_error::oom_error:
+    case datalake::writer_error::time_limit_exceeded:
+        return true;
+    case datalake::writer_error::parquet_conversion_error:
+    case datalake::writer_error::file_io_error:
+    case datalake::writer_error::no_data:
+    case datalake::writer_error::flush_error:
+    case datalake::writer_error::shutting_down:
+    case datalake::writer_error::unknown_error:
+        return false;
+    }
+}
+}; // namespace
+
 namespace datalake {
 
 namespace {
@@ -72,17 +95,19 @@ record_multiplexer::record_multiplexer(
 
 ss::future<> record_multiplexer::multiplex(
   model::record_batch_reader reader,
+  kafka::offset start_offset,
   model::timeout_clock::time_point deadline,
   ss::abort_source& as) {
     co_await std::move(reader).consume(
-      relaying_consumer{[this, &as](model::record_batch b) mutable {
-          return do_multiplex(std::move(b), as);
-      }},
+      relaying_consumer{
+        [this, start_offset, &as](model::record_batch b) mutable {
+            return do_multiplex(std::move(b), start_offset, as);
+        }},
       deadline);
 }
 
 ss::future<ss::stop_iteration> record_multiplexer::do_multiplex(
-  model::record_batch batch, ss::abort_source& as) {
+  model::record_batch batch, kafka::offset start_offset, ss::abort_source& as) {
     if (batch.compressed()) {
         batch = co_await storage::internal::decompress_batch(std::move(batch));
     }
@@ -99,6 +124,9 @@ ss::future<ss::stop_iteration> record_multiplexer::do_multiplex(
         auto timestamp = model::timestamp{
           first_timestamp + record.timestamp_delta()};
         kafka::offset offset{batch.base_offset()() + record.offset_delta()};
+        if (offset < start_offset) {
+            continue;
+        }
         int64_t estimated_size = (key ? key->size_bytes() : 0)
                                  + (val ? val->size_bytes() : 0);
         chunked_vector<std::pair<std::optional<iobuf>, std::optional<iobuf>>>
@@ -259,37 +287,36 @@ ss::future<ss::stop_iteration> record_multiplexer::do_multiplex(
             writer_iter = iter;
         }
 
-        // TODO: we want to ensure we're using an offset translating reader so
-        // that these will be Kafka offsets, not Raft offsets.
-        if (!_result.has_value()) {
-            _result = write_result{
-              .start_offset = offset,
-            };
-        }
-
-        _result.value().last_offset = offset;
-
         auto& writer = writer_iter->second;
-        auto write_result = co_await writer->add_data(
+        auto add_data_result = co_await writer->add_data(
           std::move(record_data_res.value()), estimated_size, as);
 
-        if (write_result != writer_error::ok) {
+        if (add_data_result != writer_error::ok) {
             vlog(
               _log.warn,
               "Error adding data to writer for record {}: {}",
               offset,
-              write_result);
-            _error = write_result;
+              add_data_result);
+            _error = add_data_result;
             // If a write fails, the writer is left in an indeterminate state,
             // we cannot continue in this case.
             co_return ss::stop_iteration::yes;
         }
+
+        // TODO: we want to ensure we're using an offset translating reader so
+        // that these will be Kafka offsets, not Raft offsets.
+        if (!_result.has_value()) {
+            _result = write_result{
+              .start_offset = offset,
+            };
+        }
+        _result.value().last_offset = offset;
     }
     co_return ss::stop_iteration::no;
 }
 
 ss::future<writer_error> record_multiplexer::flush_writers() {
-    if (_error) {
+    if (_error && !is_recoverable_error(_error.value())) {
         co_return *_error;
     }
     auto result = co_await ss::coroutine::as_future(ss::max_concurrent_for_each(
@@ -304,37 +331,41 @@ ss::future<writer_error> record_multiplexer::flush_writers() {
 
 ss::future<result<record_multiplexer::write_result, writer_error>>
 record_multiplexer::finish() && {
-    if (!_result) {
-        // no batches were processed.
-        co_return writer_error::no_data;
-    }
     auto writers = std::move(_writers);
     for (auto& [id, writer] : writers) {
         auto res = co_await std::move(*writer).finish();
         if (res.has_error()) {
             _error = res.error();
             continue;
         }
-        auto& files = res.value();
-        std::move(
-          files.begin(), files.end(), std::back_inserter(_result->data_files));
+        if (_result) {
+            auto& files = res.value();
+            std::move(
+              files.begin(),
+              files.end(),
+              std::back_inserter(_result->data_files));
+        }
     }
     if (_invalid_record_writer) {
         auto writer = std::move(_invalid_record_writer);
         auto res = co_await std::move(*writer).finish();
         if (res.has_error()) {
             _error = res.error();
-        } else {
+        } else if (_result) {
             auto& files = res.value();
             std::move(
               files.begin(),
               files.end(),
               std::back_inserter(_result->dlq_files));
         }
     }
-    if (_error) {
+    if (_error && !is_recoverable_error(_error.value())) {
         co_return *_error;
     }
+    if (!_result) {
+        // no batches were processed.
+        co_return writer_error::no_data;
+    }
     co_return std::move(*_result);
 }
 

src/v/datalake/record_multiplexer.h

@@ -77,17 +77,23 @@ class record_multiplexer {
     /**
      * Multiplex the data from a reader into writers per schema and partition.
      * Can be called multiple times in succession before calling finish().
+     *
+     * start_offset controls minimum offset from which multiplexer can work. If
+     * the previous translation stopped in the middle of a batch, we do not want
+     * to multiplex already translated offsets in the batch, start_offset helps
+     * solve that problem.
      */
     ss::future<> multiplex(
       model::record_batch_reader reader,
+      kafka::offset start_offset,
       model::timeout_clock::time_point deadline,
       ss::abort_source& as);
 
     /**
      * Abortable multiplexing on a single batch. Visible for testing.
      */
-    ss::future<ss::stop_iteration>
-    do_multiplex(model::record_batch batch, ss::abort_source&);
+    ss::future<ss::stop_iteration> do_multiplex(
+      model::record_batch batch, kafka::offset start_offset, ss::abort_source&);
 
     /**
      * Forces a flush on all the underlying file writers resulting in freeing