libmeminfo: Add support to fiter accounting based on page flags

Bug: 114325007
Bug: 111694435
Test: libmeminfo_test 1 --gtest_filter=ValidateProcMemInfoFlags.*

Change-Id: Ifa7006de78dd909b5b4db282a8c9931ebf97a68a
Signed-off-by: Sandeep Patil <sspatil@google.com>

3 files changed, 57 insertions, 2 deletions

diff --git a/libmeminfo/include/meminfo/procmeminfo.h b/libmeminfo/include/meminfo/procmeminfo.h
index dd599128d..3b620a00b 100644
--- a/libmeminfo/include/meminfo/procmeminfo.h
+++ b/libmeminfo/include/meminfo/procmeminfo.h
@@ -29,7 +29,7 @@ namespace meminfo {
 class ProcMemInfo final {
     // Per-process memory accounting
   public:
-    ProcMemInfo(pid_t pid, bool get_wss = false);
+    ProcMemInfo(pid_t pid, bool get_wss = false, uint64_t pgflags = 0, uint64_t pgflags_mask = 0);
 
     const std::vector<Vma>& Maps();
     const MemUsage& Usage();
@@ -45,6 +45,8 @@ class ProcMemInfo final {
 
     pid_t pid_;
     bool get_wss_;
+    uint64_t pgflags_;
+    uint64_t pgflags_mask_;
 
     std::vector<Vma> maps_;
 
diff --git a/libmeminfo/libmeminfo_test.cpp b/libmeminfo/libmeminfo_test.cpp
index 923e156d5..6d731737b 100644
--- a/libmeminfo/libmeminfo_test.cpp
+++ b/libmeminfo/libmeminfo_test.cpp
@@ -245,6 +245,55 @@ TEST_F(ValidatePageAcct, TestPageIdle) {
     }
 }
 
+TEST(ValidateProcMemInfoFlags, TestPageFlags1) {
+    // Create proc object using libpagemap
+    pm_kernel_t* ker;
+    ASSERT_EQ(0, pm_kernel_create(&ker));
+    pm_process_t* proc;
+    ASSERT_EQ(0, pm_process_create(ker, pid, &proc));
+
+    // count swapbacked pages using libpagemap
+    pm_memusage_t proc_usage;
+    pm_memusage_zero(&proc_usage);
+    ASSERT_EQ(0, pm_process_usage_flags(proc, &proc_usage, (1 << KPF_SWAPBACKED),
+                                        (1 << KPF_SWAPBACKED)));
+
+    // Create ProcMemInfo that counts swapbacked pages
+    ProcMemInfo proc_mem(pid, false, (1 << KPF_SWAPBACKED), (1 << KPF_SWAPBACKED));
+
+    EXPECT_EQ(proc_usage.vss, proc_mem.Usage().vss);
+    EXPECT_EQ(proc_usage.rss, proc_mem.Usage().rss);
+    EXPECT_EQ(proc_usage.pss, proc_mem.Usage().pss);
+    EXPECT_EQ(proc_usage.uss, proc_mem.Usage().uss);
+
+    pm_process_destroy(proc);
+    pm_kernel_destroy(ker);
+}
+
+TEST(ValidateProcMemInfoFlags, TestPageFlags2) {
+    // Create proc object using libpagemap
+    pm_kernel_t* ker;
+    ASSERT_EQ(0, pm_kernel_create(&ker));
+    pm_process_t* proc;
+    ASSERT_EQ(0, pm_process_create(ker, pid, &proc));
+
+    // count non-swapbacked pages using libpagemap
+    pm_memusage_t proc_usage;
+    pm_memusage_zero(&proc_usage);
+    ASSERT_EQ(0, pm_process_usage_flags(proc, &proc_usage, (1 << KPF_SWAPBACKED), 0));
+
+    // Create ProcMemInfo that counts non-swapbacked pages
+    ProcMemInfo proc_mem(pid, false, 0, (1 << KPF_SWAPBACKED));
+
+    EXPECT_EQ(proc_usage.vss, proc_mem.Usage().vss);
+    EXPECT_EQ(proc_usage.rss, proc_mem.Usage().rss);
+    EXPECT_EQ(proc_usage.pss, proc_mem.Usage().pss);
+    EXPECT_EQ(proc_usage.uss, proc_mem.Usage().uss);
+
+    pm_process_destroy(proc);
+    pm_kernel_destroy(ker);
+}
+
 TEST(SysMemInfoParser, TestSysMemInfoFile) {
     std::string meminfo = R"meminfo(MemTotal:        3019740 kB
 MemFree:         1809728 kB
diff --git a/libmeminfo/procmeminfo.cpp b/libmeminfo/procmeminfo.cpp
index e3839346a..56b62e826 100644
--- a/libmeminfo/procmeminfo.cpp
+++ b/libmeminfo/procmeminfo.cpp
@@ -53,7 +53,8 @@ static void add_mem_usage(MemUsage* to, const MemUsage& from) {
     to->shared_dirty += from.shared_dirty;
 }
 
-ProcMemInfo::ProcMemInfo(pid_t pid, bool get_wss) : pid_(pid), get_wss_(get_wss) {
+ProcMemInfo::ProcMemInfo(pid_t pid, bool get_wss, uint64_t pgflags, uint64_t pgflags_mask)
+    : pid_(pid), get_wss_(get_wss), pgflags_(pgflags), pgflags_mask_(pgflags_mask) {
     if (!ReadMaps(get_wss_)) {
         LOG(ERROR) << "Failed to read maps for Process " << pid_;
     }
@@ -170,6 +171,9 @@ bool ProcMemInfo::ReadVmaStats(int pagemap_fd, Vma& vma, bool get_wss) {
             return false;
         }
 
+        // skip unwanted pages from the count
+        if ((pg_flags[i] & pgflags_mask_) != pgflags_) continue;
+
         if (!pinfo.PageMapCount(page_frame, &pg_counts[i])) {
             LOG(ERROR) << "Failed to get page count for " << page_frame << " in process " << pid_;
             return false;