1
2#include <c10/util/irange.h>
3#include "StoreTestCommon.hpp"
4
5#ifndef _WIN32
6#include <unistd.h>
7#endif
8
9#include <iostream>
10#include <thread>
11
12#include <gtest/gtest.h>
13
14#include <torch/csrc/distributed/c10d/FileStore.hpp>
15#include <torch/csrc/distributed/c10d/PrefixStore.hpp>
16
17#ifdef _WIN32
18std::string tmppath() {
19 return c10d::test::autoGenerateTmpFilePath();
20}
21#else
22std::string tmppath() {
23 const char* tmpdir = getenv("TMPDIR");
24 if (tmpdir == nullptr) {
25 tmpdir = "/tmp";
26 }
27
28 // Create template
29 std::vector<char> tmp(256);
30 auto len = snprintf(tmp.data(), tmp.size(), "%s/testXXXXXX", tmpdir);
31 tmp.resize(len);
32
33 // Create temporary file
34 auto fd = mkstemp(&tmp[0]);
35 if (fd == -1) {
36 throw std::system_error(errno, std::system_category());
37 }
38 close(fd);
39 return std::string(tmp.data(), tmp.size());
40}
41#endif
42
43void testGetSet(std::string path, std::string prefix = "") {
44 // Basic Set/Get on File Store
45 {
46 auto fileStore = c10::make_intrusive<c10d::FileStore>(path, 2);
47 c10d::PrefixStore store(prefix, fileStore);
48 c10d::test::set(store, "key0", "value0");
49 c10d::test::set(store, "key1", "value1");
50 c10d::test::set(store, "key2", "value2");
51 c10d::test::check(store, "key0", "value0");
52 c10d::test::check(store, "key1", "value1");
53 c10d::test::check(store, "key2", "value2");
54 auto numKeys = fileStore->getNumKeys();
55 EXPECT_EQ(numKeys, 4);
56
57 // Check compareSet, does not check return value
58 c10d::test::compareSet(store, "key0", "wrongExpectedValue", "newValue");
59 c10d::test::check(store, "key0", "value0");
60 c10d::test::compareSet(store, "key0", "value0", "newValue");
61 c10d::test::check(store, "key0", "newValue");
62
63 // Check deleteKey
64 c10d::test::deleteKey(store, "key1");
65 numKeys = fileStore->getNumKeys();
66 EXPECT_EQ(numKeys, 3);
67 c10d::test::check(store, "key0", "newValue");
68 c10d::test::check(store, "key2", "value2");
69
70 c10d::test::set(store, "-key0", "value-");
71 c10d::test::check(store, "key0", "newValue");
72 c10d::test::check(store, "-key0", "value-");
73 numKeys = fileStore->getNumKeys();
74 EXPECT_EQ(numKeys, 4);
75 c10d::test::deleteKey(store, "-key0");
76 numKeys = fileStore->getNumKeys();
77 EXPECT_EQ(numKeys, 3);
78 c10d::test::check(store, "key0", "newValue");
79 c10d::test::check(store, "key2", "value2");
80 }
81
82 // Perform get on new instance
83 {
84 auto fileStore = c10::make_intrusive<c10d::FileStore>(path, 2);
85 c10d::PrefixStore store(prefix, fileStore);
86 c10d::test::check(store, "key0", "newValue");
87 auto numKeys = fileStore->getNumKeys();
88 // There will be 4 keys since we still use the same underlying file as the
89 // other store above.
90 EXPECT_EQ(numKeys, 4);
91 }
92}
93
94void stressTestStore(std::string path, std::string prefix = "") {
95 // Hammer on FileStore::add
96 const auto numThreads = 4;
97 const auto numIterations = 100;
98
99 std::vector<std::thread> threads;
100 c10d::test::Semaphore sem1, sem2;
101
102 for (C10_UNUSED const auto i : c10::irange(numThreads)) {
103 threads.emplace_back(std::thread([&] {
104 auto fileStore =
105 c10::make_intrusive<c10d::FileStore>(path, numThreads + 1);
106 c10d::PrefixStore store(prefix, fileStore);
107 sem1.post();
108 sem2.wait();
109 for (C10_UNUSED const auto j : c10::irange(numIterations)) {
110 store.add("counter", 1);
111 }
112 }));
113 }
114
115 sem1.wait(numThreads);
116 sem2.post(numThreads);
117 for (auto& thread : threads) {
118 thread.join();
119 }
120
121 // Check that the counter has the expected value
122 {
123 auto fileStore = c10::make_intrusive<c10d::FileStore>(path, numThreads + 1);
124 c10d::PrefixStore store(prefix, fileStore);
125 std::string expected = std::to_string(numThreads * numIterations);
126 c10d::test::check(store, "counter", expected);
127 }
128}
129
130class FileStoreTest : public ::testing::Test {
131 protected:
132 void SetUp() override {
133 path_ = tmppath();
134 }
135
136 void TearDown() override {
137 unlink(path_.c_str());
138 }
139
140 std::string path_;
141};
142
143TEST_F(FileStoreTest, testGetAndSet) {
144 testGetSet(path_);
145}
146
147TEST_F(FileStoreTest, testGetAndSetWithPrefix) {
148 testGetSet(path_, "testPrefix");
149}
150
151TEST_F(FileStoreTest, testStressStore) {
152 stressTestStore(path_);
153}
154
155TEST_F(FileStoreTest, testStressStoreWithPrefix) {
156 stressTestStore(path_, "testPrefix");
157}
158