mm: process_reclaim: do not iterate over stale task structs

swap_fn iterates through the threads of selected tasks after
a rcu_read_unlock which is wrong. But we can't extend the
rcu_read_lock since it will result in severe performance
issues. So better avoid iterating over the threads. Just
lock the group leader and use it further.

Change-Id: I36269b1b6619315f33f6f3b49ec73571a66796f2
Signed-off-by: Vinayak Menon <vinmenon@codeaurora.org>
diff --git a/mm/process_reclaim.c b/mm/process_reclaim.c
index a92daf8..cae249f 100644
--- a/mm/process_reclaim.c
+++ b/mm/process_reclaim.c
@@ -86,14 +86,17 @@
 {
 	struct task_struct *t = p;
 
-	do {
+	rcu_read_lock();
+	for_each_thread(p, t) {
 		task_lock(t);
 		if (test_tsk_thread_flag(t, flag)) {
 			task_unlock(t);
+			rcu_read_unlock();
 			return 1;
 		}
 		task_unlock(t);
-	} while_each_thread(p, t);
+	}
+	rcu_read_unlock();
 
 	return 0;
 }
@@ -123,10 +126,6 @@
 		if (tsk->flags & PF_KTHREAD)
 			continue;
 
-		/* if task no longer has any memory ignore it */
-		if (test_task_flag(tsk, TIF_MM_RELEASED))
-			continue;
-
 		if (test_task_flag(tsk, TIF_MEMDIE))
 			continue;
 
@@ -163,20 +162,20 @@
 		}
 	}
 
-	for (i = 0; i < si; i++) {
-		get_task_struct(selected[i].p);
+	for (i = 0; i < si; i++)
 		total_sz += selected[i].tasksize;
-	}
-
-	rcu_read_unlock();
 
 	/* Skip reclaim if total size is too less */
 	if (total_sz < SWAP_CLUSTER_MAX) {
-		for (i = 0; i < si; i++)
-			put_task_struct(selected[i].p);
+		rcu_read_unlock();
 		return;
 	}
 
+	for (i = 0; i < si; i++)
+		get_task_struct(selected[i].p);
+
+	rcu_read_unlock();
+
 	while (si--) {
 		nr_to_reclaim =
 			(selected[si].tasksize * per_swap_size) / total_sz;
@@ -184,12 +183,6 @@
 		if (!nr_to_reclaim)
 			nr_to_reclaim = 1;
 
-		if ((test_task_flag(selected[si].p, TIF_MM_RELEASED))
-			|| (test_task_flag(selected[si].p, TIF_MEMDIE))) {
-			put_task_struct(selected[si].p);
-			continue;
-		}
-
 		rp = reclaim_task_anon(selected[si].p, nr_to_reclaim);
 
 		trace_process_reclaim(selected[si].tasksize,