|
14 | 14 | #include <algorithm> |
15 | 15 | #include <cmath> |
16 | 16 | #include <cstddef> |
17 | | -#include <exception> |
18 | 17 | #include <functional> |
19 | 18 | #include <future> |
20 | 19 | #include <optional> |
@@ -482,66 +481,53 @@ class parquet_field_struct_list : public parquet_field { |
482 | 481 | auto const items_per_task = n / num_tasks; |
483 | 482 | auto const remainder = n % num_tasks; |
484 | 483 |
|
485 | | - std::vector<std::future<void>> tasks; |
486 | | - tasks.reserve(num_tasks); |
487 | | - |
488 | 484 | // Size once up front, then fill by index below: workers read already-filled slots while the |
489 | 485 | // main thread fills later ones, so a size-mutating emplace_back would be a data race. |
490 | 486 | std::vector<cudf::host_span<uint8_t const>> all_ranges; |
491 | 487 | all_ranges.resize(n); |
492 | 488 |
|
493 | | - // Wait for every task before all_ranges/val die on any exit path: a throw during collection |
494 | | - // must not let a worker read freed memory (host-pool futures do not block on destruction). |
495 | | - try { |
496 | | - uint32_t struct_idx = 0; |
497 | | - for (uint32_t task_id = 0; task_id < num_tasks; ++task_id) { |
498 | | - auto const task_size = items_per_task + (task_id < remainder ? 1 : 0); |
499 | | - auto const start_idx = std::min(n, struct_idx); |
500 | | - auto const end_idx = std::min(n, start_idx + task_size); |
501 | | - |
502 | | - if (start_idx >= end_idx) { break; } |
503 | | - |
504 | | - // Collect struct ranges for the next task |
505 | | - for (auto i = start_idx; i < end_idx; ++i) { |
506 | | - uint8_t const* const start = cpr->m_cur; |
507 | | - cpr->skip_struct_field(static_cast<int>(FieldType::STRUCT)); |
508 | | - all_ranges[i] = |
509 | | - cudf::host_span<uint8_t const>{start, static_cast<size_t>(cpr->m_cur - start)}; |
510 | | - } |
| 489 | + // Declared after all_ranges so that every task is waited on before all_ranges dies on any |
| 490 | + // exit path (host-pool futures do not block on destruction) |
| 491 | + std::vector<std::future<void>> tasks; |
| 492 | + cudf::detail::futures_wait_guard tasks_guard{tasks}; |
| 493 | + tasks.reserve(num_tasks); |
511 | 494 |
|
512 | | - // Launch task immediately to parse the structs in parallel while the main thread collects |
513 | | - // the remaining ranges |
514 | | - tasks.emplace_back(cudf::detail::host_worker_pool().submit_task( |
515 | | - [&val = this->val, &all_ranges, start_idx, end_idx, mode]() { |
516 | | - CompactProtocolReader local_cpr{nullptr, 0, mode}; |
517 | | - // end_idx alone bounds the loop: all_ranges[start_idx, end_idx) is fully populated |
518 | | - // before this task launches, so no concurrent read of all_ranges.size() is needed. |
519 | | - for (size_t i = start_idx; i < end_idx; ++i) { |
520 | | - local_cpr.init(all_ranges[i].data(), all_ranges[i].size()); |
521 | | - local_cpr.read(&val[i]); |
522 | | - } |
523 | | - })); |
524 | | - |
525 | | - struct_idx = end_idx; |
526 | | - } |
527 | | - } catch (...) { |
528 | | - for (auto& task : tasks) { |
529 | | - task.wait(); |
| 495 | + uint32_t struct_idx = 0; |
| 496 | + for (uint32_t task_id = 0; task_id < num_tasks; ++task_id) { |
| 497 | + auto const task_size = items_per_task + (task_id < remainder ? 1 : 0); |
| 498 | + auto const start_idx = std::min(n, struct_idx); |
| 499 | + auto const end_idx = std::min(n, start_idx + task_size); |
| 500 | + |
| 501 | + if (start_idx >= end_idx) { break; } |
| 502 | + |
| 503 | + // Collect struct ranges for the next task |
| 504 | + for (auto i = start_idx; i < end_idx; ++i) { |
| 505 | + uint8_t const* const start = cpr->m_cur; |
| 506 | + cpr->skip_struct_field(static_cast<int>(FieldType::STRUCT)); |
| 507 | + all_ranges[i] = |
| 508 | + cudf::host_span<uint8_t const>{start, static_cast<size_t>(cpr->m_cur - start)}; |
530 | 509 | } |
531 | | - throw; |
| 510 | + |
| 511 | + // Launch task immediately to parse the structs in parallel while the main thread collects |
| 512 | + // the remaining ranges |
| 513 | + tasks.emplace_back(cudf::detail::host_worker_pool().submit_task( |
| 514 | + [&val = this->val, &all_ranges, start_idx, end_idx, mode]() { |
| 515 | + CompactProtocolReader local_cpr{nullptr, 0, mode}; |
| 516 | + // end_idx alone bounds the loop: all_ranges[start_idx, end_idx) is fully populated |
| 517 | + // before this task launches, so no concurrent read of all_ranges.size() is needed. |
| 518 | + for (size_t i = start_idx; i < end_idx; ++i) { |
| 519 | + local_cpr.init(all_ranges[i].data(), all_ranges[i].size()); |
| 520 | + local_cpr.read(&val[i]); |
| 521 | + } |
| 522 | + })); |
| 523 | + |
| 524 | + struct_idx = end_idx; |
532 | 525 | } |
533 | | - // Join EVERY task before leaving, then rethrow the first failure: task.get() throws on a |
534 | | - // strict-mode mismatch, and bailing out of this loop mid-way would abandon the remaining |
535 | | - // un-joined tasks, whose lambdas still reference all_ranges/val (use-after-free at exit). |
536 | | - std::exception_ptr first_error; |
| 526 | + // Rethrows the first failure (e.g. a strict-mode mismatch); the guard waits for the |
| 527 | + // remaining tasks, whose lambdas reference all_ranges/val, while unwinding |
537 | 528 | for (auto& task : tasks) { |
538 | | - try { |
539 | | - task.get(); |
540 | | - } catch (...) { |
541 | | - if (!first_error) { first_error = std::current_exception(); } |
542 | | - } |
| 529 | + task.get(); |
543 | 530 | } |
544 | | - if (first_error) { std::rethrow_exception(first_error); } |
545 | 531 | } else { |
546 | 532 | // For small numbers of elements, use sequential processing to avoid overhead |
547 | 533 | std::for_each(val.begin(), val.end(), [&cpr](auto& elem) { cpr->read(&elem); }); |
|
0 commit comments