libpagemap: Add APIs to mark and read idle pages of a process.

pm_kernel_init_page_idle()   - Initialize idle page tracking.
pm_kernel_has_page_idle()    -  Returns 1 if idle page tracking has been
                               successfully initialized.
pm_kernel_get_page_idle()    - Gets the idle status of the page frame.
pm_kernel_mark_page_idle()   - Marks one or many page frames idle.
pm_kernel_page_is_accessed() - Returns 1 if page is accessed.

pm_map_mark_idle()           - Marks mapped pages from a pagemap idle.

Added '-i' usage flag to procmem that tells procmem to use idle page
tracking. The flag is ignored in the absense of '-w' or '-W' options.

Bug: 111694435
Test: procmem -i -h 1
Test: procmem -h -W 1 && procmem -h -w 1
Test: procmem -h -i -W 1 && procmem -h -i -w 1

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

5 files changed, 293 insertions, 69 deletions

diff --git a/libpagemap/include/pagemap/pagemap.h b/libpagemap/include/pagemap/pagemap.h
index b03614e8..d567729f 100644
--- a/libpagemap/include/pagemap/pagemap.h
+++ b/libpagemap/include/pagemap/pagemap.h
@@ -81,6 +81,7 @@ typedef struct pm_map      pm_map_t;
 struct pm_kernel {
     int kpagecount_fd;
     int kpageflags_fd;
+    int pageidle_fd;
 
     int pagesize;
 };
@@ -116,6 +117,45 @@ int pm_kernel_create(pm_kernel_t **ker_out);
 
 #define pm_kernel_pagesize(ker) ((ker)->pagesize)
 
+#define pfn_to_page_idle_offset(x) ((x >> 6) << 3)
+
+/* Sets up the interaction with kernel's idle page tracking support
+ * by opening /sys/kernel/mm/page_idle/bitmap
+ *
+ * Return:
+ * 0 on success -errno on failure.
+ */
+int pm_kernel_init_page_idle(pm_kernel_t* ker);
+
+/* Checks if the kernel idle page tracking interface has been initialized.
+ * Return:
+ * 0 if not initialized or failed to initialize.
+ * 1 if we are ready to talk to page_idle/bitmap.
+ */
+int pm_kernel_has_page_idle(pm_kernel_t* ker);
+
+/* Get idle status of given page frame number.
+ * Returns:
+ *  0 if not idle i.e. page frame was accessed by the process.
+ *  1 if page frame continues to be idle since last time idle flag was set.
+ *  -errno on failure
+ */
+int pm_kernel_get_page_idle(pm_kernel_t* ker, uint64_t pfn);
+
+/* Mark page frames idle.
+ * Return:
+ * 0 on success -errno on failure.
+ */
+int pm_kernel_mark_page_idle(pm_kernel_t* ker, uint64_t* pfn, int n);
+
+/* Determines if the page frame has been accessed by user space since
+ * the working set was last reset
+ * 0 If page frame wasn't accessed.
+ * 1 If page frame was accessed.
+ * -errno on failure.
+ */
+int pm_kernel_page_is_accessed(pm_kernel_t* ker, uint64_t pfn, uint64_t* flags);
+
 /* Get a list of probably-existing PIDs (returned through *pids_out).
  * Length of the array (in sizeof(pid_t) units) is returned through *len.
  * The array should be freed by the caller. */
@@ -194,6 +234,9 @@ int pm_process_destroy(pm_process_t *proc);
  * caller. */
 int pm_map_pagemap(pm_map_t *map, uint64_t **pagemap_out, size_t *len);
 
+/* Mark all the present pages in process's VM as idle */
+int pm_map_mark_idle(pm_map_t* map);
+
 /* Get the memory usage of this map alone. */
 int pm_map_usage(pm_map_t *map, pm_memusage_t *usage_out);
 
diff --git a/libpagemap/pm_kernel.c b/libpagemap/pm_kernel.c
index 760f7347..ec475169 100644
--- a/libpagemap/pm_kernel.c
+++ b/libpagemap/pm_kernel.c
@@ -51,6 +51,9 @@ int pm_kernel_create(pm_kernel_t **ker_out) {
         return error;
     }
 
+    /* This must be explicitly initialized through pm_kernel_init_page_idle() */
+    ker->pageidle_fd = -1;
+
     ker->pagesize = getpagesize();
 
     *ker_out = ker;
@@ -58,6 +61,19 @@ int pm_kernel_create(pm_kernel_t **ker_out) {
     return 0;
 }
 
+int pm_kernel_init_page_idle(pm_kernel_t* ker) {
+    if (!ker || ker->pageidle_fd != -1) {
+        return -EINVAL;
+    }
+
+    ker->pageidle_fd = open("/sys/kernel/mm/page_idle/bitmap", O_RDWR | O_CLOEXEC);
+    if (ker->pageidle_fd < 0) {
+        return -errno;
+    }
+
+    return 0;
+}
+
 #define INIT_PIDS 20
 int pm_kernel_pids(pm_kernel_t *ker, pid_t **pids_out, size_t *len) {
     DIR *proc;
@@ -116,41 +132,123 @@ int pm_kernel_pids(pm_kernel_t *ker, pid_t **pids_out, size_t *len) {
 int pm_kernel_count(pm_kernel_t *ker, uint64_t pfn, uint64_t *count_out) {
     off64_t off;
 
-    if (!ker || !count_out)
-        return -1;
+    if (!ker || !count_out) return -1;
 
     off = lseek64(ker->kpagecount_fd, pfn * sizeof(uint64_t), SEEK_SET);
-    if (off == (off_t)-1)
-        return -errno;
-    if (read(ker->kpagecount_fd, count_out, sizeof(uint64_t)) <
-        (ssize_t)sizeof(uint64_t))
+    if (off == (off64_t)-1) return -errno;
+    if (TEMP_FAILURE_RETRY(read(ker->kpagecount_fd, count_out, sizeof(uint64_t))) <
+        (ssize_t)sizeof(uint64_t)) {
         return -errno;
+    }
 
     return 0;
 }
 
-int pm_kernel_flags(pm_kernel_t *ker, uint64_t pfn, uint64_t *flags_out) {
+int pm_kernel_flags(pm_kernel_t* ker, uint64_t pfn, uint64_t* flags_out) {
     off64_t off;
 
-    if (!ker || !flags_out)
-        return -1;
+    if (!ker || !flags_out) {
+        return -EINVAL;
+    }
 
     off = lseek64(ker->kpageflags_fd, pfn * sizeof(uint64_t), SEEK_SET);
-    if (off == (off_t)-1)
+    if (off == (off64_t)-1) return -errno;
+    if (TEMP_FAILURE_RETRY(read(ker->kpageflags_fd, flags_out, sizeof(uint64_t))) <
+        (ssize_t)sizeof(uint64_t)) {
         return -errno;
-    if (read(ker->kpageflags_fd, flags_out, sizeof(uint64_t)) <
-        (ssize_t)sizeof(uint64_t))
+    }
+
+    return 0;
+}
+
+int pm_kernel_has_page_idle(pm_kernel_t* ker) {
+    /* Treat error to be fallback to clear_refs */
+    if (!ker) {
+        return 0;
+    }
+
+    return !(ker->pageidle_fd < 0);
+}
+
+int pm_kernel_get_page_idle(pm_kernel_t* ker, uint64_t pfn) {
+    uint64_t bits;
+    off64_t offset;
+
+    if (!ker) {
+        return -EINVAL;
+    }
+
+    if (ker->pageidle_fd < 0) {
+        return -ENXIO;
+    }
+
+    offset = pfn_to_page_idle_offset(pfn);
+    if (pread64(ker->pageidle_fd, &bits, sizeof(uint64_t), offset) < 0) {
         return -errno;
+    }
+
+    bits >>= (pfn % 64);
+
+    return bits & 1;
+}
+
+int pm_kernel_mark_page_idle(pm_kernel_t* ker, uint64_t* pfn, int n) {
+    uint64_t bits;
+
+    if (!ker) {
+        return -EINVAL;
+    }
+
+    if (ker->pageidle_fd < 0) {
+        return -ENXIO;
+    }
+
+    for (int i = 0; i < n; i++) {
+        off64_t offset = pfn_to_page_idle_offset(pfn[i]);
+        if (pread64(ker->pageidle_fd, &bits, sizeof(uint64_t), offset) < 0) {
+            return -errno;
+        }
+
+        bits |= 1ULL << (pfn[i] % 64);
+
+        if (pwrite64(ker->pageidle_fd, &bits, sizeof(uint64_t), offset) < 0) {
+            return -errno;
+        }
+    }
 
     return 0;
 }
 
+int pm_kernel_page_is_accessed(pm_kernel_t* ker, uint64_t pfn, uint64_t* flags) {
+    int error;
+    uint64_t page_flags;
+
+    if (!ker) {
+        return -EINVAL;
+    }
+
+    if (pm_kernel_has_page_idle(ker)) {
+        return pm_kernel_get_page_idle(ker, pfn);
+    }
+
+    if (!flags) {
+        flags = &page_flags;
+        error = pm_kernel_flags(ker, pfn, flags);
+        if (error) return error;
+    }
+
+    return !!(*flags & (1 << KPF_REFERENCED));
+}
+
 int pm_kernel_destroy(pm_kernel_t *ker) {
     if (!ker)
         return -1;
 
     close(ker->kpagecount_fd);
     close(ker->kpageflags_fd);
+    if (ker->pageidle_fd >= 0) {
+        close(ker->pageidle_fd);
+    }
 
     free(ker);
 
diff --git a/libpagemap/pm_map.c b/libpagemap/pm_map.c
index 48d50678..158a4faa 100644
--- a/libpagemap/pm_map.c
+++ b/libpagemap/pm_map.c
@@ -13,20 +13,51 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
+#include <errno.h>
 #include <stdlib.h>
 #include <string.h>
 
 #include <pagemap/pagemap.h>
 
 int pm_map_pagemap(pm_map_t *map, uint64_t **pagemap_out, size_t *len) {
-    if (!map)
-        return -1;
+    if (!map) {
+        return -EINVAL;
+    }
 
     return pm_process_pagemap_range(map->proc, map->start, map->end,
                                     pagemap_out, len);
 }
 
+int pm_map_mark_idle(pm_map_t* map) {
+    uint64_t* pagemap;
+    size_t len;
+    int error;
+
+    if (!map) {
+        return -EINVAL;
+    }
+
+    error = pm_map_pagemap(map, &pagemap, &len);
+    if (error) {
+        return error;
+    }
+
+    for (size_t i = 0; i < len; i++) {
+        uint64_t pfn;
+        if (!PM_PAGEMAP_PRESENT(pagemap[i])) {
+            continue;
+        }
+        pfn = PM_PAGEMAP_PFN(pagemap[i]);
+        error = pm_kernel_mark_page_idle(map->proc->ker, &pfn, 1);
+        if (error) {
+            break;
+        }
+    }
+
+    free(pagemap);
+    return error;
+}
+
 int pm_map_usage_flags(pm_map_t *map, pm_memusage_t *usage_out,
                         uint64_t flags_mask, uint64_t required_flags) {
     uint64_t *pagemap;
@@ -47,23 +78,18 @@ int pm_map_usage_flags(pm_map_t *map, pm_memusage_t *usage_out,
     for (i = 0; i < len; i++) {
         usage.vss += map->proc->ker->pagesize;
 
-        if (!PM_PAGEMAP_PRESENT(pagemap[i]) &&
-                !PM_PAGEMAP_SWAPPED(pagemap[i]))
-            continue;
+        if (!PM_PAGEMAP_PRESENT(pagemap[i]) && !PM_PAGEMAP_SWAPPED(pagemap[i])) continue;
 
         if (!PM_PAGEMAP_SWAPPED(pagemap[i])) {
             if (flags_mask) {
                 uint64_t flags;
-                error = pm_kernel_flags(map->proc->ker, PM_PAGEMAP_PFN(pagemap[i]),
-                                        &flags);
+                error = pm_kernel_flags(map->proc->ker, PM_PAGEMAP_PFN(pagemap[i]), &flags);
                 if (error) goto out;
 
-                if ((flags & flags_mask) != required_flags)
-                    continue;
+                if ((flags & flags_mask) != required_flags) continue;
             }
 
-            error = pm_kernel_count(map->proc->ker, PM_PAGEMAP_PFN(pagemap[i]),
-                                    &count);
+            error = pm_kernel_count(map->proc->ker, PM_PAGEMAP_PFN(pagemap[i]), &count);
             if (error) goto out;
 
             usage.rss += (count >= 1) ? map->proc->ker->pagesize : (0);
@@ -85,19 +111,18 @@ out:
     return error;
 }
 
-int pm_map_usage(pm_map_t *map, pm_memusage_t *usage_out) {
+int pm_map_usage(pm_map_t* map, pm_memusage_t* usage_out) {
     return pm_map_usage_flags(map, usage_out, 0, 0);
 }
 
-int pm_map_workingset(pm_map_t *map, pm_memusage_t *ws_out) {
-    uint64_t *pagemap;
+int pm_map_workingset(pm_map_t* map, pm_memusage_t* ws_out) {
+    uint64_t* pagemap;
     size_t len, i;
-    uint64_t count, flags;
+    uint64_t count;
     pm_memusage_t ws;
     int error;
 
-    if (!map || !ws_out)
-        return -1;
+    if (!map || !ws_out) return -1;
 
     error = pm_map_pagemap(map, &pagemap, &len);
     if (error) return error;
@@ -105,18 +130,19 @@ int pm_map_workingset(pm_map_t *map, pm_memusage_t *ws_out) {
     pm_memusage_zero(&ws);
 
     for (i = 0; i < len; i++) {
-        error = pm_kernel_flags(map->proc->ker, PM_PAGEMAP_PFN(pagemap[i]),
-                                &flags);
-        if (error) goto out;
-
-        if (!(flags & (1 << KPF_REFERENCED))) continue;
+        if (!pm_kernel_page_is_accessed(map->proc->ker, PM_PAGEMAP_PFN(pagemap[i]), NULL)) {
+            continue;
+        }
 
-        error = pm_kernel_count(map->proc->ker, PM_PAGEMAP_PFN(pagemap[i]),
-                                &count);
-        if (error) goto out;
+        error = pm_kernel_count(map->proc->ker, PM_PAGEMAP_PFN(pagemap[i]), &count);
+        if (error) {
+            goto out;
+        }
 
         ws.vss += map->proc->ker->pagesize;
-        if( PM_PAGEMAP_SWAPPED(pagemap[i]) ) continue;
+        if (PM_PAGEMAP_SWAPPED(pagemap[i])) {
+            continue;
+        }
         ws.rss += (count >= 1) ? (map->proc->ker->pagesize) : (0);
         ws.pss += (count >= 1) ? (map->proc->ker->pagesize / count) : (0);
         ws.uss += (count == 1) ? (map->proc->ker->pagesize) : (0);
diff --git a/libpagemap/pm_process.c b/libpagemap/pm_process.c
index 5d47e26c..98c1c6a0 100644
--- a/libpagemap/pm_process.c
+++ b/libpagemap/pm_process.c
@@ -27,6 +27,8 @@
 #include "pm_map.h"
 
 static int read_maps(pm_process_t *proc);
+static int pm_process_clear_refs(pm_process_t* proc);
+static int pm_process_mark_idle(pm_process_t* proc);
 
 #define MAX_FILENAME 64
 
@@ -113,8 +115,9 @@ int pm_process_pagemap_range(pm_process_t *proc,
     off64_t off;
     int error;
 
-    if (!proc || (low > high) || !range_out || !len)
-        return -1;
+    if (!proc || (low > high) || !range_out || !len) {
+        return -EINVAL;
+    }
 
     if (low == high) {
         *range_out = NULL;
@@ -126,12 +129,13 @@ int pm_process_pagemap_range(pm_process_t *proc,
     numpages = (high - low) / proc->ker->pagesize;
 
     range = malloc(numpages * sizeof(uint64_t));
-    if (!range)
-        return errno;
+    if (!range) {
+        return -ENOMEM;
+    }
 
     off = lseek64(proc->pagemap_fd, firstpage * sizeof(uint64_t), SEEK_SET);
     if (off == (off_t)-1) {
-        error = errno;
+        error = -errno;
         free(range);
         return error;
     }
@@ -143,7 +147,7 @@ int pm_process_pagemap_range(pm_process_t *proc,
         *range_out = NULL;
         return 0;
     } else if (error < 0 || (error > 0 && error < (int)(numpages * sizeof(uint64_t)))) {
-        error = (error < 0) ? errno : -1;
+        error = (error < 0) ? -errno : -EIO;
         free(range);
         return error;
     }
@@ -179,8 +183,6 @@ int pm_process_maps(pm_process_t *proc, pm_map_t ***maps_out, size_t *len) {
 int pm_process_workingset(pm_process_t *proc,
                           pm_memusage_t *ws_out, int reset) {
     pm_memusage_t ws, map_ws;
-    char filename[MAX_FILENAME];
-    int fd;
     int i;
     int error;
 
@@ -200,24 +202,12 @@ int pm_process_workingset(pm_process_t *proc,
 
             pm_memusage_add(&ws, &map_ws);
         }
-        
         memcpy(ws_out, &ws, sizeof(ws));
     }
 
     if (reset) {
-        error = snprintf(filename, MAX_FILENAME, "/proc/%d/clear_refs",
-                         proc->pid);
-        if (error < 0 || error >= MAX_FILENAME) {
-            return (error < 0) ? (errno) : (-1);
-        }
-
-        fd = open(filename, O_WRONLY);
-        if (fd < 0)
-            return errno;
-
-        write(fd, "1\n", strlen("1\n"));
-
-        close(fd);
+        return pm_kernel_has_page_idle(proc->ker) ? pm_process_mark_idle(proc)
+                                                  : pm_process_clear_refs(proc);
     }
 
     return 0;
@@ -329,3 +319,49 @@ static int read_maps(pm_process_t *proc) {
 
     return 0;
 }
+
+static int pm_process_clear_refs(pm_process_t* proc) {
+    int error;
+    char filename[MAX_FILENAME];
+
+    if (!proc) {
+        return -EINVAL;
+    }
+
+    error = snprintf(filename, MAX_FILENAME, "/proc/%d/clear_refs", proc->pid);
+    if (error < 0 || error >= MAX_FILENAME) {
+        return (error < 0) ? -errno : -1;
+    }
+
+    int fd = open(filename, O_WRONLY | O_CLOEXEC);
+    if (fd < 0) {
+        return -errno;
+    }
+
+    const char* cmd = "1\n";
+    size_t cmdlen = strlen(cmd);
+
+    error = TEMP_FAILURE_RETRY(write(fd, cmd, cmdlen));
+    error = (error == cmdlen) ? 0 : -errno;
+
+    close(fd);
+
+    return error;
+}
+
+static int pm_process_mark_idle(pm_process_t* proc) {
+    int error;
+
+    if (!proc) {
+        return -EINVAL;
+    }
+
+    for (int i = 0; i < proc->num_maps; i++) {
+        error = pm_map_mark_idle(proc->maps[i]);
+        if (error) {
+            return error;
+        }
+    }
+
+    return 0;
+}
diff --git a/libpagemap/procmem.cpp b/libpagemap/procmem.cpp
index 9484a6e1..7e755ae5 100644
--- a/libpagemap/procmem.cpp
+++ b/libpagemap/procmem.cpp
@@ -70,6 +70,9 @@ int main(int argc, char *argv[]) {
     int (*compfn)(const void *a, const void *b);
     int hide_zeros;
 
+    /* Use kernel's idle page tracking interface for working set determination */
+    int use_pageidle;
+
     /* temporary variables */
     size_t i, j;
     char *endptr;
@@ -83,9 +86,14 @@ int main(int argc, char *argv[]) {
     ws = WS_OFF;
     compfn = NULL;
     hide_zeros = 0;
+    use_pageidle = 0;
     for (i = 1; i < (size_t)(argc - 1); i++) {
         if (!strcmp(argv[i], "-w")) { ws = WS_ONLY; continue; }
         if (!strcmp(argv[i], "-W")) { ws = WS_RESET; continue; }
+        if (!strcmp(argv[i], "-i")) {
+            use_pageidle = 1;
+            continue;
+        }
         if (!strcmp(argv[i], "-m")) { compfn = NULL; continue; }
         if (!strcmp(argv[i], "-p")) { compfn = &comp_pss; continue; }
         if (!strcmp(argv[i], "-h")) { hide_zeros = 1; continue; }
@@ -107,6 +115,16 @@ int main(int argc, char *argv[]) {
         exit(EXIT_FAILURE);
     }
 
+    if (ws != WS_OFF && use_pageidle) {
+        error = pm_kernel_init_page_idle(ker);
+        if (error) {
+            fprintf(stderr,
+                    "error initalizing idle page tracking -- "
+                    "enable CONFIG_IDLE_PAGE_TRACKING in kernel.\n");
+            exit(EXIT_FAILURE);
+        }
+    }
+
     pagesize = pm_kernel_pagesize(ker);
 
     error = pm_process_create(ker, pid, &proc);
@@ -206,7 +224,8 @@ int main(int argc, char *argv[]) {
                     fprintf(stderr, "error getting flags for frame.\n");
                 }
 
-                if ((ws != WS_ONLY) || (flags & (1 << KPF_REFERENCED))) {
+                if ((ws != WS_ONLY) ||
+                    pm_kernel_page_is_accessed(ker, PM_PAGEMAP_PFN(mapentry), &flags)) {
                     if (count > 1) {
                         if (flags & (1 << KPF_DIRTY)) {
                             mi->shared_dirty++;
@@ -304,13 +323,15 @@ int main(int argc, char *argv[]) {
 }
 
 static void usage(const char *cmd) {
-    fprintf(stderr, "Usage: %s [ -w | -W ] [ -p | -m ] [ -h ] pid\n"
-                    "    -w  Displays statistics for the working set only.\n"
-                    "    -W  Resets the working set of the process.\n"
-                    "    -p  Sort by PSS.\n"
-                    "    -m  Sort by mapping order (as read from /proc).\n"
-                    "    -h  Hide maps with no RSS.\n",
-        cmd);
+    fprintf(stderr,
+            "Usage: %s [-i] [ -w | -W ] [ -p | -m ] [ -h ] pid\n"
+            "    -i  Uses idle page tracking for working set statistics.\n"
+            "    -w  Displays statistics for the working set only.\n"
+            "    -W  Resets the working set of the process.\n"
+            "    -p  Sort by PSS.\n"
+            "    -m  Sort by mapping order (as read from /proc).\n"
+            "    -h  Hide maps with no RSS.\n",
+            cmd);
 }
 
 int comp_pss(const void *a, const void *b) {