YugabyteDB (2.13.0.0-b42, bfc6a6643e7399ac8a0e81d06a3ee6d6571b33ab)

// Use of this source code is governed by a BSD-style license that can be

// found in the LICENSE file. See the AUTHORS file for names of contributors.

//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.

//  This source code is licensed under the BSD-style license found in the

//  LICENSE file in the root directory of this source tree. An additional grant

//  of patent rights can be found in the PATENTS file in the same directory.

//

// The following only applies to changes made to this file as part of YugaByte development.

//

// Portions Copyright (c) YugaByte, Inc.

//

// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except

// in compliance with the License.  You may obtain a copy of the License at

//

// http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing, software distributed under the License

// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express

// or implied.  See the License for the specific language governing permissions and limitations

// under the License.

//

#include <map>

#include <string>

#include <gtest/gtest.h>

#include "yb/rocksdb/db.h"

#include "yb/rocksdb/env.h"

#include "yb/rocksdb/util/hash.h"

#include "yb/rocksdb/util/kv_map.h"

#include "yb/rocksdb/util/testharness.h"

#include "yb/rocksdb/util/testutil.h"

#include "yb/util/string_util.h"

#include "yb/util/test_macros.h"

using std::unique_ptr;

namespace rocksdb {

namespace {

static const Comparator* comparator;

class KVIter : public Iterator {

 public:

  explicit KVIter(const stl_wrappers::KVMap* map)

      : map_(map), iter_(map_->end()) {}

  bool Valid() const override { return iter_ != map_->end(); }

  void SeekToFirst() override { iter_ = map_->begin(); }

  void SeekToLast() override {

    if (map_->empty()) {

      iter_ = map_->end();

    } else {

      iter_ = map_->find(map_->rbegin()->first);

    }

  }

  void Seek(const Slice& k) override {

    iter_ = map_->lower_bound(k.ToString());

  }

  void Next() override { ++iter_; }

  void Prev() override {

    if (iter_ == map_->begin()) {

      iter_ = map_->end();

      return;

    }

    --iter_;

  }

  Slice key() const override { return iter_->first; }

  Slice value() const override { return iter_->second; }

  Status status() const override { return Status::OK(); }

 private:

  const stl_wrappers::KVMap* const map_;

  stl_wrappers::KVMap::const_iterator iter_;

};

void AssertItersEqual(Iterator* iter1, Iterator* iter2) {

  ASSERT_EQ(iter1->Valid(), iter2->Valid());

  if (iter1->Valid()) {

    auto key1 = iter1->key().ToBuffer();

    auto key2 = iter2->key().ToBuffer();

    auto value1 = iter1->value().ToBuffer();

    auto value2 = iter2->value().ToBuffer();

    ASSERT_EQ(key1, key2);

    ASSERT_EQ(value1, value2);

  }

}

// Measuring operations on DB (expect to be empty).

// source_strings are candidate keys

void DoRandomIteraratorTest(DB* db, std::vector<std::string> source_strings,

                            Random* rnd, int num_writes, int num_iter_ops,

                            int num_trigger_flush) {

  stl_wrappers::KVMap map((stl_wrappers::LessOfComparator(comparator)));

  for (int i = 0; i < num_writes; i++) {

    if (num_trigger_flush > 0 && i != 0 && i % num_trigger_flush == 0) {

      ASSERT_OK(db->Flush(FlushOptions()));

    }

    int type = rnd->Uniform(2);

    int index = rnd->Uniform(static_cast<int>(source_strings.size()));

    auto& key = source_strings[index];

    switch (type) {

      case 0:

        // put

        map[key] = key;

        ASSERT_OK(db->Put(WriteOptions(), key, key));

        break;

      case 1:

        // delete

        if (map.find(key) != map.end()) {

          map.erase(key);

        }

        ASSERT_OK(db->Delete(WriteOptions(), key));

        break;

      default:

        assert(false);

    }

  }

  std::unique_ptr<Iterator> iter(db->NewIterator(ReadOptions()));

  std::unique_ptr<Iterator> result_iter(new KVIter(&map));

  bool is_valid = false;

  for (int i = 0; i < num_iter_ops; i++) {

    // Random walk and make sure iter and result_iter returns the

    // same key and value

    int type = rnd->Uniform(6);

    ASSERT_OK(iter->status());

    switch (type) {

      case 0:

        // Seek to First

        iter->SeekToFirst();

        result_iter->SeekToFirst();

        break;

      case 1:

        // Seek to last

        iter->SeekToLast();

        result_iter->SeekToLast();

        break;

      case 2: {

        // Seek to random key

        auto key_idx = rnd->Uniform(static_cast<int>(source_strings.size()));

        auto key = source_strings[key_idx];

        iter->Seek(key);

        result_iter->Seek(key);

        break;

      }

      case 3:

        // Next

        if (is_valid) {

          iter->Next();

          result_iter->Next();

        } else {

          continue;

        }

        break;

      case 4:

        // Prev

        if (is_valid) {

          iter->Prev();

          result_iter->Prev();

        } else {

          continue;

        }

        break;

      default: {

        assert(type == 5);

        auto key_idx = rnd->Uniform(static_cast<int>(source_strings.size()));

        auto key = source_strings[key_idx];

        std::string result;

        auto status = db->Get(ReadOptions(), key, &result);

        if (map.find(key) == map.end()) {

          ASSERT_TRUE(status.IsNotFound());

        } else {

          ASSERT_EQ(map[key], result);

        }

        break;

      }

    }

    AssertItersEqual(iter.get(), result_iter.get());

    is_valid = iter->Valid();

  }

}

class DoubleComparator : public Comparator {

 public:

  DoubleComparator() {}

  const char* Name() const override { return "DoubleComparator"; }

  int Compare(const Slice& a, const Slice& b) const override {

#ifndef CYGWIN

    double da = std::stod(a.ToString());

    double db = std::stod(b.ToString());

#else

    double da = std::strtod(a.ToString().c_str(), 0 /* endptr */);

    double db = std::strtod(a.ToString().c_str(), 0 /* endptr */);

#endif

    if (da == db) {

      return a.compare(b);

    } else if (da > db) {

      return 1;

    } else {

      return -1;

    }

  }

  virtual void FindShortestSeparator(std::string* start,

                                     const Slice& limit) const override {}

  void FindShortSuccessor(std::string* key) const override {}

};

class HashComparator : public Comparator {

 public:

  HashComparator() {}

  const char* Name() const override { return "HashComparator"; }

  int Compare(const Slice& a, const Slice& b) const override {

    uint32_t ha = Hash(a.data(), a.size(), 66);

    uint32_t hb = Hash(b.data(), b.size(), 66);

    if (ha == hb) {

      return a.compare(b);

    } else if (ha > hb) {

      return 1;

    } else {

      return -1;

    }

  }

  virtual void FindShortestSeparator(std::string* start,

                                     const Slice& limit) const override {}

  void FindShortSuccessor(std::string* key) const override {}

};

class TwoStrComparator : public Comparator {

 public:

  TwoStrComparator() {}

  const char* Name() const override { return "TwoStrComparator"; }

  int Compare(const Slice& a, const Slice& b) const override {

    assert(a.size() >= 2);

    assert(b.size() >= 2);

    size_t size_a1 = static_cast<size_t>(a[0]);

    size_t size_b1 = static_cast<size_t>(b[0]);

    size_t size_a2 = static_cast<size_t>(a[1]);

    size_t size_b2 = static_cast<size_t>(b[1]);

    assert(size_a1 + size_a2 + 2 == a.size());

    assert(size_b1 + size_b2 + 2 == b.size());

    Slice a1 = Slice(a.data() + 2, size_a1);

    Slice b1 = Slice(b.data() + 2, size_b1);

    Slice a2 = Slice(a.data() + 2 + size_a1, size_a2);

    Slice b2 = Slice(b.data() + 2 + size_b1, size_b2);

    if (a1 != b1) {

      return a1.compare(b1);

    }

    return a2.compare(b2);

  }

  virtual void FindShortestSeparator(std::string* start,

                                     const Slice& limit) const override {}

  void FindShortSuccessor(std::string* key) const override {}

};

}  // namespace

class ComparatorDBTest : public RocksDBTest {

 private:

  std::string dbname_;

  Env* env_;

  DB* db_;

  Options last_options_;

  std::unique_ptr<const Comparator> comparator_guard;

 public:

  ComparatorDBTest() : env_(Env::Default()), db_(nullptr) {

    comparator = BytewiseComparator();

    dbname_ = test::TmpDir() + "/comparator_db_test";

    EXPECT_OK(DestroyDB(dbname_, last_options_));

  }

  ~ComparatorDBTest() {

    delete db_;

    EXPECT_OK(DestroyDB(dbname_, last_options_));

    comparator = BytewiseComparator();

  }

  DB* GetDB() { return db_; }

  void SetOwnedComparator(const Comparator* cmp) {

    comparator_guard.reset(cmp);

    SetComparator(cmp);

  }

  void SetComparator(const Comparator* cmp) {

    comparator = cmp;

    last_options_.comparator = cmp;

  }

  // Return the current option configuration.

  Options* GetOptions() { return &last_options_; }

  void DestroyAndReopen() {

    // Destroy using last options

    Destroy();

    ASSERT_OK(TryReopen());

  }

  void Destroy() {

    delete db_;

    db_ = nullptr;

    ASSERT_OK(DestroyDB(dbname_, last_options_));

  }

  Status TryReopen() {

    delete db_;

    db_ = nullptr;

    last_options_.create_if_missing = true;

    return DB::Open(last_options_, dbname_, &db_);

  }

};

TEST_F(ComparatorDBTest, Bytewise) {

  for (int rand_seed = 301; rand_seed < 306; rand_seed++) {

    DestroyAndReopen();

    Random rnd(rand_seed);

    DoRandomIteraratorTest(GetDB(),

                           {"a", "b", "c", "d", "e", "f", "g", "h", "i"}, &rnd,

                           8, 100, 3);

  }

}

TEST_F(ComparatorDBTest, SimpleSuffixReverseComparator) {

  SetOwnedComparator(new test::SimpleSuffixReverseComparator());

  for (int rnd_seed = 301; rnd_seed < 316; rnd_seed++) {

    Options* opt = GetOptions();

    opt->comparator = comparator;

    DestroyAndReopen();

    Random rnd(rnd_seed);

    std::vector<std::string> source_strings;

    std::vector<std::string> source_prefixes;

    // Randomly generate 5 prefixes

    for (int i = 0; i < 5; i++) {

      source_prefixes.push_back(test::RandomHumanReadableString(&rnd, 8));

    }

    for (int j = 0; j < 20; j++) {

      int prefix_index = rnd.Uniform(static_cast<int>(source_prefixes.size()));

      std::string key = source_prefixes[prefix_index] +

                        test::RandomHumanReadableString(&rnd, rnd.Uniform(8));

      source_strings.push_back(key);

    }

    DoRandomIteraratorTest(GetDB(), source_strings, &rnd, 30, 600, 66);

  }

}

TEST_F(ComparatorDBTest, Uint64Comparator) {

  SetComparator(Uint64Comparator());

  for (int rnd_seed = 301; rnd_seed < 316; rnd_seed++) {

    Options* opt = GetOptions();

    opt->comparator = comparator;

    DestroyAndReopen();

    Random rnd(rnd_seed);

    Random64 rnd64(rnd_seed);

    std::vector<std::string> source_strings;

    // Randomly generate source keys

    for (int i = 0; i < 100; i++) {

      uint64_t r = rnd64.Next();

      std::string str;

      str.resize(8);

      memcpy(&str[0], static_cast<void*>(&r), 8);

      source_strings.push_back(str);

    }

    DoRandomIteraratorTest(GetDB(), source_strings, &rnd, 200, 1000, 66);

  }

}

TEST_F(ComparatorDBTest, DoubleComparator) {

  SetOwnedComparator(new DoubleComparator());

  for (int rnd_seed = 301; rnd_seed < 316; rnd_seed++) {

    Options* opt = GetOptions();

    opt->comparator = comparator;

    DestroyAndReopen();

    Random rnd(rnd_seed);

    std::vector<std::string> source_strings;

    // Randomly generate source keys

    for (int i = 0; i < 100; i++) {

      uint32_t r = rnd.Next();

      uint32_t divide_order = rnd.Uniform(8);

      double to_divide = 1.0;

      for (uint32_t j = 0; j < divide_order; j++) {

        to_divide *= 10.0;

      }

      source_strings.push_back(ToString(r / to_divide));

    }

    DoRandomIteraratorTest(GetDB(), source_strings, &rnd, 200, 1000, 66);

  }

}

TEST_F(ComparatorDBTest, HashComparator) {

  SetOwnedComparator(new HashComparator());

  for (int rnd_seed = 301; rnd_seed < 316; rnd_seed++) {

    Options* opt = GetOptions();

    opt->comparator = comparator;

    DestroyAndReopen();

    Random rnd(rnd_seed);

    std::vector<std::string> source_strings;

    // Randomly generate source keys

    for (int i = 0; i < 100; i++) {

      source_strings.push_back(test::RandomKey(&rnd, 8));

    }

    DoRandomIteraratorTest(GetDB(), source_strings, &rnd, 200, 1000, 66);

  }

}

TEST_F(ComparatorDBTest, TwoStrComparator) {

  SetOwnedComparator(new TwoStrComparator());

  for (int rnd_seed = 301; rnd_seed < 316; rnd_seed++) {

    Options* opt = GetOptions();

    opt->comparator = comparator;

    DestroyAndReopen();

    Random rnd(rnd_seed);

    std::vector<std::string> source_strings;

    // Randomly generate source keys

    for (int i = 0; i < 100; i++) {

      std::string str;

      uint32_t size1 = rnd.Uniform(8);

      uint32_t size2 = rnd.Uniform(8);

      str.append(1, static_cast<char>(size1));

      str.append(1, static_cast<char>(size2));

      str.append(test::RandomKey(&rnd, size1));

      str.append(test::RandomKey(&rnd, size2));

      source_strings.push_back(str);

    }

    DoRandomIteraratorTest(GetDB(), source_strings, &rnd, 200, 1000, 66);

  }

}

}  // namespace rocksdb

int main(int argc, char** argv) {

  ::testing::InitGoogleTest(&argc, argv);

  return RUN_ALL_TESTS();

}