dl/translation: propagate abort errors all the way

Author: bharathv

src/v/datalake/data_writer_interface.cc

@@ -26,10 +26,29 @@ std::ostream& operator<<(std::ostream& os, const writer_error& ev) {
         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";
     }
 }
 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::oom_error;
+    case unknown:
+        return writer_error::shutting_down;
+    }
+}
+
 } // namespace datalake

src/v/datalake/data_writer_interface.h

@@ -27,6 +27,8 @@ enum class writer_error {
     no_data,
     flush_error,
     oom_error,
+    time_limit_exceeded,
+    shutting_down,
 };
 std::ostream& operator<<(std::ostream&, const writer_error&);
 
@@ -45,6 +47,15 @@ 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 writers. The reservations are
  * held until the tracker object is alive or release is explicitly called.
@@ -62,7 +73,8 @@ class writer_mem_tracker {
     /**
      * Reserves passed input bytes.
      */
-    virtual ss::future<> reserve_bytes(size_t bytes, ss::abort_source&) = 0;
+    virtual ss::future<reservation_error>
+    reserve_bytes(size_t bytes, ss::abort_source&) noexcept = 0;
 
     /**
      * Frees up passed input bytes.

src/v/datalake/serde_parquet_writer.cc

@@ -22,8 +22,11 @@ ss::future<writer_error> serde_parquet_writer::add_data_struct(
         auto stats = co_await _writer.write_row(std::move(group));
         auto new_buffered_bytes = stats.buffered_size;
         if (new_buffered_bytes > _buffered_bytes) {
-            co_await _mem_tracker.reserve_bytes(
+            auto reservation_result = co_await _mem_tracker.reserve_bytes(
               new_buffered_bytes - _buffered_bytes, as);
+            if (reservation_result != reservation_error::ok) {
+                co_return map_to_writer_error(reservation_result);
+            }
         } else if (new_buffered_bytes < _buffered_bytes) {
             co_await _mem_tracker.free_bytes(
               _buffered_bytes - new_buffered_bytes, as);

src/v/datalake/tests/test_data_writer.h

@@ -23,8 +23,9 @@
 namespace datalake {
 class noop_mem_tracker : public writer_mem_tracker {
 public:
-    ss::future<> reserve_bytes(size_t, ss::abort_source&) override {
-        return ss::make_ready_future<>();
+    ss::future<reservation_error>
+    reserve_bytes(size_t, ss::abort_source&) noexcept override {
+        return ss::make_ready_future<reservation_error>(reservation_error::ok);
     }
     ss::future<> free_bytes(size_t, ss::abort_source&) override {
         return ss::make_ready_future<>();

src/v/datalake/translation/deps.cc

@@ -38,6 +38,10 @@ map_error_code(datalake::translation_task::errc errc) {
         return translation_errc::no_data;
     case datalake::translation_task::errc::oom_error:
         return translation_errc::oom_error;
+    case datalake::translation_task::errc::time_limit_exceeded:
+        return translation_errc::time_limit_exceeded;
+    case datalake::translation_task::errc::shutting_down:
+        return translation_errc::shutting_down;
     }
 }
 } // namespace
@@ -61,26 +65,39 @@ ss::future<cluster::errc> wait_stm_translated(
 }
 } // namespace
 
-ss::future<> noop_mem_tracker::reserve_bytes(size_t, ss::abort_source&) {
-    return ss::make_ready_future<>();
+ss::future<reservation_error>
+noop_mem_tracker::reserve_bytes(size_t, ss::abort_source&) noexcept {
+    return ss::make_ready_future<reservation_error>(reservation_error::ok);
 }
 ss::future<> noop_mem_tracker::free_bytes(size_t, ss::abort_source&) {
     return ss::make_ready_future<>();
 }
 void noop_mem_tracker::release() {}
 
-ss::future<>
-translator_mem_tracker::reserve_bytes(size_t bytes, ss::abort_source& as) {
+ss::future<reservation_error> translator_mem_tracker::reserve_bytes(
+  size_t bytes, ss::abort_source& as) noexcept {
     _current_usage += bytes;
-    while (_current_usage > _reservations.count()) {
-        // reserve any deficit
-        auto reservation = co_await _reservations_tracker.reserve_memory(as);
-        if (_reservations.count()) {
-            _reservations.adopt(std::move(reservation));
-        } else {
-            _reservations = std::move(reservation);
+    try {
+        while (_current_usage > _reservations.count()) {
+            // reserve any deficit
+            auto reservation = co_await _reservations_tracker.reserve_memory(
+              as);
+            if (_reservations.count()) {
+                _reservations.adopt(std::move(reservation));
+            } else {
+                _reservations = std::move(reservation);
+            }
         }
-    }
+    } catch (const translator_out_of_memory_error&) {
+        co_return reservation_error::out_of_memory;
+    } catch (const translator_shutdown_error&) {
+        co_return reservation_error::shutting_down;
+    } catch (const translator_time_quota_exceeded_error&) {
+        co_return reservation_error::time_quota_exceeded;
+    } catch (...) {
+        co_return reservation_error::unknown;
+    }
+    co_return reservation_error::ok;
 }
 
 ss::future<>
@@ -377,6 +394,10 @@ std::ostream& operator<<(std::ostream& o, translation_errc ec) {
         return o << "translation_errc::discard_error";
     case oom_error:
         return o << "translation_errc::oom_error";
+    case time_limit_exceeded:
+        return o << "translation_errc::time_limit_exceeded";
+    case shutting_down:
+        return o << "translation_errc::shutting_down";
     }
 }
 

src/v/datalake/translation/deps.h

@@ -27,9 +27,28 @@
 
 namespace datalake::translation {
 
+class translator_out_of_memory_error final : public std::runtime_error {
+public:
+    explicit translator_out_of_memory_error()
+      : std::runtime_error("translator_out_of_memory") {}
+};
+
+class translator_shutdown_error final : public std::runtime_error {
+public:
+    explicit translator_shutdown_error()
+      : std::runtime_error("translator_shutdown") {}
+};
+
+class translator_time_quota_exceeded_error final : public std::runtime_error {
+public:
+    explicit translator_time_quota_exceeded_error()
+      : std::runtime_error("translator_time_quota_exceeded") {}
+};
+
 class noop_mem_tracker : public writer_mem_tracker {
 public:
-    ss::future<> reserve_bytes(size_t, ss::abort_source&) override;
+    ss::future<reservation_error>
+    reserve_bytes(size_t, ss::abort_source&) noexcept override;
     ss::future<> free_bytes(size_t, ss::abort_source&) override;
     void release() override;
 };
@@ -43,7 +62,8 @@ class translator_mem_tracker : public writer_mem_tracker {
       scheduling::reservations_tracker& scheduling_reservations)
       : _reservations_tracker(scheduling_reservations) {}
 
-    ss::future<> reserve_bytes(size_t, ss::abort_source&) override;
+    ss::future<reservation_error>
+    reserve_bytes(size_t, ss::abort_source&) noexcept override;
     ss::future<> free_bytes(size_t, ss::abort_source&) override;
     void release() override;
 
@@ -161,6 +181,8 @@ enum translation_errc {
     flush_error,
     discard_error,
     oom_error,
+    time_limit_exceeded,
+    shutting_down,
 };
 
 std::ostream& operator<<(std::ostream&, translation_errc);

src/v/datalake/translation/partition_translator.cc

@@ -150,7 +150,8 @@ partition_translator::translate_when_notified(kafka::offset begin_offset) {
     }
     vlog(_logger.trace, "starting translation from offset: {}", begin_offset);
     ss::timer<scheduling::clock> cancellation_timer;
-    cancellation_timer.set_callback([&as] { as.request_abort(); });
+    cancellation_timer.set_callback(
+      [&as] { as.request_abort_ex(translator_time_quota_exceeded_error{}); });
 
     auto translation_f
       = _translation_ctx
@@ -472,7 +473,8 @@ ss::future<> partition_translator::close() noexcept {
     _as.request_abort();
     _ready_to_translate.broken();
     if (_inflight_translation_state) {
-        _inflight_translation_state->as.request_abort();
+        _inflight_translation_state->as.request_abort_ex(
+          translator_shutdown_error{});
     }
     _data_source->close();
     co_await _gate.close();
@@ -511,6 +513,9 @@ void partition_translator::stop_translation() {
     if (_gate.is_closed() || !_inflight_translation_state) {
         return;
     }
-    _inflight_translation_state->as.request_abort();
+    // Currently only preempted on OOM error, if the policy changes
+    // to preempt on other errors, should be updated accordingly.
+    _inflight_translation_state->as.request_abort_ex(
+      translator_out_of_memory_error{});
 }
 } // namespace datalake::translation

src/v/datalake/translation_task.cc

@@ -44,6 +44,10 @@ translation_task::errc map_error_code(writer_error errc) {
         return translation_task::errc::flush_error;
     case writer_error::oom_error:
         return translation_task::errc::oom_error;
+    case writer_error::time_limit_exceeded:
+        return translation_task::errc::time_limit_exceeded;
+    case writer_error::shutting_down:
+        return translation_task::errc::shutting_down;
     }
 }
 
@@ -359,6 +363,10 @@ std::ostream& operator<<(std::ostream& o, translation_task::errc ec) {
         return o << "no data to translate";
     case translation_task::errc::oom_error:
         return o << "memory exhausted";
+    case translation_task::errc::time_limit_exceeded:
+        return o << "time limit exceeded";
+    case translation_task::errc::shutting_down:
+        return o << "shutting down";
     }
 }
 } // namespace datalake

src/v/datalake/translation_task.h

@@ -46,6 +46,8 @@ class translation_task {
         flush_error,
         no_data,
         oom_error,
+        time_limit_exceeded,
+        shutting_down,
     };
 
     using custom_partitioning_enabled