staging: android: lowmemorykiller: implement task's adj rbtree

Based on the current LMK implementation, LMK has to scan all processes to
select the correct task to kill during low memory.
The basic idea for the optimization is to :
queue all tasks with oom_score_adj priority, and then LMK just selects the
proper task from the queue(rbtree) to kill.

performance improvement:
the current implementation: average time to find a task to kill : 1004us
the optimized implementation: average time to find a task to kill: 43us

squished another 3 commits:

staging: android: lowmemorykiller: select a new task to kill
Under certain circumstances, a process may take time to handle a SIGKILL.
When lowmemkiller is called again shortly after, it would pick the same
process to kill over and over, so that we cann't get free memory for long time.

Solution is to check fatal_signal_pending() on the selected task, and if it's
already pending, select a new task to kill.

staging: android: lmk: skip if killed by lmk
A task can be selected to kill when it's already exiting. Sending
signal to the task at that state won't set TIF_SIGPENDING for it.
After that, every task including threads in the exiting task itself
running into lmk will compete for the mutex, then spend some time
waiting for the exiting task to release memory or try to send signal
to the same task again.
We want the exiting task to exit as early as possible but it's slowed
down greatly by the mutex and waiting. Skip lmk if current thread
group leader is exiting and is killed by lmk. Also delegate TIF_MEMDIE
to current task.

staging: android: lmk: check TIF directly
When trying to fix that lmk can't kill an exiting task
(commit 895304e
(CR) staging: android: lmk: skip if killed by lmk),
additional check was added for the TIF_MEMDIE flag.
However the task lock has already been taken when examining
the task. The same helper function which tries to grab the
task lock again will be deadlocked.
Change to check the thread flag of thread group leader directly.

Change-Id: Ibbb99cbd36e27a546195b4d4bed59b3829e6b7b0

Author: alin23

drivers/staging/android/lowmemorykiller.c

@@ -1,16 +1,16 @@
 /* drivers/misc/lowmemorykiller.c
  *
  * The lowmemorykiller driver lets user-space specify a set of memory thresholds
- * where processes with a range of oom_score_adj values will get killed. Specify
- * the minimum oom_score_adj values in
+ * where processes with a range of oom_adj values will get killed. Specify
+ * the minimum oom_adj values in
  * /sys/module/lowmemorykiller/parameters/adj and the number of free pages in
  * /sys/module/lowmemorykiller/parameters/minfree. Both files take a comma
  * separated list of numbers in ascending order.
  *
  * For example, write "0,8" to /sys/module/lowmemorykiller/parameters/adj and
  * "1024,4096" to /sys/module/lowmemorykiller/parameters/minfree to kill
- * processes with a oom_score_adj value of 8 or higher when the free memory
- * drops below 4096 pages and kill processes with a oom_score_adj value of 0 or
+ * processes with a oom_adj value of 8 or higher when the free memory
+ * drops below 4096 pages and kill processes with a oom_adj value of 0 or
  * higher when the free memory drops below 1024 pages.
  *
  * The driver considers memory used for caches to be free, but if a large
@@ -56,6 +56,7 @@ static int lowmem_adj[6] = {
 	1,
 	6,
 	12,
+	15,
 };
 static int lowmem_adj_size = 4;
 static int lowmem_minfree[6] = {
@@ -149,6 +150,10 @@ void tune_lmk_zone_param(struct zonelist *zonelist, int classzone_idx,
 	}
 }
 
+static struct task_struct *pick_next_from_adj_tree(struct task_struct *task);
+static struct task_struct *pick_first_task(void);
+static struct task_struct *pick_last_task(void);
+
 void tune_lmk_param(int *other_free, int *other_file, gfp_t gfp_mask)
 {
 	struct zone *preferred_zone;
@@ -202,13 +207,19 @@ static int lowmem_shrink(struct shrinker *s, int nr_to_scan, gfp_t gfp_mask)
 	int rem = 0;
 	int tasksize;
 	int i;
-	int min_score_adj = OOM_SCORE_ADJ_MAX + 1;
+	int min_adj = OOM_SCORE_ADJ_MAX + 1;
 	int selected_tasksize = 0;
-	int selected_oom_score_adj;
+	int selected_oom_adj;
 	int array_size = ARRAY_SIZE(lowmem_adj);
 	int other_free;
 	int other_file;
-
+	
+	tsk = current->group_leader;
+	if ((tsk->flags & PF_EXITING) && test_task_flag(tsk, TIF_MEMDIE)) {
+		set_tsk_thread_flag(current, TIF_MEMDIE);
+		return 0;
+	}
+	
 	if (nr_to_scan > 0) {
 		if (mutex_lock_interruptible(&scan_mutex) < 0)
 			return 0;
@@ -227,19 +238,19 @@ static int lowmem_shrink(struct shrinker *s, int nr_to_scan, gfp_t gfp_mask)
 	for (i = 0; i < array_size; i++) {
 		if (other_free < lowmem_minfree[i] &&
 		    other_file < lowmem_minfree[i]) {
-			min_score_adj = lowmem_adj[i];
+			min_adj = lowmem_adj[i];
 			break;
 		}
 	}
 	if (nr_to_scan > 0)
 		lowmem_print(3, "lowmem_shrink %d, %x, ofree %d %d, ma %d\n",
 				nr_to_scan, gfp_mask, other_free,
-				other_file, min_score_adj);
+				other_file, min_adj);
 	rem = global_page_state(NR_ACTIVE_ANON) +
 		global_page_state(NR_ACTIVE_FILE) +
 		global_page_state(NR_INACTIVE_ANON) +
 		global_page_state(NR_INACTIVE_FILE);
-	if (nr_to_scan <= 0 || min_score_adj == OOM_SCORE_ADJ_MAX + 1) {
+	if (nr_to_scan <= 0 || min_adj == OOM_SCORE_ADJ_MAX + 1) {
 		lowmem_print(5, "lowmem_shrink %d, %x, return %d\n",
 			     nr_to_scan, gfp_mask, rem);
 
@@ -248,12 +259,14 @@ static int lowmem_shrink(struct shrinker *s, int nr_to_scan, gfp_t gfp_mask)
 
 		return rem;
 	}
-	selected_oom_score_adj = min_score_adj;
+	selected_oom_adj = min_adj;
 
 	rcu_read_lock();
-	for_each_process(tsk) {
+	for (tsk = pick_first_task();
+		tsk != pick_last_task();
+		tsk = pick_next_from_adj_tree(tsk)) {
 		struct task_struct *p;
-		int oom_score_adj;
+		int oom_adj;
 
 		if (tsk->flags & PF_KTHREAD)
 			continue;
@@ -266,7 +279,10 @@ static int lowmem_shrink(struct shrinker *s, int nr_to_scan, gfp_t gfp_mask)
 			if (test_task_flag(tsk, TIF_MEMDIE)) {
 				rcu_read_unlock();
 				/* give the system time to free up the memory */
-				msleep_interruptible(20);
+				if (!same_thread_group(current, tsk))
+					msleep_interruptible(20);
+				else
+					set_tsk_thread_flag(current, TIF_MEMDIE);
 				mutex_unlock(&scan_mutex);
 				return 0;
 			}
@@ -276,32 +292,42 @@ static int lowmem_shrink(struct shrinker *s, int nr_to_scan, gfp_t gfp_mask)
 		if (!p)
 			continue;
 
-		oom_score_adj = p->signal->oom_adj;
-		if (oom_score_adj < min_score_adj) {
+		oom_adj = p->signal->oom_adj;
+		if (oom_adj < min_adj) {
+			task_unlock(p);
+			break;
+		}
+		
+		
+		if (fatal_signal_pending(p) ||
+			((p->flags & PF_EXITING) &&
+			test_tsk_thread_flag(p, TIF_MEMDIE))) {
+			lowmem_print(2, "skip slow dying process %d\n", p->pid);
 			task_unlock(p);
 			continue;
 		}
+		
 		tasksize = get_mm_rss(p->mm);
 		task_unlock(p);
 		if (tasksize <= 0)
 			continue;
 		if (selected) {
-			if (oom_score_adj < selected_oom_score_adj)
-				continue;
-			if (oom_score_adj == selected_oom_score_adj &&
+			if (oom_adj < selected_oom_adj)
+				break;
+			if (oom_adj == selected_oom_adj &&
 			    tasksize <= selected_tasksize)
 				continue;
 		}
 		selected = p;
 		selected_tasksize = tasksize;
-		selected_oom_score_adj = oom_score_adj;
+		selected_oom_adj = oom_adj;
 		lowmem_print(2, "select %d (%s), adj %d, size %d, to kill\n",
-			     p->pid, p->comm, oom_score_adj, tasksize);
+			     p->pid, p->comm, oom_adj, tasksize);
 	}
 	if (selected) {
 		lowmem_print(1, "send sigkill to %d (%s), adj %d, size %d\n",
 			     selected->pid, selected->comm,
-			     selected_oom_score_adj, selected_tasksize);
+			     selected_oom_adj, selected_tasksize);
 		lowmem_deathpending_timeout = jiffies + HZ;
 		send_sig(SIGKILL, selected, 0);
 		set_tsk_thread_flag(selected, TIF_MEMDIE);
@@ -334,6 +360,81 @@ static void __exit lowmem_exit(void)
 	unregister_shrinker(&lowmem_shrinker);
 }
 
+DEFINE_SPINLOCK(lmk_lock);
+struct rb_root tasksadj = RB_ROOT;
+void add_2_adj_tree(struct task_struct *task)
+{		struct rb_node **link = &tasksadj.rb_node;
+struct rb_node *parent = NULL;
+struct task_struct *task_entry;
+s64 key = task->signal->oom_adj;
+/*
+ * Find the right place in the rbtree:
+ */
+spin_lock(&lmk_lock);
+while (*link) {
+	parent = *link;
+	task_entry = rb_entry(parent, struct task_struct, adj_node);
+	if (key < task_entry->signal->oom_adj)
+		link = &parent->rb_right;
+	else
+		link = &parent->rb_left;
+}
+
+rb_link_node(&task->adj_node, parent, link);
+rb_insert_color(&task->adj_node, &tasksadj);
+spin_unlock(&lmk_lock);
+}
+
+void delete_from_adj_tree(struct task_struct *task)
+{
+	spin_lock(&lmk_lock);
+	rb_erase(&task->adj_node, &tasksadj);
+	spin_unlock(&lmk_lock);
+}
+
+
+static struct task_struct *pick_next_from_adj_tree(struct task_struct *task)
+{
+	struct rb_node *next;
+	
+	spin_lock(&lmk_lock);
+	next = rb_next(&task->adj_node);
+	spin_unlock(&lmk_lock);
+	
+	if (!next)
+		return NULL;
+	
+	return rb_entry(next, struct task_struct, adj_node);
+}
+
+static struct task_struct *pick_first_task(void)
+{
+	struct rb_node *left;
+	
+	spin_lock(&lmk_lock);
+	left = rb_first(&tasksadj);
+	spin_unlock(&lmk_lock);
+	
+	if (!left)
+		return NULL;
+	
+	return rb_entry(left, struct task_struct, adj_node);
+}
+
+static struct task_struct *pick_last_task(void)
+{
+	struct rb_node *right;
+	
+	spin_lock(&lmk_lock);
+	right = rb_last(&tasksadj);
+	spin_unlock(&lmk_lock);
+	
+	if (!right)
+		return NULL;
+	
+	return rb_entry(right, struct task_struct, adj_node);
+}
+
 module_param_named(cost, lowmem_shrinker.seeks, int, S_IRUGO | S_IWUSR);
 module_param_array_named(adj, lowmem_adj, int, &lowmem_adj_size,
 			 S_IRUGO | S_IWUSR);

fs/exec.c

@@ -849,6 +849,8 @@ static int de_thread(struct task_struct *tsk)
 		transfer_pid(leader, tsk, PIDTYPE_SID);
 
 		list_replace_rcu(&leader->tasks, &tsk->tasks);
+		delete_from_adj_tree(leader);
+		add_2_adj_tree(tsk);
 		list_replace_init(&leader->sibling, &tsk->sibling);
 
 		tsk->group_leader = tsk;

fs/proc/base.c

@@ -1050,6 +1050,9 @@ static ssize_t oom_adjust_write(struct file *file, const char __user *buf,
 
 	unlock_task_sighand(task, &flags);
 	put_task_struct(task);
+	
+	delete_from_adj_tree(task);
+	add_2_adj_tree(task);
 
 	return count;
 }

include/linux/sched.h

@@ -1224,6 +1224,7 @@ struct task_struct {
 #endif
 
 	struct list_head tasks;
+	struct rb_node adj_node;
 	struct plist_node pushable_tasks;
 
 	struct mm_struct *mm, *active_mm;
@@ -1552,6 +1553,9 @@ static inline struct pid *task_tgid(struct task_struct *task)
 	return task->group_leader->pids[PIDTYPE_PID].pid;
 }
 
+extern void add_2_adj_tree(struct task_struct *task);
+extern void delete_from_adj_tree(struct task_struct *task);
+
 /*
  * Without tasklist or rcu lock it is not safe to dereference
  * the result of task_pgrp/task_session even if task == current,

kernel/exit.c

@@ -67,6 +67,7 @@ static void __unhash_process(struct task_struct *p, bool group_dead)
 		detach_pid(p, PIDTYPE_SID);
 
 		list_del_rcu(&p->tasks);
+		delete_from_adj_tree(p);
 		list_del_init(&p->sibling);
 		__get_cpu_var(process_counts)--;
 	}

kernel/fork.c

@@ -1298,6 +1298,7 @@ static struct task_struct *copy_process(unsigned long clone_flags,
 			attach_pid(p, PIDTYPE_SID, task_session(current));
 			list_add_tail(&p->sibling, &p->real_parent->children);
 			list_add_tail_rcu(&p->tasks, &init_task.tasks);
+			add_2_adj_tree(p);
 			__get_cpu_var(process_counts)++;
 		}
 		attach_pid(p, PIDTYPE_PID, pid);