| 311 | /// approach is to use a stop token to cause the generator to exhaust early. |
| 312 | template <typename T> |
| 313 | static Iterator<T> IterateGenerator( |
| 314 | internal::FnOnce<Result<std::function<Future<T>()>>(Executor*)> initial_task) { |
| 315 | auto serial_executor = std::unique_ptr<SerialExecutor>(new SerialExecutor()); |
| 316 | auto maybe_generator = std::move(initial_task)(serial_executor.get()); |
| 317 | if (!maybe_generator.ok()) { |
| 318 | return MakeErrorIterator<T>(maybe_generator.status()); |
| 319 | } |
| 320 | auto generator = maybe_generator.MoveValueUnsafe(); |
| 321 | struct SerialIterator { |
| 322 | SerialIterator(std::unique_ptr<SerialExecutor> executor, |
| 323 | std::function<Future<T>()> generator) |
| 324 | : executor(std::move(executor)), generator(std::move(generator)) {} |
| 325 | ARROW_DISALLOW_COPY_AND_ASSIGN(SerialIterator); |
| 326 | ARROW_DEFAULT_MOVE_AND_ASSIGN(SerialIterator); |
| 327 | ~SerialIterator() { |
| 328 | // A serial iterator must be consumed before it can be destroyed. Allowing it to |
| 329 | // do otherwise would lead to resource leakage. There will likely be deadlocks at |
| 330 | // this spot in the future but these will be the result of other bugs and not the |
| 331 | // fact that we are forcing consumption here. |
| 332 | |
| 333 | // If a streaming API needs to support early abandonment then it should be done so |
| 334 | // with a cancellation token and not simply discarding the iterator and expecting |
| 335 | // the underlying work to clean up correctly. |
| 336 | if (executor && !executor->IsFinished()) { |
| 337 | while (true) { |
| 338 | Result<T> maybe_next = Next(); |
| 339 | if (!maybe_next.ok() || IsIterationEnd(*maybe_next)) { |
| 340 | break; |
| 341 | } |
| 342 | } |
| 343 | } |
| 344 | } |
| 345 | |
| 346 | Result<T> Next() { |
| 347 | executor->Unpause(); |
| 348 | // This call may lead to tasks being scheduled in the serial executor |
| 349 | Future<T> next_fut = generator(); |
| 350 | next_fut.AddCallback([this](const Result<T>& res) { |
| 351 | // If we're done iterating we should drain the rest of the tasks in the executor |
| 352 | if (!res.ok() || IsIterationEnd(*res)) { |
| 353 | executor->Finish(); |
| 354 | return; |
| 355 | } |
| 356 | // Otherwise we will break out immediately, leaving the remaining tasks for |
| 357 | // the next call. |
| 358 | executor->Pause(); |
| 359 | }); |
| 360 | #ifdef ARROW_ENABLE_THREADING |
| 361 | // future must run on this thread |
| 362 | // Borrow this thread and run tasks until the future is finished |
| 363 | executor->RunLoop(); |
| 364 | #else |
| 365 | next_fut.Wait(); |
| 366 | #endif |
| 367 | if (!next_fut.is_finished()) { |
| 368 | // Not clear this is possible since RunLoop wouldn't generally exit |
| 369 | // unless we paused/finished which would imply next_fut has been |
| 370 | // finished. |
no test coverage detected