Skip to content

bugfix: rollback shared prefix blocks on allocate failure.#1146

Open
RobbieLeung wants to merge 2 commits intojd-opensource:mainfrom
RobbieLeung:bugfix/scheduler
Open

bugfix: rollback shared prefix blocks on allocate failure.#1146
RobbieLeung wants to merge 2 commits intojd-opensource:mainfrom
RobbieLeung:bugfix/scheduler

Conversation

@RobbieLeung
Copy link
Copy Markdown
Collaborator

No description provided.

XuZhang99
XuZhang99 previously approved these changes Mar 31, 2026
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces rollback logic in BlockManagerPool::allocate to handle allocation failures by releasing shared prefix blocks and embedding IDs. It also adds a unit test to verify this functionality. Review feedback highlights a resource leak caused by the incorrect ordering of reset operations and a memory leak in the test suite resulting from the use of release() instead of reset().

later_sequence->num_tokens()));
EXPECT_EQ(later_sequence->kv_state().num_kv_blocks(), 1);

(void)engine.release();
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Using engine.release() here causes a memory leak of the FakeEngine object. More importantly, it bypasses the BlockManagerImpl destructor which contains a critical safety check (CHECK_EQ(num_free_blocks_, free_blocks_.size() - 1)) to ensure all blocks have been correctly freed. This check is essential for verifying that the rollback logic being tested actually works as expected and doesn't leak blocks. Use engine.reset() or simply let it go out of scope instead.

Suggested change
(void)engine.release();
engine.reset();

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
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.

2 participants