fix(server): guard finish_read_prefetched behind retrieve_succeeded flag#2736
Conversation
When _retrieve_loop raises an exception inside the with-block, the read_prefetched_results context manager already releases the read locks in its finally clause. The retrieve() finally block was unconditionally scheduling a GPU callback to call finish_read_prefetched, causing a double-free of read locks. Introduce retrieve_succeeded flag that is only set to True when the with-block exits normally. The GPU callback is now conditional on this flag, forming a defense-in-depth pair with the context manager re-raise fix. Signed-off-by: baoloongmao <baoloongmao@tencent.com>
Summary of ChangesHello, I'm Gemini Code Assist1! 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 a critical bug in the Highlights
🧠 New Feature in Public Preview: You can now enable Memory 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. Changelog
Activity
Using Gemini Code AssistThe 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
Customization To customize 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 Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. 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. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
There was a problem hiding this comment.
Code Review
The pull request effectively addresses a critical bug where finish_read_prefetched could be called twice, leading to a double-free of read locks. By introducing the retrieve_succeeded flag, the GPU callback to release read locks is now correctly conditional on the successful exit of the read_prefetched_results context manager. This provides a robust defense-in-depth solution, preventing potential issues related to resource management and ensuring the stability of the cache engine.
The context manager was swallowing exceptions raised inside the with-block (e.g. TypeError from _retrieve_loop). This caused the caller to continue as if the retrieve succeeded, leading to a double-free of read locks: once by the context manager's finally clause, and again by the caller's GPU callback. Add 'raise' to propagate the exception so callers correctly enter their except branch and skip the redundant finish_read_prefetched. This change was originally in PR LMCache#2737 and is now combined into this PR per reviewer's suggestion to streamline the merging process. Signed-off-by: baoloongmao <baoloongmao@tencent.com>
|
@KuntaiDu Can you take another look? |
…lag (LMCache#2736) * fix(server): guard finish_read_prefetched behind retrieve_succeeded flag When _retrieve_loop raises an exception inside the with-block, the read_prefetched_results context manager already releases the read locks in its finally clause. The retrieve() finally block was unconditionally scheduling a GPU callback to call finish_read_prefetched, causing a double-free of read locks. Introduce retrieve_succeeded flag that is only set to True when the with-block exits normally. The GPU callback is now conditional on this flag, forming a defense-in-depth pair with the context manager re-raise fix. Signed-off-by: baoloongmao <baoloongmao@tencent.com> * fix(storage_manager): re-raise exception in read_prefetched_results The context manager was swallowing exceptions raised inside the with-block (e.g. TypeError from _retrieve_loop). This caused the caller to continue as if the retrieve succeeded, leading to a double-free of read locks: once by the context manager's finally clause, and again by the caller's GPU callback. Add 'raise' to propagate the exception so callers correctly enter their except branch and skip the redundant finish_read_prefetched. This change was originally in PR LMCache#2737 and is now combined into this PR per reviewer's suggestion to streamline the merging process. Signed-off-by: baoloongmao <baoloongmao@tencent.com> --------- Signed-off-by: baoloongmao <baoloongmao@tencent.com> Signed-off-by: Aaron Wu <aaron.wu@dell.com>
…lag (LMCache#2736) * fix(server): guard finish_read_prefetched behind retrieve_succeeded flag When _retrieve_loop raises an exception inside the with-block, the read_prefetched_results context manager already releases the read locks in its finally clause. The retrieve() finally block was unconditionally scheduling a GPU callback to call finish_read_prefetched, causing a double-free of read locks. Introduce retrieve_succeeded flag that is only set to True when the with-block exits normally. The GPU callback is now conditional on this flag, forming a defense-in-depth pair with the context manager re-raise fix. Signed-off-by: baoloongmao <baoloongmao@tencent.com> * fix(storage_manager): re-raise exception in read_prefetched_results The context manager was swallowing exceptions raised inside the with-block (e.g. TypeError from _retrieve_loop). This caused the caller to continue as if the retrieve succeeded, leading to a double-free of read locks: once by the context manager's finally clause, and again by the caller's GPU callback. Add 'raise' to propagate the exception so callers correctly enter their except branch and skip the redundant finish_read_prefetched. This change was originally in PR LMCache#2737 and is now combined into this PR per reviewer's suggestion to streamline the merging process. Signed-off-by: baoloongmao <baoloongmao@tencent.com> --------- Signed-off-by: baoloongmao <baoloongmao@tencent.com>
…lag (LMCache#2736) * fix(server): guard finish_read_prefetched behind retrieve_succeeded flag When _retrieve_loop raises an exception inside the with-block, the read_prefetched_results context manager already releases the read locks in its finally clause. The retrieve() finally block was unconditionally scheduling a GPU callback to call finish_read_prefetched, causing a double-free of read locks. Introduce retrieve_succeeded flag that is only set to True when the with-block exits normally. The GPU callback is now conditional on this flag, forming a defense-in-depth pair with the context manager re-raise fix. Signed-off-by: baoloongmao <baoloongmao@tencent.com> * fix(storage_manager): re-raise exception in read_prefetched_results The context manager was swallowing exceptions raised inside the with-block (e.g. TypeError from _retrieve_loop). This caused the caller to continue as if the retrieve succeeded, leading to a double-free of read locks: once by the context manager's finally clause, and again by the caller's GPU callback. Add 'raise' to propagate the exception so callers correctly enter their except branch and skip the redundant finish_read_prefetched. This change was originally in PR LMCache#2737 and is now combined into this PR per reviewer's suggestion to streamline the merging process. Signed-off-by: baoloongmao <baoloongmao@tencent.com> --------- Signed-off-by: baoloongmao <baoloongmao@tencent.com>
When _retrieve_loop raises an exception inside the with-block, the read_prefetched_results context manager already releases the read locks in its finally clause. The retrieve() finally block was unconditionally scheduling a GPU callback to call
finish_read_prefetched, causing a double-free of read locks.
Introduce retrieve_succeeded flag that is only set to True when the with-block exits normally. The GPU callback is now conditional on this flag, forming a defense-in-depth pair with the context manager re-raise fix.
What this PR does / why we need it:
In MPCacheEngine.retrieve(), the finally block unconditionally schedules a GPU callback to call finish_read_prefetched, which releases the read lock. When _retrieve_loop raises an exception, the read_prefetched_results context manager already releases the lock in its own finally clause. However, Python's finally block executes even after a return in an except block, so the GPU callback is still scheduled, causing a read lock double-free.
Fix: Introduce a retrieve_succeeded flag initialized to False, set to True only after the with block exits normally. The GPU callback in finally is now conditional on this flag. This forms a defense-in-depth pair with the context manager re-raise fix (PR #2737) — either fix alone prevents the double-free, but together they provide two independent safety layers.
Special notes for your reviewers:
If applicable: