dl/translator: graceful handling of OOM/scheduling timeouts

Author: bharathv

src/v/datalake/data_writer_interface.cc

@@ -45,7 +45,7 @@ writer_error map_to_writer_error(reservation_error reservation_err) {
     case out_of_memory:
         return writer_error::oom_error;
     case time_quota_exceeded:
-        return writer_error::oom_error;
+        return writer_error::time_limit_exceeded;
     case unknown:
         return writer_error::shutting_down;
     }

src/v/datalake/local_parquet_file_writer.cc

@@ -186,7 +186,11 @@ local_parquet_file_writer_factory::create_writer(
     // in the output stream which defaults to 8_KiB, which is only released
     // on output stream close().
     static constexpr size_t WRITER_RESERVATION_OVERHEAD = 10_KiB;
-    co_await _mem_tracker.reserve_bytes(WRITER_RESERVATION_OVERHEAD, as);
+    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);
 

src/v/datalake/record_multiplexer.cc

@@ -25,6 +25,28 @@
 
 #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:
+        return false;
+    }
+}
+}; // namespace
+
 namespace datalake {
 
 namespace {
@@ -289,7 +311,7 @@ ss::future<ss::stop_iteration> record_multiplexer::do_multiplex(
 }
 
 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(
@@ -332,7 +354,7 @@ record_multiplexer::finish() && {
               std::back_inserter(_result->dlq_files));
         }
     }
-    if (_error) {
+    if (_error && !is_recoverable_error(_error.value())) {
         co_return *_error;
     }
     co_return std::move(*_result);

src/v/datalake/translation/deps.cc

@@ -540,13 +540,15 @@ class partition_translation_context : public translation_context {
         if (!_in_progress_translation) {
             co_return translation_errc::no_data;
         }
+        vlog(datalake_log.debug, "[{}] finishing translation", _ntp);
         if (_discard_translated_state) {
             co_await discard().discard_result();
             co_return translation_errc::discard_error;
         }
         auto task = std::exchange(_in_progress_translation, std::nullopt);
         auto result = co_await std::move(task.value())
                         .finish(_cp_enabled, _upload_path_prefix, rcn, as);
+
         if (result.has_error()) {
             co_return map_error_code(result.error());
         }

src/v/datalake/translation/partition_translator.cc

@@ -137,33 +137,6 @@ partition_translator::checkpoint_translation_result(
       });
 }
 
-ss::future<>
-partition_translator::translate_when_notified(kafka::offset begin_offset) {
-    co_await _ready_to_translate.wait(
-      [this] { return _inflight_translation_state.has_value(); });
-
-    auto& as = _inflight_translation_state->as;
-    auto reader = co_await _data_source->make_log_reader(
-      begin_offset, datalake_priority(), as);
-    if (!reader) {
-        co_return;
-    }
-    vlog(_logger.trace, "starting translation from offset: {}", begin_offset);
-    ss::timer<scheduling::clock> cancellation_timer;
-    cancellation_timer.set_callback(
-      [&as] { as.request_abort_ex(translator_time_quota_exceeded_error{}); });
-
-    auto translation_f
-      = _translation_ctx
-          ->translate_now(
-            std::move(reader.value()), _inflight_translation_state->as)
-          .finally([this] { return _translation_ctx->flush(); });
-    cancellation_timer.arm(_inflight_translation_state->translate_for);
-    co_await std::move(translation_f).finally([&cancellation_timer] {
-        cancellation_timer.cancel();
-    });
-}
-
 bool partition_translator::should_finish_inflight_translation() const {
     auto bytes_flushed_pending_upload = _translation_ctx->flushed_bytes();
     auto lag_window_ended = _lag_tracking->should_finish_inflight_translation();
@@ -250,7 +223,7 @@ partition_translator::fetch_translation_offsets(retry_chain_node& rcn) {
     co_return offsets;
 }
 
-ss::future<> partition_translator::run_one_translation_iteration(
+ss::future<bool> partition_translator::run_one_translation_iteration(
   kafka::offset begin_offset) {
     _lag_tracking->notify_new_data_for_translation(begin_offset);
     // Notify the scheduler that there is some data to translate
@@ -259,14 +232,16 @@ ss::future<> partition_translator::run_one_translation_iteration(
     // time slice (i.e. scheduled in), then translate until the time
     // slice expires or we run out of data
     std::exception_ptr ex = nullptr;
+    bool should_discard_result = false;
+    bool force_flush_after_iteration = false;
     try {
         co_await _ready_to_translate.wait(
           [this] { return _inflight_translation_state.has_value(); });
         auto& as = _inflight_translation_state->as;
         auto reader = co_await _data_source->make_log_reader(
           begin_offset, datalake_priority(), as);
         if (!reader) {
-            co_return;
+            co_return force_flush_after_iteration;
         }
         vlog(
           _logger.trace, "starting translation from offset: {}", begin_offset);
@@ -282,13 +257,22 @@ ss::future<> partition_translator::run_one_translation_iteration(
         co_await std::move(translation_f).finally([&cancellation_timer] {
             cancellation_timer.cancel();
         });
-    } catch (...) {
-        ex = std::current_exception();
+        _inflight_translation_state->as.check();
+    } catch (const translator_out_of_memory_error&) {
         vlog(
           _logger.warn,
-          "Translation attempt failed: {}, discarding state to reset "
-          "translation",
-          ex);
+          "Translation attempt ran into OOM, result will be flushed "
+          "immediately");
+        force_flush_after_iteration = true;
+    } catch (const translator_time_quota_exceeded_error&) {
+        vlog(
+          _logger.debug,
+          "Translation attempt exceeded scheduler time limit quota");
+    } catch (...) {
+        // unknown exception or shutdown exception.
+        should_discard_result = true;
+        ex = std::current_exception();
+        vlog(_logger.warn, "Translation attempt ran into an exception: {}", ex);
     }
     // inflight_translation_state tracks a single scheduled chunk of
     // work, so we reset it to nullopt for the next time we're scheduled
@@ -297,10 +281,13 @@ ss::future<> partition_translator::run_one_translation_iteration(
     // Let the scheduler know we are done
     _scheduler->notify_done(id());
 
-    if (ex) {
+    if (should_discard_result) {
         co_await _translation_ctx->discard();
+    }
+    if (ex) {
         std::rethrow_exception(ex);
     }
+    co_return force_flush_after_iteration;
 }
 
 ss::future<bool> partition_translator::finish_inflight_translation(
@@ -404,6 +391,7 @@ ss::future<> partition_translator::translate_until_stopped() {
         if (!offsets) {
             continue;
         }
+        auto finish_immediately = false;
         if (offsets->next_translation_begin_offset) {
             // new data is available to translate
             auto translate_f = co_await ss::coroutine::as_future(
@@ -413,8 +401,9 @@ ss::future<> partition_translator::translate_until_stopped() {
                 translate_f.ignore_ready_future();
                 continue;
             }
+            finish_immediately = translate_f.get();
         }
-        if (should_finish_inflight_translation()) {
+        if (finish_immediately || should_finish_inflight_translation()) {
             auto success = co_await finish_inflight_translation(
               offsets->coordinator_lto, rcn);
             if (!success) {
@@ -513,6 +502,7 @@ void partition_translator::stop_translation() {
     if (_gate.is_closed() || !_inflight_translation_state) {
         return;
     }
+
     // 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(

src/v/datalake/translation/partition_translator.h

@@ -167,7 +167,7 @@ class partition_translator : public scheduling::translator {
     ss::future<std::optional<translation_offsets>>
     fetch_translation_offsets(retry_chain_node&);
 
-    ss::future<>
+    ss::future<bool>
     run_one_translation_iteration(kafka::offset translation_begin_offset);
 
     /**

src/v/datalake/translation/scheduling.cc

@@ -60,7 +60,14 @@ class default_reservations_tracker : public reservations_tracker {
             co_return std::move(opt_units.value());
         }
         _notifier.notify_memory_exhausted();
-        co_return co_await ss::get_units(_available_memory, nr, as);
+        try {
+            co_return co_await ss::get_units(_available_memory, nr, as);
+        } catch (...) {
+            // propagate the exception in abort source, if any
+            as.check();
+            // else, rethrow generic exception
+            std::rethrow_exception(std::current_exception());
+        }
     }
 
     size_t allocated_memory() const override {