ceph/src/rocksdb/db/comparator_db_test.cc

   1 // Use of this source code is governed by a BSD-style license that can be
   2 // found in the LICENSE file. See the AUTHORS file for names of contributors.
   3 //  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
   4 //  This source code is licensed under the BSD-style license found in the
   5 //  LICENSE file in the root directory of this source tree. An additional grant
   6 //  of patent rights can be found in the PATENTS file in the same directory.
   7 #include <map>
   8 #include <string>
   9
  10 #include "memtable/stl_wrappers.h"
  11 #include "rocksdb/db.h"
  12 #include "rocksdb/env.h"
  13 #include "util/hash.h"
  14 #include "util/kv_map.h"
  15 #include "util/string_util.h"
  16 #include "util/testharness.h"
  17 #include "util/testutil.h"
  18 #include "utilities/merge_operators.h"
  19
  20 using std::unique_ptr;
  21
  22 namespace rocksdb {
  23 namespace {
  24
  25 static const Comparator* comparator;
  26
  27 class KVIter : public Iterator {
  28  public:
  29   explicit KVIter(const stl_wrappers::KVMap* map)
  30       : map_(map), iter_(map_->end()) {}
  31   virtual bool Valid() const override { return iter_ != map_->end(); }
  32   virtual void SeekToFirst() override { iter_ = map_->begin(); }
  33   virtual void SeekToLast() override {
  34     if (map_->empty()) {
  35       iter_ = map_->end();
  36     } else {
  37       iter_ = map_->find(map_->rbegin()->first);
  38     }
  39   }
  40   virtual void Seek(const Slice& k) override {
  41     iter_ = map_->lower_bound(k.ToString());
  42   }
  43   virtual void SeekForPrev(const Slice& k) override {
  44     iter_ = map_->upper_bound(k.ToString());
  45     Prev();
  46   }
  47   virtual void Next() override { ++iter_; }
  48   virtual void Prev() override {
  49     if (iter_ == map_->begin()) {
  50       iter_ = map_->end();
  51       return;
  52     }
  53     --iter_;
  54   }
  55
  56   virtual Slice key() const override { return iter_->first; }
  57   virtual Slice value() const override { return iter_->second; }
  58   virtual Status status() const override { return Status::OK(); }
  59
  60  private:
  61   const stl_wrappers::KVMap* const map_;
  62   stl_wrappers::KVMap::const_iterator iter_;
  63 };
  64
  65 void AssertItersEqual(Iterator* iter1, Iterator* iter2) {
  66   ASSERT_EQ(iter1->Valid(), iter2->Valid());
  67   if (iter1->Valid()) {
  68     ASSERT_EQ(iter1->key().ToString(), iter2->key().ToString());
  69     ASSERT_EQ(iter1->value().ToString(), iter2->value().ToString());
  70   }
  71 }
  72
  73 // Measuring operations on DB (expect to be empty).
  74 // source_strings are candidate keys
  75 void DoRandomIteraratorTest(DB* db, std::vector<std::string> source_strings,
  76                             Random* rnd, int num_writes, int num_iter_ops,
  77                             int num_trigger_flush) {
  78   stl_wrappers::KVMap map((stl_wrappers::LessOfComparator(comparator)));
  79
  80   for (int i = 0; i < num_writes; i++) {
  81     if (num_trigger_flush > 0 && i != 0 && i % num_trigger_flush == 0) {
  82       db->Flush(FlushOptions());
  83     }
  84
  85     int type = rnd->Uniform(2);
  86     int index = rnd->Uniform(static_cast<int>(source_strings.size()));
  87     auto& key = source_strings[index];
  88     switch (type) {
  89       case 0:
  90         // put
  91         map[key] = key;
  92         ASSERT_OK(db->Put(WriteOptions(), key, key));
  93         break;
  94       case 1:
  95         // delete
  96         if (map.find(key) != map.end()) {
  97           map.erase(key);
  98         }
  99         ASSERT_OK(db->Delete(WriteOptions(), key));
 100         break;
 101       default:
 102         assert(false);
 103     }
 104   }
 105
 106   std::unique_ptr<Iterator> iter(db->NewIterator(ReadOptions()));
 107   std::unique_ptr<Iterator> result_iter(new KVIter(&map));
 108
 109   bool is_valid = false;
 110   for (int i = 0; i < num_iter_ops; i++) {
 111     // Random walk and make sure iter and result_iter returns the
 112     // same key and value
 113     int type = rnd->Uniform(6);
 114     ASSERT_OK(iter->status());
 115     switch (type) {
 116       case 0:
 117         // Seek to First
 118         iter->SeekToFirst();
 119         result_iter->SeekToFirst();
 120         break;
 121       case 1:
 122         // Seek to last
 123         iter->SeekToLast();
 124         result_iter->SeekToLast();
 125         break;
 126       case 2: {
 127         // Seek to random key
 128         auto key_idx = rnd->Uniform(static_cast<int>(source_strings.size()));
 129         auto key = source_strings[key_idx];
 130         iter->Seek(key);
 131         result_iter->Seek(key);
 132         break;
 133       }
 134       case 3:
 135         // Next
 136         if (is_valid) {
 137           iter->Next();
 138           result_iter->Next();
 139         } else {
 140           continue;
 141         }
 142         break;
 143       case 4:
 144         // Prev
 145         if (is_valid) {
 146           iter->Prev();
 147           result_iter->Prev();
 148         } else {
 149           continue;
 150         }
 151         break;
 152       default: {
 153         assert(type == 5);
 154         auto key_idx = rnd->Uniform(static_cast<int>(source_strings.size()));
 155         auto key = source_strings[key_idx];
 156         std::string result;
 157         auto status = db->Get(ReadOptions(), key, &result);
 158         if (map.find(key) == map.end()) {
 159           ASSERT_TRUE(status.IsNotFound());
 160         } else {
 161           ASSERT_EQ(map[key], result);
 162         }
 163         break;
 164       }
 165     }
 166     AssertItersEqual(iter.get(), result_iter.get());
 167     is_valid = iter->Valid();
 168   }
 169 }
 170
 171 class DoubleComparator : public Comparator {
 172  public:
 173   DoubleComparator() {}
 174
 175   virtual const char* Name() const override { return "DoubleComparator"; }
 176
 177   virtual int Compare(const Slice& a, const Slice& b) const override {
 178 #ifndef CYGWIN
 179     double da = std::stod(a.ToString());
 180     double db = std::stod(b.ToString());
 181 #else
 182     double da = std::strtod(a.ToString().c_str(), 0 /* endptr */);
 183     double db = std::strtod(a.ToString().c_str(), 0 /* endptr */);
 184 #endif
 185     if (da == db) {
 186       return a.compare(b);
 187     } else if (da > db) {
 188       return 1;
 189     } else {
 190       return -1;
 191     }
 192   }
 193   virtual void FindShortestSeparator(std::string* start,
 194                                      const Slice& limit) const override {}
 195
 196   virtual void FindShortSuccessor(std::string* key) const override {}
 197 };
 198
 199 class HashComparator : public Comparator {
 200  public:
 201   HashComparator() {}
 202
 203   virtual const char* Name() const override { return "HashComparator"; }
 204
 205   virtual int Compare(const Slice& a, const Slice& b) const override {
 206     uint32_t ha = Hash(a.data(), a.size(), 66);
 207     uint32_t hb = Hash(b.data(), b.size(), 66);
 208     if (ha == hb) {
 209       return a.compare(b);
 210     } else if (ha > hb) {
 211       return 1;
 212     } else {
 213       return -1;
 214     }
 215   }
 216   virtual void FindShortestSeparator(std::string* start,
 217                                      const Slice& limit) const override {}
 218
 219   virtual void FindShortSuccessor(std::string* key) const override {}
 220 };
 221
 222 class TwoStrComparator : public Comparator {
 223  public:
 224   TwoStrComparator() {}
 225
 226   virtual const char* Name() const override { return "TwoStrComparator"; }
 227
 228   virtual int Compare(const Slice& a, const Slice& b) const override {
 229     assert(a.size() >= 2);
 230     assert(b.size() >= 2);
 231     size_t size_a1 = static_cast<size_t>(a[0]);
 232     size_t size_b1 = static_cast<size_t>(b[0]);
 233     size_t size_a2 = static_cast<size_t>(a[1]);
 234     size_t size_b2 = static_cast<size_t>(b[1]);
 235     assert(size_a1 + size_a2 + 2 == a.size());
 236     assert(size_b1 + size_b2 + 2 == b.size());
 237
 238     Slice a1 = Slice(a.data() + 2, size_a1);
 239     Slice b1 = Slice(b.data() + 2, size_b1);
 240     Slice a2 = Slice(a.data() + 2 + size_a1, size_a2);
 241     Slice b2 = Slice(b.data() + 2 + size_b1, size_b2);
 242
 243     if (a1 != b1) {
 244       return a1.compare(b1);
 245     }
 246     return a2.compare(b2);
 247   }
 248   virtual void FindShortestSeparator(std::string* start,
 249                                      const Slice& limit) const override {}
 250
 251   virtual void FindShortSuccessor(std::string* key) const override {}
 252 };
 253 }  // namespace
 254
 255 class ComparatorDBTest : public testing::Test {
 256  private:
 257   std::string dbname_;
 258   Env* env_;
 259   DB* db_;
 260   Options last_options_;
 261   std::unique_ptr<const Comparator> comparator_guard;
 262
 263  public:
 264   ComparatorDBTest() : env_(Env::Default()), db_(nullptr) {
 265     comparator = BytewiseComparator();
 266     dbname_ = test::TmpDir() + "/comparator_db_test";
 267     EXPECT_OK(DestroyDB(dbname_, last_options_));
 268   }
 269
 270   ~ComparatorDBTest() {
 271     delete db_;
 272     EXPECT_OK(DestroyDB(dbname_, last_options_));
 273     comparator = BytewiseComparator();
 274   }
 275
 276   DB* GetDB() { return db_; }
 277
 278   void SetOwnedComparator(const Comparator* cmp) {
 279     comparator_guard.reset(cmp);
 280     comparator = cmp;
 281     last_options_.comparator = cmp;
 282   }
 283
 284   // Return the current option configuration.
 285   Options* GetOptions() { return &last_options_; }
 286
 287   void DestroyAndReopen() {
 288     // Destroy using last options
 289     Destroy();
 290     ASSERT_OK(TryReopen());
 291   }
 292
 293   void Destroy() {
 294     delete db_;
 295     db_ = nullptr;
 296     ASSERT_OK(DestroyDB(dbname_, last_options_));
 297   }
 298
 299   Status TryReopen() {
 300     delete db_;
 301     db_ = nullptr;
 302     last_options_.create_if_missing = true;
 303
 304     return DB::Open(last_options_, dbname_, &db_);
 305   }
 306 };
 307
 308 TEST_F(ComparatorDBTest, Bytewise) {
 309   for (int rand_seed = 301; rand_seed < 306; rand_seed++) {
 310     DestroyAndReopen();
 311     Random rnd(rand_seed);
 312     DoRandomIteraratorTest(GetDB(),
 313                            {"a", "b", "c", "d", "e", "f", "g", "h", "i"}, &rnd,
 314                            8, 100, 3);
 315   }
 316 }
 317
 318 TEST_F(ComparatorDBTest, SimpleSuffixReverseComparator) {
 319   SetOwnedComparator(new test::SimpleSuffixReverseComparator());
 320
 321   for (int rnd_seed = 301; rnd_seed < 316; rnd_seed++) {
 322     Options* opt = GetOptions();
 323     opt->comparator = comparator;
 324     DestroyAndReopen();
 325     Random rnd(rnd_seed);
 326
 327     std::vector<std::string> source_strings;
 328     std::vector<std::string> source_prefixes;
 329     // Randomly generate 5 prefixes
 330     for (int i = 0; i < 5; i++) {
 331       source_prefixes.push_back(test::RandomHumanReadableString(&rnd, 8));
 332     }
 333     for (int j = 0; j < 20; j++) {
 334       int prefix_index = rnd.Uniform(static_cast<int>(source_prefixes.size()));
 335       std::string key = source_prefixes[prefix_index] +
 336                         test::RandomHumanReadableString(&rnd, rnd.Uniform(8));
 337       source_strings.push_back(key);
 338     }
 339
 340     DoRandomIteraratorTest(GetDB(), source_strings, &rnd, 30, 600, 66);
 341   }
 342 }
 343
 344 TEST_F(ComparatorDBTest, Uint64Comparator) {
 345   SetOwnedComparator(test::Uint64Comparator());
 346
 347   for (int rnd_seed = 301; rnd_seed < 316; rnd_seed++) {
 348     Options* opt = GetOptions();
 349     opt->comparator = comparator;
 350     DestroyAndReopen();
 351     Random rnd(rnd_seed);
 352     Random64 rnd64(rnd_seed);
 353
 354     std::vector<std::string> source_strings;
 355     // Randomly generate source keys
 356     for (int i = 0; i < 100; i++) {
 357       uint64_t r = rnd64.Next();
 358       std::string str;
 359       str.resize(8);
 360       memcpy(&str[0], static_cast<void*>(&r), 8);
 361       source_strings.push_back(str);
 362     }
 363
 364     DoRandomIteraratorTest(GetDB(), source_strings, &rnd, 200, 1000, 66);
 365   }
 366 }
 367
 368 TEST_F(ComparatorDBTest, DoubleComparator) {
 369   SetOwnedComparator(new DoubleComparator());
 370
 371   for (int rnd_seed = 301; rnd_seed < 316; rnd_seed++) {
 372     Options* opt = GetOptions();
 373     opt->comparator = comparator;
 374     DestroyAndReopen();
 375     Random rnd(rnd_seed);
 376
 377     std::vector<std::string> source_strings;
 378     // Randomly generate source keys
 379     for (int i = 0; i < 100; i++) {
 380       uint32_t r = rnd.Next();
 381       uint32_t divide_order = rnd.Uniform(8);
 382       double to_divide = 1.0;
 383       for (uint32_t j = 0; j < divide_order; j++) {
 384         to_divide *= 10.0;
 385       }
 386       source_strings.push_back(ToString(r / to_divide));
 387     }
 388
 389     DoRandomIteraratorTest(GetDB(), source_strings, &rnd, 200, 1000, 66);
 390   }
 391 }
 392
 393 TEST_F(ComparatorDBTest, HashComparator) {
 394   SetOwnedComparator(new HashComparator());
 395
 396   for (int rnd_seed = 301; rnd_seed < 316; rnd_seed++) {
 397     Options* opt = GetOptions();
 398     opt->comparator = comparator;
 399     DestroyAndReopen();
 400     Random rnd(rnd_seed);
 401
 402     std::vector<std::string> source_strings;
 403     // Randomly generate source keys
 404     for (int i = 0; i < 100; i++) {
 405       source_strings.push_back(test::RandomKey(&rnd, 8));
 406     }
 407
 408     DoRandomIteraratorTest(GetDB(), source_strings, &rnd, 200, 1000, 66);
 409   }
 410 }
 411
 412 TEST_F(ComparatorDBTest, TwoStrComparator) {
 413   SetOwnedComparator(new TwoStrComparator());
 414
 415   for (int rnd_seed = 301; rnd_seed < 316; rnd_seed++) {
 416     Options* opt = GetOptions();
 417     opt->comparator = comparator;
 418     DestroyAndReopen();
 419     Random rnd(rnd_seed);
 420
 421     std::vector<std::string> source_strings;
 422     // Randomly generate source keys
 423     for (int i = 0; i < 100; i++) {
 424       std::string str;
 425       uint32_t size1 = rnd.Uniform(8);
 426       uint32_t size2 = rnd.Uniform(8);
 427       str.append(1, static_cast<char>(size1));
 428       str.append(1, static_cast<char>(size2));
 429       str.append(test::RandomKey(&rnd, size1));
 430       str.append(test::RandomKey(&rnd, size2));
 431       source_strings.push_back(str);
 432     }
 433
 434     DoRandomIteraratorTest(GetDB(), source_strings, &rnd, 200, 1000, 66);
 435   }
 436 }
 437
 438 }  // namespace rocksdb
 439
 440 int main(int argc, char** argv) {
 441   ::testing::InitGoogleTest(&argc, argv);
 442   return RUN_ALL_TESTS();
 443 }