ceph/src/rocksdb/table/mock_table.h

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