fix(quantization): guard TransformQuantizer ProcessQueryImpl against repeated allocation#2164
Conversation
…repeated allocation Add nullptr check before allocating inner_computer buf in ProcessQueryImpl, consistent with all other quantizer implementations. Without this guard, repeated SetQuery calls on the same Computer object leak query_code_size_ bytes per call. Add regression test verifying repeated SetQuery does not crash or leak. Signed-off-by: LHT129 <tianlan.lht@antgroup.com> Co-authored-by: opencode <opencode@anthropic.com>
Merge ProtectionsYour pull request matches the following merge protections and will not be merged until they are valid. 🟢 Require kind labelWonderful, this rule succeeded.
🟢 Require version labelWonderful, this rule succeeded.
🟢 Require linked issue for feature/bug PRsWonderful, this rule succeeded.
|
There was a problem hiding this comment.
Code Review
This pull request prevents memory leaks in TransformQuantizer during repeated SetQuery calls by checking if the buffer computer.inner_computer_->buf_ is already allocated before allocating memory, and adds a corresponding unit test. The reviewer suggests adding a nullptr check immediately after allocation to handle cases where the allocator returns nullptr instead of throwing std::bad_alloc, ensuring robust error handling.
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 (computer.inner_computer_->buf_ == nullptr) { | ||
| computer.inner_computer_->buf_ = | ||
| reinterpret_cast<uint8_t*>(this->allocator_->Allocate(this->query_code_size_)); | ||
| } |
There was a problem hiding this comment.
While checking for nullptr before allocation prevents repeated allocation, we should also handle the case where the allocator returns nullptr instead of throwing std::bad_alloc (which is common for some custom allocators or when exceptions are disabled). Adding a nullptr check immediately after allocation ensures robust error handling and prevents potential null pointer dereferences later in ExecuteChainTransform.
if (computer.inner_computer_->buf_ == nullptr) {
computer.inner_computer_->buf_ =
reinterpret_cast<uint8_t*>(this->allocator_->Allocate(this->query_code_size_));
if (computer.inner_computer_->buf_ == nullptr) {
throw VsagException(ErrorType::NO_ENOUGH_MEMORY, "alloc return nullptr when init computer buf");
}
}There was a problem hiding this comment.
Pull request overview
This PR fixes a memory leak in TransformQuantizer::ProcessQueryImpl when SetQuery() is invoked multiple times on the same Computer instance by guarding the internal query-buffer allocation with a nullptr check. It also adds a regression unit test intended to catch the leak under ASan/LSan CI runs (as used by this repo’s workflows).
Changes:
- Guard
computer.inner_computer_->buf_allocation to avoid overwriting an existing allocation on repeatedSetQuery()calls. - Add a regression unit test that repeatedly calls
SetQuery()on the sameComputerand performs a distance computation afterward.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
src/quantization/transform_quantization/transform_quantizer.h |
Prevents repeated allocation of inner_computer_->buf_ in ProcessQueryImpl to eliminate leaks on repeated SetQuery(). |
src/quantization/transform_quantization/transform_quantizer_test.cpp |
Adds a regression test that exercises repeated SetQuery() on the same Computer instance (useful for leak detection under sanitizers). |
Summary
inner_computer->buf_inTransformQuantizer::ProcessQueryImpl, preventing memory leak on repeatedSetQuerycallsSetQueryon the same Computer objectFixes: #2163