1 // Copyright (c) 2011-present, Facebook, Inc. All rights reserved.
2 // This source code is licensed under both the GPLv2 (found in the
3 // COPYING file in the root directory) and Apache 2.0 License
4 // (found in the LICENSE.Apache file in the root directory).
5
6 #ifndef ROCKSDB_LITE
7
8 #include <cinttypes>
9
10 #include "port/stack_trace.h"
11 #include "rocksdb/db.h"
12 #include "rocksdb/sst_file_reader.h"
13 #include "rocksdb/sst_file_writer.h"
14 #include "table/sst_file_writer_collectors.h"
15 #include "test_util/testharness.h"
16 #include "test_util/testutil.h"
17 #include "utilities/merge_operators.h"
18
19 namespace ROCKSDB_NAMESPACE {
20
EncodeAsString(uint64_t v)21 std::string EncodeAsString(uint64_t v) {
22 char buf[16];
23 snprintf(buf, sizeof(buf), "%08" PRIu64, v);
24 return std::string(buf);
25 }
26
EncodeAsUint64(uint64_t v)27 std::string EncodeAsUint64(uint64_t v) {
28 std::string dst;
29 PutFixed64(&dst, v);
30 return dst;
31 }
32
33 class SstFileReaderTest : public testing::Test {
34 public:
SstFileReaderTest()35 SstFileReaderTest() {
36 options_.merge_operator = MergeOperators::CreateUInt64AddOperator();
37 sst_name_ = test::PerThreadDBPath("sst_file");
38
39 Env* base_env = Env::Default();
40 const char* test_env_uri = getenv("TEST_ENV_URI");
41 if(test_env_uri) {
42 Env* test_env = nullptr;
43 Status s = Env::LoadEnv(test_env_uri, &test_env, &env_guard_);
44 base_env = test_env;
45 EXPECT_OK(s);
46 EXPECT_NE(Env::Default(), base_env);
47 }
48 EXPECT_NE(nullptr, base_env);
49 env_ = base_env;
50 options_.env = env_;
51 }
52
~SstFileReaderTest()53 ~SstFileReaderTest() {
54 Status s = env_->DeleteFile(sst_name_);
55 EXPECT_OK(s);
56 }
57
CreateFile(const std::string & file_name,const std::vector<std::string> & keys)58 void CreateFile(const std::string& file_name,
59 const std::vector<std::string>& keys) {
60 SstFileWriter writer(soptions_, options_);
61 ASSERT_OK(writer.Open(file_name));
62 for (size_t i = 0; i + 2 < keys.size(); i += 3) {
63 ASSERT_OK(writer.Put(keys[i], keys[i]));
64 ASSERT_OK(writer.Merge(keys[i + 1], EncodeAsUint64(i + 1)));
65 ASSERT_OK(writer.Delete(keys[i + 2]));
66 }
67 ASSERT_OK(writer.Finish());
68 }
69
CheckFile(const std::string & file_name,const std::vector<std::string> & keys,bool check_global_seqno=false)70 void CheckFile(const std::string& file_name,
71 const std::vector<std::string>& keys,
72 bool check_global_seqno = false) {
73 ReadOptions ropts;
74 SstFileReader reader(options_);
75 ASSERT_OK(reader.Open(file_name));
76 ASSERT_OK(reader.VerifyChecksum());
77 std::unique_ptr<Iterator> iter(reader.NewIterator(ropts));
78 iter->SeekToFirst();
79 for (size_t i = 0; i + 2 < keys.size(); i += 3) {
80 ASSERT_TRUE(iter->Valid());
81 ASSERT_EQ(iter->key().compare(keys[i]), 0);
82 ASSERT_EQ(iter->value().compare(keys[i]), 0);
83 iter->Next();
84 ASSERT_TRUE(iter->Valid());
85 ASSERT_EQ(iter->key().compare(keys[i + 1]), 0);
86 ASSERT_EQ(iter->value().compare(EncodeAsUint64(i + 1)), 0);
87 iter->Next();
88 }
89 ASSERT_FALSE(iter->Valid());
90 if (check_global_seqno) {
91 auto properties = reader.GetTableProperties();
92 ASSERT_TRUE(properties);
93 std::string hostname;
94 ASSERT_OK(env_->GetHostNameString(&hostname));
95 ASSERT_EQ(properties->db_host_id, hostname);
96 auto& user_properties = properties->user_collected_properties;
97 ASSERT_TRUE(
98 user_properties.count(ExternalSstFilePropertyNames::kGlobalSeqno));
99 }
100 }
101
CreateFileAndCheck(const std::vector<std::string> & keys)102 void CreateFileAndCheck(const std::vector<std::string>& keys) {
103 CreateFile(sst_name_, keys);
104 CheckFile(sst_name_, keys);
105 }
106
107 protected:
108 Options options_;
109 EnvOptions soptions_;
110 std::string sst_name_;
111 std::shared_ptr<Env> env_guard_;
112 Env* env_;
113 };
114
115 const uint64_t kNumKeys = 100;
116
TEST_F(SstFileReaderTest,Basic)117 TEST_F(SstFileReaderTest, Basic) {
118 std::vector<std::string> keys;
119 for (uint64_t i = 0; i < kNumKeys; i++) {
120 keys.emplace_back(EncodeAsString(i));
121 }
122 CreateFileAndCheck(keys);
123 }
124
TEST_F(SstFileReaderTest,Uint64Comparator)125 TEST_F(SstFileReaderTest, Uint64Comparator) {
126 options_.comparator = test::Uint64Comparator();
127 std::vector<std::string> keys;
128 for (uint64_t i = 0; i < kNumKeys; i++) {
129 keys.emplace_back(EncodeAsUint64(i));
130 }
131 CreateFileAndCheck(keys);
132 }
133
TEST_F(SstFileReaderTest,ReadOptionsOutOfScope)134 TEST_F(SstFileReaderTest, ReadOptionsOutOfScope) {
135 // Repro a bug where the SstFileReader depended on its configured ReadOptions
136 // outliving it.
137 options_.comparator = test::Uint64Comparator();
138 std::vector<std::string> keys;
139 for (uint64_t i = 0; i < kNumKeys; i++) {
140 keys.emplace_back(EncodeAsUint64(i));
141 }
142 CreateFile(sst_name_, keys);
143
144 SstFileReader reader(options_);
145 ASSERT_OK(reader.Open(sst_name_));
146 std::unique_ptr<Iterator> iter;
147 {
148 // Make sure ReadOptions go out of scope ASAP so we know the iterator
149 // operations do not depend on it.
150 ReadOptions ropts;
151 iter.reset(reader.NewIterator(ropts));
152 }
153 iter->SeekToFirst();
154 while (iter->Valid()) {
155 iter->Next();
156 }
157 }
158
TEST_F(SstFileReaderTest,ReadFileWithGlobalSeqno)159 TEST_F(SstFileReaderTest, ReadFileWithGlobalSeqno) {
160 std::vector<std::string> keys;
161 for (uint64_t i = 0; i < kNumKeys; i++) {
162 keys.emplace_back(EncodeAsString(i));
163 }
164 // Generate a SST file.
165 CreateFile(sst_name_, keys);
166
167 // Ingest the file into a db, to assign it a global sequence number.
168 Options options;
169 options.create_if_missing = true;
170 std::string db_name = test::PerThreadDBPath("test_db");
171 DB* db;
172 ASSERT_OK(DB::Open(options, db_name, &db));
173 // Bump sequence number.
174 ASSERT_OK(db->Put(WriteOptions(), keys[0], "foo"));
175 ASSERT_OK(db->Flush(FlushOptions()));
176 // Ingest the file.
177 IngestExternalFileOptions ingest_options;
178 ingest_options.write_global_seqno = true;
179 ASSERT_OK(db->IngestExternalFile({sst_name_}, ingest_options));
180 std::vector<std::string> live_files;
181 uint64_t manifest_file_size = 0;
182 ASSERT_OK(db->GetLiveFiles(live_files, &manifest_file_size));
183 // Get the ingested file.
184 std::string ingested_file;
185 for (auto& live_file : live_files) {
186 if (live_file.substr(live_file.size() - 4, std::string::npos) == ".sst") {
187 if (ingested_file.empty() || ingested_file < live_file) {
188 ingested_file = live_file;
189 }
190 }
191 }
192 ASSERT_FALSE(ingested_file.empty());
193 delete db;
194
195 // Verify the file can be open and read by SstFileReader.
196 CheckFile(db_name + ingested_file, keys, true /* check_global_seqno */);
197
198 // Cleanup.
199 ASSERT_OK(DestroyDB(db_name, options));
200 }
201
202 } // namespace ROCKSDB_NAMESPACE
203
204 #ifdef ROCKSDB_UNITTESTS_WITH_CUSTOM_OBJECTS_FROM_STATIC_LIBS
205 extern "C" {
206 void RegisterCustomObjects(int argc, char** argv);
207 }
208 #else
RegisterCustomObjects(int,char **)209 void RegisterCustomObjects(int /*argc*/, char** /*argv*/) {}
210 #endif // !ROCKSDB_UNITTESTS_WITH_CUSTOM_OBJECTS_FROM_STATIC_LIBS
211
main(int argc,char ** argv)212 int main(int argc, char** argv) {
213 ROCKSDB_NAMESPACE::port::InstallStackTraceHandler();
214 ::testing::InitGoogleTest(&argc, argv);
215 RegisterCustomObjects(argc, argv);
216 return RUN_ALL_TESTS();
217 }
218
219 #else
220 #include <stdio.h>
221
main(int,char **)222 int main(int /*argc*/, char** /*argv*/) {
223 fprintf(stderr,
224 "SKIPPED as SstFileReader is not supported in ROCKSDB_LITE\n");
225 return 0;
226 }
227
228 #endif // ROCKSDB_LITE
229