1#include <torch/csrc/distributed/c10d/PrefixStore.hpp>
2#include <utility>
3
4namespace c10d {
5
6PrefixStore::PrefixStore(std::string prefix, c10::intrusive_ptr<Store> store)
7 : prefix_(std::move(prefix)), store_(std::move(store)) {}
8
9std::string PrefixStore::joinKey(const std::string& key) {
10 return prefix_ + "/" + key;
11}
12
13std::vector<std::string> PrefixStore::joinKeys(
14 const std::vector<std::string>& keys) {
15 std::vector<std::string> joinedKeys;
16 joinedKeys.reserve(keys.size());
17 for (const auto& key : keys) {
18 joinedKeys.emplace_back(joinKey(key));
19 }
20 return joinedKeys;
21}
22
23void PrefixStore::set(
24 const std::string& key,
25 const std::vector<uint8_t>& value) {
26 store_->set(joinKey(key), value);
27}
28
29std::vector<uint8_t> PrefixStore::compareSet(
30 const std::string& key,
31 const std::vector<uint8_t>& expectedValue,
32 const std::vector<uint8_t>& desiredValue) {
33 return store_->compareSet(joinKey(key), expectedValue, desiredValue);
34}
35
36std::vector<uint8_t> PrefixStore::get(const std::string& key) {
37 return store_->get(joinKey(key));
38}
39
40int64_t PrefixStore::add(const std::string& key, int64_t value) {
41 return store_->add(joinKey(key), value);
42}
43
44bool PrefixStore::deleteKey(const std::string& key) {
45 return store_->deleteKey(joinKey(key));
46}
47
48void PrefixStore::watchKey(const std::string& key, WatchKeyCallback callback) {
49 return store_->watchKey(joinKey(key), std::move(callback));
50}
51
52int64_t PrefixStore::getNumKeys() {
53 return store_->getNumKeys();
54}
55
56bool PrefixStore::check(const std::vector<std::string>& keys) {
57 auto joinedKeys = joinKeys(keys);
58 return store_->check(joinedKeys);
59}
60
61void PrefixStore::wait(const std::vector<std::string>& keys) {
62 auto joinedKeys = joinKeys(keys);
63 store_->wait(joinedKeys);
64}
65
66void PrefixStore::wait(
67 const std::vector<std::string>& keys,
68 const std::chrono::milliseconds& timeout) {
69 auto joinedKeys = joinKeys(keys);
70 store_->wait(joinedKeys, timeout);
71}
72
73const std::chrono::milliseconds& PrefixStore::getTimeout() const noexcept {
74 return store_->getTimeout();
75}
76
77void PrefixStore::setTimeout(const std::chrono::milliseconds& timeout) {
78 store_->setTimeout(timeout);
79}
80
81c10::intrusive_ptr<Store> PrefixStore::getUnderlyingStore() {
82 return store_;
83}
84
85} // namespace c10d
86