libtimeinstate: add functions to read only recently-updated stats

Add getUidsUpdatedCpuFreqTimes and getUidsUpdatedConcurrentTimes,
which skip reading stats for UIDs that haven't been updated since
before a time passed in by their caller, and pass a new lastUpdate
time back to the caller. This is implemented by querying a new map
that holds the most recent update time for each new UID.

This approach has a potential race when a UID is updated after we have
already read its stats, but before we finish iterating through the
rest of the BPF map. By not skipping UIDs updated up to 1s before
lastUpdate, we improve the chance that such an update will be picked
up the next time getUidsUpdated*Times is called. Though this doesn't
completely eliminate the risk of a race, the consequences of the race
aren't that severe - we could undercount some runtimes by ~seconds at
worst, and only until the affected UID runs again.

Extend existing tests to check that these new functions behave like
the existing getUids*Times functions when passed a last update time of
0, and add new testcases to exercise the case where a nonzero last
update time is used.

Test: libtimeinstate_test passes
Bug: 138317993
Change-Id: I06ddf8bd7ab7812d067f3a1f5b2fbedeae016bfc
Signed-off-by: Connor O'Brien <connoro@google.com>

3 files changed, 236 insertions, 79 deletions

diff --git a/libs/cputimeinstate/cputimeinstate.cpp b/libs/cputimeinstate/cputimeinstate.cpp
index 1465296bac..05a462e026 100644
--- a/libs/cputimeinstate/cputimeinstate.cpp
+++ b/libs/cputimeinstate/cputimeinstate.cpp
@@ -58,6 +58,7 @@ static std::vector<std::vector<uint32_t>> gPolicyCpus;
 static std::set<uint32_t> gAllFreqs;
 static unique_fd gTisMapFd;
 static unique_fd gConcurrentMapFd;
+static unique_fd gUidLastUpdateMapFd;
 
 static std::optional<std::vector<uint32_t>> readNumbersFromFile(const std::string &path) {
     std::string data;
@@ -144,6 +145,10 @@ static bool initGlobals() {
             unique_fd{bpf_obj_get(BPF_FS_PATH "map_time_in_state_uid_concurrent_times_map")};
     if (gConcurrentMapFd < 0) return false;
 
+    gUidLastUpdateMapFd =
+            unique_fd{bpf_obj_get(BPF_FS_PATH "map_time_in_state_uid_last_update_map")};
+    if (gUidLastUpdateMapFd < 0) return false;
+
     gInitialized = true;
     return true;
 }
@@ -263,6 +268,18 @@ std::optional<std::vector<std::vector<uint64_t>>> getUidCpuFreqTimes(uint32_t ui
     return out;
 }
 
+static std::optional<bool> uidUpdatedSince(uint32_t uid, uint64_t lastUpdate,
+                                           uint64_t *newLastUpdate) {
+    uint64_t uidLastUpdate;
+    if (findMapEntry(gUidLastUpdateMapFd, &uid, &uidLastUpdate)) return {};
+    // Updates that occurred during the previous read may have been missed. To mitigate
+    // this, don't ignore entries updated up to 1s before *lastUpdate
+    constexpr uint64_t NSEC_PER_SEC = 1000000000;
+    if (uidLastUpdate + NSEC_PER_SEC < lastUpdate) return false;
+    if (uidLastUpdate > *newLastUpdate) *newLastUpdate = uidLastUpdate;
+    return true;
+}
+
 // Retrieve the times in ns that each uid spent running at each CPU freq.
 // Return contains no value on error, otherwise it contains a map from uids to vectors of vectors
 // using the format:
@@ -271,6 +288,14 @@ std::optional<std::vector<std::vector<uint64_t>>> getUidCpuFreqTimes(uint32_t ui
 // where ti_j_k is the ns uid i spent running on the jth cluster at the cluster's kth lowest freq.
 std::optional<std::unordered_map<uint32_t, std::vector<std::vector<uint64_t>>>>
 getUidsCpuFreqTimes() {
+    return getUidsUpdatedCpuFreqTimes(nullptr);
+}
+
+// Retrieve the times in ns that each uid spent running at each CPU freq, excluding UIDs that have
+// not run since before lastUpdate.
+// Return format is the same as getUidsCpuFreqTimes()
+std::optional<std::unordered_map<uint32_t, std::vector<std::vector<uint64_t>>>>
+getUidsUpdatedCpuFreqTimes(uint64_t *lastUpdate) {
     if (!gInitialized && !initGlobals()) return {};
     time_key_t key, prevKey;
     std::unordered_map<uint32_t, std::vector<std::vector<uint64_t>>> map;
@@ -282,8 +307,14 @@ getUidsCpuFreqTimes() {
     std::vector<std::vector<uint64_t>> mapFormat;
     for (const auto &freqList : gPolicyFreqs) mapFormat.emplace_back(freqList.size(), 0);
 
+    uint64_t newLastUpdate = lastUpdate ? *lastUpdate : 0;
     std::vector<tis_val_t> vals(gNCpus);
     do {
+        if (lastUpdate) {
+            auto uidUpdated = uidUpdatedSince(key.uid, *lastUpdate, &newLastUpdate);
+            if (!uidUpdated.has_value()) return {};
+            if (!*uidUpdated) continue;
+        }
         if (findMapEntry(gTisMapFd, &key, vals.data())) return {};
         if (map.find(key.uid) == map.end()) map.emplace(key.uid, mapFormat);
 
@@ -299,8 +330,9 @@ getUidsCpuFreqTimes() {
             }
         }
         prevKey = key;
-    } while (!getNextMapKey(gTisMapFd, &prevKey, &key));
+    } while (prevKey = key, !getNextMapKey(gTisMapFd, &prevKey, &key));
     if (errno != ENOENT) return {};
+    if (lastUpdate && newLastUpdate > *lastUpdate) *lastUpdate = newLastUpdate;
     return map;
 }
 
@@ -365,6 +397,15 @@ std::optional<concurrent_time_t> getUidConcurrentTimes(uint32_t uid, bool retry)
 // where ai is the ns spent running concurrently with tasks on i other cpus and pi_j is the ns spent
 // running on the ith cluster, concurrently with tasks on j other cpus in the same cluster.
 std::optional<std::unordered_map<uint32_t, concurrent_time_t>> getUidsConcurrentTimes() {
+    return getUidsUpdatedConcurrentTimes(nullptr);
+}
+
+// Retrieve the times in ns that each uid spent running concurrently with each possible number of
+// other tasks on each cluster (policy times) and overall (active times), excluding UIDs that have
+// not run since before lastUpdate.
+// Return format is the same as getUidsConcurrentTimes()
+std::optional<std::unordered_map<uint32_t, concurrent_time_t>> getUidsUpdatedConcurrentTimes(
+        uint64_t *lastUpdate) {
     if (!gInitialized && !initGlobals()) return {};
     time_key_t key, prevKey;
     std::unordered_map<uint32_t, concurrent_time_t> ret;
@@ -379,7 +420,13 @@ std::optional<std::unordered_map<uint32_t, concurrent_time_t>> getUidsConcurrent
     std::vector<concurrent_val_t> vals(gNCpus);
     std::vector<uint64_t>::iterator activeBegin, activeEnd, policyBegin, policyEnd;
 
+    uint64_t newLastUpdate = lastUpdate ? *lastUpdate : 0;
     do {
+        if (lastUpdate) {
+            auto uidUpdated = uidUpdatedSince(key.uid, *lastUpdate, &newLastUpdate);
+            if (!uidUpdated.has_value()) return {};
+            if (!*uidUpdated) continue;
+        }
         if (findMapEntry(gConcurrentMapFd, &key, vals.data())) return {};
         if (ret.find(key.uid) == ret.end()) ret.emplace(key.uid, retFormat);
 
@@ -405,8 +452,7 @@ std::optional<std::unordered_map<uint32_t, concurrent_time_t>> getUidsConcurrent
                                std::plus<uint64_t>());
             }
         }
-        prevKey = key;
-    } while (!getNextMapKey(gConcurrentMapFd, &prevKey, &key));
+    } while (prevKey = key, !getNextMapKey(gConcurrentMapFd, &prevKey, &key));
     if (errno != ENOENT) return {};
     for (const auto &[key, value] : ret) {
         if (!verifyConcurrentTimes(value)) {
@@ -414,6 +460,7 @@ std::optional<std::unordered_map<uint32_t, concurrent_time_t>> getUidsConcurrent
             if (val.has_value()) ret[key] = val.value();
         }
     }
+    if (lastUpdate && newLastUpdate > *lastUpdate) *lastUpdate = newLastUpdate;
     return ret;
 }
 
@@ -446,6 +493,8 @@ bool clearUidTimes(uint32_t uid) {
             return false;
         if (deleteMapEntry(gConcurrentMapFd, &key) && errno != ENOENT) return false;
     }
+
+    if (deleteMapEntry(gUidLastUpdateMapFd, &uid) && errno != ENOENT) return false;
     return true;
 }
 
diff --git a/libs/cputimeinstate/cputimeinstate.h b/libs/cputimeinstate/cputimeinstate.h
index 49469d8e04..b7600f59e0 100644
--- a/libs/cputimeinstate/cputimeinstate.h
+++ b/libs/cputimeinstate/cputimeinstate.h
@@ -26,6 +26,8 @@ bool startTrackingUidTimes();
 std::optional<std::vector<std::vector<uint64_t>>> getUidCpuFreqTimes(uint32_t uid);
 std::optional<std::unordered_map<uint32_t, std::vector<std::vector<uint64_t>>>>
     getUidsCpuFreqTimes();
+std::optional<std::unordered_map<uint32_t, std::vector<std::vector<uint64_t>>>>
+    getUidsUpdatedCpuFreqTimes(uint64_t *lastUpdate);
 std::optional<std::vector<std::vector<uint32_t>>> getCpuFreqs();
 
 struct concurrent_time_t {
@@ -35,6 +37,8 @@ struct concurrent_time_t {
 
 std::optional<concurrent_time_t> getUidConcurrentTimes(uint32_t uid, bool retry = true);
 std::optional<std::unordered_map<uint32_t, concurrent_time_t>> getUidsConcurrentTimes();
+std::optional<std::unordered_map<uint32_t, concurrent_time_t>>
+    getUidsUpdatedConcurrentTimes(uint64_t *lastUpdate);
 bool clearUidTimes(unsigned int uid);
 
 } // namespace bpf
diff --git a/libs/cputimeinstate/testtimeinstate.cpp b/libs/cputimeinstate/testtimeinstate.cpp
index 23d87fd646..ea2a2008b7 100644
--- a/libs/cputimeinstate/testtimeinstate.cpp
+++ b/libs/cputimeinstate/testtimeinstate.cpp
@@ -115,71 +115,169 @@ TEST(TimeInStateTest, SingleUidTimesConsistent) {
 }
 
 TEST(TimeInStateTest, AllUidTimeInState) {
-    vector<size_t> sizes;
-    auto map = getUidsCpuFreqTimes();
-    ASSERT_TRUE(map.has_value());
+    uint64_t zero = 0;
+    auto maps = {getUidsCpuFreqTimes(), getUidsUpdatedCpuFreqTimes(&zero)};
+    for (const auto &map : maps) {
+        ASSERT_TRUE(map.has_value());
 
-    ASSERT_FALSE(map->empty());
+        ASSERT_FALSE(map->empty());
 
-    auto firstEntry = map->begin()->second;
-    for (const auto &subEntry : firstEntry) sizes.emplace_back(subEntry.size());
+        vector<size_t> sizes;
+        auto firstEntry = map->begin()->second;
+        for (const auto &subEntry : firstEntry) sizes.emplace_back(subEntry.size());
 
-    for (const auto &vec : *map) {
-        ASSERT_EQ(vec.second.size(), sizes.size());
-        for (size_t i = 0; i < vec.second.size(); ++i) ASSERT_EQ(vec.second[i].size(), sizes[i]);
+        for (const auto &vec : *map) {
+            ASSERT_EQ(vec.second.size(), sizes.size());
+            for (size_t i = 0; i < vec.second.size(); ++i) ASSERT_EQ(vec.second[i].size(), sizes[i]);
+        }
     }
 }
 
-TEST(TimeInStateTest, SingleAndAllUidTimeInStateConsistent) {
-    auto map = getUidsCpuFreqTimes();
-    ASSERT_TRUE(map.has_value());
-    ASSERT_FALSE(map->empty());
+void TestCheckUpdate(const std::vector<std::vector<uint64_t>> &before,
+                     const std::vector<std::vector<uint64_t>> &after) {
+    ASSERT_EQ(before.size(), after.size());
+    uint64_t sumBefore = 0, sumAfter = 0;
+    for (size_t i = 0; i < before.size(); ++i) {
+        ASSERT_EQ(before[i].size(), after[i].size());
+        for (size_t j = 0; j < before[i].size(); ++j) {
+            // Times should never decrease
+            ASSERT_LE(before[i][j], after[i][j]);
+        }
+        sumBefore += std::accumulate(before[i].begin(), before[i].end(), (uint64_t)0);
+        sumAfter += std::accumulate(after[i].begin(), after[i].end(), (uint64_t)0);
+    }
+    ASSERT_LE(sumBefore, sumAfter);
+    ASSERT_LE(sumAfter - sumBefore, NSEC_PER_SEC);
+}
 
-    for (const auto &kv : *map) {
-        uint32_t uid = kv.first;
-        auto times1 = kv.second;
-        auto times2 = getUidCpuFreqTimes(uid);
-        ASSERT_TRUE(times2.has_value());
-
-        ASSERT_EQ(times1.size(), times2->size());
-        for (uint32_t i = 0; i < times1.size(); ++i) {
-            ASSERT_EQ(times1[i].size(), (*times2)[i].size());
-            for (uint32_t j = 0; j < times1[i].size(); ++j) {
-                ASSERT_LE((*times2)[i][j] - times1[i][j], NSEC_PER_SEC);
+TEST(TimeInStateTest, AllUidUpdatedTimeInState) {
+    uint64_t lastUpdate = 0;
+    auto map1 = getUidsUpdatedCpuFreqTimes(&lastUpdate);
+    ASSERT_TRUE(map1.has_value());
+    ASSERT_FALSE(map1->empty());
+    ASSERT_NE(lastUpdate, (uint64_t)0);
+    uint64_t oldLastUpdate = lastUpdate;
+
+    // Sleep briefly to trigger a context switch, ensuring we see at least one update.
+    struct timespec ts;
+    ts.tv_sec = 0;
+    ts.tv_nsec = 1000000;
+    nanosleep (&ts, NULL);
+
+    auto map2 = getUidsUpdatedCpuFreqTimes(&lastUpdate);
+    ASSERT_TRUE(map2.has_value());
+    ASSERT_FALSE(map2->empty());
+    ASSERT_NE(lastUpdate, oldLastUpdate);
+
+    bool someUidsExcluded = false;
+    for (const auto &[uid, v] : *map1) {
+        if (map2->find(uid) == map2->end()) {
+            someUidsExcluded = true;
+            break;
+        }
+    }
+    ASSERT_TRUE(someUidsExcluded);
+
+    for (const auto &[uid, newTimes] : *map2) {
+        ASSERT_NE(map1->find(uid), map1->end());
+        ASSERT_NO_FATAL_FAILURE(TestCheckUpdate((*map1)[uid], newTimes));
+    }
+}
+
+TEST(TimeInStateTest, SingleAndAllUidTimeInStateConsistent) {
+    uint64_t zero = 0;
+    auto maps = {getUidsCpuFreqTimes(), getUidsUpdatedCpuFreqTimes(&zero)};
+    for (const auto &map : maps) {
+        ASSERT_TRUE(map.has_value());
+        ASSERT_FALSE(map->empty());
+
+        for (const auto &kv : *map) {
+            uint32_t uid = kv.first;
+            auto times1 = kv.second;
+            auto times2 = getUidCpuFreqTimes(uid);
+            ASSERT_TRUE(times2.has_value());
+
+            ASSERT_EQ(times1.size(), times2->size());
+            for (uint32_t i = 0; i < times1.size(); ++i) {
+                ASSERT_EQ(times1[i].size(), (*times2)[i].size());
+                for (uint32_t j = 0; j < times1[i].size(); ++j) {
+                    ASSERT_LE((*times2)[i][j] - times1[i][j], NSEC_PER_SEC);
+                }
             }
         }
     }
 }
 
 TEST(TimeInStateTest, AllUidConcurrentTimes) {
-    auto map = getUidsConcurrentTimes();
-    ASSERT_TRUE(map.has_value());
-    ASSERT_FALSE(map->empty());
-
-    auto firstEntry = map->begin()->second;
-    for (const auto &kv : *map) {
-        ASSERT_EQ(kv.second.active.size(), firstEntry.active.size());
-        ASSERT_EQ(kv.second.policy.size(), firstEntry.policy.size());
-        for (size_t i = 0; i < kv.second.policy.size(); ++i) {
-            ASSERT_EQ(kv.second.policy[i].size(), firstEntry.policy[i].size());
+    uint64_t zero = 0;
+    auto maps = {getUidsConcurrentTimes(), getUidsUpdatedConcurrentTimes(&zero)};
+    for (const auto &map : maps) {
+        ASSERT_TRUE(map.has_value());
+        ASSERT_FALSE(map->empty());
+
+        auto firstEntry = map->begin()->second;
+        for (const auto &kv : *map) {
+            ASSERT_EQ(kv.second.active.size(), firstEntry.active.size());
+            ASSERT_EQ(kv.second.policy.size(), firstEntry.policy.size());
+            for (size_t i = 0; i < kv.second.policy.size(); ++i) {
+                ASSERT_EQ(kv.second.policy[i].size(), firstEntry.policy[i].size());
+            }
         }
     }
 }
 
-TEST(TimeInStateTest, SingleAndAllUidConcurrentTimesConsistent) {
-    auto map = getUidsConcurrentTimes();
-    ASSERT_TRUE(map.has_value());
-    for (const auto &kv : *map) {
-        uint32_t uid = kv.first;
-        auto times1 = kv.second;
-        auto times2 = getUidConcurrentTimes(uid);
-        ASSERT_TRUE(times2.has_value());
-        for (uint32_t i = 0; i < times1.active.size(); ++i) {
-            ASSERT_LE(times2->active[i] - times1.active[i], NSEC_PER_SEC);
+TEST(TimeInStateTest, AllUidUpdatedConcurrentTimes) {
+    uint64_t lastUpdate = 0;
+    auto map1 = getUidsUpdatedConcurrentTimes(&lastUpdate);
+    ASSERT_TRUE(map1.has_value());
+    ASSERT_FALSE(map1->empty());
+    ASSERT_NE(lastUpdate, (uint64_t)0);
+
+    // Sleep briefly to trigger a context switch, ensuring we see at least one update.
+    struct timespec ts;
+    ts.tv_sec = 0;
+    ts.tv_nsec = 1000000;
+    nanosleep (&ts, NULL);
+
+    uint64_t oldLastUpdate = lastUpdate;
+    auto map2 = getUidsUpdatedConcurrentTimes(&lastUpdate);
+    ASSERT_TRUE(map2.has_value());
+    ASSERT_FALSE(map2->empty());
+    ASSERT_NE(lastUpdate, oldLastUpdate);
+
+    bool someUidsExcluded = false;
+    for (const auto &[uid, v] : *map1) {
+        if (map2->find(uid) == map2->end()) {
+            someUidsExcluded = true;
+            break;
         }
-        for (uint32_t i = 0; i < times1.policy.size(); ++i) {
-            for (uint32_t j = 0; j < times1.policy[i].size(); ++j) {
-                ASSERT_LE(times2->policy[i][j] - times1.policy[i][j], NSEC_PER_SEC);
+    }
+    ASSERT_TRUE(someUidsExcluded);
+
+    for (const auto &[uid, newTimes] : *map2) {
+        ASSERT_NE(map1->find(uid), map1->end());
+        ASSERT_NO_FATAL_FAILURE(TestCheckUpdate({(*map1)[uid].active},{newTimes.active}));
+        ASSERT_NO_FATAL_FAILURE(TestCheckUpdate((*map1)[uid].policy, newTimes.policy));
+    }
+}
+
+TEST(TimeInStateTest, SingleAndAllUidConcurrentTimesConsistent) {
+    uint64_t zero = 0;
+    auto maps = {getUidsConcurrentTimes(), getUidsUpdatedConcurrentTimes(&zero)};
+    for (const auto &map : maps) {
+        ASSERT_TRUE(map.has_value());
+        for (const auto &kv : *map) {
+            uint32_t uid = kv.first;
+            auto times1 = kv.second;
+            auto times2 = getUidConcurrentTimes(uid);
+            ASSERT_TRUE(times2.has_value());
+            for (uint32_t i = 0; i < times1.active.size(); ++i) {
+                ASSERT_LE(times2->active[i] - times1.active[i], NSEC_PER_SEC);
+            }
+            for (uint32_t i = 0; i < times1.policy.size(); ++i) {
+                for (uint32_t j = 0; j < times1.policy[i].size(); ++j) {
+                    ASSERT_LE(times2->policy[i][j] - times1.policy[i][j], NSEC_PER_SEC);
+                }
             }
         }
     }
@@ -242,45 +340,51 @@ TEST(TimeInStateTest, AllUidConcurrentTimesMonotonic) {
 }
 
 TEST(TimeInStateTest, AllUidTimeInStateSanityCheck) {
-    auto map = getUidsCpuFreqTimes();
-    ASSERT_TRUE(map.has_value());
-
-    bool foundLargeValue = false;
-    for (const auto &kv : *map) {
-        for (const auto &timeVec : kv.second) {
-            for (const auto &time : timeVec) {
-                ASSERT_LE(time, NSEC_PER_YEAR);
-                if (time > UINT32_MAX) foundLargeValue = true;
+    uint64_t zero = 0;
+    auto maps = {getUidsCpuFreqTimes(), getUidsUpdatedCpuFreqTimes(&zero)};
+    for (const auto &map : maps) {
+        ASSERT_TRUE(map.has_value());
+
+        bool foundLargeValue = false;
+        for (const auto &kv : *map) {
+            for (const auto &timeVec : kv.second) {
+                for (const auto &time : timeVec) {
+                    ASSERT_LE(time, NSEC_PER_YEAR);
+                    if (time > UINT32_MAX) foundLargeValue = true;
+                }
             }
         }
+        // UINT32_MAX nanoseconds is less than 5 seconds, so if every part of our pipeline is using
+        // uint64_t as expected, we should have some times higher than that.
+        ASSERT_TRUE(foundLargeValue);
     }
-    // UINT32_MAX nanoseconds is less than 5 seconds, so if every part of our pipeline is using
-    // uint64_t as expected, we should have some times higher than that.
-    ASSERT_TRUE(foundLargeValue);
 }
 
 TEST(TimeInStateTest, AllUidConcurrentTimesSanityCheck) {
-    auto concurrentMap = getUidsConcurrentTimes();
-    ASSERT_TRUE(concurrentMap);
-
-    bool activeFoundLargeValue = false;
-    bool policyFoundLargeValue = false;
-    for (const auto &kv : *concurrentMap) {
-        for (const auto &time : kv.second.active) {
-            ASSERT_LE(time, NSEC_PER_YEAR);
-            if (time > UINT32_MAX) activeFoundLargeValue = true;
-        }
-        for (const auto &policyTimeVec : kv.second.policy) {
-            for (const auto &time : policyTimeVec) {
+    uint64_t zero = 0;
+    auto maps = {getUidsConcurrentTimes(), getUidsUpdatedConcurrentTimes(&zero)};
+    for (const auto &concurrentMap : maps) {
+        ASSERT_TRUE(concurrentMap);
+
+        bool activeFoundLargeValue = false;
+        bool policyFoundLargeValue = false;
+        for (const auto &kv : *concurrentMap) {
+            for (const auto &time : kv.second.active) {
                 ASSERT_LE(time, NSEC_PER_YEAR);
-                if (time > UINT32_MAX) policyFoundLargeValue = true;
+                if (time > UINT32_MAX) activeFoundLargeValue = true;
+            }
+            for (const auto &policyTimeVec : kv.second.policy) {
+                for (const auto &time : policyTimeVec) {
+                    ASSERT_LE(time, NSEC_PER_YEAR);
+                    if (time > UINT32_MAX) policyFoundLargeValue = true;
+                }
             }
         }
+        // UINT32_MAX nanoseconds is less than 5 seconds, so if every part of our pipeline is using
+        // uint64_t as expected, we should have some times higher than that.
+        ASSERT_TRUE(activeFoundLargeValue);
+        ASSERT_TRUE(policyFoundLargeValue);
     }
-    // UINT32_MAX nanoseconds is less than 5 seconds, so if every part of our pipeline is using
-    // uint64_t as expected, we should have some times higher than that.
-    ASSERT_TRUE(activeFoundLargeValue);
-    ASSERT_TRUE(policyFoundLargeValue);
 }
 
 TEST(TimeInStateTest, AllUidTimesConsistent) {