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 | 8 | if (j.contains("value_size")) { | |
| 24 | 8 | return j.at("value_size").get<size_t>(); | |
| 25 | } | ||
| 26 | ✗ | if (j.contains("value")) { | |
| 27 | ✗ | 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 | 8 | if (!config.json_config_.contains("detable")) { | |
| 34 | ✗ | 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 | 6 | if (value_size == 0 || value_size % sizeof(float) != 0) { | |
| 41 | ✗ | throw std::invalid_argument( | |
| 42 | ✗ | 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 | ✗ | error[0] == '\0' ? "unknown provider error" : std::string(error.data()); | |
| 51 | return std::runtime_error( | ||
| 52 | ✗ | std::string("KVEngineDETable ") + operation + " failed: " + detail); | |
| 53 | ✗ | } | |
| 54 | |||
| 55 | 6 | bool HasRequiredFunctions(const DynamicKVPluginApiV1& api) { | |
| 56 | 6 | return api.create != nullptr && api.destroy != nullptr && | |
| 57 | 6 | api.get != nullptr && api.exists != nullptr && api.put != nullptr && | |
| 58 | 18 | api.batch_get != nullptr && api.batch_put != nullptr && | |
| 59 | 12 | api.clear != nullptr; | |
| 60 | } | ||
| 61 | |||
| 62 | } // namespace | ||
| 63 | |||
| 64 | class KVEngineDETable : public BaseKV { | ||
| 65 | public: | ||
| 66 | 8 | explicit KVEngineDETable(const BaseKVConfig& config) | |
| 67 | 8 | : BaseKV(config), value_size_(ConfigValueSize(config)) { | |
| 68 | 8 | if (value_size_ == 0) { | |
| 69 | ✗ | throw std::invalid_argument("KVEngineDETable requires value_size"); | |
| 70 | } | ||
| 71 | |||
| 72 | 8 | const auto& plugin = PluginConfig(config); | |
| 73 | const std::string library_path = | ||
| 74 | 8 | plugin.value("library_path", std::string()); | |
| 75 | 8 | if (library_path.empty()) { | |
| 76 | 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 | 6 | if (library_ == nullptr) { | |
| 82 | ✗ | throw std::runtime_error( | |
| 83 | ✗ | "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 | 6 | if (symbol_error != nullptr || entry == nullptr) { | |
| 93 | ✗ | CloseLibrary(); | |
| 94 | ✗ | throw std::runtime_error( | |
| 95 | ✗ | "KVEngineDETable provider entry point is unavailable: " + | |
| 96 | ✗ | std::string( | |
| 97 | ✗ | symbol_error == nullptr ? "unknown error" : symbol_error)); | |
| 98 | } | ||
| 99 | |||
| 100 | 6 | api_ = entry(); | |
| 101 | 18 | if (api_ == nullptr || | |
| 102 | 6 | api_->abi_version != recstore::storage::plugin::kDynamicKVPluginAbiV1 || | |
| 103 | 18 | api_->struct_size < sizeof(DynamicKVPluginApiV1) || | |
| 104 | 6 | !HasRequiredFunctions(*api_)) { | |
| 105 | ✗ | CloseLibrary(); | |
| 106 | ✗ | throw std::runtime_error("KVEngineDETable provider ABI mismatch"); | |
| 107 | } | ||
| 108 | |||
| 109 | 6 | DynamicKVPluginConfigV1 plugin_config{ | |
| 110 | recstore::storage::plugin::kDynamicKVPluginAbiV1, | ||
| 111 | sizeof(DynamicKVPluginConfigV1), | ||
| 112 | 6 | config.json_config_.at("capacity").get<uint64_t>(), | |
| 113 | 6 | value_size_, | |
| 114 | 6 | plugin.value("block_size", uint64_t{10240})}; | |
| 115 | 6 | std::array<char, kErrorBufferSize> error{}; | |
| 116 | 6 | if (api_->create(&plugin_config, &handle_, error.data(), error.size()) != | |
| 117 | 12 | 0 || | |
| 118 | 6 | handle_ == nullptr) { | |
| 119 | ✗ | const auto exception = PluginError("create", error); | |
| 120 | ✗ | if (handle_ != nullptr) { | |
| 121 | ✗ | api_->destroy(handle_); | |
| 122 | ✗ | handle_ = nullptr; | |
| 123 | } | ||
| 124 | ✗ | CloseLibrary(); | |
| 125 | ✗ | throw exception; | |
| 126 | ✗ | } | |
| 127 | 10 | } | |
| 128 | |||
| 129 | 12 | ~KVEngineDETable() override { | |
| 130 | 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 | 8 | value.resize(value_size_); | |
| 139 | 8 | uint8_t found = 0; | |
| 140 | 8 | std::array<char, kErrorBufferSize> error{}; | |
| 141 | 16 | if (api_->get(handle_, | |
| 142 | key, | ||
| 143 | 8 | value.data(), | |
| 144 | value_size_, | ||
| 145 | &found, | ||
| 146 | error.data(), | ||
| 147 | 8 | error.size()) != 0) { | |
| 148 | ✗ | throw PluginError("get", error); | |
| 149 | } | ||
| 150 | 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 | if (api_->exists(handle_, key, &found, error.data(), error.size()) != 0) { | |
| 160 | ✗ | 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 | 4 | if (value.size() != value_size_) { | |
| 168 | ✗ | throw std::invalid_argument("KVEngineDETable requires fixed-size Put"); | |
| 169 | } | ||
| 170 | 4 | std::array<char, kErrorBufferSize> error{}; | |
| 171 | 8 | if (api_->put(handle_, | |
| 172 | key, | ||
| 173 | 4 | value.data(), | |
| 174 | value.size(), | ||
| 175 | error.data(), | ||
| 176 | 4 | error.size()) != 0) { | |
| 177 | ✗ | 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 | 4 | ValidateFloatAligned(value_size_, "KVEngineDETable::BatchPut"); | |
| 186 | 4 | if (values == nullptr || keys.Size() != static_cast<int>(values->size())) { | |
| 187 | ✗ | throw std::invalid_argument("KVEngineDETable::BatchPut size mismatch"); | |
| 188 | } | ||
| 189 | 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 | 14 | for (int i = 0; i < keys.Size(); ++i) { | |
| 192 | 10 | if ((*values)[i].Size() != floats_per_row) { | |
| 193 | ✗ | 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 | 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 | 2 | ValidateFloatAligned(value_size_, "KVEngineDETable::BatchGet"); | |
| 208 | 2 | if (values == nullptr) { | |
| 209 | ✗ | throw std::invalid_argument("KVEngineDETable::BatchGet values is null"); | |
| 210 | } | ||
| 211 | 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 | buffers.assign(keys.Size(), std::vector<float>(floats_per_row)); | |
| 214 | 2 | std::vector<char> flat(static_cast<size_t>(keys.Size()) * value_size_); | |
| 215 | 2 | std::vector<uint8_t> found(keys.Size(), 0); | |
| 216 | 2 | BatchGetRaw(keys, flat.data(), found.data()); | |
| 217 | |||
| 218 | 2 | values->resize(keys.Size()); | |
| 219 | 8 | for (int i = 0; i < keys.Size(); ++i) { | |
| 220 | 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 | 2 | if (values == nullptr || num_rows != keys.Size() || num_rows < 0 || | |
| 240 | 4 | embedding_dim <= 0 || | |
| 241 | 2 | static_cast<size_t>(embedding_dim) * sizeof(float) != value_size_) { | |
| 242 | ✗ | throw std::invalid_argument( | |
| 243 | ✗ | "KVEngineDETable::BatchGetFlat shape mismatch"); | |
| 244 | } | ||
| 245 | 2 | std::vector<uint8_t> found(keys.Size(), 0); | |
| 246 | 2 | BatchGetRaw(keys, values, found.data()); | |
| 247 | 2 | uint64_t missing = 0; | |
| 248 | 6 | for (int i = 0; i < keys.Size(); ++i) { | |
| 249 | 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 | 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 | ✗ | if (value == nullptr && keys.Size() > 0) { | |
| 263 | ✗ | 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 | if (api_->clear(handle_, error.data(), error.size()) != 0) { | |
| 271 | ✗ | 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 | 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 | 4 | error.size()) != 0) { | |
| 288 | ✗ | 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 | 8 | if (api_->batch_put( | |
| 295 | handle_, | ||
| 296 | keys.Data(), | ||
| 297 | 4 | keys.Size(), | |
| 298 | values, | ||
| 299 | value_size_, | ||
| 300 | error.data(), | ||
| 301 | 4 | error.size()) != 0) { | |
| 302 | ✗ | throw PluginError("batch_put", error); | |
| 303 | } | ||
| 304 | 4 | } | |
| 305 | |||
| 306 | 6 | void CloseLibrary() { | |
| 307 | 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 |