ceph/src/rocksdb/table/mock_table.h

   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 #pragma once
   8
   9 #include <algorithm>
  10 #include <atomic>
  11 #include <map>
  12 #include <memory>
  13 #include <set>
  14 #include <string>
  15 #include <utility>
  16
  17 #include "util/kv_map.h"
  18 #include "port/port.h"
  19 #include "rocksdb/comparator.h"
  20 #include "rocksdb/table.h"
  21 #include "table/internal_iterator.h"
  22 #include "table/table_builder.h"
  23 #include "table/table_reader.h"
  24 #include "util/mutexlock.h"
  25 #include "util/testharness.h"
  26 #include "util/testutil.h"
  27
  28 namespace rocksdb {
  29 namespace mock {
  30
  31 stl_wrappers::KVMap MakeMockFile(
  32     std::initializer_list<std::pair<const std::string, std::string>> l = {});
  33
  34 struct MockTableFileSystem {
  35   port::Mutex mutex;
  36   std::map<uint32_t, stl_wrappers::KVMap> files;
  37 };
  38
  39 class MockTableReader : public TableReader {
  40  public:
  41   explicit MockTableReader(const stl_wrappers::KVMap& table) : table_(table) {}
  42
  43   InternalIterator* NewIterator(const ReadOptions&, Arena* arena,
  44                                 bool skip_filters = false) override;
  45
  46   Status Get(const ReadOptions&, const Slice& key, GetContext* get_context,
  47              bool skip_filters = false) override;
  48
  49   uint64_t ApproximateOffsetOf(const Slice& key) override { return 0; }
  50
  51   virtual size_t ApproximateMemoryUsage() const override { return 0; }
  52
  53   void SetupForCompaction() override {}
  54
  55   std::shared_ptr<const TableProperties> GetTableProperties() const override;
  56
  57   ~MockTableReader() {}
  58
  59  private:
  60   const stl_wrappers::KVMap& table_;
  61 };
  62
  63 class MockTableIterator : public InternalIterator {
  64  public:
  65   explicit MockTableIterator(const stl_wrappers::KVMap& table) : table_(table) {
  66     itr_ = table_.end();
  67   }
  68
  69   bool Valid() const override { return itr_ != table_.end(); }
  70
  71   void SeekToFirst() override { itr_ = table_.begin(); }
  72
  73   void SeekToLast() override {
  74     itr_ = table_.end();
  75     --itr_;
  76   }
  77
  78   void Seek(const Slice& target) override {
  79     std::string str_target(target.data(), target.size());
  80     itr_ = table_.lower_bound(str_target);
  81   }
  82
  83   void SeekForPrev(const Slice& target) override {
  84     std::string str_target(target.data(), target.size());
  85     itr_ = table_.upper_bound(str_target);
  86     Prev();
  87   }
  88
  89   void Next() override { ++itr_; }
  90
  91   void Prev() override {
  92     if (itr_ == table_.begin()) {
  93       itr_ = table_.end();
  94     } else {
  95       --itr_;
  96     }
  97   }
  98
  99   Slice key() const override { return Slice(itr_->first); }
 100
 101   Slice value() const override { return Slice(itr_->second); }
 102
 103   Status status() const override { return Status::OK(); }
 104
 105  private:
 106   const stl_wrappers::KVMap& table_;
 107   stl_wrappers::KVMap::const_iterator itr_;
 108 };
 109
 110 class MockTableBuilder : public TableBuilder {
 111  public:
 112   MockTableBuilder(uint32_t id, MockTableFileSystem* file_system)
 113       : id_(id), file_system_(file_system) {
 114     table_ = MakeMockFile({});
 115   }
 116
 117   // REQUIRES: Either Finish() or Abandon() has been called.
 118   ~MockTableBuilder() {}
 119
 120   // Add key,value to the table being constructed.
 121   // REQUIRES: key is after any previously added key according to comparator.
 122   // REQUIRES: Finish(), Abandon() have not been called
 123   void Add(const Slice& key, const Slice& value) override {
 124     table_.insert({key.ToString(), value.ToString()});
 125   }
 126
 127   // Return non-ok iff some error has been detected.
 128   Status status() const override { return Status::OK(); }
 129
 130   Status Finish() override {
 131     MutexLock lock_guard(&file_system_->mutex);
 132     file_system_->files.insert({id_, table_});
 133     return Status::OK();
 134   }
 135
 136   void Abandon() override {}
 137
 138   uint64_t NumEntries() const override { return table_.size(); }
 139
 140   uint64_t FileSize() const override { return table_.size(); }
 141
 142   TableProperties GetTableProperties() const override {
 143     return TableProperties();
 144   }
 145
 146  private:
 147   uint32_t id_;
 148   MockTableFileSystem* file_system_;
 149   stl_wrappers::KVMap table_;
 150 };
 151
 152 class MockTableFactory : public TableFactory {
 153  public:
 154   MockTableFactory();
 155   const char* Name() const override { return "MockTable"; }
 156   Status NewTableReader(
 157       const TableReaderOptions& table_reader_options,
 158       unique_ptr<RandomAccessFileReader>&& file, uint64_t file_size,
 159       unique_ptr<TableReader>* table_reader,
 160       bool prefetch_index_and_filter_in_cache = true) const override;
 161   TableBuilder* NewTableBuilder(
 162       const TableBuilderOptions& table_builder_options,
 163       uint32_t column_familly_id, WritableFileWriter* file) const override;
 164
 165   // This function will directly create mock table instead of going through
 166   // MockTableBuilder. file_contents has to have a format of <internal_key,
 167   // value>. Those key-value pairs will then be inserted into the mock table.
 168   Status CreateMockTable(Env* env, const std::string& fname,
 169                          stl_wrappers::KVMap file_contents);
 170
 171   virtual Status SanitizeOptions(
 172       const DBOptions& db_opts,
 173       const ColumnFamilyOptions& cf_opts) const override {
 174     return Status::OK();
 175   }
 176
 177   virtual std::string GetPrintableTableOptions() const override {
 178     return std::string();
 179   }
 180
 181   // This function will assert that only a single file exists and that the
 182   // contents are equal to file_contents
 183   void AssertSingleFile(const stl_wrappers::KVMap& file_contents);
 184   void AssertLatestFile(const stl_wrappers::KVMap& file_contents);
 185
 186  private:
 187   uint32_t GetAndWriteNextID(WritableFileWriter* file) const;
 188   uint32_t GetIDFromFile(RandomAccessFileReader* file) const;
 189
 190   mutable MockTableFileSystem file_system_;
 191   mutable std::atomic<uint32_t> next_id_;
 192 };
 193
 194 }  // namespace mock
 195 }  // namespace rocksdb