Fix race conditions, error recovery, and exit handlers in job servers

[shunping]

The following error occurs in the python test workflows (e.g. https://github.com/apache/beam/actions/runs/25133968815/job/73667329011?pr=38135) from time to time. It happens during the shutdown of a job server, but looking at the code path, it could be a problem in other servers like expansion service.

Exception ignored in atexit callback <bound method StopOnExitJobServer.stop of <apache_beam.runners.portability.job_server.StopOnExitJobServer object at 0x7c8f2326c590>>:
Traceback (most recent call last):
  File "/runner/_work/beam/beam/sdks/python/test-suites/tox/py313/build/srcs/sdks/python/target/.tox-py313/py313/lib/python3.13/site-packages/apache_beam/runners/portability/job_server.py", line 97, in stop
    self._job_server.stop()
  File "/runner/_work/beam/beam/sdks/python/test-suites/tox/py313/build/srcs/sdks/python/target/.tox-py313/py313/lib/python3.13/site-packages/apache_beam/runners/portability/job_server.py", line 127, in stop
    return self._server.stop()
  File "/runner/_work/beam/beam/sdks/python/test-suites/tox/py313/build/srcs/sdks/python/target/.tox-py313/py313/lib/python3.13/site-packages/apache_beam/utils/subprocess_server.py", line 255, in stop
    self.stop_process()
  File "/runner/_work/beam/beam/sdks/python/test-suites/tox/py313/build/srcs/sdks/python/target/.tox-py313/py313/lib/python3.13/site-packages/apache_beam/utils/subprocess_server.py", line 259, in stop_process
    self._cache.purge(self._owner_id)
  File "/runner/_work/beam/beam/sdks/python/test-suites/tox/py313/build/srcs/sdks/python/target/.tox-py313/py313/lib/python3.13/site-packages/apache_beam/utils/subprocess_server.py", line 93, in purge
    raise ValueError(f"{owner} not in {self._live_owners}")
ValueError: 10 not in set()

We identified three potential issues and implemented fixes to harden the server shutdown process.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses intermittent shutdown errors in job servers by hardening the cleanup process. It introduces better state management, ensures thread safety during cache operations, and prevents resource leaks or duplicate registrations by properly managing atexit handlers and ensuring cleanup routines complete reliably.

Highlights

• Atexit Handler Management: Added logic to unregister atexit handlers when a job server is stopped, preventing duplicate registrations during server restarts.
• Thread Safety in Cache Purging: Improved thread safety in the shared cache by wrapping the owner membership check and removal within a lock, preventing race conditions.
• Robust Error Recovery: Implemented try-finally blocks in server shutdown processes to ensure internal state, such as owner IDs, is correctly cleared even if errors occur during cleanup.

New Features

🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[shunping]

r: @derrickaw

[github-actions]

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

[gemini-code-assist]

Code Review

This pull request improves the robustness of job and subprocess server cleanup by using try-finally blocks to ensure state variables are reset during failures. It also fixes a race condition in the purge method by moving a membership check inside a lock and prevents duplicate atexit registrations. Review feedback recommends removing a redundant hasattr check for atexit.unregister, as it is supported in all targeted Python versions.

[shunping]

/gemini review

[gemini-code-assist]

Code Review

This pull request improves the robustness of job and subprocess server management by ensuring consistent state cleanup and thread safety. Key changes include moving the _live_owners membership check and removal inside a lock to prevent race conditions, and wrapping cleanup calls in try...finally blocks to ensure _started and _owner_id are reset even if errors occur. It also ensures atexit handlers are unregistered to avoid duplicates on restart. Comprehensive tests were added to cover these scenarios, including concurrent purge attempts and destructor failures. I have no feedback to provide.

[shunping]

cc'ed @tvalentyn

This seems promising to fix the "grpc deadline exceeds" error in our python test suites.