134 lines
4.2 KiB
C++
134 lines
4.2 KiB
C++
// Copyright 2025-present the zvec project
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#include <cstdint>
|
|
#include <thread>
|
|
#include <gtest/gtest.h>
|
|
#include <zvec/ailego/buffer/buffer_manager.h>
|
|
#include <zvec/ailego/logger/logger.h>
|
|
#include "tests/test_util.h"
|
|
|
|
#if defined(__GNUC__) || defined(__GNUG__)
|
|
#pragma GCC diagnostic push
|
|
#pragma GCC diagnostic ignored "-Wunused-result"
|
|
#endif
|
|
|
|
using namespace zvec::ailego;
|
|
|
|
|
|
const std::string working_dir{"./buffer_manager_dir/"};
|
|
const std::string file_path_forward{working_dir + "test.forward_index"};
|
|
const std::string file_path_vector{working_dir + "test.vector_index"};
|
|
|
|
|
|
class BufferManagerTest : public testing::Test {
|
|
/***** Global initialization and cleanup - Start *****/
|
|
public:
|
|
static void SetUpTestCase() {
|
|
zvec::test_util::RemoveTestPath(working_dir);
|
|
|
|
if (!File::MakePath(working_dir)) {
|
|
LOG_ERROR("Failed to create working directory.");
|
|
return;
|
|
}
|
|
|
|
File file_vector_index;
|
|
size_t file_vector_size = 16 * 1024 * 1024;
|
|
if (!file_vector_index.create(file_path_vector, file_vector_size)) {
|
|
LOG_ERROR("Failed to create vector index file.");
|
|
return;
|
|
}
|
|
// Populate vector file with number series
|
|
for (uint32_t i = 0; i < file_vector_size / sizeof(uint32_t); ++i) {
|
|
file_vector_index.write((void *)&i, sizeof(i));
|
|
}
|
|
file_vector_index.close();
|
|
|
|
BufferManager::Instance().init(4 * 1024 * 1024, 1);
|
|
}
|
|
|
|
static void TearDownTestCase() {
|
|
BufferManager::Instance().cleanup();
|
|
zvec::test_util::RemoveTestPath(working_dir);
|
|
}
|
|
/***** Global initialization and cleanup - End *****/
|
|
;
|
|
};
|
|
|
|
|
|
TEST_F(BufferManagerTest, READ_VECTOR_FILE) {
|
|
uint32_t size_4KB = 4 * 1024;
|
|
|
|
auto read_and_verify_numbers = [&](uint32_t offset) {
|
|
BufferID id = BufferID::VectorID(file_path_vector, offset, size_4KB);
|
|
auto handle = BufferManager::Instance().acquire(id);
|
|
uint32_t *vector_data = (uint32_t *)handle.pin_vector_data();
|
|
uint32_t num_start = offset / sizeof(uint32_t);
|
|
for (uint32_t i = 0; i < size_4KB / sizeof(uint32_t); i++) {
|
|
ASSERT_EQ(*(vector_data + i), num_start + i);
|
|
}
|
|
handle.unpin_vector_data();
|
|
};
|
|
|
|
std::vector<std::thread> threads;
|
|
|
|
// Read the same part concurrently
|
|
for (int i = 0; i < 10; ++i) {
|
|
threads.emplace_back(read_and_verify_numbers, 3 * size_4KB);
|
|
}
|
|
for (auto &thread : threads) {
|
|
thread.join();
|
|
}
|
|
|
|
{ // Verify the reference count
|
|
BufferID id = BufferID::VectorID(file_path_vector, 3 * size_4KB, size_4KB);
|
|
auto handle = BufferManager::Instance().acquire(id);
|
|
handle.pin_vector_data();
|
|
ASSERT_EQ(handle.references(), 1);
|
|
handle.unpin_vector_data();
|
|
ASSERT_EQ(handle.references(), 0);
|
|
}
|
|
|
|
threads.clear();
|
|
// Read different parts concurrently
|
|
for (int i = 0; i < 30; ++i) {
|
|
threads.emplace_back(read_and_verify_numbers, i * size_4KB);
|
|
}
|
|
for (auto &thread : threads) {
|
|
thread.join();
|
|
}
|
|
ASSERT_EQ(BufferManager::Instance().total_size_in_bytes(), 30 * 4 * 1024);
|
|
|
|
{ // Read a large chunk so that the buffer is full
|
|
BufferID id =
|
|
BufferID::VectorID(file_path_vector, 4 * 1024 * 1024, 4 * 1024 * 1024);
|
|
auto handle = BufferManager::Instance().acquire(id);
|
|
handle.pin_vector_data();
|
|
handle.unpin_vector_data();
|
|
}
|
|
|
|
{ // Trigger eviction
|
|
BufferID id =
|
|
BufferID::VectorID(file_path_vector, 8 * 1024 * 1024, 4 * 1024 * 1024);
|
|
auto handle = BufferManager::Instance().acquire(id);
|
|
handle.pin_vector_data();
|
|
ASSERT_EQ(BufferManager::Instance().total_size_in_bytes(), 4 * 1024 * 1024);
|
|
handle.unpin_vector_data();
|
|
ASSERT_EQ(handle.references(), 0);
|
|
}
|
|
}
|
|
|
|
#if defined(__GNUC__) || defined(__GNUG__)
|
|
#pragma GCC diagnostic pop
|
|
#endif |