From e6e00e7f5fb03cf4397496e2f7a0e9dc4da5126d Mon Sep 17 00:00:00 2001
From: Sutou Kouhei <kou@clear-code.com>
Date: Wed, 29 May 2024 00:02:19 +0900
Subject: [PATCH] GH-41771: [C++] Iterator releases its resource immediately
 when it reads all values (#41824)

### Rationale for this change

`Iterator` keeps its resource (`ptr_`) until it's deleted but we can release its resource immediately when it reads all values. If `Iterator` keeps its resource until it's deleted, it may block closing a file. See GH-41771 for this case.

### What changes are included in this PR?

Releases `ptr_` when `Next()` returns the end.

### Are these changes tested?

Yes.

### Are there any user-facing changes?

Yes.
* GitHub Issue: #41771

Authored-by: Sutou Kouhei <kou@clear-code.com>
Signed-off-by: Benjamin Kietzman <bengilgit@gmail.com>
---
 cpp/src/arrow/util/iterator.h       | 15 ++++++++--
 cpp/src/arrow/util/iterator_test.cc | 43 +++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/cpp/src/arrow/util/iterator.h b/cpp/src/arrow/util/iterator.h
index 4da8394a0299c..5025799b9a372 100644
--- a/cpp/src/arrow/util/iterator.h
+++ b/cpp/src/arrow/util/iterator.h
@@ -105,9 +105,18 @@ class Iterator : public util::EqualityComparable<Iterator<T>> {
   Iterator() : ptr_(NULLPTR, [](void*) {}) {}
 
   /// \brief Return the next element of the sequence, IterationTraits<T>::End() when the
-  /// iteration is completed. Calling this on a default constructed Iterator
-  /// will result in undefined behavior.
-  Result<T> Next() { return next_(ptr_.get()); }
+  /// iteration is completed.
+  Result<T> Next() {
+    if (ptr_) {
+      auto next_result = next_(ptr_.get());
+      if (next_result.ok() && IsIterationEnd(next_result.ValueUnsafe())) {
+        ptr_.reset(NULLPTR);
+      }
+      return next_result;
+    } else {
+      return IterationTraits<T>::End();
+    }
+  }
 
   /// Pass each element of the sequence to a visitor. Will return any error status
   /// returned by the visitor, terminating iteration.
diff --git a/cpp/src/arrow/util/iterator_test.cc b/cpp/src/arrow/util/iterator_test.cc
index ba21ddcced209..a247ba13aef73 100644
--- a/cpp/src/arrow/util/iterator_test.cc
+++ b/cpp/src/arrow/util/iterator_test.cc
@@ -146,6 +146,49 @@ void AssertIteratorNext(T expected, Iterator<T>& it) {
   ASSERT_EQ(expected, actual);
 }
 
+template <typename T>
+class DeleteDetectableIterator {
+ public:
+  explicit DeleteDetectableIterator(std::vector<T> values, bool* deleted)
+      : values_(std::move(values)), i_(0), deleted_(deleted) {}
+
+  DeleteDetectableIterator(DeleteDetectableIterator&& source)
+      : values_(std::move(source.values_)), i_(source.i_), deleted_(source.deleted_) {
+    source.deleted_ = nullptr;
+  }
+
+  ~DeleteDetectableIterator() {
+    if (deleted_) {
+      *deleted_ = true;
+    }
+  }
+
+  Result<T> Next() {
+    if (i_ == values_.size()) {
+      return IterationTraits<T>::End();
+    }
+    return std::move(values_[i_++]);
+  }
+
+ private:
+  std::vector<T> values_;
+  size_t i_;
+  bool* deleted_;
+};
+
+// Generic iterator tests
+
+TEST(TestIterator, DeleteOnEnd) {
+  bool deleted = false;
+  Iterator<TestInt> it(DeleteDetectableIterator<TestInt>({1}, &deleted));
+  ASSERT_FALSE(deleted);
+  AssertIteratorNext({1}, it);
+  ASSERT_FALSE(deleted);
+  ASSERT_OK_AND_ASSIGN(auto value, it.Next());
+  ASSERT_TRUE(IsIterationEnd(value));
+  ASSERT_TRUE(deleted);
+}
+
 // --------------------------------------------------------------------
 // Synchronous iterator tests