1 // Copyright (C) 2019 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/result/result-state-manager.h"
16
17 #include "gmock/gmock.h"
18 #include "gtest/gtest.h"
19 #include "icing/document-builder.h"
20 #include "icing/file/filesystem.h"
21 #include "icing/portable/equals-proto.h"
22 #include "icing/result/page-result.h"
23 #include "icing/result/result-adjustment-info.h"
24 #include "icing/result/result-retriever-v2.h"
25 #include "icing/schema/schema-store.h"
26 #include "icing/scoring/priority-queue-scored-document-hits-ranker.h"
27 #include "icing/scoring/scored-document-hits-ranker.h"
28 #include "icing/store/document-store.h"
29 #include "icing/testing/common-matchers.h"
30 #include "icing/testing/fake-clock.h"
31 #include "icing/testing/icu-data-file-helper.h"
32 #include "icing/testing/test-data.h"
33 #include "icing/testing/tmp-directory.h"
34 #include "icing/tokenization/language-segmenter-factory.h"
35 #include "icing/transform/normalizer-factory.h"
36 #include "icing/transform/normalizer.h"
37 #include "icing/util/clock.h"
38 #include "unicode/uloc.h"
39
40 namespace icing {
41 namespace lib {
42 namespace {
43
44 using ::icing::lib::portable_equals_proto::EqualsProto;
45 using ::testing::Eq;
46 using ::testing::IsEmpty;
47 using ::testing::Not;
48 using ::testing::SizeIs;
49 using PageResultInfo = std::pair<uint64_t, PageResult>;
50
CreateScoringSpec()51 ScoringSpecProto CreateScoringSpec() {
52 ScoringSpecProto scoring_spec;
53 scoring_spec.set_rank_by(ScoringSpecProto::RankingStrategy::DOCUMENT_SCORE);
54 return scoring_spec;
55 }
56
CreateResultSpec(int num_per_page,ResultSpecProto::ResultGroupingType result_group_type)57 ResultSpecProto CreateResultSpec(
58 int num_per_page, ResultSpecProto::ResultGroupingType result_group_type) {
59 ResultSpecProto result_spec;
60 result_spec.set_result_group_type(result_group_type);
61 result_spec.set_num_per_page(num_per_page);
62 return result_spec;
63 }
64
CreateDocument(int id)65 DocumentProto CreateDocument(int id) {
66 return DocumentBuilder()
67 .SetNamespace("namespace")
68 .SetUri(std::to_string(id))
69 .SetSchema("Document")
70 .SetCreationTimestampMs(1574365086666 + id)
71 .SetScore(1)
72 .Build();
73 }
74
75 class ResultStateManagerTest : public testing::Test {
76 protected:
ResultStateManagerTest()77 ResultStateManagerTest() : test_dir_(GetTestTempDir() + "/icing") {
78 filesystem_.CreateDirectoryRecursively(test_dir_.c_str());
79 }
80
SetUp()81 void SetUp() override {
82 if (!IsCfStringTokenization() && !IsReverseJniTokenization()) {
83 ICING_ASSERT_OK(
84 // File generated via icu_data_file rule in //icing/BUILD.
85 icu_data_file_helper::SetUpICUDataFile(
86 GetTestFilePath("icing/icu.dat")));
87 }
88
89 clock_ = std::make_unique<FakeClock>();
90
91 language_segmenter_factory::SegmenterOptions options(ULOC_US);
92 ICING_ASSERT_OK_AND_ASSIGN(
93 language_segmenter_,
94 language_segmenter_factory::Create(std::move(options)));
95
96 ICING_ASSERT_OK_AND_ASSIGN(
97 schema_store_,
98 SchemaStore::Create(&filesystem_, test_dir_, clock_.get()));
99 SchemaProto schema;
100 schema.add_types()->set_schema_type("Document");
101 ICING_ASSERT_OK(schema_store_->SetSchema(
102 std::move(schema), /*ignore_errors_and_delete_documents=*/false,
103 /*allow_circular_schema_definitions=*/false));
104
105 ICING_ASSERT_OK_AND_ASSIGN(normalizer_, normalizer_factory::Create(
106 /*max_term_byte_size=*/10000));
107
108 ICING_ASSERT_OK_AND_ASSIGN(
109 DocumentStore::CreateResult result,
110 DocumentStore::Create(&filesystem_, test_dir_, clock_.get(),
111 schema_store_.get(),
112 /*force_recovery_and_revalidate_documents=*/false,
113 /*namespace_id_fingerprint=*/false,
114 PortableFileBackedProtoLog<
115 DocumentWrapper>::kDeflateCompressionLevel,
116 /*initialize_stats=*/nullptr));
117 document_store_ = std::move(result.document_store);
118
119 ICING_ASSERT_OK_AND_ASSIGN(
120 result_retriever_, ResultRetrieverV2::Create(
121 document_store_.get(), schema_store_.get(),
122 language_segmenter_.get(), normalizer_.get()));
123 }
124
TearDown()125 void TearDown() override {
126 filesystem_.DeleteDirectoryRecursively(test_dir_.c_str());
127 clock_.reset();
128 }
129
AddScoredDocument(DocumentId document_id)130 std::pair<ScoredDocumentHit, DocumentProto> AddScoredDocument(
131 DocumentId document_id) {
132 DocumentProto document;
133 document.set_namespace_("namespace");
134 document.set_uri(std::to_string(document_id));
135 document.set_schema("Document");
136 document.set_creation_timestamp_ms(1574365086666 + document_id);
137 document_store_->Put(document);
138 return std::make_pair(
139 ScoredDocumentHit(document_id, kSectionIdMaskNone, /*score=*/1),
140 std::move(document));
141 }
142
143 std::pair<std::vector<ScoredDocumentHit>, std::vector<DocumentProto>>
AddScoredDocuments(const std::vector<DocumentId> & document_ids)144 AddScoredDocuments(const std::vector<DocumentId>& document_ids) {
145 std::vector<ScoredDocumentHit> scored_document_hits;
146 std::vector<DocumentProto> document_protos;
147
148 for (DocumentId document_id : document_ids) {
149 std::pair<ScoredDocumentHit, DocumentProto> pair =
150 AddScoredDocument(document_id);
151 scored_document_hits.emplace_back(std::move(pair.first));
152 document_protos.emplace_back(std::move(pair.second));
153 }
154
155 std::reverse(document_protos.begin(), document_protos.end());
156
157 return std::make_pair(std::move(scored_document_hits),
158 std::move(document_protos));
159 }
160
clock()161 FakeClock* clock() { return clock_.get(); }
clock() const162 const FakeClock* clock() const { return clock_.get(); }
163
document_store()164 DocumentStore& document_store() { return *document_store_; }
document_store() const165 const DocumentStore& document_store() const { return *document_store_; }
166
schema_store()167 SchemaStore& schema_store() { return *schema_store_; }
schema_store() const168 const SchemaStore& schema_store() const { return *schema_store_; }
169
result_retriever() const170 const ResultRetrieverV2& result_retriever() const {
171 return *result_retriever_;
172 }
173
174 private:
175 Filesystem filesystem_;
176 const std::string test_dir_;
177 std::unique_ptr<FakeClock> clock_;
178 std::unique_ptr<LanguageSegmenter> language_segmenter_;
179 std::unique_ptr<SchemaStore> schema_store_;
180 std::unique_ptr<Normalizer> normalizer_;
181 std::unique_ptr<DocumentStore> document_store_;
182 std::unique_ptr<ResultRetrieverV2> result_retriever_;
183 };
184
TEST_F(ResultStateManagerTest,ShouldCacheAndRetrieveFirstPageOnePage)185 TEST_F(ResultStateManagerTest, ShouldCacheAndRetrieveFirstPageOnePage) {
186 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id1,
187 document_store().Put(CreateDocument(/*id=*/1)));
188 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id2,
189 document_store().Put(CreateDocument(/*id=*/2)));
190 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id3,
191 document_store().Put(CreateDocument(/*id=*/3)));
192 std::vector<ScoredDocumentHit> scored_document_hits = {
193 {document_id1, kSectionIdMaskNone, /*score=*/1},
194 {document_id2, kSectionIdMaskNone, /*score=*/1},
195 {document_id3, kSectionIdMaskNone, /*score=*/1}};
196 std::unique_ptr<ScoredDocumentHitsRanker> ranker = std::make_unique<
197 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
198 std::move(scored_document_hits), /*is_descending=*/true);
199
200 ResultStateManager result_state_manager(
201 /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
202
203 ICING_ASSERT_OK_AND_ASSIGN(
204 PageResultInfo page_result_info,
205 result_state_manager.CacheAndRetrieveFirstPage(
206 std::move(ranker), /*parent_adjustment_info=*/nullptr,
207 /*child_adjustment_info=*/nullptr,
208 CreateResultSpec(/*num_per_page=*/10, ResultSpecProto::NAMESPACE),
209 document_store(), result_retriever(),
210 clock()->GetSystemTimeMilliseconds()));
211
212 EXPECT_THAT(page_result_info.first, Eq(kInvalidNextPageToken));
213
214 // Should get docs.
215 ASSERT_THAT(page_result_info.second.results, SizeIs(3));
216 EXPECT_THAT(page_result_info.second.results.at(0).document(),
217 EqualsProto(CreateDocument(/*id=*/3)));
218 EXPECT_THAT(page_result_info.second.results.at(1).document(),
219 EqualsProto(CreateDocument(/*id=*/2)));
220 EXPECT_THAT(page_result_info.second.results.at(2).document(),
221 EqualsProto(CreateDocument(/*id=*/1)));
222 }
223
TEST_F(ResultStateManagerTest,ShouldCacheAndRetrieveFirstPageMultiplePages)224 TEST_F(ResultStateManagerTest, ShouldCacheAndRetrieveFirstPageMultiplePages) {
225 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id1,
226 document_store().Put(CreateDocument(/*id=*/1)));
227 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id2,
228 document_store().Put(CreateDocument(/*id=*/2)));
229 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id3,
230 document_store().Put(CreateDocument(/*id=*/3)));
231 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id4,
232 document_store().Put(CreateDocument(/*id=*/4)));
233 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id5,
234 document_store().Put(CreateDocument(/*id=*/5)));
235 std::vector<ScoredDocumentHit> scored_document_hits = {
236 {document_id1, kSectionIdMaskNone, /*score=*/1},
237 {document_id2, kSectionIdMaskNone, /*score=*/1},
238 {document_id3, kSectionIdMaskNone, /*score=*/1},
239 {document_id4, kSectionIdMaskNone, /*score=*/1},
240 {document_id5, kSectionIdMaskNone, /*score=*/1}};
241 std::unique_ptr<ScoredDocumentHitsRanker> ranker = std::make_unique<
242 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
243 std::move(scored_document_hits), /*is_descending=*/true);
244
245 ResultStateManager result_state_manager(
246 /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
247
248 // First page, 2 results
249 ICING_ASSERT_OK_AND_ASSIGN(
250 PageResultInfo page_result_info1,
251 result_state_manager.CacheAndRetrieveFirstPage(
252 std::move(ranker), /*parent_adjustment_info=*/nullptr,
253 /*child_adjustment_info=*/nullptr,
254 CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE),
255 document_store(), result_retriever(),
256 clock()->GetSystemTimeMilliseconds()));
257 EXPECT_THAT(page_result_info1.first, Not(Eq(kInvalidNextPageToken)));
258 ASSERT_THAT(page_result_info1.second.results, SizeIs(2));
259 EXPECT_THAT(page_result_info1.second.results.at(0).document(),
260 EqualsProto(CreateDocument(/*id=*/5)));
261 EXPECT_THAT(page_result_info1.second.results.at(1).document(),
262 EqualsProto(CreateDocument(/*id=*/4)));
263
264 uint64_t next_page_token = page_result_info1.first;
265
266 // Second page, 2 results
267 ICING_ASSERT_OK_AND_ASSIGN(
268 PageResultInfo page_result_info2,
269 result_state_manager.GetNextPage(next_page_token, result_retriever(),
270 clock()->GetSystemTimeMilliseconds()));
271 EXPECT_THAT(page_result_info2.first, Eq(next_page_token));
272 ASSERT_THAT(page_result_info2.second.results, SizeIs(2));
273 EXPECT_THAT(page_result_info2.second.results.at(0).document(),
274 EqualsProto(CreateDocument(/*id=*/3)));
275 EXPECT_THAT(page_result_info2.second.results.at(1).document(),
276 EqualsProto(CreateDocument(/*id=*/2)));
277
278 // Third page, 1 result
279 ICING_ASSERT_OK_AND_ASSIGN(
280 PageResultInfo page_result_info3,
281 result_state_manager.GetNextPage(next_page_token, result_retriever(),
282 clock()->GetSystemTimeMilliseconds()));
283 EXPECT_THAT(page_result_info3.first, Eq(kInvalidNextPageToken));
284 ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
285 EXPECT_THAT(page_result_info3.second.results.at(0).document(),
286 EqualsProto(CreateDocument(/*id=*/1)));
287
288 // No results
289 EXPECT_THAT(
290 result_state_manager.GetNextPage(next_page_token, result_retriever(),
291 clock()->GetSystemTimeMilliseconds()),
292 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
293 }
294
TEST_F(ResultStateManagerTest,NullRankerShouldReturnError)295 TEST_F(ResultStateManagerTest, NullRankerShouldReturnError) {
296 ResultStateManager result_state_manager(
297 /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
298
299 EXPECT_THAT(
300 result_state_manager.CacheAndRetrieveFirstPage(
301 /*ranker=*/nullptr, /*parent_adjustment_info=*/nullptr,
302 /*child_adjustment_info=*/nullptr,
303 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
304 document_store(), result_retriever(),
305 clock()->GetSystemTimeMilliseconds()),
306 StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
307 }
308
TEST_F(ResultStateManagerTest,EmptyRankerShouldReturnEmptyFirstPage)309 TEST_F(ResultStateManagerTest, EmptyRankerShouldReturnEmptyFirstPage) {
310 ResultStateManager result_state_manager(
311 /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
312 ICING_ASSERT_OK_AND_ASSIGN(
313 PageResultInfo page_result_info,
314 result_state_manager.CacheAndRetrieveFirstPage(
315 std::make_unique<
316 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
317 std::vector<ScoredDocumentHit>(), /*is_descending=*/true),
318 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
319 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
320 document_store(), result_retriever(),
321 clock()->GetSystemTimeMilliseconds()));
322
323 EXPECT_THAT(page_result_info.first, Eq(kInvalidNextPageToken));
324 EXPECT_THAT(page_result_info.second.results, IsEmpty());
325 }
326
TEST_F(ResultStateManagerTest,ShouldAllowEmptyFirstPage)327 TEST_F(ResultStateManagerTest, ShouldAllowEmptyFirstPage) {
328 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id1,
329 document_store().Put(CreateDocument(/*id=*/1)));
330 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id2,
331 document_store().Put(CreateDocument(/*id=*/2)));
332 std::vector<ScoredDocumentHit> scored_document_hits = {
333 {document_id1, kSectionIdMaskNone, /*score=*/1},
334 {document_id2, kSectionIdMaskNone, /*score=*/1}};
335
336 ResultStateManager result_state_manager(
337 /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
338
339 // Create a ResultSpec that limits "namespace" to 0 results.
340 ResultSpecProto result_spec =
341 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE);
342 ResultSpecProto::ResultGrouping* result_grouping =
343 result_spec.add_result_groupings();
344 ResultSpecProto::ResultGrouping::Entry* entry =
345 result_grouping->add_entry_groupings();
346 result_grouping->set_max_results(0);
347 entry->set_namespace_("namespace");
348
349 // First page, no result.
350 ICING_ASSERT_OK_AND_ASSIGN(
351 PageResultInfo page_result_info,
352 result_state_manager.CacheAndRetrieveFirstPage(
353 std::make_unique<
354 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
355 std::move(scored_document_hits), /*is_descending=*/true),
356 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
357 result_spec, document_store(), result_retriever(),
358 clock()->GetSystemTimeMilliseconds()));
359 // If the first page has no result, then it should be the last page.
360 EXPECT_THAT(page_result_info.first, Eq(kInvalidNextPageToken));
361 EXPECT_THAT(page_result_info.second.results, IsEmpty());
362 }
363
TEST_F(ResultStateManagerTest,ShouldAllowEmptyLastPage)364 TEST_F(ResultStateManagerTest, ShouldAllowEmptyLastPage) {
365 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id1,
366 document_store().Put(CreateDocument(/*id=*/1)));
367 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id2,
368 document_store().Put(CreateDocument(/*id=*/2)));
369 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id3,
370 document_store().Put(CreateDocument(/*id=*/3)));
371 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id4,
372 document_store().Put(CreateDocument(/*id=*/4)));
373 std::vector<ScoredDocumentHit> scored_document_hits = {
374 {document_id1, kSectionIdMaskNone, /*score=*/1},
375 {document_id2, kSectionIdMaskNone, /*score=*/1},
376 {document_id3, kSectionIdMaskNone, /*score=*/1},
377 {document_id4, kSectionIdMaskNone, /*score=*/1}};
378
379 ResultStateManager result_state_manager(
380 /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
381
382 // Create a ResultSpec that limits "namespace" to 2 results.
383 ResultSpecProto result_spec =
384 CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE);
385 ResultSpecProto::ResultGrouping* result_grouping =
386 result_spec.add_result_groupings();
387 ResultSpecProto::ResultGrouping::Entry* entry =
388 result_grouping->add_entry_groupings();
389 result_grouping->set_max_results(2);
390 entry->set_namespace_("namespace");
391
392 // First page, 2 results.
393 ICING_ASSERT_OK_AND_ASSIGN(
394 PageResultInfo page_result_info1,
395 result_state_manager.CacheAndRetrieveFirstPage(
396 std::make_unique<
397 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
398 std::move(scored_document_hits), /*is_descending=*/true),
399 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
400 result_spec, document_store(), result_retriever(),
401 clock()->GetSystemTimeMilliseconds()));
402 EXPECT_THAT(page_result_info1.first, Not(Eq(kInvalidNextPageToken)));
403 ASSERT_THAT(page_result_info1.second.results, SizeIs(2));
404 EXPECT_THAT(page_result_info1.second.results.at(0).document(),
405 EqualsProto(CreateDocument(/*id=*/4)));
406 EXPECT_THAT(page_result_info1.second.results.at(1).document(),
407 EqualsProto(CreateDocument(/*id=*/3)));
408
409 uint64_t next_page_token = page_result_info1.first;
410
411 // Second page, all remaining documents will be filtered out by group result
412 // limiter, so we should get an empty page.
413 ICING_ASSERT_OK_AND_ASSIGN(
414 PageResultInfo page_result_info2,
415 result_state_manager.GetNextPage(next_page_token, result_retriever(),
416 clock()->GetSystemTimeMilliseconds()));
417 EXPECT_THAT(page_result_info2.first, Eq(kInvalidNextPageToken));
418 EXPECT_THAT(page_result_info2.second.results, IsEmpty());
419 }
420
TEST_F(ResultStateManagerTest,ShouldInvalidateExpiredTokensWhenCacheAndRetrieveFirstPage)421 TEST_F(ResultStateManagerTest,
422 ShouldInvalidateExpiredTokensWhenCacheAndRetrieveFirstPage) {
423 auto [scored_document_hits1, document_protos1] = AddScoredDocuments(
424 {/*document_id=*/0, /*document_id=*/1, /*document_id=*/2});
425 auto [scored_document_hits2, document_protos2] = AddScoredDocuments(
426 {/*document_id=*/3, /*document_id=*/4, /*document_id=*/5});
427
428 ResultStateManager result_state_manager(
429 /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
430
431 SectionRestrictQueryTermsMap query_terms;
432 SearchSpecProto search_spec;
433 ScoringSpecProto scoring_spec = CreateScoringSpec();
434 ResultSpecProto result_spec =
435 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE);
436
437 // Set time as 1s and add state 1.
438 clock()->SetSystemTimeMilliseconds(1000);
439 ICING_ASSERT_OK_AND_ASSIGN(
440 PageResultInfo page_result_info1,
441 result_state_manager.CacheAndRetrieveFirstPage(
442 std::make_unique<
443 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
444 std::move(scored_document_hits1), /*is_descending=*/true),
445 /*parent_adjustment_info=*/
446 std::make_unique<ResultAdjustmentInfo>(search_spec, scoring_spec,
447 result_spec, &schema_store(),
448 query_terms),
449 /*child_adjustment_info=*/nullptr, result_spec, document_store(),
450 result_retriever(), clock()->GetSystemTimeMilliseconds()));
451 ASSERT_THAT(page_result_info1.first, Not(Eq(kInvalidNextPageToken)));
452
453 // Set time as 1hr1s and add state 2.
454 clock()->SetSystemTimeMilliseconds(kDefaultResultStateTtlInMs + 1000);
455 ICING_ASSERT_OK_AND_ASSIGN(
456 PageResultInfo page_result_info2,
457 result_state_manager.CacheAndRetrieveFirstPage(
458 std::make_unique<
459 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
460 std::move(scored_document_hits2), /*is_descending=*/true),
461 /*parent_adjustment_info=*/
462 std::make_unique<ResultAdjustmentInfo>(search_spec, scoring_spec,
463 result_spec, &schema_store(),
464 query_terms),
465 /*child_adjustment_info=*/nullptr, result_spec, document_store(),
466 result_retriever(), clock()->GetSystemTimeMilliseconds()));
467
468 // Calling CacheAndRetrieveFirstPage() on state 2 should invalidate the
469 // expired state 1 internally.
470 //
471 // We test the behavior by setting time back to 1s, to make sure the
472 // invalidation of state 1 was done by the previous
473 // CacheAndRetrieveFirstPage() instead of the following GetNextPage().
474 clock()->SetSystemTimeMilliseconds(1000);
475 // page_result_info1's token (page_result_info1.first) shouldn't be found.
476 EXPECT_THAT(result_state_manager.GetNextPage(
477 page_result_info1.first, result_retriever(),
478 clock()->GetSystemTimeMilliseconds()),
479 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
480 }
481
TEST_F(ResultStateManagerTest,ShouldInvalidateExpiredTokensWhenGetNextPageOnOthers)482 TEST_F(ResultStateManagerTest,
483 ShouldInvalidateExpiredTokensWhenGetNextPageOnOthers) {
484 auto [scored_document_hits1, document_protos1] = AddScoredDocuments(
485 {/*document_id=*/0, /*document_id=*/1, /*document_id=*/2});
486 auto [scored_document_hits2, document_protos2] = AddScoredDocuments(
487 {/*document_id=*/3, /*document_id=*/4, /*document_id=*/5});
488
489 ResultStateManager result_state_manager(
490 /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
491
492 // Set time as 1s and add state 1.
493 clock()->SetSystemTimeMilliseconds(1000);
494 ICING_ASSERT_OK_AND_ASSIGN(
495 PageResultInfo page_result_info1,
496 result_state_manager.CacheAndRetrieveFirstPage(
497 std::make_unique<
498 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
499 std::move(scored_document_hits1), /*is_descending=*/true),
500 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
501 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
502 document_store(), result_retriever(),
503 clock()->GetSystemTimeMilliseconds()));
504 ASSERT_THAT(page_result_info1.first, Not(Eq(kInvalidNextPageToken)));
505
506 // Set time as 2s and add state 2.
507 clock()->SetSystemTimeMilliseconds(2000);
508 ICING_ASSERT_OK_AND_ASSIGN(
509 PageResultInfo page_result_info2,
510 result_state_manager.CacheAndRetrieveFirstPage(
511 std::make_unique<
512 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
513 std::move(scored_document_hits2), /*is_descending=*/true),
514 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
515 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
516 document_store(), result_retriever(),
517 clock()->GetSystemTimeMilliseconds()));
518 ASSERT_THAT(page_result_info2.first, Not(Eq(kInvalidNextPageToken)));
519
520 // 1. Set time as 1hr1s.
521 // 2. Call GetNextPage() on state 2. It should correctly invalidate the
522 // expired state 1.
523 // 3. Then calling GetNextPage() on state 1 shouldn't get anything.
524 clock()->SetSystemTimeMilliseconds(kDefaultResultStateTtlInMs + 1000);
525 // page_result_info2's token (page_result_info2.first) should be found
526 ICING_ASSERT_OK_AND_ASSIGN(page_result_info2,
527 result_state_manager.GetNextPage(
528 page_result_info2.first, result_retriever(),
529 clock()->GetSystemTimeMilliseconds()));
530 // We test the behavior by setting time back to 2s, to make sure the
531 // invalidation of state 1 was done by the previous GetNextPage() instead of
532 // the following GetNextPage().
533 clock()->SetSystemTimeMilliseconds(2000);
534 // page_result_info1's token (page_result_info1.first) shouldn't be found.
535 EXPECT_THAT(result_state_manager.GetNextPage(
536 page_result_info1.first, result_retriever(),
537 clock()->GetSystemTimeMilliseconds()),
538 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
539 }
540
TEST_F(ResultStateManagerTest,ShouldInvalidateExpiredTokensWhenGetNextPageOnItself)541 TEST_F(ResultStateManagerTest,
542 ShouldInvalidateExpiredTokensWhenGetNextPageOnItself) {
543 auto [scored_document_hits1, document_protos1] = AddScoredDocuments(
544 {/*document_id=*/0, /*document_id=*/1, /*document_id=*/2});
545 auto [scored_document_hits2, document_protos2] = AddScoredDocuments(
546 {/*document_id=*/3, /*document_id=*/4, /*document_id=*/5});
547
548 ResultStateManager result_state_manager(
549 /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
550
551 // Set time as 1s and add state.
552 clock()->SetSystemTimeMilliseconds(1000);
553 ICING_ASSERT_OK_AND_ASSIGN(
554 PageResultInfo page_result_info,
555 result_state_manager.CacheAndRetrieveFirstPage(
556 std::make_unique<
557 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
558 std::move(scored_document_hits1), /*is_descending=*/true),
559 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
560 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
561 document_store(), result_retriever(),
562 clock()->GetSystemTimeMilliseconds()));
563 ASSERT_THAT(page_result_info.first, Not(Eq(kInvalidNextPageToken)));
564
565 // 1. Set time as 1hr1s.
566 // 2. Then calling GetNextPage() on the state shouldn't get anything.
567 clock()->SetSystemTimeMilliseconds(kDefaultResultStateTtlInMs + 1000);
568 // page_result_info's token (page_result_info.first) shouldn't be found.
569 EXPECT_THAT(result_state_manager.GetNextPage(
570 page_result_info.first, result_retriever(),
571 clock()->GetSystemTimeMilliseconds()),
572 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
573 }
574
TEST_F(ResultStateManagerTest,ShouldInvalidateOneToken)575 TEST_F(ResultStateManagerTest, ShouldInvalidateOneToken) {
576 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id1,
577 document_store().Put(CreateDocument(/*id=*/1)));
578 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id2,
579 document_store().Put(CreateDocument(/*id=*/2)));
580 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id3,
581 document_store().Put(CreateDocument(/*id=*/3)));
582 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id4,
583 document_store().Put(CreateDocument(/*id=*/4)));
584 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id5,
585 document_store().Put(CreateDocument(/*id=*/5)));
586 ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id6,
587 document_store().Put(CreateDocument(/*id=*/6)));
588 std::vector<ScoredDocumentHit> scored_document_hits1 = {
589 {document_id1, kSectionIdMaskNone, /*score=*/1},
590 {document_id2, kSectionIdMaskNone, /*score=*/1},
591 {document_id3, kSectionIdMaskNone, /*score=*/1}};
592 std::vector<ScoredDocumentHit> scored_document_hits2 = {
593 {document_id4, kSectionIdMaskNone, /*score=*/1},
594 {document_id5, kSectionIdMaskNone, /*score=*/1},
595 {document_id6, kSectionIdMaskNone, /*score=*/1}};
596
597 ResultStateManager result_state_manager(
598 /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
599
600 ICING_ASSERT_OK_AND_ASSIGN(
601 PageResultInfo page_result_info1,
602 result_state_manager.CacheAndRetrieveFirstPage(
603 std::make_unique<
604 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
605 std::move(scored_document_hits1), /*is_descending=*/true),
606 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
607 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
608 document_store(), result_retriever(),
609 clock()->GetSystemTimeMilliseconds()));
610
611 ICING_ASSERT_OK_AND_ASSIGN(
612 PageResultInfo page_result_info2,
613 result_state_manager.CacheAndRetrieveFirstPage(
614 std::make_unique<
615 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
616 std::move(scored_document_hits2), /*is_descending=*/true),
617 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
618 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
619 document_store(), result_retriever(),
620 clock()->GetSystemTimeMilliseconds()));
621
622 // Invalidate first result state by the token.
623 result_state_manager.InvalidateResultState(page_result_info1.first);
624
625 // page_result_info1's token (page_result_info1.first) shouldn't be found
626 EXPECT_THAT(result_state_manager.GetNextPage(
627 page_result_info1.first, result_retriever(),
628 clock()->GetSystemTimeMilliseconds()),
629 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
630
631 // page_result_info2's token (page_result_info2.first) should still exist
632 ICING_ASSERT_OK_AND_ASSIGN(page_result_info2,
633 result_state_manager.GetNextPage(
634 page_result_info2.first, result_retriever(),
635 clock()->GetSystemTimeMilliseconds()));
636 // Should get docs.
637 ASSERT_THAT(page_result_info2.second.results, SizeIs(1));
638 EXPECT_THAT(page_result_info2.second.results.at(0).document(),
639 EqualsProto(CreateDocument(/*id=*/5)));
640 }
641
TEST_F(ResultStateManagerTest,ShouldInvalidateAllTokens)642 TEST_F(ResultStateManagerTest, ShouldInvalidateAllTokens) {
643 auto [scored_document_hits1, document_protos1] = AddScoredDocuments(
644 {/*document_id=*/0, /*document_id=*/1, /*document_id=*/2});
645 auto [scored_document_hits2, document_protos2] = AddScoredDocuments(
646 {/*document_id=*/3, /*document_id=*/4, /*document_id=*/5});
647
648 ResultStateManager result_state_manager(
649 /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
650
651 ICING_ASSERT_OK_AND_ASSIGN(
652 PageResultInfo page_result_info1,
653 result_state_manager.CacheAndRetrieveFirstPage(
654 std::make_unique<
655 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
656 std::move(scored_document_hits1), /*is_descending=*/true),
657 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
658 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
659 document_store(), result_retriever(),
660 clock()->GetSystemTimeMilliseconds()));
661
662 ICING_ASSERT_OK_AND_ASSIGN(
663 PageResultInfo page_result_info2,
664 result_state_manager.CacheAndRetrieveFirstPage(
665 std::make_unique<
666 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
667 std::move(scored_document_hits2), /*is_descending=*/true),
668 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
669 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
670 document_store(), result_retriever(),
671 clock()->GetSystemTimeMilliseconds()));
672
673 result_state_manager.InvalidateAllResultStates();
674
675 // page_result_info1's token (page_result_info1.first) shouldn't be found
676 EXPECT_THAT(result_state_manager.GetNextPage(
677 page_result_info1.first, result_retriever(),
678 clock()->GetSystemTimeMilliseconds()),
679 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
680
681 // page_result_info2's token (page_result_info2.first) shouldn't be found
682 EXPECT_THAT(result_state_manager.GetNextPage(
683 page_result_info2.first, result_retriever(),
684 clock()->GetSystemTimeMilliseconds()),
685 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
686 }
687
TEST_F(ResultStateManagerTest,ShouldRemoveOldestResultState)688 TEST_F(ResultStateManagerTest, ShouldRemoveOldestResultState) {
689 auto [scored_document_hits1, document_protos1] =
690 AddScoredDocuments({/*document_id=*/0, /*document_id=*/1});
691 auto [scored_document_hits2, document_protos2] =
692 AddScoredDocuments({/*document_id=*/2, /*document_id=*/3});
693 auto [scored_document_hits3, document_protos3] =
694 AddScoredDocuments({/*document_id=*/4, /*document_id=*/5});
695
696 ResultStateManager result_state_manager(/*max_total_hits=*/2,
697 document_store());
698
699 ICING_ASSERT_OK_AND_ASSIGN(
700 PageResultInfo page_result_info1,
701 result_state_manager.CacheAndRetrieveFirstPage(
702 std::make_unique<
703 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
704 std::move(scored_document_hits1), /*is_descending=*/true),
705 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
706 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
707 document_store(), result_retriever(),
708 clock()->GetSystemTimeMilliseconds()));
709
710 ICING_ASSERT_OK_AND_ASSIGN(
711 PageResultInfo page_result_info2,
712 result_state_manager.CacheAndRetrieveFirstPage(
713 std::make_unique<
714 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
715 std::move(scored_document_hits2), /*is_descending=*/true),
716 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
717 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
718 document_store(), result_retriever(),
719 clock()->GetSystemTimeMilliseconds()));
720
721 // Adding state 3 should cause state 1 to be removed.
722 ICING_ASSERT_OK_AND_ASSIGN(
723 PageResultInfo page_result_info3,
724 result_state_manager.CacheAndRetrieveFirstPage(
725 std::make_unique<
726 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
727 std::move(scored_document_hits3), /*is_descending=*/true),
728 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
729 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
730 document_store(), result_retriever(),
731 clock()->GetSystemTimeMilliseconds()));
732
733 EXPECT_THAT(result_state_manager.GetNextPage(
734 page_result_info1.first, result_retriever(),
735 clock()->GetSystemTimeMilliseconds()),
736 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
737
738 ICING_ASSERT_OK_AND_ASSIGN(page_result_info2,
739 result_state_manager.GetNextPage(
740 page_result_info2.first, result_retriever(),
741 clock()->GetSystemTimeMilliseconds()));
742 ASSERT_THAT(page_result_info2.second.results, SizeIs(1));
743 EXPECT_THAT(page_result_info2.second.results.at(0).document(),
744 EqualsProto(document_protos2.at(1)));
745
746 ICING_ASSERT_OK_AND_ASSIGN(page_result_info3,
747 result_state_manager.GetNextPage(
748 page_result_info3.first, result_retriever(),
749 clock()->GetSystemTimeMilliseconds()));
750 ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
751 EXPECT_THAT(page_result_info3.second.results.at(0).document(),
752 EqualsProto(document_protos3.at(1)));
753 }
754
TEST_F(ResultStateManagerTest,InvalidatedResultStateShouldDecreaseCurrentHitsCount)755 TEST_F(ResultStateManagerTest,
756 InvalidatedResultStateShouldDecreaseCurrentHitsCount) {
757 auto [scored_document_hits1, document_protos1] =
758 AddScoredDocuments({/*document_id=*/0, /*document_id=*/1});
759 auto [scored_document_hits2, document_protos2] =
760 AddScoredDocuments({/*document_id=*/2, /*document_id=*/3});
761 auto [scored_document_hits3, document_protos3] =
762 AddScoredDocuments({/*document_id=*/4, /*document_id=*/5});
763
764 // Add the first three states. Remember, the first page for each result state
765 // won't be cached (since it is returned immediately from
766 // CacheAndRetrieveFirstPage). Each result state has a page size of 1 and a
767 // result set of 2 hits. So each result will take up one hit of our three hit
768 // budget.
769 ResultStateManager result_state_manager(/*max_total_hits=*/3,
770 document_store());
771
772 ICING_ASSERT_OK_AND_ASSIGN(
773 PageResultInfo page_result_info1,
774 result_state_manager.CacheAndRetrieveFirstPage(
775 std::make_unique<
776 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
777 std::move(scored_document_hits1), /*is_descending=*/true),
778 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
779 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
780 document_store(), result_retriever(),
781 clock()->GetSystemTimeMilliseconds()));
782
783 ICING_ASSERT_OK_AND_ASSIGN(
784 PageResultInfo page_result_info2,
785 result_state_manager.CacheAndRetrieveFirstPage(
786 std::make_unique<
787 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
788 std::move(scored_document_hits2), /*is_descending=*/true),
789 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
790 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
791 document_store(), result_retriever(),
792 clock()->GetSystemTimeMilliseconds()));
793
794 ICING_ASSERT_OK_AND_ASSIGN(
795 PageResultInfo page_result_info3,
796 result_state_manager.CacheAndRetrieveFirstPage(
797 std::make_unique<
798 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
799 std::move(scored_document_hits3), /*is_descending=*/true),
800 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
801 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
802 document_store(), result_retriever(),
803 clock()->GetSystemTimeMilliseconds()));
804
805 // Invalidates state 2, so that the number of hits current cached should be
806 // decremented to 2.
807 result_state_manager.InvalidateResultState(page_result_info2.first);
808
809 // If invalidating state 2 correctly decremented the current hit count to 2,
810 // then adding state 4 should still be within our budget and no other result
811 // states should be evicted.
812 auto [scored_document_hits4, document_protos4] =
813 AddScoredDocuments({/*document_id=*/6, /*document_id=*/7});
814 ICING_ASSERT_OK_AND_ASSIGN(
815 PageResultInfo page_result_info4,
816 result_state_manager.CacheAndRetrieveFirstPage(
817 std::make_unique<
818 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
819 std::move(scored_document_hits4), /*is_descending=*/true),
820 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
821 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
822 document_store(), result_retriever(),
823 clock()->GetSystemTimeMilliseconds()));
824
825 ICING_ASSERT_OK_AND_ASSIGN(page_result_info1,
826 result_state_manager.GetNextPage(
827 page_result_info1.first, result_retriever(),
828 clock()->GetSystemTimeMilliseconds()));
829 ASSERT_THAT(page_result_info1.second.results, SizeIs(1));
830 EXPECT_THAT(page_result_info1.second.results.at(0).document(),
831 EqualsProto(document_protos1.at(1)));
832
833 EXPECT_THAT(result_state_manager.GetNextPage(
834 page_result_info2.first, result_retriever(),
835 clock()->GetSystemTimeMilliseconds()),
836 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
837
838 ICING_ASSERT_OK_AND_ASSIGN(page_result_info3,
839 result_state_manager.GetNextPage(
840 page_result_info3.first, result_retriever(),
841 clock()->GetSystemTimeMilliseconds()));
842 ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
843 EXPECT_THAT(page_result_info3.second.results.at(0).document(),
844 EqualsProto(document_protos3.at(1)));
845
846 ICING_ASSERT_OK_AND_ASSIGN(page_result_info4,
847 result_state_manager.GetNextPage(
848 page_result_info4.first, result_retriever(),
849 clock()->GetSystemTimeMilliseconds()));
850 ASSERT_THAT(page_result_info4.second.results, SizeIs(1));
851 EXPECT_THAT(page_result_info4.second.results.at(0).document(),
852 EqualsProto(document_protos4.at(1)));
853 }
854
TEST_F(ResultStateManagerTest,InvalidatedAllResultStatesShouldResetCurrentHitCount)855 TEST_F(ResultStateManagerTest,
856 InvalidatedAllResultStatesShouldResetCurrentHitCount) {
857 auto [scored_document_hits1, document_protos1] =
858 AddScoredDocuments({/*document_id=*/0, /*document_id=*/1});
859 auto [scored_document_hits2, document_protos2] =
860 AddScoredDocuments({/*document_id=*/2, /*document_id=*/3});
861 auto [scored_document_hits3, document_protos3] =
862 AddScoredDocuments({/*document_id=*/4, /*document_id=*/5});
863
864 // Add the first three states. Remember, the first page for each result state
865 // won't be cached (since it is returned immediately from
866 // CacheAndRetrieveFirstPage). Each result state has a page size of 1 and a
867 // result set of 2 hits. So each result will take up one hit of our three hit
868 // budget.
869 ResultStateManager result_state_manager(/*max_total_hits=*/3,
870 document_store());
871
872 ICING_ASSERT_OK_AND_ASSIGN(
873 PageResultInfo page_result_info1,
874 result_state_manager.CacheAndRetrieveFirstPage(
875 std::make_unique<
876 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
877 std::move(scored_document_hits1), /*is_descending=*/true),
878 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
879 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
880 document_store(), result_retriever(),
881 clock()->GetSystemTimeMilliseconds()));
882
883 ICING_ASSERT_OK_AND_ASSIGN(
884 PageResultInfo page_result_info2,
885 result_state_manager.CacheAndRetrieveFirstPage(
886 std::make_unique<
887 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
888 std::move(scored_document_hits2), /*is_descending=*/true),
889 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
890 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
891 document_store(), result_retriever(),
892 clock()->GetSystemTimeMilliseconds()));
893
894 ICING_ASSERT_OK_AND_ASSIGN(
895 PageResultInfo page_result_info3,
896 result_state_manager.CacheAndRetrieveFirstPage(
897 std::make_unique<
898 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
899 std::move(scored_document_hits3), /*is_descending=*/true),
900 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
901 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
902 document_store(), result_retriever(),
903 clock()->GetSystemTimeMilliseconds()));
904
905 // Invalidates all states so that the current hit count will be 0.
906 result_state_manager.InvalidateAllResultStates();
907
908 // If invalidating all states correctly reset the current hit count to 0,
909 // then adding state 4, 5, 6 should still be within our budget and no other
910 // result states should be evicted.
911 auto [scored_document_hits4, document_protos4] =
912 AddScoredDocuments({/*document_id=*/6, /*document_id=*/7});
913 auto [scored_document_hits5, document_protos5] =
914 AddScoredDocuments({/*document_id=*/8, /*document_id=*/9});
915 auto [scored_document_hits6, document_protos6] =
916 AddScoredDocuments({/*document_id=*/10, /*document_id=*/11});
917
918 ICING_ASSERT_OK_AND_ASSIGN(
919 PageResultInfo page_result_info4,
920 result_state_manager.CacheAndRetrieveFirstPage(
921 std::make_unique<
922 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
923 std::move(scored_document_hits4), /*is_descending=*/true),
924 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
925 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
926 document_store(), result_retriever(),
927 clock()->GetSystemTimeMilliseconds()));
928
929 ICING_ASSERT_OK_AND_ASSIGN(
930 PageResultInfo page_result_info5,
931 result_state_manager.CacheAndRetrieveFirstPage(
932 std::make_unique<
933 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
934 std::move(scored_document_hits5), /*is_descending=*/true),
935 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
936 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
937 document_store(), result_retriever(),
938 clock()->GetSystemTimeMilliseconds()));
939
940 ICING_ASSERT_OK_AND_ASSIGN(
941 PageResultInfo page_result_info6,
942 result_state_manager.CacheAndRetrieveFirstPage(
943 std::make_unique<
944 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
945 std::move(scored_document_hits6), /*is_descending=*/true),
946 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
947 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
948 document_store(), result_retriever(),
949 clock()->GetSystemTimeMilliseconds()));
950
951 EXPECT_THAT(result_state_manager.GetNextPage(
952 page_result_info1.first, result_retriever(),
953 clock()->GetSystemTimeMilliseconds()),
954 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
955
956 EXPECT_THAT(result_state_manager.GetNextPage(
957 page_result_info2.first, result_retriever(),
958 clock()->GetSystemTimeMilliseconds()),
959 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
960
961 EXPECT_THAT(result_state_manager.GetNextPage(
962 page_result_info3.first, result_retriever(),
963 clock()->GetSystemTimeMilliseconds()),
964 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
965
966 ICING_ASSERT_OK_AND_ASSIGN(page_result_info4,
967 result_state_manager.GetNextPage(
968 page_result_info4.first, result_retriever(),
969 clock()->GetSystemTimeMilliseconds()));
970 ASSERT_THAT(page_result_info4.second.results, SizeIs(1));
971 EXPECT_THAT(page_result_info4.second.results.at(0).document(),
972 EqualsProto(document_protos4.at(1)));
973
974 ICING_ASSERT_OK_AND_ASSIGN(page_result_info5,
975 result_state_manager.GetNextPage(
976 page_result_info5.first, result_retriever(),
977 clock()->GetSystemTimeMilliseconds()));
978 ASSERT_THAT(page_result_info5.second.results, SizeIs(1));
979 EXPECT_THAT(page_result_info5.second.results.at(0).document(),
980 EqualsProto(document_protos5.at(1)));
981
982 ICING_ASSERT_OK_AND_ASSIGN(page_result_info6,
983 result_state_manager.GetNextPage(
984 page_result_info6.first, result_retriever(),
985 clock()->GetSystemTimeMilliseconds()));
986 ASSERT_THAT(page_result_info6.second.results, SizeIs(1));
987 EXPECT_THAT(page_result_info6.second.results.at(0).document(),
988 EqualsProto(document_protos6.at(1)));
989 }
990
TEST_F(ResultStateManagerTest,InvalidatedResultStateShouldDecreaseCurrentHitsCountByExactStateHitCount)991 TEST_F(
992 ResultStateManagerTest,
993 InvalidatedResultStateShouldDecreaseCurrentHitsCountByExactStateHitCount) {
994 auto [scored_document_hits1, document_protos1] =
995 AddScoredDocuments({/*document_id=*/0, /*document_id=*/1});
996 auto [scored_document_hits2, document_protos2] =
997 AddScoredDocuments({/*document_id=*/2, /*document_id=*/3});
998 auto [scored_document_hits3, document_protos3] =
999 AddScoredDocuments({/*document_id=*/4, /*document_id=*/5});
1000
1001 // Add the first three states. Remember, the first page for each result state
1002 // won't be cached (since it is returned immediately from
1003 // CacheAndRetrieveFirstPage). Each result state has a page size of 1 and a
1004 // result set of 2 hits. So each result will take up one hit of our three hit
1005 // budget.
1006 ResultStateManager result_state_manager(/*max_total_hits=*/3,
1007 document_store());
1008
1009 ICING_ASSERT_OK_AND_ASSIGN(
1010 PageResultInfo page_result_info1,
1011 result_state_manager.CacheAndRetrieveFirstPage(
1012 std::make_unique<
1013 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1014 std::move(scored_document_hits1), /*is_descending=*/true),
1015 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1016 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1017 document_store(), result_retriever(),
1018 clock()->GetSystemTimeMilliseconds()));
1019
1020 ICING_ASSERT_OK_AND_ASSIGN(
1021 PageResultInfo page_result_info2,
1022 result_state_manager.CacheAndRetrieveFirstPage(
1023 std::make_unique<
1024 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1025 std::move(scored_document_hits2), /*is_descending=*/true),
1026 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1027 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1028 document_store(), result_retriever(),
1029 clock()->GetSystemTimeMilliseconds()));
1030
1031 ICING_ASSERT_OK_AND_ASSIGN(
1032 PageResultInfo page_result_info3,
1033 result_state_manager.CacheAndRetrieveFirstPage(
1034 std::make_unique<
1035 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1036 std::move(scored_document_hits3), /*is_descending=*/true),
1037 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1038 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1039 document_store(), result_retriever(),
1040 clock()->GetSystemTimeMilliseconds()));
1041
1042 // Invalidates state 2, so that the number of hits current cached should be
1043 // decremented to 2.
1044 result_state_manager.InvalidateResultState(page_result_info2.first);
1045
1046 // If invalidating state 2 correctly decremented the current hit count to 2,
1047 // then adding state 4 should still be within our budget and no other result
1048 // states should be evicted.
1049 auto [scored_document_hits4, document_protos4] =
1050 AddScoredDocuments({/*document_id=*/6, /*document_id=*/7});
1051 ICING_ASSERT_OK_AND_ASSIGN(
1052 PageResultInfo page_result_info4,
1053 result_state_manager.CacheAndRetrieveFirstPage(
1054 std::make_unique<
1055 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1056 std::move(scored_document_hits4), /*is_descending=*/true),
1057 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1058 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1059 document_store(), result_retriever(),
1060 clock()->GetSystemTimeMilliseconds()));
1061
1062 // If invalidating result state 2 correctly decremented the current hit count
1063 // to 2 and adding state 4 correctly incremented it to 3, then adding this
1064 // result state should trigger the eviction of state 1.
1065 auto [scored_document_hits5, document_protos5] =
1066 AddScoredDocuments({/*document_id=*/8, /*document_id=*/9});
1067 ICING_ASSERT_OK_AND_ASSIGN(
1068 PageResultInfo page_result_info5,
1069 result_state_manager.CacheAndRetrieveFirstPage(
1070 std::make_unique<
1071 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1072 std::move(scored_document_hits5), /*is_descending=*/true),
1073 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1074 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1075 document_store(), result_retriever(),
1076 clock()->GetSystemTimeMilliseconds()));
1077
1078 EXPECT_THAT(result_state_manager.GetNextPage(
1079 page_result_info1.first, result_retriever(),
1080 clock()->GetSystemTimeMilliseconds()),
1081 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
1082
1083 EXPECT_THAT(result_state_manager.GetNextPage(
1084 page_result_info2.first, result_retriever(),
1085 clock()->GetSystemTimeMilliseconds()),
1086 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
1087
1088 ICING_ASSERT_OK_AND_ASSIGN(page_result_info3,
1089 result_state_manager.GetNextPage(
1090 page_result_info3.first, result_retriever(),
1091 clock()->GetSystemTimeMilliseconds()));
1092 ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
1093 EXPECT_THAT(page_result_info3.second.results.at(0).document(),
1094 EqualsProto(document_protos3.at(1)));
1095
1096 ICING_ASSERT_OK_AND_ASSIGN(page_result_info4,
1097 result_state_manager.GetNextPage(
1098 page_result_info4.first, result_retriever(),
1099 clock()->GetSystemTimeMilliseconds()));
1100 ASSERT_THAT(page_result_info4.second.results, SizeIs(1));
1101 EXPECT_THAT(page_result_info4.second.results.at(0).document(),
1102 EqualsProto(document_protos4.at(1)));
1103
1104 ICING_ASSERT_OK_AND_ASSIGN(page_result_info5,
1105 result_state_manager.GetNextPage(
1106 page_result_info5.first, result_retriever(),
1107 clock()->GetSystemTimeMilliseconds()));
1108 ASSERT_THAT(page_result_info5.second.results, SizeIs(1));
1109 EXPECT_THAT(page_result_info5.second.results.at(0).document(),
1110 EqualsProto(document_protos5.at(1)));
1111 }
1112
TEST_F(ResultStateManagerTest,GetNextPageShouldDecreaseCurrentHitsCount)1113 TEST_F(ResultStateManagerTest, GetNextPageShouldDecreaseCurrentHitsCount) {
1114 auto [scored_document_hits1, document_protos1] =
1115 AddScoredDocuments({/*document_id=*/0, /*document_id=*/1});
1116 auto [scored_document_hits2, document_protos2] =
1117 AddScoredDocuments({/*document_id=*/2, /*document_id=*/3});
1118 auto [scored_document_hits3, document_protos3] =
1119 AddScoredDocuments({/*document_id=*/4, /*document_id=*/5});
1120
1121 // Add the first three states. Remember, the first page for each result state
1122 // won't be cached (since it is returned immediately from
1123 // CacheAndRetrieveFirstPage). Each result state has a page size of 1 and a
1124 // result set of 2 hits. So each result will take up one hit of our three hit
1125 // budget.
1126 ResultStateManager result_state_manager(/*max_total_hits=*/3,
1127 document_store());
1128
1129 ICING_ASSERT_OK_AND_ASSIGN(
1130 PageResultInfo page_result_info1,
1131 result_state_manager.CacheAndRetrieveFirstPage(
1132 std::make_unique<
1133 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1134 std::move(scored_document_hits1), /*is_descending=*/true),
1135 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1136 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1137 document_store(), result_retriever(),
1138 clock()->GetSystemTimeMilliseconds()));
1139
1140 ICING_ASSERT_OK_AND_ASSIGN(
1141 PageResultInfo page_result_info2,
1142 result_state_manager.CacheAndRetrieveFirstPage(
1143 std::make_unique<
1144 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1145 std::move(scored_document_hits2), /*is_descending=*/true),
1146 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1147 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1148 document_store(), result_retriever(),
1149 clock()->GetSystemTimeMilliseconds()));
1150
1151 ICING_ASSERT_OK_AND_ASSIGN(
1152 PageResultInfo page_result_info3,
1153 result_state_manager.CacheAndRetrieveFirstPage(
1154 std::make_unique<
1155 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1156 std::move(scored_document_hits3), /*is_descending=*/true),
1157 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1158 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1159 document_store(), result_retriever(),
1160 clock()->GetSystemTimeMilliseconds()));
1161
1162 // GetNextPage for result state 1 should return its result and decrement the
1163 // number of cached hits to 2.
1164 ICING_ASSERT_OK_AND_ASSIGN(page_result_info1,
1165 result_state_manager.GetNextPage(
1166 page_result_info1.first, result_retriever(),
1167 clock()->GetSystemTimeMilliseconds()));
1168 ASSERT_THAT(page_result_info1.second.results, SizeIs(1));
1169 EXPECT_THAT(page_result_info1.second.results.at(0).document(),
1170 EqualsProto(document_protos1.at(1)));
1171
1172 // If retrieving the next page for result state 1 correctly decremented the
1173 // current hit count to 2, then adding state 4 should still be within our
1174 // budget and no other result states should be evicted.
1175 auto [scored_document_hits4, document_protos4] =
1176 AddScoredDocuments({/*document_id=*/6, /*document_id=*/7});
1177 ICING_ASSERT_OK_AND_ASSIGN(
1178 PageResultInfo page_result_info4,
1179 result_state_manager.CacheAndRetrieveFirstPage(
1180 std::make_unique<
1181 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1182 std::move(scored_document_hits4), /*is_descending=*/true),
1183 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1184 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1185 document_store(), result_retriever(),
1186 clock()->GetSystemTimeMilliseconds()));
1187
1188 EXPECT_THAT(result_state_manager.GetNextPage(
1189 page_result_info1.first, result_retriever(),
1190 clock()->GetSystemTimeMilliseconds()),
1191 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
1192
1193 ICING_ASSERT_OK_AND_ASSIGN(page_result_info2,
1194 result_state_manager.GetNextPage(
1195 page_result_info2.first, result_retriever(),
1196 clock()->GetSystemTimeMilliseconds()));
1197 ASSERT_THAT(page_result_info2.second.results, SizeIs(1));
1198 EXPECT_THAT(page_result_info2.second.results.at(0).document(),
1199 EqualsProto(document_protos2.at(1)));
1200
1201 ICING_ASSERT_OK_AND_ASSIGN(page_result_info3,
1202 result_state_manager.GetNextPage(
1203 page_result_info3.first, result_retriever(),
1204 clock()->GetSystemTimeMilliseconds()));
1205 ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
1206 EXPECT_THAT(page_result_info3.second.results.at(0).document(),
1207 EqualsProto(document_protos3.at(1)));
1208
1209 ICING_ASSERT_OK_AND_ASSIGN(page_result_info4,
1210 result_state_manager.GetNextPage(
1211 page_result_info4.first, result_retriever(),
1212 clock()->GetSystemTimeMilliseconds()));
1213 ASSERT_THAT(page_result_info4.second.results, SizeIs(1));
1214 EXPECT_THAT(page_result_info4.second.results.at(0).document(),
1215 EqualsProto(document_protos4.at(1)));
1216 }
1217
TEST_F(ResultStateManagerTest,GetNextPageShouldDecreaseCurrentHitsCountByExactlyOnePage)1218 TEST_F(ResultStateManagerTest,
1219 GetNextPageShouldDecreaseCurrentHitsCountByExactlyOnePage) {
1220 auto [scored_document_hits1, document_protos1] =
1221 AddScoredDocuments({/*document_id=*/0, /*document_id=*/1});
1222 auto [scored_document_hits2, document_protos2] =
1223 AddScoredDocuments({/*document_id=*/2, /*document_id=*/3});
1224 auto [scored_document_hits3, document_protos3] =
1225 AddScoredDocuments({/*document_id=*/4, /*document_id=*/5});
1226
1227 // Add the first three states. Remember, the first page for each result state
1228 // won't be cached (since it is returned immediately from
1229 // CacheAndRetrieveFirstPage). Each result state has a page size of 1 and a
1230 // result set of 2 hits. So each result will take up one hit of our three hit
1231 // budget.
1232 ResultStateManager result_state_manager(/*max_total_hits=*/3,
1233 document_store());
1234
1235 ICING_ASSERT_OK_AND_ASSIGN(
1236 PageResultInfo page_result_info1,
1237 result_state_manager.CacheAndRetrieveFirstPage(
1238 std::make_unique<
1239 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1240 std::move(scored_document_hits1), /*is_descending=*/true),
1241 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1242 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1243 document_store(), result_retriever(),
1244 clock()->GetSystemTimeMilliseconds()));
1245
1246 ICING_ASSERT_OK_AND_ASSIGN(
1247 PageResultInfo page_result_info2,
1248 result_state_manager.CacheAndRetrieveFirstPage(
1249 std::make_unique<
1250 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1251 std::move(scored_document_hits2), /*is_descending=*/true),
1252 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1253 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1254 document_store(), result_retriever(),
1255 clock()->GetSystemTimeMilliseconds()));
1256
1257 ICING_ASSERT_OK_AND_ASSIGN(
1258 PageResultInfo page_result_info3,
1259 result_state_manager.CacheAndRetrieveFirstPage(
1260 std::make_unique<
1261 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1262 std::move(scored_document_hits3), /*is_descending=*/true),
1263 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1264 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1265 document_store(), result_retriever(),
1266 clock()->GetSystemTimeMilliseconds()));
1267
1268 // GetNextPage for result state 1 should return its result and decrement the
1269 // number of cached hits to 2.
1270 ICING_ASSERT_OK_AND_ASSIGN(page_result_info1,
1271 result_state_manager.GetNextPage(
1272 page_result_info1.first, result_retriever(),
1273 clock()->GetSystemTimeMilliseconds()));
1274 ASSERT_THAT(page_result_info1.second.results, SizeIs(1));
1275 EXPECT_THAT(page_result_info1.second.results.at(0).document(),
1276 EqualsProto(document_protos1.at(1)));
1277
1278 // If retrieving the next page for result state 1 correctly decremented the
1279 // current hit count to 2, then adding state 4 should still be within our
1280 // budget and no other result states should be evicted.
1281 auto [scored_document_hits4, document_protos4] =
1282 AddScoredDocuments({/*document_id=*/6, /*document_id=*/7});
1283 ICING_ASSERT_OK_AND_ASSIGN(
1284 PageResultInfo page_result_info4,
1285 result_state_manager.CacheAndRetrieveFirstPage(
1286 std::make_unique<
1287 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1288 std::move(scored_document_hits4), /*is_descending=*/true),
1289 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1290 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1291 document_store(), result_retriever(),
1292 clock()->GetSystemTimeMilliseconds()));
1293
1294 // If retrieving the next page for result state 1 correctly decremented the
1295 // current hit count to 2 and adding state 4 correctly incremented it to 3,
1296 // then adding this result state should trigger the eviction of state 2.
1297 auto [scored_document_hits5, document_protos5] =
1298 AddScoredDocuments({/*document_id=*/8, /*document_id=*/9});
1299 ICING_ASSERT_OK_AND_ASSIGN(
1300 PageResultInfo page_result_info5,
1301 result_state_manager.CacheAndRetrieveFirstPage(
1302 std::make_unique<
1303 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1304 std::move(scored_document_hits5), /*is_descending=*/true),
1305 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1306 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1307 document_store(), result_retriever(),
1308 clock()->GetSystemTimeMilliseconds()));
1309
1310 EXPECT_THAT(result_state_manager.GetNextPage(
1311 page_result_info1.first, result_retriever(),
1312 clock()->GetSystemTimeMilliseconds()),
1313 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
1314
1315 EXPECT_THAT(result_state_manager.GetNextPage(
1316 page_result_info2.first, result_retriever(),
1317 clock()->GetSystemTimeMilliseconds()),
1318 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
1319
1320 ICING_ASSERT_OK_AND_ASSIGN(page_result_info3,
1321 result_state_manager.GetNextPage(
1322 page_result_info3.first, result_retriever(),
1323 clock()->GetSystemTimeMilliseconds()));
1324 ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
1325 EXPECT_THAT(page_result_info3.second.results.at(0).document(),
1326 EqualsProto(document_protos3.at(1)));
1327
1328 ICING_ASSERT_OK_AND_ASSIGN(page_result_info4,
1329 result_state_manager.GetNextPage(
1330 page_result_info4.first, result_retriever(),
1331 clock()->GetSystemTimeMilliseconds()));
1332 ASSERT_THAT(page_result_info4.second.results, SizeIs(1));
1333 EXPECT_THAT(page_result_info4.second.results.at(0).document(),
1334 EqualsProto(document_protos4.at(1)));
1335
1336 ICING_ASSERT_OK_AND_ASSIGN(page_result_info5,
1337 result_state_manager.GetNextPage(
1338 page_result_info5.first, result_retriever(),
1339 clock()->GetSystemTimeMilliseconds()));
1340 ASSERT_THAT(page_result_info5.second.results, SizeIs(1));
1341 EXPECT_THAT(page_result_info5.second.results.at(0).document(),
1342 EqualsProto(document_protos5.at(1)));
1343 }
1344
TEST_F(ResultStateManagerTest,AddingOverBudgetResultStateShouldEvictAllStates)1345 TEST_F(ResultStateManagerTest,
1346 AddingOverBudgetResultStateShouldEvictAllStates) {
1347 auto [scored_document_hits1, document_protos1] = AddScoredDocuments(
1348 {/*document_id=*/0, /*document_id=*/1, /*document_id=*/2});
1349 auto [scored_document_hits2, document_protos2] =
1350 AddScoredDocuments({/*document_id=*/3, /*document_id=*/4});
1351
1352 // Add the first two states. Remember, the first page for each result state
1353 // won't be cached (since it is returned immediately from
1354 // CacheAndRetrieveFirstPage). Each result state has a page size of 1. So 3
1355 // hits will remain cached.
1356 ResultStateManager result_state_manager(/*max_total_hits=*/4,
1357 document_store());
1358
1359 ICING_ASSERT_OK_AND_ASSIGN(
1360 PageResultInfo page_result_info1,
1361 result_state_manager.CacheAndRetrieveFirstPage(
1362 std::make_unique<
1363 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1364 std::move(scored_document_hits1), /*is_descending=*/true),
1365 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1366 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1367 document_store(), result_retriever(),
1368 clock()->GetSystemTimeMilliseconds()));
1369
1370 ICING_ASSERT_OK_AND_ASSIGN(
1371 PageResultInfo page_result_info2,
1372 result_state_manager.CacheAndRetrieveFirstPage(
1373 std::make_unique<
1374 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1375 std::move(scored_document_hits2), /*is_descending=*/true),
1376 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1377 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1378 document_store(), result_retriever(),
1379 clock()->GetSystemTimeMilliseconds()));
1380
1381 // Add a result state that is larger than the entire budget. This should
1382 // result in all previous result states being evicted, the first hit from
1383 // result state 3 being returned and the next four hits being cached (the last
1384 // hit should be dropped because it exceeds the max).
1385 auto [scored_document_hits3, document_protos3] = AddScoredDocuments(
1386 {/*document_id=*/5, /*document_id=*/6, /*document_id=*/7,
1387 /*document_id=*/8, /*document_id=*/9, /*document_id=*/10});
1388 ICING_ASSERT_OK_AND_ASSIGN(
1389 PageResultInfo page_result_info3,
1390 result_state_manager.CacheAndRetrieveFirstPage(
1391 std::make_unique<
1392 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1393 std::move(scored_document_hits3), /*is_descending=*/true),
1394 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1395 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1396 document_store(), result_retriever(),
1397 clock()->GetSystemTimeMilliseconds()));
1398 EXPECT_THAT(page_result_info3.first, Not(Eq(kInvalidNextPageToken)));
1399
1400 // GetNextPage for result state 1 and 2 should return NOT_FOUND.
1401 EXPECT_THAT(result_state_manager.GetNextPage(
1402 page_result_info1.first, result_retriever(),
1403 clock()->GetSystemTimeMilliseconds()),
1404 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
1405
1406 EXPECT_THAT(result_state_manager.GetNextPage(
1407 page_result_info2.first, result_retriever(),
1408 clock()->GetSystemTimeMilliseconds()),
1409 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
1410
1411 // Only the next four results in state 3 should be retrievable.
1412 uint64_t next_page_token3 = page_result_info3.first;
1413 ICING_ASSERT_OK_AND_ASSIGN(
1414 page_result_info3,
1415 result_state_manager.GetNextPage(next_page_token3, result_retriever(),
1416 clock()->GetSystemTimeMilliseconds()));
1417 EXPECT_THAT(page_result_info3.first, Eq(next_page_token3));
1418 ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
1419 EXPECT_THAT(page_result_info3.second.results.at(0).document(),
1420 EqualsProto(document_protos3.at(1)));
1421
1422 ICING_ASSERT_OK_AND_ASSIGN(
1423 page_result_info3,
1424 result_state_manager.GetNextPage(next_page_token3, result_retriever(),
1425 clock()->GetSystemTimeMilliseconds()));
1426 EXPECT_THAT(page_result_info3.first, Eq(next_page_token3));
1427 ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
1428 EXPECT_THAT(page_result_info3.second.results.at(0).document(),
1429 EqualsProto(document_protos3.at(2)));
1430
1431 ICING_ASSERT_OK_AND_ASSIGN(
1432 page_result_info3,
1433 result_state_manager.GetNextPage(next_page_token3, result_retriever(),
1434 clock()->GetSystemTimeMilliseconds()));
1435 EXPECT_THAT(page_result_info3.first, Eq(next_page_token3));
1436 ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
1437 EXPECT_THAT(page_result_info3.second.results.at(0).document(),
1438 EqualsProto(document_protos3.at(3)));
1439
1440 ICING_ASSERT_OK_AND_ASSIGN(
1441 page_result_info3,
1442 result_state_manager.GetNextPage(next_page_token3, result_retriever(),
1443 clock()->GetSystemTimeMilliseconds()));
1444 // The final document should have been dropped because it exceeded the budget,
1445 // so the next page token of the second last round should be
1446 // kInvalidNextPageToken.
1447 EXPECT_THAT(page_result_info3.first, Eq(kInvalidNextPageToken));
1448 ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
1449 EXPECT_THAT(page_result_info3.second.results.at(0).document(),
1450 EqualsProto(document_protos3.at(4)));
1451
1452 // Double check that next_page_token3 is not retrievable anymore.
1453 EXPECT_THAT(
1454 result_state_manager.GetNextPage(next_page_token3, result_retriever(),
1455 clock()->GetSystemTimeMilliseconds()),
1456 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
1457 }
1458
TEST_F(ResultStateManagerTest,AddingResultStateShouldEvictOverBudgetResultState)1459 TEST_F(ResultStateManagerTest,
1460 AddingResultStateShouldEvictOverBudgetResultState) {
1461 // Add a result state that is larger than the entire budget. The entire result
1462 // state will still be cached
1463 auto [scored_document_hits1, document_protos1] = AddScoredDocuments(
1464 {/*document_id=*/0, /*document_id=*/1, /*document_id=*/2,
1465 /*document_id=*/3, /*document_id=*/4, /*document_id=*/5});
1466
1467 ResultStateManager result_state_manager(/*max_total_hits=*/4,
1468 document_store());
1469
1470 ICING_ASSERT_OK_AND_ASSIGN(
1471 PageResultInfo page_result_info1,
1472 result_state_manager.CacheAndRetrieveFirstPage(
1473 std::make_unique<
1474 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1475 std::move(scored_document_hits1), /*is_descending=*/true),
1476 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1477 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1478 document_store(), result_retriever(),
1479 clock()->GetSystemTimeMilliseconds()));
1480
1481 // Add a result state. Because state2 + state1 is larger than the budget,
1482 // state1 should be evicted.
1483 auto [scored_document_hits2, document_protos2] =
1484 AddScoredDocuments({/*document_id=*/6, /*document_id=*/7});
1485 ICING_ASSERT_OK_AND_ASSIGN(
1486 PageResultInfo page_result_info2,
1487 result_state_manager.CacheAndRetrieveFirstPage(
1488 std::make_unique<
1489 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1490 std::move(scored_document_hits2), /*is_descending=*/true),
1491 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1492 CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
1493 document_store(), result_retriever(),
1494 clock()->GetSystemTimeMilliseconds()));
1495
1496 // state1 should have been evicted and state2 should still be retrievable.
1497 EXPECT_THAT(result_state_manager.GetNextPage(
1498 page_result_info1.first, result_retriever(),
1499 clock()->GetSystemTimeMilliseconds()),
1500 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
1501
1502 ICING_ASSERT_OK_AND_ASSIGN(page_result_info2,
1503 result_state_manager.GetNextPage(
1504 page_result_info2.first, result_retriever(),
1505 clock()->GetSystemTimeMilliseconds()));
1506 ASSERT_THAT(page_result_info2.second.results, SizeIs(1));
1507 EXPECT_THAT(page_result_info2.second.results.at(0).document(),
1508 EqualsProto(document_protos2.at(1)));
1509 }
1510
TEST_F(ResultStateManagerTest,AddingResultStateShouldNotTruncatedAfterFirstPage)1511 TEST_F(ResultStateManagerTest,
1512 AddingResultStateShouldNotTruncatedAfterFirstPage) {
1513 // Add a result state that is larger than the entire budget, but within the
1514 // entire budget after the first page. The entire result state will still be
1515 // cached and not truncated.
1516 auto [scored_document_hits, document_protos] = AddScoredDocuments(
1517 {/*document_id=*/0, /*document_id=*/1, /*document_id=*/2,
1518 /*document_id=*/3, /*document_id=*/4});
1519
1520 ResultStateManager result_state_manager(/*max_total_hits=*/4,
1521 document_store());
1522
1523 // The 5 input scored document hits will not be truncated. The first page of
1524 // two hits will be returned immediately and the other three hits will fit
1525 // within our caching budget.
1526 ICING_ASSERT_OK_AND_ASSIGN(
1527 PageResultInfo page_result_info1,
1528 result_state_manager.CacheAndRetrieveFirstPage(
1529 std::make_unique<
1530 PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1531 std::move(scored_document_hits), /*is_descending=*/true),
1532 /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1533 CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE),
1534 document_store(), result_retriever(),
1535 clock()->GetSystemTimeMilliseconds()));
1536
1537 // First page, 2 results
1538 ASSERT_THAT(page_result_info1.second.results, SizeIs(2));
1539 EXPECT_THAT(page_result_info1.second.results.at(0).document(),
1540 EqualsProto(document_protos.at(0)));
1541 EXPECT_THAT(page_result_info1.second.results.at(1).document(),
1542 EqualsProto(document_protos.at(1)));
1543
1544 uint64_t next_page_token = page_result_info1.first;
1545
1546 // Second page, 2 results.
1547 ICING_ASSERT_OK_AND_ASSIGN(
1548 PageResultInfo page_result_info2,
1549 result_state_manager.GetNextPage(next_page_token, result_retriever(),
1550 clock()->GetSystemTimeMilliseconds()));
1551 ASSERT_THAT(page_result_info2.second.results, SizeIs(2));
1552 EXPECT_THAT(page_result_info2.second.results.at(0).document(),
1553 EqualsProto(document_protos.at(2)));
1554 EXPECT_THAT(page_result_info2.second.results.at(1).document(),
1555 EqualsProto(document_protos.at(3)));
1556
1557 // Third page, 1 result.
1558 ICING_ASSERT_OK_AND_ASSIGN(
1559 PageResultInfo page_result_info3,
1560 result_state_manager.GetNextPage(next_page_token, result_retriever(),
1561 clock()->GetSystemTimeMilliseconds()));
1562 ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
1563 EXPECT_THAT(page_result_info3.second.results.at(0).document(),
1564 EqualsProto(document_protos.at(4)));
1565
1566 // Fourth page, 0 results.
1567 EXPECT_THAT(
1568 result_state_manager.GetNextPage(next_page_token, result_retriever(),
1569 clock()->GetSystemTimeMilliseconds()),
1570 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
1571 }
1572
1573 } // namespace
1574 } // namespace lib
1575 } // namespace icing
1576