From e59287ec6887c7f9bf4aba384b6501db5d9821a0 Mon Sep 17 00:00:00 2001 From: LHT129 Date: Mon, 8 Jun 2026 16:36:34 +0800 Subject: [PATCH] fix(quantization): guard TransformQuantizer ProcessQueryImpl against 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 Co-authored-by: opencode --- .../transform_quantizer.h | 6 ++- .../transform_quantizer_test.cpp | 40 +++++++++++++++++++ 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/src/quantization/transform_quantization/transform_quantizer.h b/src/quantization/transform_quantization/transform_quantizer.h index f53fdbbe0..3ee545279 100644 --- a/src/quantization/transform_quantization/transform_quantizer.h +++ b/src/quantization/transform_quantization/transform_quantizer.h @@ -284,8 +284,10 @@ TransformQuantizer::ProcessQueryImpl( const float* query, Computer& computer) const { // 0. allocate try { - computer.inner_computer_->buf_ = - reinterpret_cast(this->allocator_->Allocate(this->query_code_size_)); + if (computer.inner_computer_->buf_ == nullptr) { + computer.inner_computer_->buf_ = + reinterpret_cast(this->allocator_->Allocate(this->query_code_size_)); + } } catch (const std::bad_alloc& e) { computer.inner_computer_->buf_ = nullptr; throw VsagException(ErrorType::NO_ENOUGH_MEMORY, "bad alloc when init computer buf"); diff --git a/src/quantization/transform_quantization/transform_quantizer_test.cpp b/src/quantization/transform_quantization/transform_quantizer_test.cpp index d9d7ddb91..4050d8f8a 100644 --- a/src/quantization/transform_quantization/transform_quantizer_test.cpp +++ b/src/quantization/transform_quantization/transform_quantizer_test.cpp @@ -124,3 +124,43 @@ TEST_CASE("TQ Serialize and Deserialize", "[ut][TransformQuantizer]") { } } } + +TEST_CASE("TQ Repeated SetQuery No Leak", "[ut][TransformQuantizer]") { + constexpr MetricType metric = MetricType::METRIC_TYPE_L2SQR; + uint64_t dim = 128; + uint64_t count = 200; + + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + auto param = std::make_shared(); + static constexpr const char* param_template = R"( + {{ + "tq_chain": "rom, fp32", + "pca_dim": {}, + "mrle_dim": {} + }} + )"; + auto param_str = fmt::format(param_template, dim, dim - 1); + auto param_json = vsag::JsonType::Parse(param_str); + param->FromJson(param_json); + + IndexCommonParam common_param; + common_param.allocator_ = allocator; + common_param.dim_ = dim; + TransformQuantizer, metric> quantizer(param, common_param); + + auto vecs = fixtures::generate_vectors(count, dim); + quantizer.Train(vecs.data(), count); + + auto computer = quantizer.FactoryComputer(); + auto queries = fixtures::generate_vectors(10, dim, false, 42); + + for (int i = 0; i < 10; ++i) { + computer->SetQuery(queries.data() + i * dim); + } + + std::vector codes(quantizer.GetCodeSize()); + quantizer.EncodeOne(vecs.data(), codes.data()); + float dist = 0; + quantizer.ComputeDist(*computer, codes.data(), &dist); + REQUIRE(dist >= 0); +}