1 // Copyright (C) 2022 Google LLC
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 // http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14
15 #include "icing/store/key-mapper.h"
16
17 #include <memory>
18 #include <string>
19 #include <unordered_map>
20
21 #include "icing/text_classifier/lib3/utils/base/status.h"
22 #include "icing/text_classifier/lib3/utils/base/statusor.h"
23 #include "gmock/gmock.h"
24 #include "gtest/gtest.h"
25 #include "icing/file/filesystem.h"
26 #include "icing/store/document-id.h"
27 #include "icing/store/dynamic-trie-key-mapper.h"
28 #include "icing/store/persistent-hash-map-key-mapper.h"
29 #include "icing/testing/common-matchers.h"
30 #include "icing/testing/tmp-directory.h"
31
32 using ::testing::IsEmpty;
33 using ::testing::IsTrue;
34 using ::testing::Pair;
35 using ::testing::UnorderedElementsAre;
36
37 namespace icing {
38 namespace lib {
39
40 namespace {
41
42 constexpr int kMaxDynamicTrieKeyMapperSize = 3 * 1024 * 1024; // 3 MiB
43
44 enum class KeyMapperType {
45 kDynamicTrie,
46 kPersistentHashMap,
47 };
48
49 struct KeyMapperTestParam {
50 KeyMapperType key_mapper_type;
51 bool pre_mapping_fbv;
52
KeyMapperTestParamicing::lib::__anon4e5a97c10111::KeyMapperTestParam53 explicit KeyMapperTestParam(KeyMapperType key_mapper_type_in,
54 bool pre_mapping_fbv_in)
55 : key_mapper_type(key_mapper_type_in),
56 pre_mapping_fbv(pre_mapping_fbv_in) {}
57 };
58
59 class KeyMapperTest : public ::testing::TestWithParam<KeyMapperTestParam> {
60 protected:
SetUp()61 void SetUp() override {
62 base_dir_ = GetTestTempDir() + "/icing";
63 ASSERT_THAT(filesystem_.CreateDirectoryRecursively(base_dir_.c_str()),
64 IsTrue());
65
66 working_dir_ = base_dir_ + "/key_mapper";
67 }
68
TearDown()69 void TearDown() override {
70 filesystem_.DeleteDirectoryRecursively(base_dir_.c_str());
71 }
72
73 libtextclassifier3::StatusOr<std::unique_ptr<KeyMapper<DocumentId>>>
CreateKeyMapper()74 CreateKeyMapper() {
75 const KeyMapperTestParam& param = GetParam();
76 switch (param.key_mapper_type) {
77 case KeyMapperType::kDynamicTrie:
78 return DynamicTrieKeyMapper<DocumentId>::Create(
79 filesystem_, working_dir_, kMaxDynamicTrieKeyMapperSize);
80 case KeyMapperType::kPersistentHashMap:
81 return PersistentHashMapKeyMapper<DocumentId>::Create(
82 filesystem_, working_dir_, param.pre_mapping_fbv);
83 }
84 }
85
DeleteKeyMapper()86 libtextclassifier3::Status DeleteKeyMapper() {
87 const KeyMapperTestParam& param = GetParam();
88 switch (param.key_mapper_type) {
89 case KeyMapperType::kDynamicTrie:
90 return DynamicTrieKeyMapper<DocumentId>::Delete(filesystem_,
91 working_dir_);
92 case KeyMapperType::kPersistentHashMap:
93 return PersistentHashMapKeyMapper<DocumentId>::Delete(filesystem_,
94 working_dir_);
95 }
96 }
97
98 std::string base_dir_;
99 std::string working_dir_;
100 Filesystem filesystem_;
101 };
102
GetAllKeyValuePairs(const KeyMapper<DocumentId> * key_mapper)103 std::unordered_map<std::string, DocumentId> GetAllKeyValuePairs(
104 const KeyMapper<DocumentId>* key_mapper) {
105 std::unordered_map<std::string, DocumentId> ret;
106
107 std::unique_ptr<typename KeyMapper<DocumentId>::Iterator> itr =
108 key_mapper->GetIterator();
109 while (itr->Advance()) {
110 ret.emplace(itr->GetKey(), itr->GetValue());
111 }
112 return ret;
113 }
114
TEST_P(KeyMapperTest,CreateNewKeyMapper)115 TEST_P(KeyMapperTest, CreateNewKeyMapper) {
116 ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<KeyMapper<DocumentId>> key_mapper,
117 CreateKeyMapper());
118 EXPECT_THAT(key_mapper->num_keys(), 0);
119 }
120
TEST_P(KeyMapperTest,CanUpdateSameKeyMultipleTimes)121 TEST_P(KeyMapperTest, CanUpdateSameKeyMultipleTimes) {
122 ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<KeyMapper<DocumentId>> key_mapper,
123 CreateKeyMapper());
124
125 ICING_EXPECT_OK(key_mapper->Put("default-google.com", 100));
126 ICING_EXPECT_OK(key_mapper->Put("default-youtube.com", 50));
127
128 EXPECT_THAT(key_mapper->Get("default-google.com"), IsOkAndHolds(100));
129
130 ICING_EXPECT_OK(key_mapper->Put("default-google.com", 200));
131 EXPECT_THAT(key_mapper->Get("default-google.com"), IsOkAndHolds(200));
132 EXPECT_THAT(key_mapper->num_keys(), 2);
133
134 ICING_EXPECT_OK(key_mapper->Put("default-google.com", 300));
135 EXPECT_THAT(key_mapper->Get("default-google.com"), IsOkAndHolds(300));
136 EXPECT_THAT(key_mapper->num_keys(), 2);
137 }
138
TEST_P(KeyMapperTest,GetOrPutOk)139 TEST_P(KeyMapperTest, GetOrPutOk) {
140 ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<KeyMapper<DocumentId>> key_mapper,
141 CreateKeyMapper());
142
143 EXPECT_THAT(key_mapper->Get("foo"),
144 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
145 EXPECT_THAT(key_mapper->GetOrPut("foo", 1), IsOkAndHolds(1));
146 EXPECT_THAT(key_mapper->Get("foo"), IsOkAndHolds(1));
147 }
148
TEST_P(KeyMapperTest,CanPersistToDiskRegularly)149 TEST_P(KeyMapperTest, CanPersistToDiskRegularly) {
150 ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<KeyMapper<DocumentId>> key_mapper,
151 CreateKeyMapper());
152
153 // Can persist an empty DynamicTrieKeyMapper.
154 ICING_EXPECT_OK(key_mapper->PersistToDisk());
155 EXPECT_THAT(key_mapper->num_keys(), 0);
156
157 // Can persist the smallest DynamicTrieKeyMapper.
158 ICING_EXPECT_OK(key_mapper->Put("default-google.com", 100));
159 ICING_EXPECT_OK(key_mapper->PersistToDisk());
160 EXPECT_THAT(key_mapper->num_keys(), 1);
161 EXPECT_THAT(key_mapper->Get("default-google.com"), IsOkAndHolds(100));
162
163 // Can continue to add keys after PersistToDisk().
164 ICING_EXPECT_OK(key_mapper->Put("default-youtube.com", 200));
165 EXPECT_THAT(key_mapper->num_keys(), 2);
166 EXPECT_THAT(key_mapper->Get("default-youtube.com"), IsOkAndHolds(200));
167
168 // Can continue to update the same key after PersistToDisk().
169 ICING_EXPECT_OK(key_mapper->Put("default-google.com", 300));
170 EXPECT_THAT(key_mapper->Get("default-google.com"), IsOkAndHolds(300));
171 EXPECT_THAT(key_mapper->num_keys(), 2);
172 }
173
TEST_P(KeyMapperTest,CanUseAcrossMultipleInstances)174 TEST_P(KeyMapperTest, CanUseAcrossMultipleInstances) {
175 ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<KeyMapper<DocumentId>> key_mapper,
176 CreateKeyMapper());
177 ICING_EXPECT_OK(key_mapper->Put("default-google.com", 100));
178 ICING_EXPECT_OK(key_mapper->PersistToDisk());
179
180 key_mapper.reset();
181
182 ICING_ASSERT_OK_AND_ASSIGN(key_mapper, CreateKeyMapper());
183 EXPECT_THAT(key_mapper->num_keys(), 1);
184 EXPECT_THAT(key_mapper->Get("default-google.com"), IsOkAndHolds(100));
185
186 // Can continue to read/write to the KeyMapper.
187 ICING_EXPECT_OK(key_mapper->Put("default-youtube.com", 200));
188 ICING_EXPECT_OK(key_mapper->Put("default-google.com", 300));
189 EXPECT_THAT(key_mapper->num_keys(), 2);
190 EXPECT_THAT(key_mapper->Get("default-youtube.com"), IsOkAndHolds(200));
191 EXPECT_THAT(key_mapper->Get("default-google.com"), IsOkAndHolds(300));
192 }
193
TEST_P(KeyMapperTest,CanDeleteAndRestartKeyMapping)194 TEST_P(KeyMapperTest, CanDeleteAndRestartKeyMapping) {
195 // Can delete even if there's nothing there
196 ICING_EXPECT_OK(DeleteKeyMapper());
197
198 ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<KeyMapper<DocumentId>> key_mapper,
199 CreateKeyMapper());
200 ICING_EXPECT_OK(key_mapper->Put("default-google.com", 100));
201 ICING_EXPECT_OK(key_mapper->PersistToDisk());
202 ICING_EXPECT_OK(DeleteKeyMapper());
203
204 key_mapper.reset();
205 ICING_ASSERT_OK_AND_ASSIGN(key_mapper, CreateKeyMapper());
206 EXPECT_THAT(key_mapper->num_keys(), 0);
207 ICING_EXPECT_OK(key_mapper->Put("default-google.com", 100));
208 EXPECT_THAT(key_mapper->num_keys(), 1);
209 }
210
TEST_P(KeyMapperTest,Iterator)211 TEST_P(KeyMapperTest, Iterator) {
212 ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<KeyMapper<DocumentId>> key_mapper,
213 CreateKeyMapper());
214 EXPECT_THAT(GetAllKeyValuePairs(key_mapper.get()), IsEmpty());
215
216 ICING_EXPECT_OK(key_mapper->Put("foo", /*value=*/1));
217 ICING_EXPECT_OK(key_mapper->Put("bar", /*value=*/2));
218 EXPECT_THAT(GetAllKeyValuePairs(key_mapper.get()),
219 UnorderedElementsAre(Pair("foo", 1), Pair("bar", 2)));
220
221 ICING_EXPECT_OK(key_mapper->Put("baz", /*value=*/3));
222 EXPECT_THAT(
223 GetAllKeyValuePairs(key_mapper.get()),
224 UnorderedElementsAre(Pair("foo", 1), Pair("bar", 2), Pair("baz", 3)));
225 }
226
227 INSTANTIATE_TEST_SUITE_P(
228 KeyMapperTest, KeyMapperTest,
229 testing::Values(KeyMapperTestParam(KeyMapperType::kDynamicTrie,
230 /*pre_mapping_fbv_in=*/true),
231 KeyMapperTestParam(KeyMapperType::kPersistentHashMap,
232 /*pre_mapping_fbv_in=*/true),
233 KeyMapperTestParam(KeyMapperType::kPersistentHashMap,
234 /*pre_mapping_fbv_in=*/false)));
235
236 } // namespace
237
238 } // namespace lib
239 } // namespace icing
240