1 files changed, 8 insertions, 5 deletions
diff --git a/gd/btaa/linux_generic/attribution_processor.cc b/gd/btaa/linux_generic/attribution_processor.cc
index a55243a06..4c0eac780 100644
--- a/gd/btaa/linux_generic/attribution_processor.cc
+++ b/gd/btaa/linux_generic/attribution_processor.cc
@@ -64,7 +64,7 @@ void AttributionProcessor::OnWakelockReleased(uint32_t duration_ms) {
   }
 
   ms_per_byte = duration_ms / total_byte_count;
-  auto cur_time = std::chrono::system_clock::now();
+  auto cur_time = now_func_();
   for (auto& it : wakelock_duration_aggregator_) {
     it.second.wakelock_duration_ms = ms_per_byte * it.second.byte_count;
     if (btaa_aggregator_.find(it.first) == btaa_aggregator_.end()) {
@@ -91,12 +91,15 @@ void AttributionProcessor::OnWakelockReleased(uint32_t duration_ms) {
     return;
   }
   // Trim down the transient entries in the aggregator to avoid that it overgrows
-  for (auto& it : btaa_aggregator_) {
+  auto it = btaa_aggregator_.begin();
+  while (it != btaa_aggregator_.end()) {
     auto elapsed_time_sec =
-        std::chrono::duration_cast<std::chrono::seconds>(cur_time - it.second.creation_time).count();
+        std::chrono::duration_cast<std::chrono::seconds>(cur_time - it->second.creation_time).count();
     if (elapsed_time_sec > kDurationTransientDeviceActivityEntrySecs &&
-        it.second.byte_count < kByteCountTransientDeviceActivityEntry) {
-      btaa_aggregator_.erase(it.first);
+        it->second.byte_count < kByteCountTransientDeviceActivityEntry) {
+      it = btaa_aggregator_.erase(it);
+    } else {
+      it++;
     }
   }
 }