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 GFLAGS
7 #include <cstdio>
main()8 int main() {
9   fprintf(stderr, "Please install gflags to run rocksdb tools\n");
10   return 1;
11 }
12 #else
13 
14 #include "db/db_impl/db_impl.h"
15 #include "db/dbformat.h"
16 #include "file/random_access_file_reader.h"
17 #include "monitoring/histogram.h"
18 #include "rocksdb/db.h"
19 #include "rocksdb/file_system.h"
20 #include "rocksdb/slice_transform.h"
21 #include "rocksdb/system_clock.h"
22 #include "rocksdb/table.h"
23 #include "table/block_based/block_based_table_factory.h"
24 #include "table/get_context.h"
25 #include "table/internal_iterator.h"
26 #include "table/plain/plain_table_factory.h"
27 #include "table/table_builder.h"
28 #include "test_util/testharness.h"
29 #include "test_util/testutil.h"
30 #include "util/gflags_compat.h"
31 
32 using GFLAGS_NAMESPACE::ParseCommandLineFlags;
33 using GFLAGS_NAMESPACE::SetUsageMessage;
34 
35 namespace ROCKSDB_NAMESPACE {
36 
37 namespace {
38 // Make a key that i determines the first 4 characters and j determines the
39 // last 4 characters.
MakeKey(int i,int j,bool through_db)40 static std::string MakeKey(int i, int j, bool through_db) {
41   char buf[100];
42   snprintf(buf, sizeof(buf), "%04d__key___%04d", i, j);
43   if (through_db) {
44     return std::string(buf);
45   }
46   // If we directly query table, which operates on internal keys
47   // instead of user keys, we need to add 8 bytes of internal
48   // information (row type etc) to user key to make an internal
49   // key.
50   InternalKey key(std::string(buf), 0, ValueType::kTypeValue);
51   return key.Encode().ToString();
52 }
53 
Now(SystemClock * clock,bool measured_by_nanosecond)54 uint64_t Now(SystemClock* clock, bool measured_by_nanosecond) {
55   return measured_by_nanosecond ? clock->NowNanos() : clock->NowMicros();
56 }
57 }  // namespace
58 
59 // A very simple benchmark that.
60 // Create a table with roughly numKey1 * numKey2 keys,
61 // where there are numKey1 prefixes of the key, each has numKey2 number of
62 // distinguished key, differing in the suffix part.
63 // If if_query_empty_keys = false, query the existing keys numKey1 * numKey2
64 // times randomly.
65 // If if_query_empty_keys = true, query numKey1 * numKey2 random empty keys.
66 // Print out the total time.
67 // If through_db=true, a full DB will be created and queries will be against
68 // it. Otherwise, operations will be directly through table level.
69 //
70 // If for_terator=true, instead of just query one key each time, it queries
71 // a range sharing the same prefix.
72 namespace {
TableReaderBenchmark(Options & opts,EnvOptions & env_options,ReadOptions & read_options,int num_keys1,int num_keys2,int num_iter,int,bool if_query_empty_keys,bool for_iterator,bool through_db,bool measured_by_nanosecond)73 void TableReaderBenchmark(Options& opts, EnvOptions& env_options,
74                           ReadOptions& read_options, int num_keys1,
75                           int num_keys2, int num_iter, int /*prefix_len*/,
76                           bool if_query_empty_keys, bool for_iterator,
77                           bool through_db, bool measured_by_nanosecond) {
78   ROCKSDB_NAMESPACE::InternalKeyComparator ikc(opts.comparator);
79 
80   std::string file_name =
81       test::PerThreadDBPath("rocksdb_table_reader_benchmark");
82   std::string dbname = test::PerThreadDBPath("rocksdb_table_reader_bench_db");
83   WriteOptions wo;
84   Env* env = Env::Default();
85   auto* clock = env->GetSystemClock().get();
86   TableBuilder* tb = nullptr;
87   DB* db = nullptr;
88   Status s;
89   const ImmutableOptions ioptions(opts);
90   const ColumnFamilyOptions cfo(opts);
91   const MutableCFOptions moptions(cfo);
92   std::unique_ptr<WritableFileWriter> file_writer;
93   if (!through_db) {
94     ASSERT_OK(WritableFileWriter::Create(env->GetFileSystem(), file_name,
95                                          FileOptions(env_options), &file_writer,
96                                          nullptr));
97 
98     IntTblPropCollectorFactories int_tbl_prop_collector_factories;
99 
100     int unknown_level = -1;
101     tb = opts.table_factory->NewTableBuilder(
102         TableBuilderOptions(
103             ioptions, moptions, ikc, &int_tbl_prop_collector_factories,
104             CompressionType::kNoCompression, CompressionOptions(),
105             0 /* column_family_id */, kDefaultColumnFamilyName, unknown_level),
106         file_writer.get());
107   } else {
108     s = DB::Open(opts, dbname, &db);
109     ASSERT_OK(s);
110     ASSERT_TRUE(db != nullptr);
111   }
112   // Populate slightly more than 1M keys
113   for (int i = 0; i < num_keys1; i++) {
114     for (int j = 0; j < num_keys2; j++) {
115       std::string key = MakeKey(i * 2, j, through_db);
116       if (!through_db) {
117         tb->Add(key, key);
118       } else {
119         db->Put(wo, key, key);
120       }
121     }
122   }
123   if (!through_db) {
124     tb->Finish();
125     file_writer->Close();
126   } else {
127     db->Flush(FlushOptions());
128   }
129 
130   std::unique_ptr<TableReader> table_reader;
131   if (!through_db) {
132     const auto& fs = env->GetFileSystem();
133     FileOptions fopts(env_options);
134 
135     std::unique_ptr<FSRandomAccessFile> raf;
136     s = fs->NewRandomAccessFile(file_name, fopts, &raf, nullptr);
137     if (!s.ok()) {
138       fprintf(stderr, "Create File Error: %s\n", s.ToString().c_str());
139       exit(1);
140     }
141     uint64_t file_size;
142     fs->GetFileSize(file_name, fopts.io_options, &file_size, nullptr);
143     std::unique_ptr<RandomAccessFileReader> file_reader(
144         new RandomAccessFileReader(std::move(raf), file_name));
145     s = opts.table_factory->NewTableReader(
146         TableReaderOptions(ioptions, moptions.prefix_extractor.get(),
147                            env_options, ikc),
148         std::move(file_reader), file_size, &table_reader);
149     if (!s.ok()) {
150       fprintf(stderr, "Open Table Error: %s\n", s.ToString().c_str());
151       exit(1);
152     }
153   }
154 
155   Random rnd(301);
156   std::string result;
157   HistogramImpl hist;
158 
159   for (int it = 0; it < num_iter; it++) {
160     for (int i = 0; i < num_keys1; i++) {
161       for (int j = 0; j < num_keys2; j++) {
162         int r1 = rnd.Uniform(num_keys1) * 2;
163         int r2 = rnd.Uniform(num_keys2);
164         if (if_query_empty_keys) {
165           r1++;
166           r2 = num_keys2 * 2 - r2;
167         }
168 
169         if (!for_iterator) {
170           // Query one existing key;
171           std::string key = MakeKey(r1, r2, through_db);
172           uint64_t start_time = Now(clock, measured_by_nanosecond);
173           if (!through_db) {
174             PinnableSlice value;
175             MergeContext merge_context;
176             SequenceNumber max_covering_tombstone_seq = 0;
177             GetContext get_context(
178                 ioptions.user_comparator, ioptions.merge_operator.get(),
179                 ioptions.logger, ioptions.stats, GetContext::kNotFound,
180                 Slice(key), &value, nullptr, &merge_context, true,
181                 &max_covering_tombstone_seq, clock);
182             s = table_reader->Get(read_options, key, &get_context, nullptr);
183           } else {
184             s = db->Get(read_options, key, &result);
185           }
186           hist.Add(Now(clock, measured_by_nanosecond) - start_time);
187         } else {
188           int r2_len;
189           if (if_query_empty_keys) {
190             r2_len = 0;
191           } else {
192             r2_len = rnd.Uniform(num_keys2) + 1;
193             if (r2_len + r2 > num_keys2) {
194               r2_len = num_keys2 - r2;
195             }
196           }
197           std::string start_key = MakeKey(r1, r2, through_db);
198           std::string end_key = MakeKey(r1, r2 + r2_len, through_db);
199           uint64_t total_time = 0;
200           uint64_t start_time = Now(clock, measured_by_nanosecond);
201           Iterator* iter = nullptr;
202           InternalIterator* iiter = nullptr;
203           if (!through_db) {
204             iiter = table_reader->NewIterator(
205                 read_options, /*prefix_extractor=*/nullptr, /*arena=*/nullptr,
206                 /*skip_filters=*/false, TableReaderCaller::kUncategorized);
207           } else {
208             iter = db->NewIterator(read_options);
209           }
210           int count = 0;
211           for (through_db ? iter->Seek(start_key) : iiter->Seek(start_key);
212                through_db ? iter->Valid() : iiter->Valid();
213                through_db ? iter->Next() : iiter->Next()) {
214             if (if_query_empty_keys) {
215               break;
216             }
217             // verify key;
218             total_time += Now(clock, measured_by_nanosecond) - start_time;
219             assert(Slice(MakeKey(r1, r2 + count, through_db)) ==
220                    (through_db ? iter->key() : iiter->key()));
221             start_time = Now(clock, measured_by_nanosecond);
222             if (++count >= r2_len) {
223               break;
224             }
225           }
226           if (count != r2_len) {
227             fprintf(
228                 stderr, "Iterator cannot iterate expected number of entries. "
229                 "Expected %d but got %d\n", r2_len, count);
230             assert(false);
231           }
232           delete iter;
233           total_time += Now(clock, measured_by_nanosecond) - start_time;
234           hist.Add(total_time);
235         }
236       }
237     }
238   }
239 
240   fprintf(
241       stderr,
242       "==================================================="
243       "====================================================\n"
244       "InMemoryTableSimpleBenchmark: %20s   num_key1:  %5d   "
245       "num_key2: %5d  %10s\n"
246       "==================================================="
247       "===================================================="
248       "\nHistogram (unit: %s): \n%s",
249       opts.table_factory->Name(), num_keys1, num_keys2,
250       for_iterator ? "iterator" : (if_query_empty_keys ? "empty" : "non_empty"),
251       measured_by_nanosecond ? "nanosecond" : "microsecond",
252       hist.ToString().c_str());
253   if (!through_db) {
254     env->DeleteFile(file_name);
255   } else {
256     delete db;
257     db = nullptr;
258     DestroyDB(dbname, opts);
259   }
260 }
261 }  // namespace
262 }  // namespace ROCKSDB_NAMESPACE
263 
264 DEFINE_bool(query_empty, false, "query non-existing keys instead of existing "
265             "ones.");
266 DEFINE_int32(num_keys1, 4096, "number of distinguish prefix of keys");
267 DEFINE_int32(num_keys2, 512, "number of distinguish keys for each prefix");
268 DEFINE_int32(iter, 3, "query non-existing keys instead of existing ones");
269 DEFINE_int32(prefix_len, 16, "Prefix length used for iterators and indexes");
270 DEFINE_bool(iterator, false, "For test iterator");
271 DEFINE_bool(through_db, false, "If enable, a DB instance will be created and "
272             "the query will be against DB. Otherwise, will be directly against "
273             "a table reader.");
274 DEFINE_bool(mmap_read, true, "Whether use mmap read");
275 DEFINE_string(table_factory, "block_based",
276               "Table factory to use: `block_based` (default), `plain_table` or "
277               "`cuckoo_hash`.");
278 DEFINE_string(time_unit, "microsecond",
279               "The time unit used for measuring performance. User can specify "
280               "`microsecond` (default) or `nanosecond`");
281 
main(int argc,char ** argv)282 int main(int argc, char** argv) {
283   SetUsageMessage(std::string("\nUSAGE:\n") + std::string(argv[0]) +
284                   " [OPTIONS]...");
285   ParseCommandLineFlags(&argc, &argv, true);
286 
287   std::shared_ptr<ROCKSDB_NAMESPACE::TableFactory> tf;
288   ROCKSDB_NAMESPACE::Options options;
289   if (FLAGS_prefix_len < 16) {
290     options.prefix_extractor.reset(
291         ROCKSDB_NAMESPACE::NewFixedPrefixTransform(FLAGS_prefix_len));
292   }
293   ROCKSDB_NAMESPACE::ReadOptions ro;
294   ROCKSDB_NAMESPACE::EnvOptions env_options;
295   options.create_if_missing = true;
296   options.compression = ROCKSDB_NAMESPACE::CompressionType::kNoCompression;
297 
298   if (FLAGS_table_factory == "cuckoo_hash") {
299 #ifndef ROCKSDB_LITE
300     options.allow_mmap_reads = FLAGS_mmap_read;
301     env_options.use_mmap_reads = FLAGS_mmap_read;
302     ROCKSDB_NAMESPACE::CuckooTableOptions table_options;
303     table_options.hash_table_ratio = 0.75;
304     tf.reset(ROCKSDB_NAMESPACE::NewCuckooTableFactory(table_options));
305 #else
306     fprintf(stderr, "Plain table is not supported in lite mode\n");
307     exit(1);
308 #endif  // ROCKSDB_LITE
309   } else if (FLAGS_table_factory == "plain_table") {
310 #ifndef ROCKSDB_LITE
311     options.allow_mmap_reads = FLAGS_mmap_read;
312     env_options.use_mmap_reads = FLAGS_mmap_read;
313 
314     ROCKSDB_NAMESPACE::PlainTableOptions plain_table_options;
315     plain_table_options.user_key_len = 16;
316     plain_table_options.bloom_bits_per_key = (FLAGS_prefix_len == 16) ? 0 : 8;
317     plain_table_options.hash_table_ratio = 0.75;
318 
319     tf.reset(new ROCKSDB_NAMESPACE::PlainTableFactory(plain_table_options));
320     options.prefix_extractor.reset(
321         ROCKSDB_NAMESPACE::NewFixedPrefixTransform(FLAGS_prefix_len));
322 #else
323     fprintf(stderr, "Cuckoo table is not supported in lite mode\n");
324     exit(1);
325 #endif  // ROCKSDB_LITE
326   } else if (FLAGS_table_factory == "block_based") {
327     tf.reset(new ROCKSDB_NAMESPACE::BlockBasedTableFactory());
328   } else {
329     fprintf(stderr, "Invalid table type %s\n", FLAGS_table_factory.c_str());
330   }
331 
332   if (tf) {
333     // if user provides invalid options, just fall back to microsecond.
334     bool measured_by_nanosecond = FLAGS_time_unit == "nanosecond";
335 
336     options.table_factory = tf;
337     ROCKSDB_NAMESPACE::TableReaderBenchmark(
338         options, env_options, ro, FLAGS_num_keys1, FLAGS_num_keys2, FLAGS_iter,
339         FLAGS_prefix_len, FLAGS_query_empty, FLAGS_iterator, FLAGS_through_db,
340         measured_by_nanosecond);
341   } else {
342     return 1;
343   }
344 
345   return 0;
346 }
347 
348 #endif  // GFLAGS
349