From fafd6a3c9c04737a7c9b4bed1b270c81cb77435e Mon Sep 17 00:00:00 2001
From: Felix Chern <fchern@google.com>
Date: Tue, 30 Jan 2024 15:49:53 -0800
Subject: [PATCH] internal changes.

PiperOrigin-RevId: 602863918
---
 cpp/BUILD                  |  1 +
 cpp/array_record_writer.cc | 21 ++++++++++++++++++++-
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/cpp/BUILD b/cpp/BUILD
index 1c77e47..ea17ec8 100644
--- a/cpp/BUILD
+++ b/cpp/BUILD
@@ -114,6 +114,7 @@ cc_library(
         "@com_google_protobuf//:protobuf_lite",
         "@com_google_riegeli//riegeli/base:object",
         "@com_google_riegeli//riegeli/base:options_parser",
+        "@com_google_riegeli//riegeli/base:status",
         "@com_google_riegeli//riegeli/bytes:chain_writer",
         "@com_google_riegeli//riegeli/chunk_encoding:chunk",
         "@com_google_riegeli//riegeli/chunk_encoding:chunk_encoder",
diff --git a/cpp/array_record_writer.cc b/cpp/array_record_writer.cc
index 3ef250b..4f22221 100644
--- a/cpp/array_record_writer.cc
+++ b/cpp/array_record_writer.cc
@@ -43,6 +43,7 @@ limitations under the License.
 #include "cpp/thread_pool.h"
 #include "riegeli/base/object.h"
 #include "riegeli/base/options_parser.h"
+#include "riegeli/base/status.h"
 #include "riegeli/bytes/chain_writer.h"
 #include "riegeli/chunk_encoding/chunk.h"
 #include "riegeli/chunk_encoding/chunk_encoder.h"
@@ -365,10 +366,28 @@ void ArrayRecordWriterBase::Initialize() {
 }
 
 void ArrayRecordWriterBase::Done() {
+  if (!ok()) {
+    return;
+  }
   auto writer = get_writer();
+  if (writer == nullptr) {
+    Fail(absl::InternalError("writer should not be a nullptr."));
+    return;
+  }
+  if (!writer->ok()) {
+    Fail(riegeli::Annotate(writer->status(), "SequencedChunkWriter failure"));
+    return;
+  }
+  if (chunk_encoder_ == nullptr) {
+    Fail(absl::InternalError("chunk_encoder_ should not be a nullptr."));
+    return;
+  }
   if (chunk_encoder_->num_records() > 0) {
     std::promise<absl::StatusOr<Chunk>> chunk_promise;
-    writer->CommitFutureChunk(chunk_promise.get_future());
+    if (!writer->CommitFutureChunk(chunk_promise.get_future())) {
+      Fail(riegeli::Annotate(writer->status(), "Cannot commit future chunk"));
+      return;
+    }
     chunk_promise.set_value(EncodeChunk(chunk_encoder_.get()));
   }
   submit_chunk_callback_->WriteFooterAndPostscript(writer.get());