[Wheel] Gracefully stop Store REST service on SIGTERM#2874
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces graceful shutdown capabilities to the Mooncake Store Service by integrating signal handling and an asyncio shutdown event during both startup and execution. It also adds comprehensive unit tests to verify the shutdown behavior. The review feedback recommends ensuring that the store is properly closed and cleared if a shutdown is requested immediately after a successful setup to prevent resource leaks, setting self.store = None after closing it in the stop method to ensure idempotency, and adding a comment to explain the second call to _unblock_shutdown_signals() in main.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if shutdown_event is not None: | ||
| # setup() is synchronous. Give asyncio signal callbacks a | ||
| # chance to publish a shutdown requested while it ran. | ||
| await asyncio.sleep(0) | ||
| if shutdown_event.is_set(): | ||
| logging.info( | ||
| "Store startup cancelled by shutdown request" | ||
| ) | ||
| return False |
There was a problem hiding this comment.
If shutdown_event is set immediately after self.store.setup() succeeds, the method returns False but leaves self.store initialized and active. This can cause resource leaks if the caller does not call stop(). To prevent this, we should close the store and set self.store = None before returning False.
| if shutdown_event is not None: | |
| # setup() is synchronous. Give asyncio signal callbacks a | |
| # chance to publish a shutdown requested while it ran. | |
| await asyncio.sleep(0) | |
| if shutdown_event.is_set(): | |
| logging.info( | |
| "Store startup cancelled by shutdown request" | |
| ) | |
| return False | |
| if shutdown_event is not None: | |
| # setup() is synchronous. Give asyncio signal callbacks a | |
| # chance to publish a shutdown requested while it ran. | |
| await asyncio.sleep(0) | |
| if shutdown_event.is_set(): | |
| logging.info( | |
| "Store startup cancelled by shutdown request" | |
| ) | |
| if hasattr(self.store, "close"): | |
| self.store.close() | |
| self.store = None | |
| return False |
| ret = self.store.close() | ||
| if ret != 0: | ||
| logging.warning("Mooncake service close returned %s", ret) | ||
| else: | ||
| logging.info("Mooncake service stopped") |
There was a problem hiding this comment.
To make the stop() method idempotent and prevent potential double-close issues or operations on a closed store, set self.store = None immediately after calling self.store.close().
| ret = self.store.close() | |
| if ret != 0: | |
| logging.warning("Mooncake service close returned %s", ret) | |
| else: | |
| logging.info("Mooncake service stopped") | |
| ret = self.store.close() | |
| self.store = None | |
| if ret != 0: | |
| logging.warning("Mooncake service close returned %s", ret) | |
| else: | |
| logging.info("Mooncake service stopped") |
| _unblock_shutdown_signals() | ||
| await asyncio.sleep(0) | ||
| if shutdown_event.is_set(): | ||
| return |
There was a problem hiding this comment.
The call to _unblock_shutdown_signals() is repeated here after start_store_service. While this is likely intended to restore the signal mask in case the synchronous C++ store setup modified it, adding a brief comment explaining this non-obvious behavior would greatly improve maintainability.
| _unblock_shutdown_signals() | |
| await asyncio.sleep(0) | |
| if shutdown_event.is_set(): | |
| return | |
| # Unblock signals again in case the C++ store setup modified the signal mask. | |
| _unblock_shutdown_signals() | |
| await asyncio.sleep(0) | |
| if shutdown_event.is_set(): | |
| return |
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
* fix(wheel): handle store service shutdown signals * fix(wheel): address shutdown review feedback
4 participants
Description
The Store REST service currently relies on
KeyboardInterruptaround a polling loop for shutdown. Process managers commonly terminate services withSIGTERM, and a shutdown request received during startup can remain pending while the service is retrying or proceed into HTTP startup before the asyncio signal callback runs. Cleanup is also duplicated across exception paths instead of being guaranteed for every exit path.This PR:
SIGINTandSIGTERMand unblocks inherited signal masks;asyncio.Eventto drive the service lifetime;finallyblock and reports nonzero close results.MooncakeDistributedStore.setup()remains synchronous and is not cancelled while it is executing. A pending shutdown request is honored as soon as setup returns.Module
mooncake-transfer-engine)mooncake-store)mooncake-ep)mooncake-pg)mooncake-integration)mooncake-p2p-store)mooncake-wheel)mooncake-common)mooncake-rl)Type of Change
How Has This Been Tested?
The shutdown tests cover signal-handler dispatch, cancellation before and after Store setup, interruption of retry sleeps, guaranteed cleanup from
main(), and nonzero close results.Test commands:
Test results:
Checklist
./scripts/code_format.sh(no C/C++ files changed; Python files pass Ruff format check)pre-commit run --all-filesand all hooks passAI Assistance Disclosure
Codex assisted with implementation, test design, and validation. The submitter reviewed the resulting changes.