From cd27987b5d4153ad75586cc3ed67cf879e53d046 Mon Sep 17 00:00:00 2001
From: Marius Tolzmann <tolzmann@molgen.mpg.de>
Date: Thu, 22 Oct 2015 21:28:45 +0200
Subject: [PATCH] mx_proc: Add process tree related functions

---
 mx_proc.c | 350 +++++++++++++++++++++++++++++++++++++++++++++++++++++-
 mx_proc.h |  31 +++++
 2 files changed, 377 insertions(+), 4 deletions(-)

diff --git a/mx_proc.c b/mx_proc.c
index 2af1a5da..611cda8b 100644
--- a/mx_proc.c
+++ b/mx_proc.c
@@ -2,6 +2,8 @@
 #include <stdio.h>
 #include <errno.h>
 #include <assert.h>
+#include <dirent.h>
+#include <ctype.h>
 
 #include "mx_util.h"
 #include "mx_proc.h"
@@ -81,21 +83,57 @@ static int _mx_proc_pid_stat_strscan(char *str, struct mx_proc_pid_stat *pps)
 }
 
 int mx_proc_pid_stat(struct mx_proc_pid_stat **pps, pid_t pid)
+{
+    struct mx_proc_pid_stat *pstat;
+    int res;
+
+    pstat = *pps;
+    if (!pstat)
+        pstat = mx_calloc_forever(1, sizeof(*pstat));
+
+    res = mx_proc_pid_stat_read(pstat, "/proc/%d/stat", pid);
+    if (res < 0)
+        return res;
+
+    *pps = pstat;
+    return 0;
+}
+
+int mx_proc_pid_task_tid_stat(struct mx_proc_pid_stat **pps, pid_t pid, pid_t tid)
+{
+    struct mx_proc_pid_stat *pstat;
+    int res;
+
+    pstat = *pps;
+    if (!pstat)
+        pstat = mx_calloc_forever(1, sizeof(*pstat));
+
+    res = mx_proc_pid_stat_read(pstat, "/proc/%d/task/%d/stat", pid, tid);
+    if (res < 0)
+        return res;
+
+    *pps = pstat;
+    return 0;
+}
+
+int mx_proc_pid_stat_read(struct mx_proc_pid_stat *pps, char *fmt, ...)
 {
     _mx_cleanup_free_ char *fname = NULL;
     _mx_cleanup_free_ char *line = NULL;
+    va_list ap;
     int res;
 
-    mx_asprintf_forever(&fname, "/proc/%d/stat", pid);
+    assert(pps);
 
-    if (!*pps)
-        *pps = mx_calloc_forever(1, sizeof(**pps));
+    va_start(ap, fmt);
+    mx_vasprintf_forever(&fname, fmt, ap);
+    va_end(ap);
 
     res = mx_read_first_line_from_file(fname, &line);
     if (res < 0)
         return res;
 
-    res = _mx_proc_pid_stat_strscan(line, *pps);
+    res = _mx_proc_pid_stat_strscan(line, pps);
     if (res < 0)
         return res;
 
@@ -109,3 +147,307 @@ void mx_proc_pid_stat_free_content(struct mx_proc_pid_stat *pps)
 
     mx_free_null(pps->comm);
 }
+
+static void mx_proc_tree_update_parent_pinfo(struct mx_proc_tree_node *this, struct mx_proc_info *pinfo)
+{
+    if (!this)
+        return;
+
+    this->pinfo.sum_rss += pinfo->sum_rss;
+
+    mx_proc_tree_update_parent_pinfo(this->parent, pinfo);
+}
+
+static void mx_proc_tree_add_to_list_sorted(struct mx_proc_tree_node **ptn_ptr, struct mx_proc_tree_node *new)
+{
+    struct mx_proc_tree_node *current;
+
+    assert(new);
+    assert(new->pinfo.pstat);
+    assert(!new->next);
+    assert(new->pinfo.pstat->pid > 0);
+
+    current = *ptn_ptr;
+
+    /* update stats */
+    if (new->parent) {
+        new->parent->nchilds++;
+        mx_proc_tree_update_parent_pinfo(new->parent, &(new->pinfo));
+    }
+
+    /* empty list? -> start new list */
+    if (!current) {
+        *ptn_ptr = new;
+        return;
+    }
+
+    /* new is first entry */
+    if (new->pinfo.pstat->pid < current->pinfo.pstat->pid) {
+        new->next = current;
+        *ptn_ptr = new;
+        return;
+    }
+
+    /* find position */
+    while (1) {
+        assert(new->pinfo.pstat->pid > current->pinfo.pstat->pid);
+
+        /* new is last entry */
+        if (!current->next) {
+            current->next = new;
+            break;
+        }
+
+        assert(current->next->pinfo.pstat->pid > current->pinfo.pstat->pid);
+
+        /* add new between current and current->next */
+        if (new->pinfo.pstat->pid < current->next->pinfo.pstat->pid) {
+            new->next = current->next;
+            current->next = new;
+            break;
+        }
+
+        current = current->next;
+    }
+
+    return;
+}
+
+static struct mx_proc_tree_node *mx_proc_tree_find_by_pid(struct mx_proc_tree_node *ptn, long long int pid)
+{
+    assert(ptn);
+    assert(pid >= 0);
+
+    struct mx_proc_tree_node *current;
+    struct mx_proc_tree_node *node;
+
+    if (pid == 0)
+        return NULL;
+
+    current = ptn;
+
+    for (current = ptn; current; current=current->next) {
+        if (current->pinfo.pstat->pid == pid)
+            return current;
+
+        if (!current->childs)
+            continue;
+
+        node = mx_proc_tree_find_by_pid(current->childs, pid);
+        if (node)
+            return node;
+    }
+
+    return NULL;
+}
+
+#define ppid_or_pgrp(x) (((x)->ppid != 1 || (x)->pid == (x)->pgrp) ? (x)->ppid : (x)->pgrp)
+
+static struct mx_proc_tree_node *mx_proc_tree_add(struct mx_proc_tree *pt, struct mx_proc_pid_stat *pps)
+{
+    assert(pps);
+    assert(pt);
+    struct mx_proc_tree_node *new;
+    struct mx_proc_tree_node *current;
+    struct mx_proc_tree_node *next;
+    struct mx_proc_tree_node *parent;
+
+    new = mx_calloc_forever(1, sizeof(*new));
+
+    pt->nentries++;
+
+    new->pinfo.pstat   = pps;
+    new->pinfo.sum_rss = pps->rss;
+
+    if (!(pt->root)) {
+        pt->root = new;
+        return new;
+    }
+
+    assert(pt->root);
+
+    /* new is second to last roots parent? -> collect */
+    current = pt->root;
+    while (current->next) {
+        if (current->next->pinfo.pstat->ppid != new->pinfo.pstat->pid) {
+            current = current->next;
+            continue;
+        }
+        assert(current->next->pinfo.pstat->ppid == new->pinfo.pstat->pid);
+
+        /* disconnect next */
+        next          = current->next;
+        current->next = current->next->next;
+        next->next    = NULL;
+
+        /* add as child of new */
+        next->parent = new;
+        mx_proc_tree_add_to_list_sorted(&new->childs, next);
+    }
+
+    /* new is first roots parent? -> new is new root */
+    if (pt->root->pinfo.pstat->ppid == new->pinfo.pstat->pid) {
+        assert(!new->next);
+
+        current = pt->root;
+        pt->root = pt->root->next;
+
+        current->next = NULL;
+        current->parent = new;
+
+        mx_proc_tree_add_to_list_sorted(&new->childs, current);
+
+        if (!(pt->root)) {
+            pt->root = new;
+            return new;
+        }
+    }
+
+
+    parent = mx_proc_tree_find_by_pid(pt->root, new->pinfo.pstat->ppid);
+    if (parent) {
+        new->parent = parent;
+        mx_proc_tree_add_to_list_sorted(&parent->childs, new);
+    } else {
+        mx_proc_tree_add_to_list_sorted(&pt->root, new);
+    }
+
+    return new;
+}
+
+static void mx_proc_tree_reorder_roots(struct mx_proc_tree *pt)
+{
+    struct mx_proc_tree_node *current;
+    struct mx_proc_tree_node *pid1;
+    struct mx_proc_tree_node *last = NULL;
+    struct mx_proc_tree_node *next = NULL;
+
+    for (current = pt->root; current; current = current->next) {
+        if (current->pinfo.pstat->pid == 1) {
+            pid1 = current;
+            break;
+        }
+    }
+
+    if (!pid1)
+        return;
+
+    for (current = pt->root; current; current = next) {
+        next = current->next;
+
+        if (current->pinfo.pstat->ppid != 1) {
+            last = current;
+            continue;
+        }
+
+        if (!last) {
+            if (!current->next)
+                return;
+            pt->root = current->next;
+        } else {
+            last->next = current->next;
+        }
+        current->next = NULL;
+        current->parent = pid1;
+        mx_proc_tree_add_to_list_sorted(&pid1->childs, current);
+    }
+}
+
+static int _mx_filter_numbers(const struct dirent *d)
+{
+    if (!isdigit(d->d_name[0]))
+        return 0;
+
+    return 1;
+}
+
+int mx_proc_tree(struct mx_proc_tree **newtree)
+{
+    struct mx_proc_tree *pt;
+    struct dirent **namelist = NULL;
+    struct mx_proc_pid_stat *pps;
+    int n;
+    int i;
+    int res;
+    unsigned long long int pid;
+
+    assert(*newtree == NULL);
+
+    pt = mx_calloc_forever(1, sizeof(*pt));
+
+    n = scandir("/proc", &namelist, _mx_filter_numbers, NULL);
+    if (n < 0)
+        return -errno;
+
+    if (n == 0)
+        return -(errno=ENOENT);
+
+    for (i=0; i < n; i++) {
+        res = mx_strtoull(namelist[i]->d_name, &pid);
+        free(namelist[i]);
+        if (res < 0)
+            continue;
+
+        pps = NULL;
+        res = mx_proc_pid_stat(&pps, pid);
+        if (res < 0)
+            continue;
+
+        mx_proc_tree_add(pt, pps);
+    }
+    free(namelist);
+
+    mx_proc_tree_reorder_roots(pt);
+
+    *newtree = pt;
+    return 0;
+}
+
+static void _mx_proc_tree_node_free_recursive(struct mx_proc_tree_node *ptn)
+{
+    assert(ptn);
+
+    struct mx_proc_tree_node *current;
+    struct mx_proc_tree_node *next;
+
+    for (current = ptn; current; current=next) {
+
+        if (current->childs)
+            _mx_proc_tree_node_free_recursive(current->childs);
+
+        next = current->next;
+
+        mx_proc_pid_stat_free_content(current->pinfo.pstat);
+        mx_free_null(current->pinfo.pstat);
+        mx_free_null(current);
+    }
+
+    return;
+}
+
+int mx_proc_tree_free(struct mx_proc_tree **tree)
+{
+    struct mx_proc_tree *pt;
+
+    pt = *tree;
+
+    _mx_proc_tree_node_free_recursive(pt->root);
+
+    mx_free_null(*tree);
+
+    return 0;
+}
+
+struct mx_proc_info *mx_proc_tree_proc_info(struct mx_proc_tree *tree, pid_t pid)
+{
+    struct mx_proc_tree_node *ptn;
+
+    assert(tree);
+
+    ptn = mx_proc_tree_find_by_pid(tree->root, pid);
+
+    if (!ptn)
+        return NULL;
+
+    return &(ptn->pinfo);
+}
diff --git a/mx_proc.h b/mx_proc.h
index 4a0efb8e..9f16806e 100644
--- a/mx_proc.h
+++ b/mx_proc.h
@@ -3,6 +3,29 @@
 
 #include <sys/types.h>
 
+struct mx_proc_info {
+    struct mx_proc_pid_stat *pstat;
+
+    unsigned long long int sum_rss;
+
+    char **environment;
+};
+
+struct mx_proc_tree {
+    struct mx_proc_tree_node *root;
+    int nentries;
+};
+
+struct mx_proc_tree_node {
+    struct mx_proc_tree_node *parent;
+    struct mx_proc_tree_node *next;
+
+    struct mx_proc_info pinfo;
+
+    unsigned long long int nchilds;
+    struct mx_proc_tree_node *childs;
+};
+
 struct mx_proc_pid_stat {
     long long int pid;     /* 1 */
     char *comm;            /* 2 (comm) */
@@ -51,8 +74,16 @@ struct mx_proc_pid_stat {
 };
 
 int mx_proc_pid_stat_read(struct mx_proc_pid_stat *pps, char *fmt, ...);
+
 int mx_proc_pid_stat(struct mx_proc_pid_stat **pps, pid_t pid);
+int mx_proc_pid_task_tid_stat(struct mx_proc_pid_stat **pps, pid_t pid, pid_t tid);
 
 void mx_proc_pid_stat_free_content(struct mx_proc_pid_stat *pps);
 
+int mx_proc_tree(struct mx_proc_tree **newtree);
+
+int mx_proc_tree_free(struct mx_proc_tree **tree);
+
+struct mx_proc_info *mx_proc_tree_proc_info(struct mx_proc_tree *tree, pid_t pid);
+
 #endif