pitrou commented on a change in pull request #9533:
URL: https://github.com/apache/arrow/pull/9533#discussion_r592421631
##########
File path: cpp/src/arrow/util/iterator_test.cc
##########
@@ -726,6 +730,138 @@ TEST(TestAsyncUtil, CompleteBackgroundStressTest) {
}
}
+template <typename T>
+class ReentrantChecker {
+ public:
+ explicit ReentrantChecker(AsyncGenerator<T> source)
+ : state_(std::make_shared<State>(std::move(source))) {}
+
+ Future<T> operator()() {
+ if (state_->in.load()) {
+ state_->valid.store(false);
+ }
+ state_->in.store(true);
+ auto result = state_->source();
+ return result.Then(Callback{state_});
+ }
+
+ void AssertValid() {
+ EXPECT_EQ(true, state_->valid.load())
+ << "The generator was accessed in a reentrant manner";
+ }
+
+ private:
+ struct State {
+ explicit State(AsyncGenerator<T> source_)
+ : source(std::move(source_)), in(false), valid(true) {}
+
+ AsyncGenerator<T> source;
+ std::atomic<bool> in;
+ std::atomic<bool> valid;
+ };
+ struct Callback {
+ Future<T> operator()(const Result<T>& result) {
+ state_->in.store(false);
+ return result;
+ }
+ std::shared_ptr<State> state_;
+ };
+
+ std::shared_ptr<State> state_;
+};
+
+TEST(TestAsyncUtil, SerialReadaheadSlowProducer) {
+ AsyncGenerator<TestInt> it = BackgroundAsyncVectorIt({1, 2, 3, 4, 5});
+ ReentrantChecker<TestInt> checker(std::move(it));
+ SerialReadaheadGenerator<TestInt> serial_readahead(checker, 2);
+ AssertAsyncGeneratorMatch({1, 2, 3, 4, 5},
+
static_cast<AsyncGenerator<TestInt>>(serial_readahead));
+ checker.AssertValid();
+}
+
+TEST(TestAsyncUtil, SerialReadaheadSlowConsumer) {
+ int num_delivered = 0;
+ auto source = [&num_delivered]() {
+ if (num_delivered < 5) {
+ return Future<TestInt>::MakeFinished(num_delivered++);
+ } else {
+ return Future<TestInt>::MakeFinished(IterationTraits<TestInt>::End());
+ }
+ };
+ SerialReadaheadGenerator<TestInt> serial_readahead(std::move(source), 3);
+ SleepABit();
+ ASSERT_EQ(0, num_delivered);
Review comment:
Fair enough. I suppose both behaviours make sense in different
situations.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
