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/bpfloader/BpfLoader.cpp b/bpfloader/BpfLoader.cpp
index 47dbb3d..7ca044e 100644
--- a/bpfloader/BpfLoader.cpp
+++ b/bpfloader/BpfLoader.cpp
@@ -264,6 +264,7 @@
 }  // 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_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 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 @@
     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 e37f1a1..a59cb6d 100644
--- a/bpfloader/bpf_kern.h
+++ b/bpfloader/bpf_kern.h
@@ -194,5 +194,6 @@
 
     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 aa604ce..eafbb5e 100644
--- a/libbpf/BpfNetworkStats.cpp
+++ b/libbpf/BpfNetworkStats.cpp
@@ -42,31 +42,27 @@
 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 e804c40..8fd4943 100644
--- a/libbpf/BpfNetworkStatsTest.cpp
+++ b/libbpf/BpfNetworkStatsTest.cpp
@@ -65,15 +65,20 @@
 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 @@
             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 @@
                          .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 @@
     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 @@
     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 dd02a39..80dff88 100644
--- a/libbpf/include/bpf/BpfNetworkStats.h
+++ b/libbpf/include/bpf/BpfNetworkStats.h
@@ -46,7 +46,7 @@
 };
 // 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 cfee6a9..ac107df 100644
--- a/libbpf/include/bpf/BpfUtils.h
+++ b/libbpf/include/bpf/BpfUtils.h
@@ -75,19 +75,29 @@
 #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* 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 4a4904d..217b76f 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 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) {
diff --git a/tests/bpf_base_test.cpp b/tests/bpf_base_test.cpp
index fa7de0d..e023052 100644
--- a/tests/bpf_base_test.cpp
+++ b/tests/bpf_base_test.cpp
@@ -102,6 +102,33 @@
     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 @@
     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 @@
     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));