Use a separate map to store per app stats

To avoid iterating through the eBPF map to get the total stats of a
specific uid. A new bpf map called appUidStatsMap is added to the
trafficController so that TrafficStats API can directly read that map
for per uid total stats regardless of tag, counterSet and iface
information. This could make this call more efficient and solve the
possible racing problem.

Bug: 79171384
Test: netd_unit_test, libbpf_test, netd_integration_test
Change-Id: I47a4ac3466caa729c5730a498a2de226303d6b77
diff --git a/server/TrafficController.cpp b/server/TrafficController.cpp
index d8eb6b6..d6a6480 100644
--- a/server/TrafficController.cpp
+++ b/server/TrafficController.cpp
@@ -149,6 +149,11 @@
                                         "UidCounterSetMap", false));
 
     RETURN_IF_NOT_OK(
+        mAppUidStatsMap.getOrCreate(UID_STATS_MAP_SIZE, APP_UID_STATS_MAP_PATH, BPF_MAP_TYPE_HASH));
+    RETURN_IF_NOT_OK(
+        changeOwnerAndMode(APP_UID_STATS_MAP_PATH, AID_NET_BW_STATS, "AppUidStatsMap", false));
+
+    RETURN_IF_NOT_OK(
         mUidStatsMap.getOrCreate(UID_STATS_MAP_SIZE, UID_STATS_MAP_PATH, BPF_MAP_TYPE_HASH));
     RETURN_IF_NOT_OK(changeOwnerAndMode(UID_STATS_MAP_PATH, AID_NET_BW_STATS, "UidStatsMap",
                                         false));
@@ -391,6 +396,18 @@
               strerror(res.code()));
     }
     mUidStatsMap.iterate(deleteMatchedUidTagEntries);
+
+    auto deleteAppUidStatsEntry = [uid](const uint32_t& key, BpfMap<uint32_t, StatsValue>& map) {
+        if (key == uid) {
+            Status res = map.deleteValue(key);
+            if (isOk(res) || (res.code() == ENOENT)) {
+                return netdutils::status::ok;
+            }
+            ALOGE("Failed to delete data(uid=%u): %s", key, strerror(res.code()));
+        }
+        return netdutils::status::ok;
+    };
+    mAppUidStatsMap.iterate(deleteAppUidStatsEntry);
     return 0;
 }
 
@@ -589,6 +606,8 @@
                getMapStatus(mCookieTagMap.getMap(), COOKIE_TAG_MAP_PATH).c_str());
     dw.println("mUidCounterSetMap status: %s",
                getMapStatus(mUidCounterSetMap.getMap(), UID_COUNTERSET_MAP_PATH).c_str());
+    dw.println("mAppUidStatsMap status: %s",
+               getMapStatus(mAppUidStatsMap.getMap(), APP_UID_STATS_MAP_PATH).c_str());
     dw.println("mUidStatsMap status: %s",
                getMapStatus(mUidStatsMap.getMap(), UID_STATS_MAP_PATH).c_str());
     dw.println("mTagStatsMap status: %s",
@@ -644,6 +663,20 @@
         dw.println("mUidCounterSetMap print end with error: %s", res.msg().c_str());
     }
 
+    // Print AppUidStatsMap content
+    std::string appUidStatsHeader = StringPrintf("uid rxBytes rxPackets txBytes txPackets");
+    dumpBpfMap("mAppUidStatsMap:", dw, appUidStatsHeader);
+    auto printAppUidStatsInfo = [&dw](const uint32_t& key, const StatsValue& value,
+                                      const BpfMap<uint32_t, StatsValue>&) {
+        dw.println("%u %" PRIu64 " %" PRIu64 " %" PRIu64 " %" PRIu64, key, value.rxBytes,
+                   value.rxPackets, value.txBytes, value.txPackets);
+        return netdutils::status::ok;
+    };
+    res = mAppUidStatsMap.iterateWithValue(printAppUidStatsInfo);
+    if (!res.ok()) {
+        dw.println("mAppUidStatsMap print end with error: %s", res.msg().c_str());
+    }
+
     // Print uidStatsMap content
     std::string statsHeader = StringPrintf("ifaceIndex ifaceName tag_hex uid_int cnt_set rxBytes"
                                            " rxPackets txBytes txPackets");
diff --git a/server/TrafficController.h b/server/TrafficController.h
index 05d91df..79f7d14 100644
--- a/server/TrafficController.h
+++ b/server/TrafficController.h
@@ -132,6 +132,13 @@
     BpfMap<uint32_t, uint8_t> mUidCounterSetMap;
 
     /*
+     * mAppUidStatsMap: Store the total traffic stats for a uid regardless of
+     * tag, counterSet and iface. The stats is used by TrafficStats.getUidStats
+     * API to return persistent stats for a specific uid since device boot.
+     */
+    BpfMap<uint32_t, StatsValue> mAppUidStatsMap;
+
+    /*
      * mUidStatsMap: Store the traffic statistics for a specific combination of
      * uid, iface and counterSet. We maintain this map in addition to
      * mTagStatsMap because we want to be able to track per-UID data usage even
diff --git a/server/TrafficControllerTest.cpp b/server/TrafficControllerTest.cpp
index ca4b163..a354f83 100644
--- a/server/TrafficControllerTest.cpp
+++ b/server/TrafficControllerTest.cpp
@@ -72,6 +72,7 @@
     TrafficController mTc;
     BpfMap<uint64_t, UidTag> mFakeCookieTagMap;
     BpfMap<uint32_t, uint8_t> mFakeUidCounterSetMap;
+    BpfMap<uint32_t, StatsValue> mFakeAppUidStatsMap;
     BpfMap<StatsKey, StatsValue> mFakeUidStatsMap;
     BpfMap<StatsKey, StatsValue> mFakeTagStatsMap;
     BpfMap<uint32_t, uint8_t> mFakeDozableUidMap;
@@ -90,6 +91,10 @@
             createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(uint8_t), TEST_MAP_SIZE, 0));
         ASSERT_LE(0, mFakeUidCounterSetMap.getMap());
 
+        mFakeAppUidStatsMap.reset(createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t),
+                                            sizeof(struct StatsValue), TEST_MAP_SIZE, 0));
+        ASSERT_LE(0, mFakeAppUidStatsMap.getMap());
+
         mFakeUidStatsMap.reset(createMap(BPF_MAP_TYPE_HASH, sizeof(struct StatsKey),
                                          sizeof(struct StatsValue), TEST_MAP_SIZE, 0));
         ASSERT_LE(0, mFakeUidStatsMap.getMap());
@@ -114,6 +119,7 @@
 
         mTc.mCookieTagMap.reset(mFakeCookieTagMap.getMap());
         mTc.mUidCounterSetMap.reset(mFakeUidCounterSetMap.getMap());
+        mTc.mAppUidStatsMap.reset(mFakeAppUidStatsMap.getMap());
         mTc.mUidStatsMap.reset(mFakeUidStatsMap.getMap());
         mTc.mTagStatsMap.reset(mFakeTagStatsMap.getMap());
         mTc.mDozableUidMap.reset(mFakeDozableUidMap.getMap());
@@ -151,6 +157,7 @@
         EXPECT_TRUE(isOk(mFakeTagStatsMap.writeValue(*key, statsMapValue, BPF_ANY)));
         key->tag = 0;
         EXPECT_TRUE(isOk(mFakeUidStatsMap.writeValue(*key, statsMapValue, BPF_ANY)));
+        EXPECT_TRUE(isOk(mFakeAppUidStatsMap.writeValue(uid, statsMapValue, BPF_ANY)));
         // put tag information back to statsKey
         key->tag = tag;
     }
@@ -217,6 +224,7 @@
         std::lock_guard<std::mutex> ownerGuard(mTc.mOwnerMatchMutex);
         mFakeCookieTagMap.reset();
         mFakeUidCounterSetMap.reset();
+        mFakeAppUidStatsMap.reset();
         mFakeUidStatsMap.reset();
         mFakeTagStatsMap.reset();
         mTc.mDozableUidMap.reset();
@@ -331,6 +339,10 @@
     ASSERT_TRUE(isOk(statsMapResult));
     ASSERT_EQ((uint64_t)1, statsMapResult.value().rxPackets);
     ASSERT_EQ((uint64_t)100, statsMapResult.value().rxBytes);
+    auto appStatsResult = mFakeAppUidStatsMap.readValue(TEST_UID);
+    ASSERT_TRUE(isOk(appStatsResult));
+    ASSERT_EQ((uint64_t)1, appStatsResult.value().rxPackets);
+    ASSERT_EQ((uint64_t)100, appStatsResult.value().rxBytes);
 }
 
 TEST_F(TrafficControllerTest, TestDeleteAllUidData) {
@@ -347,6 +359,7 @@
     ASSERT_FALSE(isOk(mFakeTagStatsMap.readValue(tagStatsMapKey)));
     tagStatsMapKey.tag = 0;
     ASSERT_FALSE(isOk(mFakeUidStatsMap.readValue(tagStatsMapKey)));
+    ASSERT_FALSE(isOk(mFakeAppUidStatsMap.readValue(TEST_UID)));
 }
 
 TEST_F(TrafficControllerTest, TestDeleteDataWithTwoTags) {
@@ -401,15 +414,21 @@
     ASSERT_FALSE(isOk(mFakeTagStatsMap.readValue(tagStatsMapKey2)));
     tagStatsMapKey2.tag = 0;
     ASSERT_FALSE(isOk(mFakeUidStatsMap.readValue(tagStatsMapKey2)));
+    ASSERT_FALSE(isOk(mFakeAppUidStatsMap.readValue(uid2)));
     tagStatsMapKey1.tag = 0;
     StatusOr<StatsValue> statsMapResult = mFakeUidStatsMap.readValue(tagStatsMapKey1);
     ASSERT_TRUE(isOk(statsMapResult));
     ASSERT_EQ((uint64_t)1, statsMapResult.value().rxPackets);
     ASSERT_EQ((uint64_t)100, statsMapResult.value().rxBytes);
+    auto appStatsResult = mFakeAppUidStatsMap.readValue(uid1);
+    ASSERT_TRUE(isOk(appStatsResult));
+    ASSERT_EQ((uint64_t)1, appStatsResult.value().rxPackets);
+    ASSERT_EQ((uint64_t)100, appStatsResult.value().rxBytes);
 
     // Delete the stats of the other uid.
     ASSERT_EQ(0, mTc.deleteTagData(0, uid1));
     ASSERT_FALSE(isOk(mFakeUidStatsMap.readValue(tagStatsMapKey1)));
+    ASSERT_FALSE(isOk(mFakeAppUidStatsMap.readValue(uid1)));
 }
 
 TEST_F(TrafficControllerTest, TestUpdateOwnerMapEntry) {