Skip to content

[TSL] Wait for queued and running work in default UnboundedWorkQueue::~UnboundedWorkQueue - #3653

Merged
copybara-service[bot] merged 1 commit into
mainfrom
test_993490972
Oct 5, 2026
Merged

copybara-service[bot] merged 1 commit into
mainfrom
test_993490972

Conversation

@copybara-service

Copy link
Copy Markdown

[TSL] Wait for queued and running work in default UnboundedWorkQueue::~UnboundedWorkQueue

In the default/OSS implementation of tsl::UnboundedWorkQueue (xla/tsl/platform/default/unbounded_work_queue.{h,cc}), ~UnboundedWorkQueue() previously set cancelled_ = true immediately without waiting for pending or in-flight work items to finish, and then acquired thread_pool_mu_ while joining thread_pool_ threads.

This differed from the Google-internal implementation (which waits for all work to complete in ~UnboundedWorkQueue()) and caused a lock-order inversion deadlock when an in-flight worker thread called UnboundedWorkQueue::Schedule() concurrently with ~UnboundedWorkQueue() (e.g., in //xla/pjrt/c:pjrt_c_api_gpu_test_nvgpu_any during PjrtCApiTest.BufferTransferImmutableUntilTransferCompletes, where CopyRawHostToDeviceAndReturnEvent calls ExecuteWhenReady -> Schedule() on the H2D Dispatch worker thread after on_done_with_host_buffer has already unblocked client destruction):

  • Main thread in ~UnboundedWorkQueue() holds thread_pool_mu_ and blocks in pthread_join waiting for the worker thread.
  • Worker thread in Schedule() holds work_queue_mu_ and blocks trying to acquire thread_pool_mu_.

Update default/unbounded_work_queue.{h,cc} to track num_running_functions_ (guarded by work_queue_mu_) and wait in ~UnboundedWorkQueue() until work_queue_.empty() && num_running_functions_ == 0 before setting cancelled_ = true and joining thread_pool_. Also update unbounded_work_queue_test.cc to expect all closures to finish in RacyDestructor and add NestedClosureDuringDestructor to cover scheduling nested work during destruction.

…::~UnboundedWorkQueue`

In the default/OSS implementation of `tsl::UnboundedWorkQueue` (`xla/tsl/platform/default/unbounded_work_queue.{h,cc}`), `~UnboundedWorkQueue()` previously set `cancelled_ = true` immediately without waiting for pending or in-flight work items to finish, and then acquired `thread_pool_mu_` while joining `thread_pool_` threads.

This differed from the Google-internal implementation (which waits for all work to complete in `~UnboundedWorkQueue()`) and caused a lock-order inversion deadlock when an in-flight worker thread called `UnboundedWorkQueue::Schedule()` concurrently with `~UnboundedWorkQueue()` (e.g., in `//xla/pjrt/c:pjrt_c_api_gpu_test_nvgpu_any` during `PjrtCApiTest.BufferTransferImmutableUntilTransferCompletes`, where `CopyRawHostToDeviceAndReturnEvent` calls `ExecuteWhenReady` -> `Schedule()` on the `H2D Dispatch` worker thread after `on_done_with_host_buffer` has already unblocked client destruction):
- Main thread in `~UnboundedWorkQueue()` holds `thread_pool_mu_` and blocks in `pthread_join` waiting for the worker thread.
- Worker thread in `Schedule()` holds `work_queue_mu_` and blocks trying to acquire `thread_pool_mu_`.

Update `default/unbounded_work_queue.{h,cc}` to track `num_running_functions_` (guarded by `work_queue_mu_`) and wait in `~UnboundedWorkQueue()` until `work_queue_.empty() && num_running_functions_ == 0` before setting `cancelled_ = true` and joining `thread_pool_`. Also update `unbounded_work_queue_test.cc` to expect all closures to finish in `RacyDestructor` and add `NestedClosureDuringDestructor` to cover scheduling nested work during destruction.

PiperOrigin-RevId: 993570944
@copybara-service
copybara-service Bot merged commit 32b347f into main Oct 5, 2026
2 checks passed
@copybara-service
copybara-service Bot deleted the test_993490972 branch October 5, 2026 10:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant