wehub-resource-sync
258 lines
8.2 KiB
C++
258 lines
8.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 <filesystem>
|
|
#include <thread>
|
|
#ifdef __ANDROID__
|
|
#include <unistd.h> // _exit()
|
|
#endif
|
|
#include <zvec/ailego/internal/platform.h>
|
|
#include <zvec/db/collection.h>
|
|
#include "zvec/ailego/logger/logger.h"
|
|
#include "utility.h"
|
|
|
|
|
|
constexpr int kBatchSize = 20;
|
|
constexpr int kBatchDelayMs = 10;
|
|
|
|
|
|
struct Config {
|
|
std::string path;
|
|
int start_id = 0;
|
|
int end_id = 0;
|
|
std::string operation; // "insert", "upsert", "update", "delete"
|
|
int version = 999999;
|
|
};
|
|
|
|
|
|
bool ParseArgs(int argc, char **argv, Config &config) {
|
|
for (int i = 1; i < argc; i++) {
|
|
std::string arg = argv[i];
|
|
|
|
if (arg == "--path" && i + 1 < argc) {
|
|
config.path = argv[++i];
|
|
} else if (arg == "--start" && i + 1 < argc) {
|
|
config.start_id = std::stoi(argv[++i]);
|
|
} else if (arg == "--end" && i + 1 < argc) {
|
|
config.end_id = std::stoi(argv[++i]);
|
|
} else if (arg == "--op" && i + 1 < argc) {
|
|
config.operation = argv[++i];
|
|
} else if (arg == "--version" && i + 1 < argc) {
|
|
config.version = std::stoi(argv[++i]);
|
|
} else if (arg == "--help" || arg == "-h") {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
// Validate required arguments
|
|
if (config.path.empty() || config.operation.empty() ||
|
|
config.start_id >= config.end_id || config.version == 999999) {
|
|
return false;
|
|
}
|
|
|
|
// Validate operation
|
|
if (config.operation != "insert" && config.operation != "upsert" &&
|
|
config.operation != "update" && config.operation != "delete") {
|
|
std::cerr << "Error: Invalid operation '" << config.operation
|
|
<< "'. Must be 'insert', 'upsert', 'update', or 'delete'."
|
|
<< std::endl;
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
|
|
void PrintUsage(const char *program) {
|
|
std::cout << "Usage: " << program
|
|
<< " --path <collection_path> --start <start_id> --end <end_id> "
|
|
"--op <operation>"
|
|
<< std::endl;
|
|
std::cout << std::endl;
|
|
std::cout << "Arguments:" << std::endl;
|
|
std::cout << " --path Path to the collection (required)" << std::endl;
|
|
std::cout << " --start Starting document ID (inclusive, required)"
|
|
<< std::endl;
|
|
std::cout << " --end Ending document ID (exclusive, required)"
|
|
<< std::endl;
|
|
std::cout
|
|
<< " --op Operation: insert, upsert, update, or delete (required)"
|
|
<< std::endl;
|
|
std::cout << " --version Operation: version (required)" << std::endl;
|
|
std::cout << std::endl;
|
|
std::cout << "Examples:" << std::endl;
|
|
std::cout << " # Insert 1000 documents (pk_0 to pk_999)" << std::endl;
|
|
std::cout << " " << program
|
|
<< " --path ./test_db --start 0 --end 1000 --op insert --version 0"
|
|
<< std::endl;
|
|
std::cout << std::endl;
|
|
std::cout << " # Update documents 1000-1999" << std::endl;
|
|
std::cout
|
|
<< " " << program
|
|
<< " --path ./test_db --start 1000 --end 2000 --op update --version 1"
|
|
<< std::endl;
|
|
std::cout << std::endl;
|
|
std::cout << " # Upsert documents 0-499" << std::endl;
|
|
std::cout << " " << program
|
|
<< " --path ./test_db --start 0 --end 500 --op upsert --version 2"
|
|
<< std::endl;
|
|
}
|
|
|
|
|
|
int main(int argc, char **argv) {
|
|
Config config;
|
|
|
|
// Parse arguments
|
|
if (!ParseArgs(argc, argv, config)) {
|
|
PrintUsage(argv[0]);
|
|
#ifdef __ANDROID__
|
|
_exit(1);
|
|
#endif
|
|
return 1;
|
|
}
|
|
|
|
try {
|
|
std::filesystem::path cwd = std::filesystem::current_path();
|
|
std::cout << "[data_generator] Current Working Directory: " << cwd.string()
|
|
<< std::endl;
|
|
} catch (const std::filesystem::filesystem_error &e) {
|
|
std::cout
|
|
<< "[data_generator] Failed to get the current working directory: "
|
|
<< e.what() << std::endl;
|
|
}
|
|
|
|
std::cout << "Configuration:" << std::endl;
|
|
std::cout << " Path: " << config.path << std::endl;
|
|
std::cout << " Range: [" << config.start_id << ", " << config.end_id
|
|
<< ")" << std::endl;
|
|
std::cout << " Operation: " << config.operation << std::endl;
|
|
std::cout << " BatchSize: " << kBatchSize << std::endl;
|
|
std::cout << " BatchDelay: " << kBatchDelayMs << "ms" << std::endl;
|
|
std::cout << std::endl;
|
|
|
|
// Scope 'result' so its shared_ptr is released before we call _exit().
|
|
// Without scoping, result.value() returns a reference and the copy means
|
|
// result still owns a second shared_ptr; collection.reset() would only
|
|
// drop the refcount to 1, and _exit() would skip the Collection destructor,
|
|
// leaving the WAL dirty.
|
|
zvec::Collection::Ptr collection;
|
|
{
|
|
auto result = zvec::Collection::Open(
|
|
config.path, zvec::CollectionOptions{false, true, 4 * 1024 * 1024});
|
|
if (!result) {
|
|
LOG_ERROR("Failed to open collection[%s]: %s", config.path.c_str(),
|
|
result.error().c_str());
|
|
#ifdef __ANDROID__
|
|
_exit(1);
|
|
#endif
|
|
return -1;
|
|
}
|
|
collection = result.value();
|
|
}
|
|
LOG_INFO("Collection[%s] opened successfully", config.path.c_str());
|
|
|
|
// Process documents in batches
|
|
int total_docs = config.end_id - config.start_id;
|
|
int processed = 0;
|
|
int batch_num = 0;
|
|
int next_progress_threshold = total_docs / 10; // 10% increments
|
|
int progress_percent = 0;
|
|
|
|
while (config.start_id < config.end_id) {
|
|
int batch_end = std::min(config.start_id + kBatchSize, config.end_id);
|
|
int batch_count = batch_end - config.start_id;
|
|
|
|
std::vector<zvec::Doc> docs;
|
|
docs.reserve(batch_count);
|
|
for (int i = config.start_id; i < batch_end; i++) {
|
|
docs.push_back(zvec::CreateTestDoc(i, config.version));
|
|
}
|
|
|
|
zvec::Result<zvec::WriteResults> results;
|
|
if (config.operation == "insert") {
|
|
results = collection->Insert(docs);
|
|
} else if (config.operation == "upsert") {
|
|
results = collection->Upsert(docs);
|
|
} else if (config.operation == "update") {
|
|
results = collection->Update(docs);
|
|
} else if (config.operation == "delete") {
|
|
std::vector<std::string> pks{};
|
|
for (const auto &doc : docs) {
|
|
pks.emplace_back(doc.pk());
|
|
}
|
|
results = collection->Delete(pks);
|
|
}
|
|
if (!results) {
|
|
LOG_ERROR("Failed to perform operation[%s], reason: %s",
|
|
config.operation.c_str(), results.error().message().c_str());
|
|
#ifdef __ANDROID__
|
|
collection.reset();
|
|
sync();
|
|
_exit(1);
|
|
#endif
|
|
return 1;
|
|
}
|
|
for (auto &s : results.value()) {
|
|
if (!s.ok()) {
|
|
LOG_ERROR("Failed to perform operation[%s], reason: %s",
|
|
config.operation.c_str(), s.message().c_str());
|
|
#ifdef __ANDROID__
|
|
collection.reset();
|
|
sync();
|
|
_exit(1);
|
|
#endif
|
|
return 1;
|
|
}
|
|
}
|
|
|
|
processed += batch_count;
|
|
config.start_id = batch_end;
|
|
batch_num++;
|
|
|
|
// Print progress every 10%
|
|
if (processed >= next_progress_threshold) {
|
|
progress_percent++;
|
|
LOG_INFO("Progress: %d (%d/%d documents)", progress_percent * 10,
|
|
processed, total_docs);
|
|
next_progress_threshold = (progress_percent + 1) * total_docs / 10;
|
|
}
|
|
|
|
// Sleep between batches
|
|
if (config.start_id < config.end_id) {
|
|
std::this_thread::sleep_for(std::chrono::milliseconds(kBatchDelayMs));
|
|
}
|
|
}
|
|
|
|
std::cout << std::endl;
|
|
std::cout << "Success! Processed " << processed << " documents in "
|
|
<< batch_num << " batches." << std::endl;
|
|
|
|
#ifdef __ANDROID__
|
|
// On Android with c++_static STL, static destructors of glog/gflags
|
|
// crash during process teardown. Use _exit() to skip them.
|
|
// Flush + close the collection to persist all buffered data (WAL, memtable),
|
|
// then sync() to ensure kernel buffers are written to disk before _exit().
|
|
collection->Flush();
|
|
collection.reset();
|
|
sync();
|
|
std::cout.flush();
|
|
std::cerr.flush();
|
|
fflush(stdout);
|
|
fflush(stderr);
|
|
_exit(0);
|
|
#endif
|
|
return 0;
|
|
}
|