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