2 files changed, 66 insertions, 17 deletions

diff --git a/include/linux/cgroup.h b/include/linux/cgroup.h
index b7dd23040cd5..79faa6467f76 100644
--- a/include/linux/cgroup.h
+++ b/include/linux/cgroup.h
@@ -38,6 +38,8 @@
 
 /* walk only threadgroup leaders */
 #define CSS_TASK_ITER_PROCS             (1U << 0)
+/* walk all threaded css_sets in the domain */
+#define CSS_TASK_ITER_THREADED          (1U << 1)
 
 /* a css_task_iter should be treated as an opaque object */
 struct css_task_iter {
@@ -47,11 +49,15 @@ struct css_task_iter {
         struct list_head                *cset_pos;
         struct list_head                *cset_head;
 
+        struct list_head                *tcset_pos;
+        struct list_head                *tcset_head;
+
         struct list_head                *task_pos;
         struct list_head                *tasks_head;
         struct list_head                *mg_tasks_head;
 
         struct css_set                  *cur_cset;
+        struct css_set                  *cur_dcset;
         struct task_struct              *cur_task;
         struct list_head                iters_node;     /* css_set->task_iters */
 };
diff --git a/kernel/cgroup/cgroup.c b/kernel/cgroup/cgroup.c
index c7e1c243b77d..a1d59af274a9 100644
--- a/kernel/cgroup/cgroup.c
+++ b/kernel/cgroup/cgroup.c
@@ -3629,6 +3629,58 @@ bool css_has_online_children(struct cgroup_subsys_state *css)
         return ret;
 }
 
+static struct css_set *css_task_iter_next_css_set(struct css_task_iter *it)
+{
+        struct list_head *l;
+        struct cgrp_cset_link *link;
+        struct css_set *cset;
+
+        lockdep_assert_held(&css_set_lock);
+
+        /* find the next threaded cset */
+        if (it->tcset_pos) {
+                l = it->tcset_pos->next;
+
+                if (l != it->tcset_head) {
+                        it->tcset_pos = l;
+                        return container_of(l, struct css_set,
+                                            threaded_csets_node);
+                }
+
+                it->tcset_pos = NULL;
+        }
+
+        /* find the next cset */
+        l = it->cset_pos;
+        l = l->next;
+        if (l == it->cset_head) {
+                it->cset_pos = NULL;
+                return NULL;
+        }
+
+        if (it->ss) {
+                cset = container_of(l, struct css_set, e_cset_node[it->ss->id]);
+        } else {
+                link = list_entry(l, struct cgrp_cset_link, cset_link);
+                cset = link->cset;
+        }
+
+        it->cset_pos = l;
+
+        /* initialize threaded css_set walking */
+        if (it->flags & CSS_TASK_ITER_THREADED) {
+                if (it->cur_dcset)
+                        put_css_set_locked(it->cur_dcset);
+                it->cur_dcset = cset;
+                get_css_set(cset);
+
+                it->tcset_head = &cset->threaded_csets;
+                it->tcset_pos = &cset->threaded_csets;
+        }
+
+        return cset;
+}
+
 /**
  * css_task_iter_advance_css_set - advance a task itererator to the next css_set
  * @it: the iterator to advance
@@ -3637,32 +3689,19 @@ bool css_has_online_children(struct cgroup_subsys_state *css)
  */
 static void css_task_iter_advance_css_set(struct css_task_iter *it)
 {
-        struct list_head *l = it->cset_pos;
-        struct cgrp_cset_link *link;
         struct css_set *cset;
 
         lockdep_assert_held(&css_set_lock);
 
         /* Advance to the next non-empty css_set */
         do {
-                l = l->next;
-                if (l == it->cset_head) {
-                        it->cset_pos = NULL;
+                cset = css_task_iter_next_css_set(it);
+                if (!cset) {
                         it->task_pos = NULL;
                         return;
                 }
-
-                if (it->ss) {
-                        cset = container_of(l, struct css_set,
-                                            e_cset_node[it->ss->id]);
-                } else {
-                        link = list_entry(l, struct cgrp_cset_link, cset_link);
-                        cset = link->cset;
-                }
         } while (!css_set_populated(cset));
 
-        it->cset_pos = l;
-
         if (!list_empty(&cset->tasks))
                 it->task_pos = cset->tasks.next;
         else
@@ -3805,6 +3844,9 @@ void css_task_iter_end(struct css_task_iter *it)
                 spin_unlock_irq(&css_set_lock);
         }
 
+        if (it->cur_dcset)
+                put_css_set(it->cur_dcset);
+
         if (it->cur_task)
                 put_task_struct(it->cur_task);
 }
@@ -3830,6 +3872,7 @@ static void *cgroup_procs_start(struct seq_file *s, loff_t *pos)
         struct kernfs_open_file *of = s->private;
         struct cgroup *cgrp = seq_css(s)->cgroup;
         struct css_task_iter *it = of->priv;
+        unsigned iter_flags = CSS_TASK_ITER_PROCS | CSS_TASK_ITER_THREADED;
 
         /*
          * When a seq_file is seeked, it's always traversed sequentially
@@ -3843,10 +3886,10 @@ static void *cgroup_procs_start(struct seq_file *s, loff_t *pos)
                 if (!it)
                         return ERR_PTR(-ENOMEM);
                 of->priv = it;
-                css_task_iter_start(&cgrp->self, CSS_TASK_ITER_PROCS, it);
+                css_task_iter_start(&cgrp->self, iter_flags, it);
         } else if (!(*pos)++) {
                 css_task_iter_end(it);
-                css_task_iter_start(&cgrp->self, CSS_TASK_ITER_PROCS, it);
+                css_task_iter_start(&cgrp->self, iter_flags, it);
         }
 
         return cgroup_procs_next(s, NULL, NULL);