Function CheckScannerBackpressure

cpp/src/arrow/dataset/scanner_test.cc:585–636

                         });

void CheckScannerBackpressure(std::shared_ptr<MockDataset> dataset, ScanV2Options options,
                              int maxConcurrentFragments, int maxConcurrentBatches,
                              ::arrow::internal::ThreadPool* thread_pool) {
  // Start scanning
  acero::Declaration scan_decl = acero::Declaration("scan2", std::move(options));
  Future<RecordBatchVector> batches_fut =
      acero::DeclarationToBatchesAsync(std::move(scan_decl));

  auto get_num_inspected = [&] {
    int num_inspected = 0;
    for (const auto& frag : dataset->fragments_) {
      if (frag->has_inspected()) {
        num_inspected++;
      }
    }
    return num_inspected;
  };
  BusyWait(10, [&] {
    return get_num_inspected() == static_cast<int>(maxConcurrentFragments);
  });
  SleepABit();
  ASSERT_EQ(get_num_inspected(), static_cast<int>(maxConcurrentFragments));

  int total_batches = 0;
  for (const auto& frag : dataset->fragments_) {
    total_batches += frag->fragment_scanner_->NumBatches();
    frag->FinishInspection();
    frag->FinishScanBegin();
  }

  int batches_scanned = 0;
  while (batches_scanned < total_batches) {
    MockScanTask* next_task_to_deliver = nullptr;
    thread_pool->WaitForIdle();
    int batches_started = 0;
    for (const auto& frag : dataset->fragments_) {
      for (int i = 0; i < frag->fragment_scanner_->NumBatches(); i++) {
        if (frag->HasBatchStarted(i)) {
          batches_started++;
          if (next_task_to_deliver == nullptr && !frag->HasBatchDelivered(i)) {
            next_task_to_deliver = &frag->fragment_scanner_->scan_tasks_[i];
          }
        }
      }
    }
    ASSERT_LE(batches_started - batches_scanned, maxConcurrentBatches)
        << " too many scan tasks were allowed to run";
    ASSERT_NE(next_task_to_deliver, nullptr);
    next_task_to_deliver->batch_future.MarkFinished(next_task_to_deliver->batch);
    batches_scanned++;
  }
}

TEST(TestNewScanner, Backpressure) {
  constexpr int kNumFragments = 4;