1 files changed, 17 insertions, 17 deletions

diff --git a/icing/result/result-state-manager.cc b/icing/result/result-state-manager.cc
index f2fb94f..382f7db 100644
--- a/icing/result/result-state-manager.cc
+++ b/icing/result/result-state-manager.cc
@@ -31,13 +31,11 @@ namespace icing {
 namespace lib {
 
 ResultStateManager::ResultStateManager(int max_total_hits,
-                                       const DocumentStore& document_store,
-                                       const Clock* clock)
+                                       const DocumentStore& document_store)
     : document_store_(document_store),
       max_total_hits_(max_total_hits),
       num_total_hits_(0),
-      random_generator_(GetSteadyTimeNanoseconds()),
-      clock_(*clock) {}
+      random_generator_(GetSteadyTimeNanoseconds()) {}
 
 libtextclassifier3::StatusOr<std::pair<uint64_t, PageResult>>
 ResultStateManager::CacheAndRetrieveFirstPage(
@@ -45,7 +43,7 @@ ResultStateManager::CacheAndRetrieveFirstPage(
     std::unique_ptr<ResultAdjustmentInfo> parent_adjustment_info,
     std::unique_ptr<ResultAdjustmentInfo> child_adjustment_info,
     const ResultSpecProto& result_spec, const DocumentStore& document_store,
-    const ResultRetrieverV2& result_retriever) {
+    const ResultRetrieverV2& result_retriever, int64_t current_time_ms) {
   if (ranker == nullptr) {
     return absl_ports::InvalidArgumentError("Should not provide null ranker");
   }
@@ -59,7 +57,7 @@ ResultStateManager::CacheAndRetrieveFirstPage(
   // Retrieve docs outside of ResultStateManager critical section.
   // Will enter ResultState critical section inside ResultRetriever.
   auto [page_result, has_more_results] =
-      result_retriever.RetrieveNextPage(*result_state);
+      result_retriever.RetrieveNextPage(*result_state, current_time_ms);
   if (!has_more_results) {
     // No more pages, won't store ResultState, returns directly
     return std::make_pair(kInvalidNextPageToken, std::move(page_result));
@@ -86,37 +84,40 @@ ResultStateManager::CacheAndRetrieveFirstPage(
     absl_ports::unique_lock l(&mutex_);
 
     // Remove expired result states first.
-    InternalInvalidateExpiredResultStates(kDefaultResultStateTtlInMs);
+    InternalInvalidateExpiredResultStates(kDefaultResultStateTtlInMs,
+                                          current_time_ms);
     // Remove states to make room for this new state.
     RemoveStatesIfNeeded(num_hits_to_add);
     // Generate a new unique token and add it into result_state_map_.
-    next_page_token = Add(std::move(result_state));
+    next_page_token = Add(std::move(result_state), current_time_ms);
   }
 
   return std::make_pair(next_page_token, std::move(page_result));
 }
 
-uint64_t ResultStateManager::Add(std::shared_ptr<ResultStateV2> result_state) {
+uint64_t ResultStateManager::Add(std::shared_ptr<ResultStateV2> result_state,
+                                 int64_t current_time_ms) {
   uint64_t new_token = GetUniqueToken();
 
   result_state_map_.emplace(new_token, std::move(result_state));
   // Tracks the insertion order
-  token_queue_.push(
-      std::make_pair(new_token, clock_.GetSystemTimeMilliseconds()));
+  token_queue_.push(std::make_pair(new_token, current_time_ms));
 
   return new_token;
 }
 
 libtextclassifier3::StatusOr<std::pair<uint64_t, PageResult>>
 ResultStateManager::GetNextPage(uint64_t next_page_token,
-                                const ResultRetrieverV2& result_retriever) {
+                                const ResultRetrieverV2& result_retriever,
+                                int64_t current_time_ms) {
   std::shared_ptr<ResultStateV2> result_state = nullptr;
   {
     // ResultStateManager critical section
     absl_ports::unique_lock l(&mutex_);
 
     // Remove expired result states before fetching
-    InternalInvalidateExpiredResultStates(kDefaultResultStateTtlInMs);
+    InternalInvalidateExpiredResultStates(kDefaultResultStateTtlInMs,
+                                          current_time_ms);
 
     const auto& state_iterator = result_state_map_.find(next_page_token);
     if (state_iterator == result_state_map_.end()) {
@@ -128,7 +129,7 @@ ResultStateManager::GetNextPage(uint64_t next_page_token,
   // Retrieve docs outside of ResultStateManager critical section.
   // Will enter ResultState critical section inside ResultRetriever.
   auto [page_result, has_more_results] =
-      result_retriever.RetrieveNextPage(*result_state);
+      result_retriever.RetrieveNextPage(*result_state, current_time_ms);
 
   if (!has_more_results) {
     {
@@ -233,10 +234,9 @@ void ResultStateManager::InternalInvalidateResultState(uint64_t token) {
 }
 
 void ResultStateManager::InternalInvalidateExpiredResultStates(
-    int64_t result_state_ttl) {
-  int64_t current_time = clock_.GetSystemTimeMilliseconds();
+    int64_t result_state_ttl, int64_t current_time_ms) {
   while (!token_queue_.empty() &&
-         current_time - token_queue_.front().second >= result_state_ttl) {
+         current_time_ms - token_queue_.front().second >= result_state_ttl) {
     auto itr = result_state_map_.find(token_queue_.front().first);
     if (itr != result_state_map_.end()) {
       // We don't have to decrement num_total_hits_ here, since erasing the