Skip to content

fix(quantization): guard TransformQuantizer ProcessQueryImpl against repeated allocation#2164

Open
LHT129 wants to merge 1 commit into
antgroup:mainfrom
LHT129:fix/transform-quantizer-processquery-memleak
Open

fix(quantization): guard TransformQuantizer ProcessQueryImpl against repeated allocation#2164
LHT129 wants to merge 1 commit into
antgroup:mainfrom
LHT129:fix/transform-quantizer-processquery-memleak

Conversation

@LHT129

@LHT129 LHT129 commented Jun 8, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add nullptr check before allocating inner_computer->buf_ in TransformQuantizer::ProcessQueryImpl, preventing memory leak on repeated SetQuery calls
  • Consistent with all other quantizer implementations (FP32, INT8, SQ8Uniform, etc.)
  • Add regression test exercising repeated SetQuery on the same Computer object

Fixes: #2163

…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>
Copilot AI review requested due to automatic review settings June 8, 2026 08:39
@LHT129 LHT129 self-assigned this Jun 8, 2026
@mergify

mergify Bot commented Jun 8, 2026

Copy link
Copy Markdown
Contributor

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 Require kind label

Wonderful, this rule succeeded.
  • label~=^kind/

🟢 Require version label

Wonderful, this rule succeeded.
  • label~=^version/

🟢 Require linked issue for feature/bug PRs

Wonderful, this rule succeeded.
  • body~=(?im)(?:^|[\s\-\*])(?:close[sd]?|fix(?:e[sd])?|resolve[sd]?)\s*:?\s+(?:#\d+|[\w.\-]+/[\w.\-]+#\d+|https?://github\.com/[\w.\-]+/[\w.\-]+/issues/\d+)

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

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.

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.

Comment on lines +287 to +290
if (computer.inner_computer_->buf_ == nullptr) {
computer.inner_computer_->buf_ =
reinterpret_cast<uint8_t*>(this->allocator_->Allocate(this->query_code_size_));
}

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.

medium

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");
            }
        }

Copilot AI left a comment

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.

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 repeated SetQuery() calls.
  • Add a regression unit test that repeatedly calls SetQuery() on the same Computer and 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).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Memory leak in TransformQuantizer::ProcessQueryImpl on repeated SetQuery calls

2 participants