dl/translation: refactor translation mem tracking

Two main changes

- Centralizes mem tracking in translation_ctx. This memory tracker is
  shared by all parquet writers from the same translator to reserve
  memory. Having access to the memory tracker within translation_ctx
  makes the call to buffered_bytes() inexpensive (doesn't need to loop
  through every writer to compute buffered bytes as they are rolled up
  into the mem tracker)

- Changes API to do reserve(bytes, as) and free(bytes, as). The existing
  API turned out to be buggy in the presence of multiple writers as we
  overwrite current usage from the latest writer. This change is also a
  preparatory work for the next commit that wires up a fixed cost per
  writer.

Author: bharathv

src/v/datalake/data_writer_interface.h

@@ -46,7 +46,7 @@ inline std::error_code make_error_code(writer_error e) noexcept {
 }
 
 /**
- * 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 {
@@ -60,14 +60,14 @@ 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<> reserve_bytes(size_t bytes, ss::abort_source&) = 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

src/v/datalake/local_parquet_file_writer.cc

@@ -167,17 +167,17 @@ 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, ss::abort_source&) {
     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,7 +65,7 @@ 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, ss::abort_source&) final;
@@ -76,7 +76,7 @@ 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> _writer_factory;
-    std::unique_ptr<writer_mem_tracker> _mem_tracker;
+    writer_mem_tracker& _mem_tracker;
 };
 
 } // namespace datalake

src/v/datalake/serde_parquet_writer.cc

@@ -20,9 +20,16 @@ ss::future<writer_error> serde_parquet_writer::add_data_struct(
       std::move(conversion_result.value()));
     try {
         auto stats = co_await _writer.write_row(std::move(group));
-        _buffered_bytes = stats.buffered_size;
+        auto new_buffered_bytes = stats.buffered_size;
+        if (new_buffered_bytes > _buffered_bytes) {
+            co_await _mem_tracker.reserve_bytes(
+              new_buffered_bytes - _buffered_bytes, as);
+        } else if (new_buffered_bytes < _buffered_bytes) {
+            co_await _mem_tracker.free_bytes(
+              _buffered_bytes - new_buffered_bytes, as);
+        }
+        _buffered_bytes = new_buffered_bytes;
         _flushed_bytes = stats.flushed_size;
-        co_await _mem_tracker.update_current_memory_usage(_buffered_bytes, as);
     } catch (...) {
         vlog(
           datalake_log.warn,
@@ -45,12 +52,10 @@ ss::future<> serde_parquet_writer::flush() {
       _buffered_bytes == 0,
       "Memory buffered in the writer after flush: {}",
       _buffered_bytes);
-    _mem_tracker.release();
 }
 
 ss::future<writer_error> serde_parquet_writer::finish() {
     co_await _writer.close();
-    _mem_tracker.release();
     _buffered_bytes = _flushed_bytes = 0;
     co_return writer_error::ok;
 }

src/v/datalake/serde_parquet_writer.h

@@ -25,8 +25,8 @@ class serde_parquet_writer : public parquet_ostream {
 private:
     serde::parquet::writer _writer;
     writer_mem_tracker& _mem_tracker;
-    size_t _buffered_bytes{0};
-    size_t _flushed_bytes{0};
+    int64_t _buffered_bytes{0};
+    int64_t _flushed_bytes{0};
 };
 
 class serde_parquet_writer_factory : public parquet_ostream_factory {

src/v/datalake/tests/test_data_writer.h

@@ -23,8 +23,10 @@
 namespace datalake {
 class noop_mem_tracker : public writer_mem_tracker {
 public:
-    ss::future<>
-    update_current_memory_usage(size_t, ss::abort_source&) override {
+    ss::future<> reserve_bytes(size_t, ss::abort_source&) override {
+        return ss::make_ready_future<>();
+    }
+    ss::future<> free_bytes(size_t, ss::abort_source&) override {
         return ss::make_ready_future<>();
     }
     void release() override {}

src/v/datalake/translation/deps.cc

@@ -61,18 +61,18 @@ ss::future<cluster::errc> wait_stm_translated(
 }
 } // namespace
 
-ss::future<>
-noop_mem_tracker::update_current_memory_usage(size_t, ss::abort_source&) {
+ss::future<> noop_mem_tracker::reserve_bytes(size_t, ss::abort_source&) {
+    return ss::make_ready_future<>();
+}
+ss::future<> noop_mem_tracker::free_bytes(size_t, ss::abort_source&) {
     return ss::make_ready_future<>();
 }
 void noop_mem_tracker::release() {}
 
-ss::future<> writer_reservations_impl::update_current_memory_usage(
-  size_t new_used_bytes, ss::abort_source& as) {
-    _current_used_bytes = new_used_bytes;
-    // todo: as a potential optimization we can choose to release bytes if the
-    // new usage is less than current
-    while (_current_used_bytes > _reservations.count()) {
+ss::future<>
+translator_mem_tracker::reserve_bytes(size_t bytes, ss::abort_source& as) {
+    _current_usage += bytes;
+    while (_current_usage > _reservations.count()) {
         // reserve any deficit
         auto reservation = co_await _reservations_tracker.reserve_memory(as);
         if (_reservations.count()) {
@@ -83,16 +83,20 @@ ss::future<> writer_reservations_impl::update_current_memory_usage(
     }
 }
 
-void writer_reservations_impl::release() {
-    _current_used_bytes = 0;
-    _reservations.return_all();
+ss::future<>
+translator_mem_tracker::free_bytes(size_t bytes, ss::abort_source&) {
+    _current_usage -= std::min(_current_usage, bytes);
+    return ss::now();
 }
 
-size_t writer_reservations_impl::current_usage() const {
-    return _current_used_bytes;
+void translator_mem_tracker::release() {
+    _current_usage = 0;
+    _reservations.return_all();
 }
 
-size_t writer_reservations_impl::total_reserved() const {
+size_t translator_mem_tracker::current_usage() const { return _current_usage; }
+
+size_t translator_mem_tracker::total_reserved() const {
     return _reservations.count();
 }
 
@@ -409,7 +413,8 @@ class partition_translation_context : public translation_context {
       , _probe(std::move(probe))
       , _invalid_record_action(compute_invalid_record_action())
       , _cp_enabled(translation_task::custom_partitioning_enabled{
-          _features.is_active(features::feature::datalake_iceberg_ga)}) {}
+          _features.is_active(features::feature::datalake_iceberg_ga)})
+      , _mem_tracker(translator_mem_tracker{_reservations}) {}
 
     ss::future<> translate_now(
       model::record_batch_reader reader, ss::abort_source& as) final {
@@ -480,8 +485,9 @@ class partition_translation_context : public translation_context {
 
     ss::future<> flush() final {
         if (_in_progress_translation) {
-            return _in_progress_translation->flush().then_wrapped(
-              [](auto result_f) {
+            vlog(datalake_log.trace, "[{}] flushing writers", _ntp);
+            return _in_progress_translation->flush()
+              .then_wrapped([](auto result_f) {
                   if (result_f.failed()) {
                       return ss::make_exception_future(
                         result_f.get_exception());
@@ -491,13 +497,15 @@ class partition_translation_context : public translation_context {
                       return ss::make_exception_future(result.error());
                   }
                   return ss::now();
-              });
+              })
+              .finally([this]() { _mem_tracker.release(); });
         }
         return ss::now();
     }
 
     ss::future<checked<coordinator::translated_offset_range, translation_errc>>
     finish(retry_chain_node& rcn, ss::abort_source& as) final {
+        _mem_tracker.release();
         if (!_in_progress_translation) {
             co_return translation_errc::no_data;
         }
@@ -522,11 +530,7 @@ class partition_translation_context : public translation_context {
         co_await std::move(task.value()).discard().discard_result();
     }
 
-    size_t buffered_bytes() const final {
-        return _in_progress_translation
-                 ? _in_progress_translation->buffered_bytes()
-                 : 0;
-    }
+    size_t buffered_bytes() const final { return _mem_tracker.current_usage(); }
 
 private:
     scheduling::clock::duration compute_target_lag() const {
@@ -546,7 +550,7 @@ class partition_translation_context : public translation_context {
           _writer_scratch_space, // storage temp files are written to
           "",                    // file prefix
           ss::make_shared<serde_parquet_writer_factory>(),
-          std::make_unique<writer_reservations_impl>(_reservations));
+          _mem_tracker);
     }
 
     model::iceberg_invalid_record_action compute_invalid_record_action() const {
@@ -577,7 +581,7 @@ class partition_translation_context : public translation_context {
     ss::lw_shared_ptr<translation_probe> _probe;
     model::iceberg_invalid_record_action _invalid_record_action;
     translation_task::custom_partitioning_enabled _cp_enabled;
-
+    translator_mem_tracker _mem_tracker;
     std::optional<translation_task> _in_progress_translation;
     bool _discard_translated_state{false};
 };

src/v/datalake/translation/deps.h

@@ -29,26 +29,29 @@ namespace datalake::translation {
 
 class noop_mem_tracker : public writer_mem_tracker {
 public:
-    ss::future<>
-    update_current_memory_usage(size_t, ss::abort_source&) override;
+    ss::future<> reserve_bytes(size_t, ss::abort_source&) override;
+    ss::future<> free_bytes(size_t, ss::abort_source&) override;
     void release() override;
 };
 
-class writer_reservations_impl : public writer_mem_tracker {
+/**
+ * Tracks memory usage across all writers in a single translator.
+ */
+class translator_mem_tracker : public writer_mem_tracker {
 public:
-    explicit writer_reservations_impl(
+    explicit translator_mem_tracker(
       scheduling::reservations_tracker& scheduling_reservations)
       : _reservations_tracker(scheduling_reservations) {}
 
-    ss::future<>
-    update_current_memory_usage(size_t, ss::abort_source&) override;
+    ss::future<> reserve_bytes(size_t, ss::abort_source&) override;
+    ss::future<> free_bytes(size_t, ss::abort_source&) override;
     void release() override;
 
     size_t current_usage() const;
     size_t total_reserved() const;
 
 private:
-    size_t _current_used_bytes{0};
+    size_t _current_usage{0};
     scheduling::reservations_tracker& _reservations_tracker;
     ssx::semaphore_units _reservations;
 };

src/v/datalake/translation/tests/partition_translator_tests.cc

@@ -517,51 +517,53 @@ using scheduling::scheduler_fixture;
 TEST_F_CORO(scheduler_fixture, test_writer_reservations_accounting) {
     auto& reservations = *_scheduler->reservations();
     {
-        writer_reservations_impl writer{reservations};
+        translator_mem_tracker writer{reservations};
 
         auto cleanup = ss::defer([&writer] { writer.release(); });
 
         ss::abort_source as;
 
-        ASSERT_EQ_CORO(writer.current_usage(), 0);
+        size_t total_usage = 0;
+        ASSERT_EQ_CORO(writer.current_usage(), total_usage);
         ASSERT_EQ_CORO(writer.total_reserved(), 0);
 
-        auto current_usage = 1_MiB;
+        auto bytes = 1_MiB;
         // update initial usage, should result in a reservation
-        co_await writer.update_current_memory_usage(current_usage, as);
+        co_await writer.reserve_bytes(bytes, as);
+        total_usage += bytes;
 
-        ASSERT_EQ_CORO(writer.current_usage(), current_usage);
+        ASSERT_EQ_CORO(writer.current_usage(), total_usage);
         ASSERT_EQ_CORO(writer.total_reserved(), block_size);
 
         // try 3 more times, we are still within the block limit;
-        while (current_usage <= block_size) {
-            co_await writer.update_current_memory_usage(current_usage, as);
-            current_usage += 1_MiB;
+        while (total_usage < block_size) {
+            co_await writer.reserve_bytes(bytes, as);
+            total_usage += bytes;
         }
 
         ASSERT_EQ_CORO(writer.current_usage(), block_size);
         ASSERT_EQ_CORO(writer.total_reserved(), block_size);
 
-        current_usage += 1_MiB;
         // update again, should reserve a new block.
-        co_await writer.update_current_memory_usage(current_usage, as);
+        co_await writer.reserve_bytes(bytes, as);
+        total_usage += bytes;
 
-        ASSERT_EQ_CORO(writer.current_usage(), current_usage);
+        ASSERT_EQ_CORO(writer.current_usage(), total_usage);
         ASSERT_EQ_CORO(writer.total_reserved(), 2 * block_size);
 
         // exhaust all memory
         while (writer.current_usage() != total_memory) {
-            current_usage += 1_MiB;
-            co_await writer.update_current_memory_usage(current_usage, as);
+            co_await writer.reserve_bytes(bytes, as);
+            total_usage += 1_MiB;
         }
 
         ss::promise<> done;
         // update again, should block on reservations due to memory limit
         // exhaustion.
         auto f = ss::with_timeout(
           ss::steady_clock_type::now() + 500ms,
-          writer.update_current_memory_usage(current_usage + 1_MiB, as)
-            .finally([&done] { done.set_value(); }));
+          writer.reserve_bytes(bytes, as).finally(
+            [&done] { done.set_value(); }));
 
         ASSERT_THROW_CORO(co_await std::move(f), ss::timed_out_error);