Security: Thread Safety Issue in Asyncio Utils

[tomaioo]

Summary

Security: Thread Safety Issue in Asyncio Utils

Problem

Severity: Medium | File: checkpoint/orbax/checkpoint/_src/asyncio_utils.py:L58

The run_sync function in asyncio_utils.py creates a new event loop in a daemon thread when uvloop is available. The finally block calls event_loop.call_soon_threadsafe(event_loop.stop) and then thread.join(), but there's a race condition where the thread might not stop cleanly. Additionally, if nest_asyncio is used, it patches the running event loop which can have global side effects.

Solution

Add proper exception handling for thread cleanup, ensure thread.join() has a timeout, and consider the implications of globally patching asyncio with nest_asyncio. Document that this function may have side effects on the global event loop state.

Changes

• checkpoint/orbax/checkpoint/_src/asyncio_utils.py (modified)

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.