From a1fc1fba460be48dff2484b438888ede557930d4 Mon Sep 17 00:00:00 2001
From: Yun Hsiang <yun.hsiang@mediatek.com>
Date: Wed, 20 May 2020 10:31:35 +0800
Subject: [PATCH] ANDROID: sched: add restrict vendor hook to modify task
 placement policy in EAS

For modifying task placement policy, we add the hook on the top of
select_task_rq_fair(). It allows us to modify wakeup/fork/exec balance
paths.

Bug: 158263250
Signed-off-by: Yun Hsiang <yun.hsiang@mediatek.com>
Change-Id: I80ff870453472e183ab2aae7381bff91e49ae296
---
 drivers/android/vendor_hooks.c |  3 ++-
 include/trace/hooks/sched.h    | 23 +++++++++++++++++++++++
 kernel/sched/fair.c            |  7 +++++++
 3 files changed, 32 insertions(+), 1 deletion(-)
 create mode 100644 include/trace/hooks/sched.h

diff --git a/drivers/android/vendor_hooks.c b/drivers/android/vendor_hooks.c
index cf75759e4c3f..bb6c969475ee 100644
--- a/drivers/android/vendor_hooks.c
+++ b/drivers/android/vendor_hooks.c
@@ -8,9 +8,10 @@
 
 #define CREATE_TRACE_POINTS
 #include <trace/hooks/vendor_hooks.h>
+#include <trace/hooks/sched.h>
 
 /*
  * Export tracepoints that act as a bare tracehook (ie: have no trace event
  * associated with them) to allow external modules to probe them.
  */
-
+EXPORT_TRACEPOINT_SYMBOL_GPL(android_rvh_select_task_rq_fair);
diff --git a/include/trace/hooks/sched.h b/include/trace/hooks/sched.h
new file mode 100644
index 000000000000..e4ef56686a76
--- /dev/null
+++ b/include/trace/hooks/sched.h
@@ -0,0 +1,23 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+#undef TRACE_SYSTEM
+#define TRACE_SYSTEM sched
+#define TRACE_INCLUDE_PATH trace/hooks
+#if !defined(_TRACE_HOOK_SCHED_H) || defined(TRACE_HEADER_MULTI_READ)
+#define _TRACE_HOOK_SCHED_H
+#include <linux/tracepoint.h>
+#include <trace/hooks/vendor_hooks.h>
+/*
+ * Following tracepoints are not exported in tracefs and provide a
+ * mechanism for vendor modules to hook and extend functionality
+ */
+#if defined(CONFIG_TRACEPOINTS) && defined(CONFIG_ANDROID_VENDOR_HOOKS)
+struct task_struct;
+DECLARE_RESTRICTED_HOOK(android_rvh_select_task_rq_fair,
+	TP_PROTO(struct task_struct *p, int prev_cpu, int sd_flag, int wake_flags, int *new_cpu),
+	TP_ARGS(p, prev_cpu, sd_flag, wake_flags, new_cpu), 1);
+#else
+#define trace_android_rvh_select_task_rq_fair(p, prev_cpu, sd_flag, wake_flags, new_cpu)
+#endif
+#endif /* _TRACE_HOOK_SCHED_H */
+/* This part must be outside protection */
+#include <trace/define_trace.h>
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index d7bb300126fc..2cd5558180ad 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -23,6 +23,7 @@
 #include "sched.h"
 
 #include <trace/events/sched.h>
+#include <trace/hooks/sched.h>
 
 /*
  * Targeted preemption latency for CPU-bound tasks:
@@ -6706,6 +6707,12 @@ select_task_rq_fair(struct task_struct *p, int prev_cpu, int sd_flag, int wake_f
 	int new_cpu = prev_cpu;
 	int want_affine = 0;
 	int sync = (wake_flags & WF_SYNC) && !(current->flags & PF_EXITING);
+	int target_cpu = -1;
+
+	trace_android_rvh_select_task_rq_fair(p, prev_cpu, sd_flag,
+			wake_flags, &target_cpu);
+	if (target_cpu >= 0)
+		return target_cpu;
 
 	if (sd_flag & SD_BALANCE_WAKE) {
 		record_wakee(p);