fix some fail path

Author: lxy-9602

include/paimon/file_store_write.h

@@ -64,6 +64,9 @@ class PAIMON_EXPORT FileStoreWrite {
     /// independently. A future implementation will support spilling sealed segments to a
     /// temporary directory; currently this method only creates the in-memory segment boundary.
     /// Calling this method on a non-real-time writer returns an error.
+    /// If sealing fails, the caller must recreate both the `RealtimeContext` and writer. The
+    /// upstream must then recover input from the durable recovery offset persisted in the
+    /// snapshot. Reusing the failed writer is unsupported.
     virtual Status Seal();
 
     /// Compact data stored in given partition and bucket. Note that compaction process is only
@@ -112,6 +115,10 @@ class PAIMON_EXPORT FileStoreWrite {
     /// @param commit_identifier Identifier of this prepare-commit operation in streaming mode.
     /// @return Real-time commit messages with their partition-bucket offset ranges.
     /// @note Calling this method on a non-real-time writer or in batch mode returns an error.
+    /// @note If preparation fails, the caller must recreate both the `RealtimeContext` and writer.
+    ///       The upstream must then recover input from the durable recovery offset persisted in the
+    ///       snapshot. The failed writer may contain partially prepared bucket state and must not
+    ///       be reused.
     virtual Result<std::vector<RealtimeCommitProgress>> PrepareCommitWithProgress(
         int64_t commit_identifier);
 
@@ -128,6 +135,12 @@ class PAIMON_EXPORT FileStoreWrite {
     virtual Status RefreshCommittedSnapshot(int64_t snapshot_id);
 
     virtual std::shared_ptr<Metrics> GetMetrics() const = 0;
+
+    /// Releases resources owned by this writer.
+    ///
+    /// Closing a real-time writer with data not covered by a successful
+    /// `PrepareCommitWithProgress()` returns an error and invalidates its `RealtimeContext`.
+    /// The caller must rebuild both objects and recover input from the durable snapshot offset.
     virtual Status Close() = 0;
 };
 

src/paimon/core/operation/abstract_file_store_write.cpp

@@ -31,6 +31,7 @@
 #include "paimon/core/operation/file_system_write_restore.h"
 #include "paimon/core/operation/metrics/compaction_metrics.h"
 #include "paimon/core/operation/restore_files.h"
+#include "paimon/core/realtime/realtime_context_impl.h"
 #include "paimon/core/realtime/realtime_schema_layout.h"
 #include "paimon/core/schema/table_schema.h"
 #include "paimon/core/snapshot.h"
@@ -97,6 +98,7 @@ Status AbstractFileStoreWrite::Write(std::unique_ptr<RecordBatch>&& batch) {
     if (PAIMON_UNLIKELY(batch == nullptr)) {
         return Status::Invalid("batch is null pointer");
     }
+    PAIMON_RETURN_NOT_OK(CheckRealtimeWriteUsable());
     // in FileStoreWrite::Create() we have checked the table kind and bucket mode, here we only
     // check the bucket id in batch
     if (options_.GetBucket() == -1) {
@@ -149,6 +151,8 @@ Status AbstractFileStoreWrite::Seal() {
     if (!IsRealtimeWrite()) {
         return FileStoreWrite::Seal();
     }
+    std::lock_guard<std::mutex> seal_prepare_lock(realtime_seal_prepare_mutex_);
+    PAIMON_RETURN_NOT_OK(CheckRealtimeWriteUsable());
     std::vector<std::shared_ptr<BatchWriter>> writers;
     {
         std::lock_guard<std::mutex> lock(writers_mutex_);
@@ -159,13 +163,18 @@ Status AbstractFileStoreWrite::Seal() {
         }
     }
     for (const std::shared_ptr<BatchWriter>& writer : writers) {
-        PAIMON_RETURN_NOT_OK(writer->Seal());
+        Status status = writer->Seal();
+        if (!status.ok()) {
+            // Earlier buckets may already be sealed, so the writer cannot safely continue.
+            return PoisonRealtimeWrite(status);
+        }
     }
     return Status::OK();
 }
 
 Status AbstractFileStoreWrite::Compact(const std::map<std::string, std::string>& partition,
                                        int32_t bucket, bool full_compaction) {
+    PAIMON_RETURN_NOT_OK(CheckRealtimeWriteUsable());
     PAIMON_ASSIGN_OR_RAISE(BinaryRow part, file_store_path_factory_->ToBinaryRow(partition));
     PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<BatchWriter> writer, GetWriter(part, bucket));
     assert(writer);
@@ -300,10 +309,18 @@ Result<std::vector<RealtimeCommitProgress>> AbstractFileStoreWrite::PrepareCommi
     if (is_streaming_mode_ == false) {
         return Status::Invalid("PrepareCommitWithProgress requires streaming mode");
     }
+    // Keep one seal/prepare attempt in flight so a failure fences every later checkpoint.
+    std::lock_guard<std::mutex> seal_prepare_lock(realtime_seal_prepare_mutex_);
+    PAIMON_RETURN_NOT_OK(CheckRealtimeWriteUsable());
     // Real-time prepare snapshots writers under a short lock so writes can continue while sealed
     // segments are flushed. The normal path iterates and may erase writers in place, which would
     // either race with concurrent writer creation or hold writers_mutex_ for the whole prepare.
-    return PrepareRealtimeCommit();
+    Result<std::vector<RealtimeCommitProgress>> result = PrepareRealtimeCommit();
+    if (!result.ok()) {
+        // Prepare may have drained an earlier bucket or partially written the failing bucket.
+        return PoisonRealtimeWrite(result.status());
+    }
+    return result;
 }
 
 Result<std::vector<RealtimeCommitProgress>> AbstractFileStoreWrite::PrepareRealtimeCommit() {
@@ -362,22 +379,91 @@ Result<std::vector<RealtimeCommitProgress>> AbstractFileStoreWrite::PrepareRealt
 }
 
 Status AbstractFileStoreWrite::Close() {
+    std::unique_lock<std::mutex> seal_prepare_lock(realtime_seal_prepare_mutex_, std::defer_lock);
+    if (IsRealtimeWrite()) {
+        seal_prepare_lock.lock();
+    }
     std::lock_guard<std::mutex> lock(writers_mutex_);
+    bool has_unprepared_realtime_data = false;
+    Status unprepared_error = Status::OK();
+    if (IsRealtimeWrite()) {
+        for (const auto& [_, bucket_writers] : writers_) {
+            for (const auto& [_, writer_container] : bucket_writers) {
+                if (writer_container.writer->HasUnpreparedRealtimeData()) {
+                    has_unprepared_realtime_data = true;
+                    break;
+                }
+            }
+            if (has_unprepared_realtime_data) {
+                break;
+            }
+        }
+        if (has_unprepared_realtime_data) {
+            // Close cannot publish pending segments for its caller.
+            unprepared_error = PoisonRealtimeWrite(
+                Status::Invalid("real-time writer closed with data not covered by a successful "
+                                "PrepareCommitWithProgress"));
+        }
+    }
+    Status first_error = Status::OK();
     for (auto& [_, bucket_writers] : writers_) {
         for (auto& [_, writer_container] : bucket_writers) {
             writer_memory_manager_->UnregisterWriter(writer_container.writer.get());
-            PAIMON_RETURN_NOT_OK(writer_container.writer->Close());
+            Status status = writer_container.writer->Close();
+            if (!status.ok() && first_error.ok()) {
+                first_error = std::move(status);
+            }
         }
     }
     writers_.clear();
     compact_executor_->ShutdownNow();
-    return Status::OK();
+    if (!first_error.ok() && IsRealtimeWrite()) {
+        return PoisonRealtimeWrite(first_error);
+    }
+    if (has_unprepared_realtime_data) {
+        return unprepared_error;
+    }
+    return first_error;
 }
 
 std::shared_ptr<Metrics> AbstractFileStoreWrite::GetMetrics() const {
     return metrics_;
 }
 
+Status AbstractFileStoreWrite::CheckRealtimeWriteUsable() const {
+    if (IsRealtimeWrite() && realtime_write_poisoned_.load()) {
+        return Status::Invalid(
+            "real-time writer cannot be reused after Seal or PrepareCommitWithProgress failed, "
+            "or after it closed with unprepared data; recreate RealtimeContext and "
+            "FileStoreWrite, then let upstream recover input from the durable recovery offset "
+            "persisted in the snapshot");
+    }
+    std::shared_ptr<RealtimeContext> context = GetRealtimeContext();
+    if (!context) {
+        return Status::OK();
+    }
+    PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<RealtimeContextImpl> context_impl,
+                           RealtimeContextImpl::Cast(context));
+    return context_impl->CheckUsable();
+}
+
+void AbstractFileStoreWrite::InvalidateRealtimeContext() {
+    std::shared_ptr<RealtimeContextImpl> context =
+        std::dynamic_pointer_cast<RealtimeContextImpl>(GetRealtimeContext());
+    if (context) {
+        context->Invalidate();
+    }
+}
+
+Status AbstractFileStoreWrite::PoisonRealtimeWrite(const Status& cause) {
+    realtime_write_poisoned_.store(true);
+    InvalidateRealtimeContext();
+    return cause.WithMessage(cause.message(),
+                             "; recreate RealtimeContext and FileStoreWrite, then let upstream "
+                             "recover input from the durable recovery offset persisted in the "
+                             "snapshot");
+}
+
 int32_t AbstractFileStoreWrite::GetDefaultBucketNum() const {
     return options_.GetBucket();
 }

src/paimon/core/operation/abstract_file_store_write.h

@@ -18,6 +18,7 @@
 
 #pragma once
 
+#include <atomic>
 #include <cstdint>
 #include <limits>
 #include <memory>
@@ -65,6 +66,7 @@ class MemoryPool;
 class RecordBatch;
 class RestoreFiles;
 class IOManager;
+class RealtimeContext;
 class RealtimeSchemaLayout;
 
 class AbstractFileStoreWrite : public FileStoreWrite {
@@ -126,6 +128,10 @@ class AbstractFileStoreWrite : public FileStoreWrite {
         return false;
     }
 
+    virtual std::shared_ptr<RealtimeContext> GetRealtimeContext() const {
+        return nullptr;
+    }
+
     Result<std::shared_ptr<RestoreFiles>> ScanExistingFileMetas(const BinaryRow& partition,
                                                                 int32_t bucket) const;
     int32_t GetDefaultBucketNum() const;
@@ -152,18 +158,24 @@ class AbstractFileStoreWrite : public FileStoreWrite {
     std::shared_ptr<CompactionMetrics> compaction_metrics_;
 
  private:
+    Status CheckRealtimeWriteUsable() const;
+    void InvalidateRealtimeContext();
+    Status PoisonRealtimeWrite(const Status& cause);
+
     Result<std::shared_ptr<BatchWriter>> GetWriter(const BinaryRow& partition, int32_t bucket);
     Result<std::vector<RealtimeCommitProgress>> PrepareRealtimeCommit();
 
  private:
     std::unordered_map<BinaryRow, std::unordered_map<int32_t, WriterContainer<BatchWriter>>>
         writers_;
     std::mutex writers_mutex_;
+    std::mutex realtime_seal_prepare_mutex_;
     std::mutex realtime_metrics_mutex_;
     bool ignore_previous_files_ = false;
     bool is_streaming_mode_ = false;
     bool ignore_num_bucket_check_ = false;
     bool batch_committed_ = false;
+    std::atomic<bool> realtime_write_poisoned_{false};
 
     std::shared_ptr<MetricsImpl> metrics_;
     std::unique_ptr<Logger> logger_;

src/paimon/core/operation/append_only_file_store_write.h

@@ -120,6 +120,10 @@ class AppendOnlyFileStoreWrite : public AbstractFileStoreWrite {
         return realtime_context_ != nullptr;
     }
 
+    std::shared_ptr<RealtimeContext> GetRealtimeContext() const override {
+        return realtime_context_;
+    }
+
     /// @param plan_factory The active shredding write plan, or nullptr when the rewrite stays a
     ///                     plain passthrough. Resolved by the caller because
     ///                     `GetDictionaryPassthroughVetoReason` needs the same answer.

src/paimon/core/operation/file_store_write.cpp

@@ -63,14 +63,15 @@ namespace {
 Status RestoreRealtimeCommittedProgress(const std::shared_ptr<RealtimeContext>& realtime_context,
                                         const std::shared_ptr<SnapshotManager>& snapshot_manager,
                                         const CoreOptions& options) {
+    PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<RealtimeContextImpl> realtime_context_impl,
+                           RealtimeContextImpl::Cast(realtime_context));
+    PAIMON_RETURN_NOT_OK(realtime_context_impl->CheckUsable());
     PAIMON_ASSIGN_OR_RAISE(std::optional<Snapshot> latest_snapshot,
                            snapshot_manager->LatestSnapshot());
     if (latest_snapshot) {
         PAIMON_ASSIGN_OR_RAISE(
             RealtimeOffsetMap realtime_committed_offsets,
             RealtimeCommitProperties::ReadOffsets(latest_snapshot, options.GetFileSystem()));
-        PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<RealtimeContextImpl> realtime_context_impl,
-                               RealtimeContextImpl::Cast(realtime_context));
         PAIMON_RETURN_NOT_OK(realtime_context_impl->AdvanceCommittedProgress(
             latest_snapshot->Id(), realtime_committed_offsets));
     }

src/paimon/core/operation/key_value_file_store_write.cpp

@@ -188,9 +188,9 @@ Status KeyValueFileStoreWrite::RefreshCommittedSnapshot(int64_t snapshot_id) {
 }
 
 Status KeyValueFileStoreWrite::Close() {
-    PAIMON_RETURN_NOT_OK(AbstractFileStoreWrite::Close());
+    Status writer_status = AbstractFileStoreWrite::Close();
     compact_manager_factory_->Close();
-    return Status::OK();
+    return writer_status;
 }
 
 }  // namespace paimon

src/paimon/core/operation/key_value_file_store_write.h

@@ -87,6 +87,10 @@ class KeyValueFileStoreWrite : public AbstractFileStoreWrite {
         return realtime_context_ != nullptr;
     }
 
+    std::shared_ptr<RealtimeContext> GetRealtimeContext() const override {
+        return realtime_context_;
+    }
+
  private:
     bool enable_multi_thread_spill_;
     std::shared_ptr<RealtimeContext> realtime_context_;

src/paimon/core/operation/key_value_file_store_write_test.cpp

@@ -519,6 +519,9 @@ TEST_F(KeyValueFileStoreWriteTest, TestRealtimePool) {
     const int64_t allocations_before_write = pool->allocation_count;
     ASSERT_OK(writer->Write(MakeBatch(realtime_schema, R"([[0, 1, "one"]])")));
     ASSERT_GT(pool->allocation_count, allocations_before_write);
+    ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> progress,
+                         writer->PrepareCommitWithProgress(/*commit_identifier=*/0));
+    ASSERT_EQ(1, progress.size());
     ASSERT_OK(writer->Close());
     writer.reset();
     ASSERT_OK_AND_ASSIGN(auto retained_rows, ReadRealtimePrimaryKeyStoreRows(realtime_context));
@@ -615,6 +618,9 @@ TEST_F(KeyValueFileStoreWriteTest, TestRealtimeLimits) {
     ASSERT_EQ(1, read_state.views.size());
     ASSERT_EQ(std::optional<OffsetRange>(OffsetRange(max - 1, max)),
               read_state.views[0].read_view->GetOffsetRange());
+    ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> boundary_progress,
+                         writer->PrepareCommitWithProgress(/*commit_identifier=*/1));
+    ASSERT_EQ(1, boundary_progress.size());
     ASSERT_OK(writer->Close());
     ASSERT_GE(snapshot_id, 1);
 }

src/paimon/core/realtime/realtime_append_only_writer.cpp

@@ -110,6 +110,7 @@ Status RealtimeAppendOnlyWriter::Write(std::unique_ptr<RecordBatch>&& batch) {
     PAIMON_RETURN_NOT_OK(
         realtime_store_->Write(RealtimeWriteBatch{std::move(batch), validated.offset_range}));
     next_offset_ = validated.offset_range.end;
+    has_building_data_ = true;
     return Status::OK();
 }
 
@@ -121,6 +122,7 @@ Status RealtimeAppendOnlyWriter::SealCurrentSegment() {
             return Status::Invalid("append real-time store sealed a null segment");
         }
         sealed_segments_.push_back(std::move(segment.value()));
+        has_building_data_ = false;
     }
     return Status::OK();
 }
@@ -226,6 +228,11 @@ Status RealtimeAppendOnlyWriter::Close() {
     return file_writer_->Close();
 }
 
+bool RealtimeAppendOnlyWriter::HasUnpreparedRealtimeData() const {
+    std::lock_guard<std::mutex> lock(realtime_store_mutex_);
+    return has_building_data_ || !sealed_segments_.empty();
+}
+
 std::shared_ptr<Metrics> RealtimeAppendOnlyWriter::GetMetrics() const {
     return file_writer_->GetMetrics();
 }

src/paimon/core/realtime/realtime_append_only_writer.h

@@ -69,6 +69,8 @@ class RealtimeAppendOnlyWriter : public BatchWriter {
 
     Status Close() override;
 
+    bool HasUnpreparedRealtimeData() const override;
+
     std::shared_ptr<Metrics> GetMetrics() const override;
 
  private:
@@ -86,7 +88,8 @@ class RealtimeAppendOnlyWriter : public BatchWriter {
     std::shared_ptr<RealtimeSchemaLayout> schema_layout_;
     std::vector<std::shared_ptr<RealtimeSegmentHandle>> sealed_segments_;
     int64_t next_offset_;
-    std::mutex realtime_store_mutex_;
+    bool has_building_data_ = false;
+    mutable std::mutex realtime_store_mutex_;
     std::mutex prepare_mutex_;
 };