GCC Code Coverage Report

Directory: src/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 80.9% 419 / 0 / 518
Functions: 83.3% 25 / 0 / 30
Branches: 42.9% 374 / 0 / 872
storage/kv_engine/engine_composite.h
Line Branch Exec Source
1 #pragma once
2
3 #include <algorithm>
4 #include <array>
5 #include <chrono>
6 #include <cstring>
7 #include <filesystem>
8 #include <fstream>
9 #include <limits>
10 #include <memory>
11 #include <mutex>
12 #include <shared_mutex>
13 #include <stdexcept>
14 #include <string>
15 #include <unordered_set>
16 #include <utility>
17 #include <vector>
18
19 #include "base/factory.h"
20 #include "storage/index/dram/extendible_hash_index.h"
21 #include "storage/index/dram/pet_hash_index.h"
22 #include "storage/index/dram/unordered_map_index.h"
23 #include "storage/index/utils/hash.h"
24 #include "storage/kv_engine/base_kv.h"
25 #include "storage/value_store/dram_value_store.h"
26 #include "storage/value_store/hybrid_value_store.h"
27 #include "storage/value_store/ssd_value_store.h"
28
29 class KVEngineComposite : public BaseKV {
30 public:
31 KVEngineComposite(std::unique_ptr<Index> index,
32 std::unique_ptr<ValueStore> value_store,
33 int num_threads = 0)
34 : BaseKV(BaseKVConfig{}),
35 index_(std::move(index)),
36 value_store_(std::move(value_store)),
37 num_threads_(num_threads) {}
38
39 1100 explicit KVEngineComposite(const BaseKVConfig& config) : BaseKV(config) {
40 1100 config_ = config;
41 1100 const auto& j = config.json_config_;
42
3/6
✓ Branch 1 taken 1100 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 1100 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 1100 times.
✗ Branch 8 not taken.
 1100 const std::string index_type = j.at("index").at("type").get<std::string>();
43
3/6
✓ Branch 1 taken 1100 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 1100 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 1100 times.
✗ Branch 8 not taken.
 1100 const std::string value_type = j.at("value").at("type").get<std::string>();
44 using IF = base::Factory<Index, const BaseKVConfig&>;
45 using VF = base::Factory<ValueStore, const BaseKVConfig&>;
46
1/2
✓ Branch 1 taken 1100 times.
✗ Branch 2 not taken.
 1100 index_.reset(IF::NewInstance(index_type, config));
47
2/2
✓ Branch 1 taken 1098 times.
✓ Branch 2 taken 2 times.
 1100 value_store_.reset(VF::NewInstance(value_type, config));
48 1098 num_threads_ = config.num_threads_;
49 1098 default_value_size_hint_ =
50
2/4
✓ Branch 1 taken 1098 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 1098 times.
✗ Branch 5 not taken.
 1098 j.at("value").value("default_value_size_hint", 0);
51
3/6
✓ Branch 1 taken 1098 times.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 1098 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 1098 times.
 1098 if (!index_ || !value_store_) {
52 ✗ throw std::runtime_error("failed to create KVEngine components");
53 }
54 1112 }
55
56 200902 void Get(const uint64_t key, std::string& value, unsigned tid) override {
57 (void)tid;
58 200902 Value_t handle = kValueHandleNone;
59
1/2
✓ Branch 2 taken 200902 times.
✗ Branch 3 not taken.
 200902 index_->Get(key, handle);
60
2/2
✓ Branch 0 taken 5996 times.
✓ Branch 1 taken 194906 times.
 200902 if (handle == kValueHandleNone) {
61 5996 value.clear();
62 148798 return;
63 }
64
3/4
✓ Branch 2 taken 194906 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 142802 times.
✓ Branch 5 taken 52104 times.
 194906 if (const char* ptr = value_store_->DirectPtr(handle)) {
65
2/4
✓ Branch 2 taken 142802 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 142802 times.
✗ Branch 6 not taken.
 142802 value.resize(value_store_->SlotCapacity(handle));
66 142802 std::memcpy(value.data(), ptr, value.size());
67 142802 return;
68 }
69
2/4
✓ Branch 2 taken 52104 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 52104 times.
✗ Branch 6 not taken.
 52104 value.resize(value_store_->SlotCapacity(handle));
70 const size_t actual =
71
1/2
✓ Branch 4 taken 52104 times.
✗ Branch 5 not taken.
 52104 value_store_->Read(handle, value.data(), value.size());
72
1/2
✓ Branch 1 taken 52104 times.
✗ Branch 2 not taken.
 52104 value.resize(actual);
73 }
74
75 8 bool Exists(const uint64_t key, unsigned tid) override {
76 (void)tid;
77 8 Value_t handle = kValueHandleNone;
78
1/2
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 index_->Get(key, handle);
79 8 return handle != kValueHandleNone;
80 }
81
82 354558 void Put(const uint64_t key,
83 const std::string_view& value,
84 unsigned tid) override {
85
1/2
✓ Branch 1 taken 354558 times.
✗ Branch 2 not taken.
 354558 std::shared_lock<std::shared_mutex> checkpoint_lock(checkpoint_mu_);
86
1/2
✓ Branch 3 taken 354558 times.
✗ Branch 4 not taken.
 354558 PutInternal(key, value.data(), value.size(), tid, true);
87 354558 }
88
89 294 void BatchPut(base::ConstArray<uint64_t> keys,
90 std::vector<base::ConstArray<float>>* values,
91 unsigned tid) override {
92
3/6
✓ Branch 0 taken 294 times.
✗ Branch 1 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 294 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 294 times.
 294 if (values == nullptr || keys.Size() != static_cast<int>(values->size())) {
93 ✗ LOG(FATAL) << "KVEngine::BatchPut size mismatch";
94 }
95 (void)tid;
96
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 294 times.
 294 if (keys.Size() == 0) {
97 ✗ return;
98 }
99
1/2
✓ Branch 1 taken 294 times.
✗ Branch 2 not taken.
 294 std::shared_lock<std::shared_mutex> checkpoint_lock(checkpoint_mu_);
100
101 294 std::unordered_set<uint64_t> seen_keys;
102
1/2
✓ Branch 2 taken 294 times.
✗ Branch 3 not taken.
 294 seen_keys.reserve(static_cast<size_t>(keys.Size()));
103 294 bool has_duplicate_key = false;
104
2/2
✓ Branch 1 taken 742 times.
✓ Branch 2 taken 294 times.
 1036 for (int i = 0; i < keys.Size(); ++i) {
105
2/4
✓ Branch 2 taken 742 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 742 times.
 742 if (!seen_keys.insert(keys[i]).second) {
106 ✗ has_duplicate_key = true;
107 ✗ break;
108 }
109 }
110
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 294 times.
 294 if (has_duplicate_key) {
111 ✗ for (int i = 0; i < keys.Size(); ++i) {
112 ✗ const auto& item = (*values)[i];
113 ✗ PutInternal(keys[i],
114 ✗ item.Data(),
115 ✗ static_cast<size_t>(item.Size()) * sizeof(float),
116 tid,
117 false);
118 }
119 ✗ return;
120 }
121
122 struct PutItem {
123 uint64_t key = 0;
124 ValueStore::WriteSpec spec{};
125 };
126 294 std::vector<PutItem> items;
127
1/2
✓ Branch 2 taken 294 times.
✗ Branch 3 not taken.
 294 items.reserve(static_cast<size_t>(keys.Size()));
128
129
2/2
✓ Branch 1 taken 742 times.
✓ Branch 2 taken 294 times.
 1036 for (int i = 0; i < keys.Size(); ++i) {
130 742 const auto& item = (*values)[i];
131 742 const void* data = item.Data();
132 742 const size_t size = static_cast<size_t>(item.Size()) * sizeof(float);
133
1/2
✓ Branch 2 taken 742 times.
✗ Branch 3 not taken.
 742 items.push_back(PutItem{keys[i], ValueStore::WriteSpec{data, size}});
134 }
135
136 294 std::vector<ValueStore::WriteSpec> specs;
137
1/2
✓ Branch 2 taken 294 times.
✗ Branch 3 not taken.
 294 specs.reserve(items.size());
138
2/2
✓ Branch 5 taken 742 times.
✓ Branch 6 taken 294 times.
 1036 for (const auto& item : items) {
139
1/2
✓ Branch 1 taken 742 times.
✗ Branch 2 not taken.
 742 specs.push_back(item.spec);
140 }
141
1/2
✓ Branch 2 taken 294 times.
✗ Branch 3 not taken.
 294 const auto new_handles = value_store_->BatchAllocAndWrite(specs);
142
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 294 times.
 294 if (new_handles.size() != items.size()) {
143 ✗ LOG(FATAL) << "KVEngine::BatchPut allocation result size mismatch";
144 }
145
2/2
✓ Branch 1 taken 742 times.
✓ Branch 2 taken 294 times.
 1036 for (size_t i = 0; i < items.size(); ++i) {
146
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 742 times.
 742 if (new_handles[i] == kValueHandleNone) {
147 ✗ LOG(FATAL) << "KVEngine batch value allocation failed, key="
148 ✗ << items[i].key << " size=" << items[i].spec.size;
149 }
150 }
151
152
2/2
✓ Branch 1 taken 742 times.
✓ Branch 2 taken 294 times.
 1036 for (size_t i = 0; i < items.size(); ++i) {
153
1/2
✓ Branch 4 taken 742 times.
✗ Branch 5 not taken.
 742 Value_t old_handle = index_->Put(items[i].key, new_handles[i], tid);
154
2/2
✓ Branch 0 taken 216 times.
✓ Branch 1 taken 526 times.
 742 if (old_handle != kValueHandleNone) {
155
1/2
✓ Branch 2 taken 216 times.
✗ Branch 3 not taken.
 216 value_store_->Retire(old_handle);
156 }
157 }
158
1/2
✓ Branch 1 taken 294 times.
✗ Branch 2 not taken.
 294 TrackKeys(keys);
159
2/4
✓ Branch 4 taken 294 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 294 times.
✗ Branch 8 not taken.
 294 }
160
161 1536 void BatchGet(base::ConstArray<uint64_t> keys,
162 std::vector<base::ConstArray<float>>* values,
163 unsigned tid) override {
164 (void)tid;
165
1/2
✓ Branch 2 taken 1536 times.
✗ Branch 3 not taken.
 1536 values->resize(keys.Size());
166
2/2
✓ Branch 0 taken 252 times.
✓ Branch 1 taken 1284 times.
 1536 thread_local std::vector<Value_t> handles;
167
2/2
✓ Branch 0 taken 252 times.
✓ Branch 1 taken 1284 times.
 1536 thread_local std::vector<std::vector<float>> buffers;
168
1/2
✓ Branch 2 taken 1536 times.
✗ Branch 3 not taken.
 1536 handles.assign(keys.Size(), kValueHandleNone);
169 1536 buffers.clear();
170
1/2
✓ Branch 2 taken 1536 times.
✗ Branch 3 not taken.
 1536 buffers.resize(keys.Size());
171
172
1/2
✓ Branch 1 taken 1536 times.
✗ Branch 2 not taken.
 1536 if (keys.Size() > 0) {
173
1/2
✓ Branch 3 taken 1536 times.
✗ Branch 4 not taken.
 1536 index_->BatchGet(keys, handles.data(), tid);
174 }
175 1536 std::vector<uint64_t> batch_handles;
176 1536 std::vector<size_t> batch_indices;
177
1/2
✓ Branch 2 taken 1536 times.
✗ Branch 3 not taken.
 1536 batch_handles.reserve(static_cast<size_t>(keys.Size()));
178
1/2
✓ Branch 2 taken 1536 times.
✗ Branch 3 not taken.
 1536 batch_indices.reserve(static_cast<size_t>(keys.Size()));
179
2/2
✓ Branch 1 taken 67426 times.
✓ Branch 2 taken 1536 times.
 68962 for (int i = 0; i < keys.Size(); ++i) {
180
2/2
✓ Branch 1 taken 324 times.
✓ Branch 2 taken 67102 times.
 67426 if (handles[i] == kValueHandleNone) {
181 324 (*values)[i] = base::ConstArray<float>();
182 324 continue;
183 }
184
3/4
✓ Branch 3 taken 67102 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 45542 times.
✓ Branch 6 taken 21560 times.
 67102 if (const char* ptr = value_store_->DirectPtr(handles[i])) {
185
1/2
✓ Branch 3 taken 45542 times.
✗ Branch 4 not taken.
 45542 const size_t bytes = value_store_->SlotCapacity(handles[i]);
186 45542 (*values)[i] = base::ConstArray<float>(
187 reinterpret_cast<float*>(const_cast<char*>(ptr)),
188 45542 bytes / sizeof(float));
189 45542 continue;
190 45542 }
191
1/2
✓ Branch 2 taken 21560 times.
✗ Branch 3 not taken.
 21560 batch_handles.push_back(handles[i]);
192
1/2
✓ Branch 1 taken 21560 times.
✗ Branch 2 not taken.
 21560 batch_indices.push_back(static_cast<size_t>(i));
193 }
194
195 1536 std::vector<ValueStore::ReadResult> batch_results;
196
2/2
✓ Branch 1 taken 410 times.
✓ Branch 2 taken 1126 times.
 1536 if (!batch_handles.empty()) {
197
1/2
✓ Branch 2 taken 410 times.
✗ Branch 3 not taken.
 410 value_store_->BatchRead(batch_handles, batch_results);
198
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 410 times.
 410 if (batch_results.size() != batch_indices.size()) {
199 ✗ LOG(FATAL) << "KVEngine::BatchGet read result size mismatch";
200 }
201
2/2
✓ Branch 1 taken 21560 times.
✓ Branch 2 taken 410 times.
 21970 for (size_t i = 0; i < batch_indices.size(); ++i) {
202 21560 const size_t idx = batch_indices[i];
203 21560 const auto& result = batch_results[i];
204
1/2
✓ Branch 3 taken 21560 times.
✗ Branch 4 not taken.
 21560 buffers[idx].resize(result.data.size() / sizeof(float));
205
1/2
✓ Branch 1 taken 21560 times.
✗ Branch 2 not taken.
 21560 if (!result.data.empty()) {
206 43120 std::memcpy(
207 21560 buffers[idx].data(), result.data.data(), result.data.size());
208 }
209 21560 (*values)[idx] =
210 43120 base::ConstArray<float>(buffers[idx].data(), buffers[idx].size());
211 }
212 }
213 1536 }
214
215 122 bool BatchGetFlat(base::ConstArray<uint64_t> keys,
216 float* values,
217 int64_t num_rows,
218 int64_t embedding_dim,
219 unsigned tid,
220 BatchGetFlatStats* stats = nullptr) override {
221
4/8
✓ Branch 0 taken 122 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 122 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 122 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 122 times.
 244 if (values == nullptr || num_rows < 0 || embedding_dim <= 0 ||
222
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 122 times.
 122 keys.Size() != static_cast<size_t>(num_rows)) {
223 ✗ return false;
224 }
225 122 const size_t row_bytes = static_cast<size_t>(embedding_dim) * sizeof(float);
226
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 106 times.
 122 thread_local std::vector<Value_t> handles;
227
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 106 times.
 122 thread_local std::vector<char> read_buffer;
228
1/2
✓ Branch 2 taken 122 times.
✗ Branch 3 not taken.
 122 handles.assign(keys.Size(), kValueHandleNone);
229 const auto index_lookup_start =
230 54 stats != nullptr ? std::chrono::steady_clock::now()
231
2/2
✓ Branch 0 taken 54 times.
✓ Branch 1 taken 68 times.
 122 : std::chrono::steady_clock::time_point{};
232
1/2
✓ Branch 1 taken 122 times.
✗ Branch 2 not taken.
 122 if (keys.Size() > 0) {
233
1/2
✓ Branch 3 taken 122 times.
✗ Branch 4 not taken.
 122 index_->BatchGet(keys, handles.data(), tid);
234 }
235
2/2
✓ Branch 0 taken 54 times.
✓ Branch 1 taken 68 times.
 122 if (stats != nullptr) {
236 54 stats->index_lookup_ns = static_cast<std::uint64_t>(
237
1/2
✓ Branch 1 taken 54 times.
✗ Branch 2 not taken.
 54 std::chrono::duration_cast< std::chrono::nanoseconds>(
238
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 108 std::chrono::steady_clock::now() - index_lookup_start)
239 54 .count());
240 }
241
242 122 std::uint64_t missing_zero_fill_ns = 0;
243 122 std::uint64_t missing_rows = 0;
244 const auto row_copy_start =
245 54 stats != nullptr ? std::chrono::steady_clock::now()
246
2/2
✓ Branch 0 taken 54 times.
✓ Branch 1 taken 68 times.
 122 : std::chrono::steady_clock::time_point{};
247
4/4
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 110 times.
✓ Branch 2 taken 12 times.
✓ Branch 3 taken 110 times.
 134 if (default_value_size_hint_ == row_bytes &&
248
2/4
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 12 times.
✗ Branch 5 not taken.
 24 value_store_->ReadFlatFixedRows(
249 12 handles.data(),
250 static_cast<size_t>(num_rows),
251 values,
252 row_bytes,
253 &missing_rows)) {
254
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 12 times.
 12 if (stats != nullptr) {
255 ✗ stats->zero_fill_ns = 0;
256 ✗ stats->row_copy_ns = static_cast<std::uint64_t>(
257 ✗ std::chrono::duration_cast< std::chrono::nanoseconds>(
258 ✗ std::chrono::steady_clock::now() - row_copy_start)
259 ✗ .count());
260 ✗ stats->missing_rows = missing_rows;
261 }
262 12 return true;
263 }
264
2/2
✓ Branch 0 taken 218 times.
✓ Branch 1 taken 54 times.
 272 for (int64_t row = 0; row < num_rows; ++row) {
265 218 const Value_t handle = handles[static_cast<size_t>(row)];
266 218 float* dst = values + row * embedding_dim;
267
2/2
✓ Branch 0 taken 54 times.
✓ Branch 1 taken 164 times.
 218 if (handle == kValueHandleNone) {
268 const auto missing_zero_start =
269 54 stats != nullptr ? std::chrono::steady_clock::now()
270
1/2
✓ Branch 0 taken 54 times.
✗ Branch 1 not taken.
 54 : std::chrono::steady_clock::time_point{};
271 54 std::memset(dst, 0, row_bytes);
272
1/2
✓ Branch 0 taken 54 times.
✗ Branch 1 not taken.
 54 if (stats != nullptr) {
273 54 missing_zero_fill_ns += static_cast<std::uint64_t>(
274
1/2
✓ Branch 1 taken 54 times.
✗ Branch 2 not taken.
 54 std::chrono::duration_cast< std::chrono::nanoseconds>(
275
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 108 std::chrono::steady_clock::now() - missing_zero_start)
276 54 .count());
277 }
278 54 ++missing_rows;
279 54 continue;
280 54 }
281
3/4
✓ Branch 2 taken 164 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 110 times.
✓ Branch 5 taken 54 times.
 164 if (const char* ptr = value_store_->DirectPtr(handle)) {
282
1/2
✓ Branch 0 taken 110 times.
✗ Branch 1 not taken.
 110 if (default_value_size_hint_ != row_bytes) {
283
1/2
✓ Branch 2 taken 110 times.
✗ Branch 3 not taken.
 110 const size_t slot_bytes = value_store_->SlotCapacity(handle);
284
2/2
✓ Branch 0 taken 38 times.
✓ Branch 1 taken 72 times.
 110 if (slot_bytes != row_bytes) {
285
4/8
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 38 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 38 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 38 times.
✗ Branch 11 not taken.
 76 LOG(ERROR) << "KVEngine::BatchGetFlat row size mismatch row=" << row
286
2/4
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
✓ Branch 5 taken 38 times.
✗ Branch 6 not taken.
 38 << " key=" << keys[static_cast<int>(row)]
287
2/4
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 38 times.
✗ Branch 5 not taken.
 38 << " expected_bytes=" << row_bytes
288
2/4
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 38 times.
✗ Branch 5 not taken.
 38 << " actual_bytes=" << slot_bytes;
289 38 return false;
290 }
291 }
292 72 std::memcpy(dst, ptr, row_bytes);
293 } else {
294
1/2
✓ Branch 1 taken 54 times.
✗ Branch 2 not taken.
 54 read_buffer.resize(row_bytes + sizeof(float));
295 const size_t actual =
296
1/2
✓ Branch 4 taken 54 times.
✗ Branch 5 not taken.
 54 value_store_->Read(handle, read_buffer.data(), read_buffer.size());
297
2/2
✓ Branch 0 taken 18 times.
✓ Branch 1 taken 36 times.
 54 if (actual != row_bytes) {
298
4/8
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 18 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 18 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 18 times.
✗ Branch 11 not taken.
 36 LOG(ERROR) << "KVEngine::BatchGetFlat read size mismatch row=" << row
299
2/4
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
✓ Branch 5 taken 18 times.
✗ Branch 6 not taken.
 18 << " key=" << keys[static_cast<int>(row)]
300
2/4
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 18 times.
✗ Branch 5 not taken.
 18 << " expected_bytes=" << row_bytes
301
2/4
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 18 times.
✗ Branch 5 not taken.
 18 << " actual_bytes=" << actual;
302 18 return false;
303 }
304 36 std::memcpy(dst, read_buffer.data(), row_bytes);
305 }
306 }
307
1/2
✓ Branch 0 taken 54 times.
✗ Branch 1 not taken.
 54 if (stats != nullptr) {
308 54 stats->zero_fill_ns = missing_zero_fill_ns;
309 54 stats->row_copy_ns = static_cast<std::uint64_t>(
310
1/2
✓ Branch 1 taken 54 times.
✗ Branch 2 not taken.
 54 std::chrono::duration_cast< std::chrono::nanoseconds>(
311
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 54 std::chrono::steady_clock::now() - row_copy_start)
312 54 .count());
313 54 stats->missing_rows = missing_rows;
314 }
315 54 return true;
316 }
317
318 ✗ bool BatchGetIndexOnly(base::ConstArray<uint64_t> keys,
319 unsigned tid,
320 BatchGetFlatStats* stats = nullptr) override {
321 ✗ thread_local std::vector<Value_t> handles;
322 ✗ handles.assign(keys.Size(), kValueHandleNone);
323 ✗ if (keys.Size() > 0) {
324 ✗ index_->BatchGet(keys, handles.data(), tid);
325 }
326 ✗ if (stats != nullptr) {
327 ✗ std::uint64_t missing_rows = 0;
328 ✗ for (const Value_t handle : handles) {
329 ✗ if (handle == kValueHandleNone) {
330 ✗ ++missing_rows;
331 }
332 }
333 ✗ stats->missing_rows = missing_rows;
334 }
335 ✗ return true;
336 }
337
338 24 bool BatchGetDirectFixedRows(
339 base::ConstArray<uint64_t> keys,
340 int64_t num_rows,
341 int64_t embedding_dim,
342 unsigned tid,
343 std::vector<DirectFixedRow>* rows,
344 BatchGetFlatStats* stats = nullptr) override {
345
4/8
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 24 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 24 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 24 times.
 48 if (rows == nullptr || num_rows < 0 || embedding_dim <= 0 ||
346
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 24 times.
 24 keys.Size() != static_cast<size_t>(num_rows)) {
347 ✗ return false;
348 }
349 24 const size_t row_bytes = static_cast<size_t>(embedding_dim) * sizeof(float);
350
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 24 times.
 24 if (default_value_size_hint_ != row_bytes) {
351 ✗ return false;
352 }
353
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 16 times.
 24 thread_local std::vector<Value_t> handles;
354
1/2
✓ Branch 2 taken 24 times.
✗ Branch 3 not taken.
 24 handles.assign(keys.Size(), kValueHandleNone);
355 const auto index_lookup_start =
356 ✗ stats != nullptr ? std::chrono::steady_clock::now()
357
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 24 times.
 24 : std::chrono::steady_clock::time_point{};
358
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 if (keys.Size() > 0) {
359
1/2
✓ Branch 3 taken 24 times.
✗ Branch 4 not taken.
 24 index_->BatchGet(keys, handles.data(), tid);
360 }
361
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 24 times.
 24 if (stats != nullptr) {
362 ✗ stats->index_lookup_ns = static_cast<std::uint64_t>(
363 ✗ std::chrono::duration_cast< std::chrono::nanoseconds>(
364 ✗ std::chrono::steady_clock::now() - index_lookup_start)
365 ✗ .count());
366 }
367
368
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 16 times.
 24 thread_local std::vector<ValueStore::DirectFixedRow> store_rows;
369
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 store_rows.resize(static_cast<size_t>(num_rows));
370 24 uint64_t missing_rows = 0;
371
2/4
✓ Branch 3 taken 24 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 24 times.
 48 if (!value_store_->GetDirectFixedRows(
372 24 handles.data(),
373 static_cast<size_t>(num_rows),
374 row_bytes,
375 store_rows.data(),
376 &missing_rows)) {
377 ✗ return false;
378 }
379
1/2
✓ Branch 2 taken 24 times.
✗ Branch 3 not taken.
 24 rows->resize(store_rows.size());
380
2/2
✓ Branch 1 taken 46 times.
✓ Branch 2 taken 24 times.
 70 for (size_t i = 0; i < store_rows.size(); ++i) {
381 46 (*rows)[i] = DirectFixedRow{
382 46 store_rows[i].data, store_rows[i].size, store_rows[i].missing};
383 }
384
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 24 times.
 24 if (stats != nullptr) {
385 ✗ stats->missing_rows = missing_rows;
386 }
387 24 return true;
388 }
389
390 ✗ RDMABackingRegion GetRDMABackingRegion() const override {
391 ✗ if (!value_store_) {
392 ✗ return {};
393 }
394 return RDMABackingRegion{
395 ✗ value_store_->RDMABackingData(), value_store_->RDMABackingSize()};
396 }
397
398 24 bool ApplySgdUpdateFlat(
399 base::ConstArray<uint64_t> keys,
400 const float* grads,
401 int64_t num_rows,
402 int64_t embedding_dim,
403 float learning_rate,
404 uint8_t tag,
405 unsigned tid) override {
406
4/8
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
✓ Branch 3 taken 24 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 24 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 24 times.
 24 if (grads == nullptr || keys.Size() != static_cast<size_t>(num_rows) ||
407 embedding_dim <= 0) {
408 ✗ return false;
409 }
410
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 std::shared_lock<std::shared_mutex> checkpoint_lock(checkpoint_mu_);
411 24 const int tag_bits = static_cast<int>(sizeof(tag) * 8);
412 24 const int shift = static_cast<int>(sizeof(uint64_t) * 8) - tag_bits;
413 24 const uint64_t key_mask = ~0ULL >> tag_bits;
414 24 const size_t row_bytes = static_cast<size_t>(embedding_dim) * sizeof(float);
415
416
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 16 times.
 24 thread_local std::vector<uint64_t> tagged_keys;
417
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 tagged_keys.resize(static_cast<size_t>(num_rows));
418
2/2
✓ Branch 0 taken 46 times.
✓ Branch 1 taken 24 times.
 70 for (int64_t r = 0; r < num_rows; ++r) {
419 46 tagged_keys[static_cast<size_t>(r)] =
420 92 (static_cast<uint64_t>(tag) << shift) |
421 46 (keys[static_cast<size_t>(r)] & key_mask);
422 }
423
424
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 16 times.
 24 thread_local std::vector<DirectFixedRow> rows;
425
2/4
✓ Branch 2 taken 24 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 24 times.
 24 if (!BatchGetDirectFixedRows(
426 base::ConstArray<uint64_t>(tagged_keys),
427 num_rows,
428 embedding_dim,
429 tid,
430 &rows)) {
431 ✗ return false;
432 }
433
434
1/2
✓ Branch 2 taken 24 times.
✗ Branch 3 not taken.
 24 std::vector<float> missing_row(static_cast<size_t>(embedding_dim));
435
2/2
✓ Branch 0 taken 46 times.
✓ Branch 1 taken 24 times.
 70 for (int64_t r = 0; r < num_rows; ++r) {
436 46 const auto& row = rows[static_cast<size_t>(r)];
437 46 const float* grad = grads + r * embedding_dim;
438
2/2
✓ Branch 0 taken 20 times.
✓ Branch 1 taken 26 times.
 46 if (row.missing) {
439
2/2
✓ Branch 0 taken 584 times.
✓ Branch 1 taken 20 times.
 604 for (int64_t c = 0; c < embedding_dim; ++c) {
440 584 missing_row[static_cast<size_t>(c)] = -learning_rate * grad[c];
441 }
442
1/2
✓ Branch 2 taken 20 times.
✗ Branch 3 not taken.
 20 PutInternal(tagged_keys[static_cast<size_t>(r)],
443 20 missing_row.data(),
444 row_bytes,
445 tid,
446 false);
447 20 continue;
448 20 }
449
450 26 float* value = reinterpret_cast<float*>(const_cast<char*>(row.data));
451 26 #pragma omp simd
452 for (int64_t c = 0; c < embedding_dim; ++c) {
453 608 value[c] -= learning_rate * grad[c];
454 }
455 }
456 24 return true;
457 24 }
458
459 56 void BulkLoad(base::ConstArray<uint64_t> keys, const void* value) override {
460 56 const auto& j = config_.json_config_;
461
2/4
✓ Branch 1 taken 56 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 56 times.
✗ Branch 5 not taken.
 56 const size_t value_size = j.at("value").value("default_value_size_hint", 0);
462
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 56 times.
 56 if (value_size == 0) {
463 ✗ LOG(FATAL) << "KVEngine::BulkLoad requires value_size hint";
464 }
465
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 56 times.
 56 if (keys.Size() == 0) {
466 ✗ return;
467 }
468
1/2
✓ Branch 1 taken 56 times.
✗ Branch 2 not taken.
 56 std::shared_lock<std::shared_mutex> checkpoint_lock(checkpoint_mu_);
469 56 const char* data = reinterpret_cast<const char*>(value);
470 56 std::vector<ValueStore::WriteSpec> specs;
471
1/2
✓ Branch 2 taken 56 times.
✗ Branch 3 not taken.
 56 specs.reserve(static_cast<size_t>(keys.Size()));
472
2/2
✓ Branch 1 taken 434 times.
✓ Branch 2 taken 56 times.
 490 for (int i = 0; i < keys.Size(); ++i) {
473
1/2
✓ Branch 1 taken 434 times.
✗ Branch 2 not taken.
 434 specs.push_back(ValueStore::WriteSpec{data + i * value_size, value_size});
474 }
475
1/2
✓ Branch 2 taken 56 times.
✗ Branch 3 not taken.
 56 std::vector<uint64_t> handles = value_store_->BatchAllocAndWrite(specs);
476
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 56 times.
 56 if (handles.size() != static_cast<size_t>(keys.Size())) {
477 ✗ LOG(FATAL) << "KVEngine::BulkLoad allocation result size mismatch";
478 }
479
2/2
✓ Branch 1 taken 434 times.
✓ Branch 2 taken 56 times.
 490 for (int i = 0; i < keys.Size(); ++i) {
480
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 434 times.
 434 if (handles[static_cast<size_t>(i)] == kValueHandleNone) {
481 ✗ LOG(FATAL) << "KVEngine bulk value allocation failed, key=" << keys[i]
482 ✗ << " size=" << value_size;
483 }
484 }
485
1/2
✓ Branch 3 taken 56 times.
✗ Branch 4 not taken.
 56 index_->BatchPut(keys, handles.data(), 0);
486
1/2
✓ Branch 1 taken 56 times.
✗ Branch 2 not taken.
 56 TrackKeys(keys);
487 56 }
488
489 4 bool SaveCheckpoint(const std::string& file,
490 const std::string& metadata) override {
491
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 std::unique_lock<std::shared_mutex> checkpoint_lock(checkpoint_mu_);
492
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
 4 if (file.empty()) {
493 ✗ LOG(ERROR) << "KVEngine checkpoint path is empty";
494 ✗ return false;
495 }
496
497 4 std::vector<uint64_t> keys;
498 {
499
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 std::lock_guard<std::mutex> lock(checkpoint_keys_mu_);
500
1/2
✓ Branch 3 taken 4 times.
✗ Branch 4 not taken.
 4 keys.assign(checkpoint_keys_.begin(), checkpoint_keys_.end());
501 4 }
502
1/2
✓ Branch 3 taken 4 times.
✗ Branch 4 not taken.
 4 std::sort(keys.begin(), keys.end());
503
504
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 const std::filesystem::path checkpoint_path(file);
505
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 std::filesystem::path temp_path(file);
506
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 temp_path += ".tmp";
507 4 std::error_code error;
508 4 std::filesystem::remove(temp_path, error);
509
510
1/2
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 std::ofstream output(temp_path, std::ios::binary | std::ios::trunc);
511
2/4
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 4 times.
 4 if (!output) {
512 ✗ LOG(ERROR) << "Failed to open checkpoint temp file: " << temp_path;
513 ✗ return false;
514 }
515
516 4 const uint64_t metadata_size = static_cast<uint64_t>(metadata.size());
517 4 const uint64_t record_count = static_cast<uint64_t>(keys.size());
518 4 uint64_t checksum = kCheckpointChecksumSeed;
519 4 UpdateChecksum(&checksum, kCheckpointMagic.data(), kCheckpointMagic.size());
520 4 UpdateChecksumPod(&checksum, kCheckpointVersion);
521 4 UpdateChecksumPod(&checksum, metadata_size);
522 4 UpdateChecksumPod(&checksum, record_count);
523 4 UpdateChecksum(&checksum, metadata.data(), metadata.size());
524 bool write_ok =
525
1/2
✓ Branch 3 taken 4 times.
✗ Branch 4 not taken.
 4 WriteBytes(output, kCheckpointMagic.data(), kCheckpointMagic.size()) &&
526
2/4
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
✗ Branch 4 not taken.
 4 WritePod(output, kCheckpointVersion) &&
527
5/10
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 3 taken 4 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 4 times.
✗ Branch 6 not taken.
✓ Branch 8 taken 4 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 4 times.
✗ Branch 11 not taken.
 12 WritePod(output, metadata_size) && WritePod(output, record_count) &&
528
2/4
✓ Branch 3 taken 4 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 4 times.
✗ Branch 6 not taken.
 4 WriteBytes(output, metadata.data(), metadata.size());
529
2/2
✓ Branch 5 taken 8 times.
✓ Branch 6 taken 4 times.
 12 for (const uint64_t key : keys) {
530 8 std::string value;
531
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 Get(key, value, 0);
532
2/4
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 8 times.
 8 if (!Exists(key, 0)) {
533 ✗ LOG(ERROR) << "KVEngine checkpoint tracked key is missing: " << key;
534 ✗ write_ok = false;
535 ✗ break;
536 }
537 8 const uint64_t value_size = static_cast<uint64_t>(value.size());
538 8 UpdateChecksumPod(&checksum, key);
539 8 UpdateChecksumPod(&checksum, value_size);
540 8 UpdateChecksum(&checksum, value.data(), value.size());
541 8 write_ok =
542
5/10
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 3 taken 8 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 8 times.
✗ Branch 6 not taken.
✓ Branch 8 taken 8 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 8 times.
✗ Branch 11 not taken.
 16 write_ok && WritePod(output, key) && WritePod(output, value_size) &&
543
2/4
✓ Branch 3 taken 8 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 8 times.
✗ Branch 6 not taken.
 8 WriteBytes(output, value.data(), value.size());
544
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (!write_ok) {
545 ✗ break;
546 }
547
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 }
548
3/6
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 3 taken 4 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 4 times.
✗ Branch 6 not taken.
 4 write_ok = write_ok && WritePod(output, checksum);
549
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 output.flush();
550
3/6
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 3 taken 4 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 4 times.
✗ Branch 6 not taken.
 4 write_ok = write_ok && output.good();
551
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 output.close();
552
3/6
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 3 taken 4 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 4 times.
✗ Branch 6 not taken.
 4 write_ok = write_ok && !output.fail();
553
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (!write_ok) {
554 ✗ LOG(ERROR) << "Failed to write checkpoint temp file: " << temp_path;
555 ✗ std::filesystem::remove(temp_path, error);
556 ✗ return false;
557 }
558
559 4 error.clear();
560 4 std::filesystem::rename(temp_path, checkpoint_path, error);
561
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
 4 if (error) {
562 ✗ LOG(ERROR) << "Failed to publish checkpoint " << checkpoint_path << ": "
563 ✗ << error.message();
564 ✗ std::filesystem::remove(temp_path, error);
565 ✗ return false;
566 }
567 4 return true;
568 4 }
569
570 12 bool LoadCheckpoint(const std::string& file,
571 const std::string& expected_metadata) override {
572
1/2
✓ Branch 1 taken 12 times.
✗ Branch 2 not taken.
 12 std::unique_lock<std::shared_mutex> checkpoint_lock(checkpoint_mu_);
573 {
574
1/2
✓ Branch 1 taken 12 times.
✗ Branch 2 not taken.
 12 std::lock_guard<std::mutex> lock(checkpoint_keys_mu_);
575
2/2
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 10 times.
 12 if (!checkpoint_keys_.empty()) {
576
3/6
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 2 times.
✗ Branch 8 not taken.
 2 LOG(ERROR) << "KVEngine checkpoint load requires an empty engine";
577 2 return false;
578 }
579
2/2
✓ Branch 1 taken 10 times.
✓ Branch 2 taken 2 times.
 12 }
580
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 10 times.
 10 if (file.empty()) {
581 ✗ LOG(ERROR) << "KVEngine checkpoint path is empty";
582 ✗ return false;
583 }
584
585 10 std::error_code error;
586
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 const std::uintmax_t file_size = std::filesystem::file_size(file, error);
587
3/6
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 10 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 10 times.
 10 if (error || file_size > std::numeric_limits<uint64_t>::max()) {
588 ✗ LOG(ERROR) << "Failed to inspect checkpoint file: " << file;
589 ✗ return false;
590 }
591 10 uint64_t remaining = static_cast<uint64_t>(file_size);
592
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 std::ifstream input(file, std::ios::binary);
593
2/4
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 10 times.
 10 if (!input) {
594 ✗ LOG(ERROR) << "Failed to open checkpoint file: " << file;
595 ✗ return false;
596 }
597
598 10 std::array<char, kCheckpointMagic.size()> magic{};
599 10 uint32_t version = 0;
600 10 uint64_t metadata_size = 0;
601 10 uint64_t record_count = 0;
602
1/2
✓ Branch 3 taken 10 times.
✗ Branch 4 not taken.
 10 if (!ReadBytes(input, magic.data(), magic.size(), &remaining) ||
603
2/4
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 10 times.
✗ Branch 4 not taken.
 10 !ReadPod(input, &version, &remaining) ||
604
2/4
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 10 times.
✗ Branch 4 not taken.
 10 !ReadPod(input, &metadata_size, &remaining) ||
605
2/4
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 10 times.
✗ Branch 4 not taken.
 10 !ReadPod(input, &record_count, &remaining) ||
606
4/6
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 8 times.
✗ Branch 6 not taken.
 10 magic != kCheckpointMagic || version != kCheckpointVersion ||
607
4/6
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 8 times.
 28 metadata_size > remaining ||
608
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
 8 metadata_size > std::numeric_limits<size_t>::max()) {
609
4/8
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 2 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 2 times.
✗ Branch 11 not taken.
 2 LOG(ERROR) << "Invalid or truncated checkpoint header: " << file;
610 2 return false;
611 }
612
613
1/2
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 std::string metadata(static_cast<size_t>(metadata_size), '\0');
614
2/4
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 8 times.
 8 if (!ReadBytes(input, metadata.data(), metadata_size, &remaining)) {
615 ✗ LOG(ERROR) << "Truncated checkpoint metadata: " << file;
616 ✗ return false;
617 }
618
2/2
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 6 times.
 8 if (metadata != expected_metadata) {
619
4/8
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 2 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 2 times.
✗ Branch 11 not taken.
 2 LOG(ERROR) << "Checkpoint metadata mismatch: " << file;
620 2 return false;
621 }
622
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (remaining < sizeof(uint64_t)) {
623 ✗ LOG(ERROR) << "Checkpoint checksum is missing: " << file;
624 ✗ return false;
625 }
626
2/4
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
 12 if (record_count > std::numeric_limits<size_t>::max() ||
627 6 record_count >
628
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 (remaining - sizeof(uint64_t)) / (2 * sizeof(uint64_t))) {
629 ✗ LOG(ERROR) << "Invalid checkpoint record count: " << file;
630 ✗ return false;
631 }
632
633 6 uint64_t checksum = kCheckpointChecksumSeed;
634 6 UpdateChecksum(&checksum, magic.data(), magic.size());
635 6 UpdateChecksumPod(&checksum, version);
636 6 UpdateChecksumPod(&checksum, metadata_size);
637 6 UpdateChecksumPod(&checksum, record_count);
638 6 UpdateChecksum(&checksum, metadata.data(), metadata.size());
639 6 std::vector<std::pair<uint64_t, std::string>> records;
640
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 records.reserve(static_cast<size_t>(record_count));
641 6 std::unordered_set<uint64_t> seen_keys;
642
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 seen_keys.reserve(static_cast<size_t>(record_count));
643
2/2
✓ Branch 0 taken 10 times.
✓ Branch 1 taken 6 times.
 16 for (uint64_t i = 0; i < record_count; ++i) {
644 10 uint64_t key = 0;
645 10 uint64_t value_size = 0;
646
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 if (!ReadPod(input, &key, &remaining) ||
647
3/6
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 10 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 10 times.
✗ Branch 6 not taken.
 10 !ReadPod(input, &value_size, &remaining) || value_size > remaining ||
648
3/6
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✓ Branch 3 taken 10 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 10 times.
 30 value_size > std::numeric_limits<size_t>::max() ||
649
2/4
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 10 times.
 10 !seen_keys.insert(key).second) {
650 ✗ LOG(ERROR) << "Invalid or truncated checkpoint record: " << file;
651 ✗ return false;
652 }
653
1/2
✓ Branch 2 taken 10 times.
✗ Branch 3 not taken.
 10 std::string value(static_cast<size_t>(value_size), '\0');
654
2/4
✓ Branch 2 taken 10 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 10 times.
 10 if (!ReadBytes(input, value.data(), value_size, &remaining)) {
655 ✗ LOG(ERROR) << "Truncated checkpoint value: " << file;
656 ✗ return false;
657 }
658 10 UpdateChecksumPod(&checksum, key);
659 10 UpdateChecksumPod(&checksum, value_size);
660 10 UpdateChecksum(&checksum, value.data(), value.size());
661
1/2
✓ Branch 2 taken 10 times.
✗ Branch 3 not taken.
 10 records.emplace_back(key, std::move(value));
662
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 }
663 6 uint64_t saved_checksum = 0;
664
6/8
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 4 times.
✓ Branch 7 taken 2 times.
✓ Branch 8 taken 4 times.
 6 if (!ReadPod(input, &saved_checksum, &remaining) || remaining != 0) {
665
4/8
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 2 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 2 times.
✗ Branch 11 not taken.
 2 LOG(ERROR) << "Checkpoint contains trailing data: " << file;
666 2 return false;
667 }
668
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 2 times.
 4 if (saved_checksum != checksum) {
669
4/8
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 2 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 2 times.
✗ Branch 11 not taken.
 2 LOG(ERROR) << "Checkpoint checksum mismatch: " << file;
670 2 return false;
671 }
672
673
2/2
✓ Branch 5 taken 6 times.
✓ Branch 6 taken 2 times.
 8 for (const auto& record : records) {
674 6 PutInternal(
675
1/2
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
 6 record.first, record.second.data(), record.second.size(), 0, false);
676 }
677 2 return true;
678 12 }
679
680 14 uint64_t CheckpointRecordCount() const override {
681
1/2
✓ Branch 1 taken 14 times.
✗ Branch 2 not taken.
 14 std::shared_lock<std::shared_mutex> checkpoint_lock(checkpoint_mu_);
682
1/2
✓ Branch 1 taken 14 times.
✗ Branch 2 not taken.
 14 std::lock_guard<std::mutex> lock(checkpoint_keys_mu_);
683 28 return static_cast<uint64_t>(checkpoint_keys_.size());
684 14 }
685
686 ✗ void Util() override {
687 ✗ LOG(INFO) << "KVEngine index utilization=" << index_->Utilization()
688 ✗ << " value=" << value_store_->GetInfo();
689 ✗ }
690
691 ✗ void DebugInfo() const override {
692 ✗ index_->DebugInfo();
693 ✗ LOG(INFO) << value_store_->GetInfo();
694 ✗ }
695
696 ✗ std::string ExtraResultFields() const override {
697 ✗ return value_store_ ? value_store_->ExtraResultFields() : "";
698 }
699
700 private:
701 48 static bool WriteBytes(std::ofstream& output, const void* data, size_t size) {
702 48 if (size >
703
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 48 times.
 48 static_cast<size_t>((std::numeric_limits<std::streamsize>::max)())) {
704 ✗ return false;
705 }
706
1/2
✓ Branch 0 taken 48 times.
✗ Branch 1 not taken.
 48 if (size != 0) {
707 48 output.write(static_cast<const char*>(data),
708 static_cast<std::streamsize>(size));
709 }
710 48 return output.good();
711 }
712
713 template <typename T>
714 32 static bool WritePod(std::ofstream& output, const T& value) {
715 32 return WriteBytes(output, &value, sizeof(value));
716 }
717
718 84 static bool ReadBytes(
719 std::ifstream& input, void* data, uint64_t size, uint64_t* remaining) {
720
5/6
✓ Branch 0 taken 84 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 82 times.
✓ Branch 3 taken 2 times.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 82 times.
 166 if (remaining == nullptr || size > *remaining ||
721
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 82 times.
 82 size > static_cast<uint64_t>(
722 82 (std::numeric_limits<std::streamsize>::max)())) {
723 2 return false;
724 }
725
1/2
✓ Branch 0 taken 82 times.
✗ Branch 1 not taken.
 82 if (size != 0) {
726 82 input.read(static_cast<char*>(data), static_cast<std::streamsize>(size));
727
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 82 times.
 82 if (!input) {
728 ✗ return false;
729 }
730 }
731 82 *remaining -= size;
732 82 return true;
733 }
734
735 template <typename T>
736 56 static bool ReadPod(std::ifstream& input, T* value, uint64_t* remaining) {
737 56 return ReadBytes(input, value, sizeof(*value), remaining);
738 }
739
740 static void
741 104 UpdateChecksum(uint64_t* checksum, const void* data, size_t size) {
742
1/2
✓ Branch 0 taken 104 times.
✗ Branch 1 not taken.
 104 if (size != 0) {
743 104 *checksum = xxhash(data, size, *checksum);
744 }
745 104 }
746
747 template <typename T>
748 66 static void UpdateChecksumPod(uint64_t* checksum, const T& value) {
749 66 UpdateChecksum(checksum, &value, sizeof(value));
750 66 }
751
752 354584 void TrackKey(uint64_t key) {
753
1/2
✓ Branch 1 taken 354584 times.
✗ Branch 2 not taken.
 354584 std::lock_guard<std::mutex> lock(checkpoint_keys_mu_);
754
1/2
✓ Branch 1 taken 354584 times.
✗ Branch 2 not taken.
 354584 checkpoint_keys_.insert(key);
755 354584 }
756
757 350 void TrackKeys(base::ConstArray<uint64_t> keys) {
758
1/2
✓ Branch 1 taken 350 times.
✗ Branch 2 not taken.
 350 std::lock_guard<std::mutex> lock(checkpoint_keys_mu_);
759
2/2
✓ Branch 1 taken 1176 times.
✓ Branch 2 taken 350 times.
 1526 for (int i = 0; i < keys.Size(); ++i) {
760
1/2
✓ Branch 2 taken 1176 times.
✗ Branch 3 not taken.
 1176 checkpoint_keys_.insert(keys[i]);
761 }
762 350 }
763
764 354584 void PutInternal(uint64_t key,
765 const void* data,
766 size_t size,
767 unsigned tid,
768 bool emit_fence) {
769 (void)tid;
770 (void)emit_fence;
771 354584 Value_t new_handle = value_store_->AllocAndWrite(data, size);
772
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 354584 times.
 354584 if (new_handle == kValueHandleNone) {
773 ✗ LOG(FATAL) << "KVEngine value allocation failed, key=" << key
774 ✗ << " size=" << size;
775 return;
776 }
777 354584 Value_t old_handle = index_->Put(key, new_handle, tid);
778
2/2
✓ Branch 0 taken 156926 times.
✓ Branch 1 taken 197658 times.
 354584 if (old_handle != kValueHandleNone) {
779 156926 value_store_->Retire(old_handle);
780 }
781 354584 TrackKey(key);
782 }
783
784 inline static constexpr std::array<char, 8> kCheckpointMagic = {
785 'R', 'S', 'K', 'V', 'C', 'P', '0', '1'};
786 inline static constexpr uint32_t kCheckpointVersion = 2;
787 inline static constexpr uint64_t kCheckpointChecksumSeed =
788 0x9e3779b97f4a7c15ULL;
789
790 BaseKVConfig config_;
791 std::unique_ptr<Index> index_;
792 std::unique_ptr<ValueStore> value_store_;
793 int num_threads_ = 0;
794 size_t default_value_size_hint_ = 0;
795 mutable std::shared_mutex checkpoint_mu_;
796 mutable std::mutex checkpoint_keys_mu_;
797 std::unordered_set<uint64_t> checkpoint_keys_;
798 };
799
800 FACTORY_REGISTER(
801 BaseKV, KVEngineComposite, KVEngineComposite, const BaseKVConfig&);
802