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
Merged-In: I47a4ac3466caa729c5730a498a2de226303d6b77
(cherry picked from aosp commit bc4a15f91f97fbfcbfdc9dc19d73226f380bc977)

11 files changed, 163 insertions, 42 deletions

diff --git a/bpfloader/BpfLoader.cpp b/bpfloader/BpfLoader.cpp
index 375ed924..8ab34d95 100644
--- a/bpfloader/BpfLoader.cpp
+++ b/bpfloader/BpfLoader.cpp
@@ -264,6 +264,7 @@ int loadAndAttachProgram(bpf_attach_type type, const char* path, const char* nam
 }  // namespace bpf
 }  // namespace android
 
+using android::bpf::APP_UID_STATS_MAP_PATH;
 using android::bpf::BPF_EGRESS_PROG_PATH;
 using android::bpf::BPF_INGRESS_PROG_PATH;
 using android::bpf::COOKIE_TAG_MAP_PATH;
@@ -276,6 +277,7 @@ using android::bpf::UID_COUNTERSET_MAP_PATH;
 using android::bpf::UID_STATS_MAP_PATH;
 using android::bpf::XT_BPF_EGRESS_PROG_PATH;
 using android::bpf::XT_BPF_INGRESS_PROG_PATH;
+
 using android::bpf::ReplacePattern;
 using android::bpf::loadAndAttachProgram;
 
@@ -291,6 +293,7 @@ int main(int argc, char** argv) {
     int ret = 0;
     DECLARE_MAP(cookieTagMap, COOKIE_TAG_MAP_PATH);
     DECLARE_MAP(uidCounterSetMap, UID_COUNTERSET_MAP_PATH);
+    DECLARE_MAP(appUidStatsMap, APP_UID_STATS_MAP_PATH);
     DECLARE_MAP(uidStatsMap, UID_STATS_MAP_PATH);
     DECLARE_MAP(tagStatsMap, TAG_STATS_MAP_PATH);
     DECLARE_MAP(ifaceStatsMap, IFACE_STATS_MAP_PATH);
@@ -301,6 +304,7 @@ int main(int argc, char** argv) {
     const std::vector<ReplacePattern> mapPatterns = {
         ReplacePattern(COOKIE_TAG_MAP, cookieTagMap.get()),
         ReplacePattern(UID_COUNTERSET_MAP, uidCounterSetMap.get()),
+        ReplacePattern(APP_UID_STATS_MAP, appUidStatsMap.get()),
         ReplacePattern(UID_STATS_MAP, uidStatsMap.get()),
         ReplacePattern(TAG_STATS_MAP, tagStatsMap.get()),
         ReplacePattern(IFACE_STATS_MAP, ifaceStatsMap.get()),
diff --git a/bpfloader/bpf_kern.h b/bpfloader/bpf_kern.h
index e37f1a1d..a59cb6d8 100644
--- a/bpfloader/bpf_kern.h
+++ b/bpfloader/bpf_kern.h
@@ -194,5 +194,6 @@ static __always_inline inline int bpf_traffic_account(struct __sk_buff* skb, int
 
     key.tag = 0;
     bpf_update_stats(skb, UID_STATS_MAP, direction, &key);
+    bpf_update_stats(skb, APP_UID_STATS_MAP, direction, &uid);
     return match;
 }
diff --git a/libbpf/BpfNetworkStats.cpp b/libbpf/BpfNetworkStats.cpp
index aa604ce9..eafbb5e3 100644
--- a/libbpf/BpfNetworkStats.cpp
+++ b/libbpf/BpfNetworkStats.cpp
@@ -42,31 +42,27 @@ using netdutils::Status;
 static constexpr uint32_t BPF_OPEN_FLAGS = BPF_F_RDONLY;
 
 int bpfGetUidStatsInternal(uid_t uid, Stats* stats,
-                           const BpfMap<StatsKey, StatsValue>& uidStatsMap) {
-    const auto processUidStats = [stats, uid](const StatsKey& key,
-                                              const BpfMap<StatsKey, StatsValue>& uidStatsMap) {
-        if (key.uid == uid) {
-            StatsValue statsEntry;
-            ASSIGN_OR_RETURN(statsEntry, uidStatsMap.readValue(key));
-            stats->rxPackets += statsEntry.rxPackets;
-            stats->txPackets += statsEntry.txPackets;
-            stats->rxBytes += statsEntry.rxBytes;
-            stats->txBytes += statsEntry.txBytes;
-        }
-        return netdutils::status::ok;
-    };
-    return -uidStatsMap.iterate(processUidStats).code();
+                           const BpfMap<uint32_t, StatsValue>& appUidStatsMap) {
+    auto statsEntry = appUidStatsMap.readValue(uid);
+    if (isOk(statsEntry)) {
+        stats->rxPackets = statsEntry.value().rxPackets;
+        stats->txPackets = statsEntry.value().txPackets;
+        stats->rxBytes = statsEntry.value().rxBytes;
+        stats->txBytes = statsEntry.value().txBytes;
+    }
+    return -statsEntry.status().code();
 }
 
 int bpfGetUidStats(uid_t uid, Stats* stats) {
-    BpfMap<StatsKey, StatsValue> uidStatsMap(mapRetrieve(UID_STATS_MAP_PATH, BPF_OPEN_FLAGS));
+    BpfMap<uint32_t, StatsValue> appUidStatsMap(
+        mapRetrieve(APP_UID_STATS_MAP_PATH, BPF_OPEN_FLAGS));
 
-    if (!uidStatsMap.isValid()) {
+    if (!appUidStatsMap.isValid()) {
         int ret = -errno;
-        ALOGE("Opening map fd from %s failed: %s", UID_STATS_MAP_PATH, strerror(errno));
+        ALOGE("Opening appUidStatsMap(%s) failed: %s", UID_STATS_MAP_PATH, strerror(errno));
         return ret;
     }
-    return bpfGetUidStatsInternal(uid, stats, uidStatsMap);
+    return bpfGetUidStatsInternal(uid, stats, appUidStatsMap);
 }
 
 int bpfGetIfaceStatsInternal(const char* iface, Stats* stats,
diff --git a/libbpf/BpfNetworkStatsTest.cpp b/libbpf/BpfNetworkStatsTest.cpp
index e804c407..8fd49439 100644
--- a/libbpf/BpfNetworkStatsTest.cpp
+++ b/libbpf/BpfNetworkStatsTest.cpp
@@ -65,15 +65,20 @@ constexpr uint64_t TEST_PACKET1 = 200;
 constexpr const char IFACE_NAME1[] = "lo";
 constexpr const char IFACE_NAME2[] = "wlan0";
 constexpr const char IFACE_NAME3[] = "rmnet_data0";
+// A iface name that the size is bigger then IFNAMSIZ
+constexpr const char LONG_IFACE_NAME[] = "wlanWithALongName";
+constexpr const char TRUNCATED_IFACE_NAME[] = "wlanWithALongNa";
 constexpr uint32_t IFACE_INDEX1 = 1;
 constexpr uint32_t IFACE_INDEX2 = 2;
 constexpr uint32_t IFACE_INDEX3 = 3;
+constexpr uint32_t IFACE_INDEX4 = 4;
 constexpr uint32_t UNKNOWN_IFACE = 0;
 
 class BpfNetworkStatsHelperTest : public testing::Test {
   protected:
     BpfNetworkStatsHelperTest() {}
     BpfMap<uint64_t, UidTag> mFakeCookieTagMap;
+    BpfMap<uint32_t, StatsValue> mFakeAppUidStatsMap;
     BpfMap<StatsKey, StatsValue> mFakeUidStatsMap;
     BpfMap<StatsKey, StatsValue> mFakeTagStatsMap;
     BpfMap<uint32_t, IfaceValue> mFakeIfaceIndexNameMap;
@@ -84,6 +89,10 @@ class BpfNetworkStatsHelperTest : public testing::Test {
             BPF_MAP_TYPE_HASH, sizeof(uint64_t), sizeof(struct UidTag), TEST_MAP_SIZE, 0));
         ASSERT_LE(0, mFakeCookieTagMap.getMap());
 
+        mFakeAppUidStatsMap = BpfMap<uint32_t, StatsValue>(createMap(
+            BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(struct StatsValue), TEST_MAP_SIZE, 0));
+        ASSERT_LE(0, mFakeAppUidStatsMap.getMap());
+
         mFakeUidStatsMap =
             BpfMap<StatsKey, StatsValue>(createMap(BPF_MAP_TYPE_HASH, sizeof(struct StatsKey),
                                                    sizeof(struct StatsValue), TEST_MAP_SIZE, 0));
@@ -208,24 +217,26 @@ TEST_F(BpfNetworkStatsHelperTest, TestGetUidStatsTotal) {
                          .rxPackets = TEST_PACKET0,
                          .txBytes = TEST_BYTES1,
                          .txPackets = TEST_PACKET1,};
-    populateFakeStats(TEST_UID1, 0, IFACE_INDEX1, TEST_COUNTERSET0, value1, mFakeUidStatsMap);
-    populateFakeStats(TEST_UID1, 0, IFACE_INDEX2, TEST_COUNTERSET1, value1, mFakeUidStatsMap);
-    populateFakeStats(TEST_UID2, 0, IFACE_INDEX3, TEST_COUNTERSET1, value1, mFakeUidStatsMap);
-    Stats result1 = {};
-    ASSERT_EQ(0, bpfGetUidStatsInternal(TEST_UID1, &result1, mFakeUidStatsMap));
-    StatsValue uid1Value = {
+    StatsValue value2 = {
         .rxBytes = TEST_BYTES0 * 2,
         .rxPackets = TEST_PACKET0 * 2,
         .txBytes = TEST_BYTES1 * 2,
         .txPackets = TEST_PACKET1 * 2,
     };
-    expectStatsEqual(uid1Value, result1);
+    ASSERT_TRUE(isOk(mFakeAppUidStatsMap.writeValue(TEST_UID1, value1, BPF_ANY)));
+    ASSERT_TRUE(isOk(mFakeAppUidStatsMap.writeValue(TEST_UID2, value2, BPF_ANY)));
+    Stats result1 = {};
+    ASSERT_EQ(0, bpfGetUidStatsInternal(TEST_UID1, &result1, mFakeAppUidStatsMap));
+    expectStatsEqual(value1, result1);
 
     Stats result2 = {};
-    ASSERT_EQ(0, bpfGetUidStatsInternal(TEST_UID2, &result2, mFakeUidStatsMap));
-    expectStatsEqual(value1, result2);
+    ASSERT_EQ(0, bpfGetUidStatsInternal(TEST_UID2, &result2, mFakeAppUidStatsMap));
+    expectStatsEqual(value2, result2);
     std::vector<stats_line> lines;
     std::vector<std::string> ifaces;
+    populateFakeStats(TEST_UID1, 0, IFACE_INDEX1, TEST_COUNTERSET0, value1, mFakeUidStatsMap);
+    populateFakeStats(TEST_UID1, 0, IFACE_INDEX2, TEST_COUNTERSET1, value1, mFakeUidStatsMap);
+    populateFakeStats(TEST_UID2, 0, IFACE_INDEX3, TEST_COUNTERSET1, value1, mFakeUidStatsMap);
     ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(&lines, ifaces, TAG_ALL, TEST_UID1,
                                                     mFakeUidStatsMap, mFakeIfaceIndexNameMap));
     ASSERT_EQ((unsigned long)2, lines.size());
@@ -395,6 +406,7 @@ TEST_F(BpfNetworkStatsHelperTest, TestGetIfaceStatsDetail) {
     updateIfaceMap(IFACE_NAME1, IFACE_INDEX1);
     updateIfaceMap(IFACE_NAME2, IFACE_INDEX2);
     updateIfaceMap(IFACE_NAME3, IFACE_INDEX3);
+    updateIfaceMap(LONG_IFACE_NAME, IFACE_INDEX4);
     StatsValue value1 = {
         .rxBytes = TEST_BYTES0,
         .rxPackets = TEST_PACKET0,
@@ -413,16 +425,20 @@ TEST_F(BpfNetworkStatsHelperTest, TestGetIfaceStatsDetail) {
     EXPECT_TRUE(isOk(mFakeIfaceStatsMap.writeValue(ifaceStatsKey, value2, BPF_ANY)));
     ifaceStatsKey = IFACE_INDEX3;
     EXPECT_TRUE(isOk(mFakeIfaceStatsMap.writeValue(ifaceStatsKey, value1, BPF_ANY)));
+    ifaceStatsKey = IFACE_INDEX4;
+    EXPECT_TRUE(isOk(mFakeIfaceStatsMap.writeValue(ifaceStatsKey, value2, BPF_ANY)));
     std::vector<stats_line> lines;
     ASSERT_EQ(0,
               parseBpfNetworkStatsDevInternal(&lines, mFakeIfaceStatsMap, mFakeIfaceIndexNameMap));
-    ASSERT_EQ((unsigned long)3, lines.size());
+    ASSERT_EQ((unsigned long)4, lines.size());
     std::sort(lines.begin(), lines.end(), [](const auto& line1, const auto& line2)-> bool {
         return strcmp(line1.iface, line2.iface) < 0;
     });
     expectStatsLineEqual(value1, IFACE_NAME1, UID_ALL, SET_ALL, TAG_NONE, lines[0]);
     expectStatsLineEqual(value1, IFACE_NAME3, UID_ALL, SET_ALL, TAG_NONE, lines[1]);
     expectStatsLineEqual(value2, IFACE_NAME2, UID_ALL, SET_ALL, TAG_NONE, lines[2]);
+    ASSERT_EQ(0, strcmp(TRUNCATED_IFACE_NAME, lines[3].iface));
+    expectStatsLineEqual(value2, TRUNCATED_IFACE_NAME, UID_ALL, SET_ALL, TAG_NONE, lines[3]);
 }
 }  // namespace bpf
 }  // namespace android
diff --git a/libbpf/include/bpf/BpfNetworkStats.h b/libbpf/include/bpf/BpfNetworkStats.h
index dd02a392..80dff88e 100644
--- a/libbpf/include/bpf/BpfNetworkStats.h
+++ b/libbpf/include/bpf/BpfNetworkStats.h
@@ -46,7 +46,7 @@ struct stats_line {
 };
 // For test only
 int bpfGetUidStatsInternal(uid_t uid, struct Stats* stats,
-                           const BpfMap<StatsKey, StatsValue>& uidStatsMap);
+                           const BpfMap<uint32_t, StatsValue>& appUidStatsMap);
 // For test only
 int bpfGetIfaceStatsInternal(const char* iface, Stats* stats,
                              const BpfMap<uint32_t, StatsValue>& ifaceStatsMap,
diff --git a/libbpf/include/bpf/BpfUtils.h b/libbpf/include/bpf/BpfUtils.h
index cfee6a96..ac107df8 100644
--- a/libbpf/include/bpf/BpfUtils.h
+++ b/libbpf/include/bpf/BpfUtils.h
@@ -75,19 +75,29 @@ struct IfaceValue {
 #define BPF_PATH "/sys/fs/bpf"
 
 // Since we cannot garbage collect the stats map since device boot, we need to make these maps as
-// large as possible. The current rlimit of MEM_LOCK allows at most 10000 map entries for each
-// stats map. In the old qtaguid module, we don't have a total limit for data entries but only have
-// limitation of tags each uid can have. (default is 1024 in kernel);
-// cookie_tag_map:      key:  8 bytes, value:  8 bytes, total:10000*8*2 bytes         =  160Kbytes
-// uid_counter_set_map: key:  4 bytes, value:  1 bytes, total:2000*5 bytes            =   10Kbytes
-// uid_stats_map:       key: 16 bytes, value: 32 bytes, total:10000*16+10000*32 bytes =  480Kbytes
-// tag_stats_map:       key: 16 bytes, value: 32 bytes, total:10000*16+10000*32 bytes =  480Kbytes
-// iface_index_name_map:key:  4 bytes, value: 32 bytes, total:1000*36 bytes           =   36Kbytes
-// iface_stats_map:     key:  4 bytes, value: 32 bytes, total:1000*36 bytes           =   36Kbytes
-// dozable_uid_map:     key:  4 bytes, value:  1 bytes, total:2000*5 bytes            =   10Kbytes
-// standby_uid_map:     key:  4 bytes, value:  1 bytes, total:2000*5 bytes            =   10Kbytes
-// powersave_uid_map:   key:  4 bytes, value:  1 bytes, total:2000*5 bytes            =   10Kbytes
-// total:                                                                               1232Kbytes
+// large as possible. The maximum size of number of map entries we can have is depend on the rlimit
+// of MEM_LOCK granted to netd. The memory space needed by each map can be calculated by the
+// following fomula:
+//      elem_size = 40 + roundup(key_size, 8) + roundup(value_size, 8)
+//      cost = roundup_pow_of_two(max_entries) * 16 + elem_size * max_entries +
+//              elem_size * number_of_CPU
+// And the cost of each map currently used is(assume the device have 8 CPUs):
+// cookie_tag_map:      key:  8 bytes, value:  8 bytes, cost:  822592 bytes    =   823Kbytes
+// uid_counter_set_map: key:  4 bytes, value:  1 bytes, cost:  145216 bytes    =   145Kbytes
+// app_uid_stats_map:   key:  4 bytes, value: 32 bytes, cost: 1062784 bytes    =  1063Kbytes
+// uid_stats_map:       key: 16 bytes, value: 32 bytes, cost: 1142848 bytes    =  1143Kbytes
+// tag_stats_map:       key: 16 bytes, value: 32 bytes, cost: 1142848 bytes    =  1143Kbytes
+// iface_index_name_map:key:  4 bytes, value: 16 bytes, cost:   80896 bytes    =    81Kbytes
+// iface_stats_map:     key:  4 bytes, value: 32 bytes, cost:   97024 bytes    =    97Kbytes
+// dozable_uid_map:     key:  4 bytes, value:  1 bytes, cost:  145216 bytes    =   145Kbytes
+// standby_uid_map:     key:  4 bytes, value:  1 bytes, cost:  145216 bytes    =   145Kbytes
+// powersave_uid_map:   key:  4 bytes, value:  1 bytes, cost:  145216 bytes    =   145Kbytes
+// total:                                                                         4930Kbytes
+// It takes maximum 4.9MB kernel memory space if all maps are full, which requires any devices
+// running this module to have a memlock rlimit to be larger then 5MB. In the old qtaguid module,
+// we don't have a total limit for data entries but only have limitation of tags each uid can have.
+// (default is 1024 in kernel);
+
 constexpr const int COOKIE_UID_MAP_SIZE = 10000;
 constexpr const int UID_COUNTERSET_MAP_SIZE = 2000;
 constexpr const int UID_STATS_MAP_SIZE = 10000;
@@ -105,6 +115,7 @@ constexpr const char* CGROUP_ROOT_PATH = "/dev/cg2_bpf";
 
 constexpr const char* COOKIE_TAG_MAP_PATH = BPF_PATH "/traffic_cookie_tag_map";
 constexpr const char* UID_COUNTERSET_MAP_PATH = BPF_PATH "/traffic_uid_counterSet_map";
+constexpr const char* APP_UID_STATS_MAP_PATH = BPF_PATH "/traffic_app_uid_stats_map";
 constexpr const char* UID_STATS_MAP_PATH = BPF_PATH "/traffic_uid_stats_map";
 constexpr const char* TAG_STATS_MAP_PATH = BPF_PATH "/traffic_tag_stats_map";
 constexpr const char* IFACE_INDEX_NAME_MAP_PATH = BPF_PATH "/traffic_iface_index_name_map";
diff --git a/libbpf/include/bpf/bpf_shared.h b/libbpf/include/bpf/bpf_shared.h
index 4a4904d2..217b76f6 100644
--- a/libbpf/include/bpf/bpf_shared.h
+++ b/libbpf/include/bpf/bpf_shared.h
@@ -23,6 +23,7 @@
 
 #define COOKIE_TAG_MAP 0xbfceaaffffffffff
 #define UID_COUNTERSET_MAP 0xbfdceeafffffffff
+#define APP_UID_STATS_MAP 0xbfa1daafffffffff
 #define UID_STATS_MAP 0xbfdaafffffffffff
 #define TAG_STATS_MAP 0xbfaaafffffffffff
 #define IFACE_STATS_MAP 0xbf1faceaafffffff
diff --git a/server/TrafficController.cpp b/server/TrafficController.cpp
index d8eb6b6e..d6a64802 100644
--- a/server/TrafficController.cpp
+++ b/server/TrafficController.cpp
@@ -149,6 +149,11 @@ Status TrafficController::initMaps() {
                                         "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 @@ int TrafficController::deleteTagData(uint32_t tag, uid_t uid) {
               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 @@ void TrafficController::dump(DumpWriter& dw, bool verbose) {
                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 @@ void TrafficController::dump(DumpWriter& dw, bool verbose) {
         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 05d91dfd..79f7d141 100644
--- a/server/TrafficController.h
+++ b/server/TrafficController.h
@@ -132,6 +132,13 @@ class TrafficController {
     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 ca4b1634..a354f839 100644
--- a/server/TrafficControllerTest.cpp
+++ b/server/TrafficControllerTest.cpp
@@ -72,6 +72,7 @@ class TrafficControllerTest : public ::testing::Test {
     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 @@ class TrafficControllerTest : public ::testing::Test {
             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 @@ class TrafficControllerTest : public ::testing::Test {
 
         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 @@ class TrafficControllerTest : public ::testing::Test {
         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 @@ class TrafficControllerTest : public ::testing::Test {
         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 @@ TEST_F(TrafficControllerTest, TestDeleteTagData) {
     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 @@ TEST_F(TrafficControllerTest, TestDeleteAllUidData) {
     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 @@ TEST_F(TrafficControllerTest, TestDeleteDataWithTwoUids) {
     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) {
diff --git a/tests/bpf_base_test.cpp b/tests/bpf_base_test.cpp
index fa7de0db..e023052a 100644
--- a/tests/bpf_base_test.cpp
+++ b/tests/bpf_base_test.cpp
@@ -102,6 +102,33 @@ TEST_F(BpfBasicTest, TestTagSocket) {
     ASSERT_EQ(ENOENT, tagResult.status().code());
 }
 
+TEST_F(BpfBasicTest, TestCloseSocketWithoutUntag) {
+    SKIP_IF_BPF_NOT_SUPPORTED;
+
+    BpfMap<uint64_t, UidTag> cookieTagMap(mapRetrieve(COOKIE_TAG_MAP_PATH, 0));
+    ASSERT_LE(0, cookieTagMap.getMap());
+    int sock = socket(AF_INET6, SOCK_STREAM, 0);
+    ASSERT_LE(0, sock);
+    uint64_t cookie = getSocketCookie(sock);
+    ASSERT_NE(NONEXISTENT_COOKIE, cookie);
+    ASSERT_EQ(0, qtaguid_tagSocket(sock, TEST_TAG, TEST_UID));
+    StatusOr<UidTag> tagResult = cookieTagMap.readValue(cookie);
+    ASSERT_TRUE(isOk(tagResult));
+    ASSERT_EQ(TEST_UID, tagResult.value().uid);
+    ASSERT_EQ(TEST_TAG, tagResult.value().tag);
+    ASSERT_EQ(0, close(sock));
+    // Check map periodically until sk destroy handler have done its job.
+    for (int i = 0; i < 10; i++) {
+        tagResult = cookieTagMap.readValue(cookie);
+        if (!isOk(tagResult)) {
+            ASSERT_EQ(ENOENT, tagResult.status().code());
+            return;
+        }
+        usleep(50);
+    }
+    FAIL() << "socket tag still exist after 500ms";
+}
+
 TEST_F(BpfBasicTest, TestChangeCounterSet) {
     SKIP_IF_BPF_NOT_SUPPORTED;
 
@@ -125,6 +152,8 @@ TEST_F(BpfBasicTest, TestDeleteTagData) {
     ASSERT_LE(0, uidStatsMap.getMap());
     BpfMap<StatsKey, StatsValue> tagStatsMap(mapRetrieve(TAG_STATS_MAP_PATH, 0));
     ASSERT_LE(0, tagStatsMap.getMap());
+    BpfMap<uint32_t, StatsValue> appUidStatsMap(mapRetrieve(APP_UID_STATS_MAP_PATH, 0));
+    ASSERT_LE(0, appUidStatsMap.getMap());
 
     StatsKey key = {.uid = TEST_UID, .tag = TEST_TAG, .counterSet = TEST_COUNTERSET,
                     .ifaceIndex = 1};
@@ -132,10 +161,14 @@ TEST_F(BpfBasicTest, TestDeleteTagData) {
     EXPECT_TRUE(isOk(tagStatsMap.writeValue(key, statsMapValue, BPF_ANY)));
     key.tag = 0;
     EXPECT_TRUE(isOk(uidStatsMap.writeValue(key, statsMapValue, BPF_ANY)));
+    EXPECT_TRUE(isOk(appUidStatsMap.writeValue(TEST_UID, statsMapValue, BPF_ANY)));
     ASSERT_EQ(0, qtaguid_deleteTagData(0, TEST_UID));
     StatusOr<StatsValue> statsResult = uidStatsMap.readValue(key);
     ASSERT_FALSE(isOk(statsResult));
     ASSERT_EQ(ENOENT, statsResult.status().code());
+    statsResult = appUidStatsMap.readValue(TEST_UID);
+    ASSERT_FALSE(isOk(statsResult));
+    ASSERT_EQ(ENOENT, statsResult.status().code());
     key.tag = TEST_TAG;
     statsResult = tagStatsMap.readValue(key);
     ASSERT_FALSE(isOk(statsResult));