src/
Coverage:
low: ≥ 0%
medium: ≥ 75.0%
high: ≥ 90.0%
Lines:
137 / 185
74.1%
Branches:
88 / 254
34.6%
Functions:
17 / 19
89.5%

storage/external/detable/detable_engine.cc
Line Branch Exec Source
1 #include <dlfcn.h>
2
3 #include <array>
4 #include <cstring>
5 #include <stdexcept>
6 #include <string>
7 #include <string_view>
8 #include <vector>
9
10 #include "base/factory.h"
11 #include "storage/external/detable/dynamic_kv_plugin.h"
12 #include "storage/kv_engine/base_kv.h"
13
14 namespace {
15
16 using recstore::storage::plugin::DynamicKVPluginApiV1;
17 using recstore::storage::plugin::DynamicKVPluginConfigV1;
18
19 constexpr size_t kErrorBufferSize = 1024;
20
21 8 size_t ConfigValueSize(const BaseKVConfig& config) {
22 8 const auto& j = config.json_config_;
23
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
8 if (j.contains("value_size")) {
24 8 return j.at("value_size").get<size_t>();
25 }
26
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ if (j.contains("value")) {
27
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ return j.at("value").value("default_value_size_hint", 0);
28 }
29 ✗ return 0;
30 }
31
32 8 const json& PluginConfig(const BaseKVConfig& config) {
33
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
8 if (!config.json_config_.contains("detable")) {
34
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ throw std::invalid_argument("KVEngineDETable requires detable config");
35 }
36 8 return config.json_config_.at("detable");
37 }
38
39 6 void ValidateFloatAligned(size_t value_size, const char* operation) {
40
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
6 if (value_size == 0 || value_size % sizeof(float) != 0) {
41
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ throw std::invalid_argument(
42
0/4
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ std::string(operation) +
43 ✗ " requires a non-zero float-aligned value_size");
44 }
45 6 }
46
47 ✗ std::runtime_error PluginError(
48 const char* operation, const std::array<char, kErrorBufferSize>& error) {
49 const std::string detail =
50
0/14
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
✗ Branch 13 not taken.
✗ Branch 14 not taken.
✗ Branch 15 not taken.
✗ Branch 16 not taken.
✗ Branch 17 not taken.
✗ Branch 18 not taken.
✗ Branch 19 not taken.
✗ error[0] == '\0' ? "unknown provider error" : std::string(error.data());
51 return std::runtime_error(
52
0/10
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
✗ Branch 14 not taken.
✗ Branch 15 not taken.
✗ std::string("KVEngineDETable ") + operation + " failed: " + detail);
53 ✗ }
54
55 6 bool HasRequiredFunctions(const DynamicKVPluginApiV1& api) {
56
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
6 return api.create != nullptr && api.destroy != nullptr &&
57
3/6
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
6 api.get != nullptr && api.exists != nullptr && api.put != nullptr &&
58
3/6
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
18 api.batch_get != nullptr && api.batch_put != nullptr &&
59
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
12 api.clear != nullptr;
60 }
61
62 } // namespace
63
64 class KVEngineDETable : public BaseKV {
65 public:
66 8 explicit KVEngineDETable(const BaseKVConfig& config)
67
1/2
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
8 : BaseKV(config), value_size_(ConfigValueSize(config)) {
68
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
8 if (value_size_ == 0) {
69
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ throw std::invalid_argument("KVEngineDETable requires value_size");
70 }
71
72
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
8 const auto& plugin = PluginConfig(config);
73 const std::string library_path =
74
1/2
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
8 plugin.value("library_path", std::string());
75
2/2
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 6 times.
8 if (library_path.empty()) {
76
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
2 throw std::invalid_argument(
77 4 "KVEngineDETable requires detable.library_path");
78 }
79
80 6 library_ = dlopen(library_path.c_str(), RTLD_NOW | RTLD_LOCAL);
81
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
6 if (library_ == nullptr) {
82
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ throw std::runtime_error(
83
0/4
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ "KVEngineDETable failed to load provider: " + std::string(dlerror()));
84 }
85
86 6 dlerror();
87 auto entry =
88 reinterpret_cast< recstore::storage::plugin::GetDynamicKVPluginApiV1>(
89 6 dlsym(library_,
90 6 recstore::storage::plugin::kDynamicKVPluginEntryPoint));
91 6 const char* symbol_error = dlerror();
92
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
6 if (symbol_error != nullptr || entry == nullptr) {
93 ✗ CloseLibrary();
94
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ throw std::runtime_error(
95
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ "KVEngineDETable provider entry point is unavailable: " +
96
0/4
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ std::string(
97 ✗ symbol_error == nullptr ? "unknown error" : symbol_error));
98 }
99
100
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
6 api_ = entry();
101 18 if (api_ == nullptr ||
102
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
6 api_->abi_version != recstore::storage::plugin::kDynamicKVPluginAbiV1 ||
103
3/6
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 6 times.
18 api_->struct_size < sizeof(DynamicKVPluginApiV1) ||
104
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
6 !HasRequiredFunctions(*api_)) {
105 ✗ CloseLibrary();
106
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ throw std::runtime_error("KVEngineDETable provider ABI mismatch");
107 }
108
109 6 DynamicKVPluginConfigV1 plugin_config{
110 recstore::storage::plugin::kDynamicKVPluginAbiV1,
111 sizeof(DynamicKVPluginConfigV1),
112
2/4
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
6 config.json_config_.at("capacity").get<uint64_t>(),
113 6 value_size_,
114
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
6 plugin.value("block_size", uint64_t{10240})};
115 6 std::array<char, kErrorBufferSize> error{};
116
1/2
✓ Branch 3 taken 6 times.
✗ Branch 4 not taken.
6 if (api_->create(&plugin_config, &handle_, error.data(), error.size()) !=
117
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
12 0 ||
118
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
6 handle_ == nullptr) {
119
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ const auto exception = PluginError("create", error);
120
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ if (handle_ != nullptr) {
121
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ api_->destroy(handle_);
122 ✗ handle_ = nullptr;
123 }
124 ✗ CloseLibrary();
125 ✗ throw exception;
126 ✗ }
127 10 }
128
129 12 ~KVEngineDETable() override {
130
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
6 if (api_ != nullptr && handle_ != nullptr) {
131 6 api_->destroy(handle_);
132 }
133 6 CloseLibrary();
134 12 }
135
136 8 void Get(uint64_t key, std::string& value, unsigned tid) override {
137 (void)tid;
138
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
8 value.resize(value_size_);
139 8 uint8_t found = 0;
140 8 std::array<char, kErrorBufferSize> error{};
141
1/2
✓ Branch 3 taken 8 times.
✗ Branch 4 not taken.
16 if (api_->get(handle_,
142 key,
143 8 value.data(),
144 value_size_,
145 &found,
146 error.data(),
147
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
8 error.size()) != 0) {
148
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ throw PluginError("get", error);
149 }
150
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 4 times.
8 if (!found) {
151 4 value.clear();
152 }
153 8 }
154
155 2 bool Exists(uint64_t key, unsigned tid) override {
156 (void)tid;
157 2 uint8_t found = 0;
158 2 std::array<char, kErrorBufferSize> error{};
159
2/4
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 2 times.
2 if (api_->exists(handle_, key, &found, error.data(), error.size()) != 0) {
160
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ throw PluginError("exists", error);
161 }
162 2 return found != 0;
163 }
164
165 4 void Put(uint64_t key, const std::string_view& value, unsigned tid) override {
166 (void)tid;
167
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
4 if (value.size() != value_size_) {
168
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ throw std::invalid_argument("KVEngineDETable requires fixed-size Put");
169 }
170 4 std::array<char, kErrorBufferSize> error{};
171
1/2
✓ Branch 4 taken 4 times.
✗ Branch 5 not taken.
8 if (api_->put(handle_,
172 key,
173 4 value.data(),
174 value.size(),
175 error.data(),
176
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
4 error.size()) != 0) {
177
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ throw PluginError("put", error);
178 }
179 4 }
180
181 4 void BatchPut(base::ConstArray<uint64_t> keys,
182 std::vector<base::ConstArray<float>>* values,
183 unsigned tid) override {
184 (void)tid;
185
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
4 ValidateFloatAligned(value_size_, "KVEngineDETable::BatchPut");
186
3/6
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 4 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 4 times.
4 if (values == nullptr || keys.Size() != static_cast<int>(values->size())) {
187
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ throw std::invalid_argument("KVEngineDETable::BatchPut size mismatch");
188 }
189
1/2
✓ Branch 3 taken 4 times.
✗ Branch 4 not taken.
4 std::vector<char> flat(static_cast<size_t>(keys.Size()) * value_size_);
190 4 const int floats_per_row = static_cast<int>(value_size_ / sizeof(float));
191
2/2
✓ Branch 1 taken 10 times.
✓ Branch 2 taken 4 times.
14 for (int i = 0; i < keys.Size(); ++i) {
192
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 10 times.
10 if ((*values)[i].Size() != floats_per_row) {
193
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ throw std::invalid_argument(
194 ✗ "KVEngineDETable::BatchPut row size mismatch");
195 }
196 10 std::memcpy(flat.data() + static_cast<size_t>(i) * value_size_,
197 10 (*values)[i].Data(),
198 value_size_);
199 }
200
1/2
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
4 BatchPutRaw(keys, flat.data());
201 4 }
202
203 2 void BatchGet(base::ConstArray<uint64_t> keys,
204 std::vector<base::ConstArray<float>>* values,
205 unsigned tid) override {
206 (void)tid;
207
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
2 ValidateFloatAligned(value_size_, "KVEngineDETable::BatchGet");
208
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
2 if (values == nullptr) {
209
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ throw std::invalid_argument("KVEngineDETable::BatchGet values is null");
210 }
211
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
2 thread_local std::vector<std::vector<float>> buffers;
212 2 const int floats_per_row = static_cast<int>(value_size_ / sizeof(float));
213
2/4
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 6 taken 2 times.
✗ Branch 7 not taken.
2 buffers.assign(keys.Size(), std::vector<float>(floats_per_row));
214
1/2
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
2 std::vector<char> flat(static_cast<size_t>(keys.Size()) * value_size_);
215
1/2
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
2 std::vector<uint8_t> found(keys.Size(), 0);
216
1/2
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
2 BatchGetRaw(keys, flat.data(), found.data());
217
218
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
2 values->resize(keys.Size());
219
2/2
✓ Branch 1 taken 6 times.
✓ Branch 2 taken 2 times.
8 for (int i = 0; i < keys.Size(); ++i) {
220
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
6 if (!found[i]) {
221 ✗ (*values)[i] = base::ConstArray<float>();
222 ✗ continue;
223 }
224 6 std::memcpy(buffers[i].data(),
225 6 flat.data() + static_cast<size_t>(i) * value_size_,
226 value_size_);
227 6 (*values)[i] =
228 12 base::ConstArray<float>(buffers[i].data(), buffers[i].size());
229 }
230 2 }
231
232 2 bool BatchGetFlat(base::ConstArray<uint64_t> keys,
233 float* values,
234 int64_t num_rows,
235 int64_t embedding_dim,
236 unsigned tid,
237 BatchGetFlatStats* stats = nullptr) override {
238 (void)tid;
239
3/6
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
✗ Branch 6 not taken.
2 if (values == nullptr || num_rows != keys.Size() || num_rows < 0 ||
240
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
4 embedding_dim <= 0 ||
241
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
2 static_cast<size_t>(embedding_dim) * sizeof(float) != value_size_) {
242
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ throw std::invalid_argument(
243 ✗ "KVEngineDETable::BatchGetFlat shape mismatch");
244 }
245
1/2
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
2 std::vector<uint8_t> found(keys.Size(), 0);
246
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
2 BatchGetRaw(keys, values, found.data());
247 2 uint64_t missing = 0;
248
2/2
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 2 times.
6 for (int i = 0; i < keys.Size(); ++i) {
249
2/2
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 2 times.
4 if (!found[i]) {
250 2 std::memset(
251 2 values + static_cast<size_t>(i) * embedding_dim, 0, value_size_);
252 2 ++missing;
253 }
254 }
255
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
2 if (stats != nullptr) {
256 2 stats->missing_rows += missing;
257 }
258 2 return true;
259 2 }
260
261 ✗ void BulkLoad(base::ConstArray<uint64_t> keys, const void* value) override {
262
0/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ if (value == nullptr && keys.Size() > 0) {
263
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ throw std::invalid_argument("KVEngineDETable::BulkLoad value is null");
264 }
265 ✗ BatchPutRaw(keys, value);
266 ✗ }
267
268 2 void clear() override {
269 2 std::array<char, kErrorBufferSize> error{};
270
2/4
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 2 times.
2 if (api_->clear(handle_, error.data(), error.size()) != 0) {
271
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ throw PluginError("clear", error);
272 }
273 2 }
274
275 private:
276 void
277 4 BatchGetRaw(base::ConstArray<uint64_t> keys, void* values, uint8_t* found) {
278 4 std::array<char, kErrorBufferSize> error{};
279
1/2
✓ Branch 4 taken 4 times.
✗ Branch 5 not taken.
8 if (api_->batch_get(
280 handle_,
281 keys.Data(),
282 4 keys.Size(),
283 values,
284 value_size_,
285 found,
286 error.data(),
287
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
4 error.size()) != 0) {
288
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ throw PluginError("batch_get", error);
289 }
290 4 }
291
292 4 void BatchPutRaw(base::ConstArray<uint64_t> keys, const void* values) {
293 4 std::array<char, kErrorBufferSize> error{};
294
1/2
✓ Branch 4 taken 4 times.
✗ Branch 5 not taken.
8 if (api_->batch_put(
295 handle_,
296 keys.Data(),
297 4 keys.Size(),
298 values,
299 value_size_,
300 error.data(),
301
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
4 error.size()) != 0) {
302
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ throw PluginError("batch_put", error);
303 }
304 4 }
305
306 6 void CloseLibrary() {
307
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
6 if (library_ != nullptr) {
308 6 dlclose(library_);
309 6 library_ = nullptr;
310 }
311 6 }
312
313 size_t value_size_ = 0;
314 void* library_ = nullptr;
315 const DynamicKVPluginApiV1* api_ = nullptr;
316 void* handle_ = nullptr;
317 };
318
319 10 extern "C" void RecStoreForceLinkDETableEngine() {}
320
321 FACTORY_REGISTER(BaseKV, KVEngineDETable, KVEngineDETable, const BaseKVConfig&);
322