From d39f39ccb96ce1fef5b7a402c01010ebeb66134a Mon Sep 17 00:00:00 2001
From: Georgi Djakov <quic_c_gdjako@quicinc.com>
Date: Mon, 19 Apr 2021 09:55:11 -0700
Subject: [PATCH] ANDROID: psi: Add vendor hooks for PSI tracing

Add hooks to capture various per-zone memory stats when
a trigger threshold is hit.

Bug: 268290366
Change-Id: Ibe39263ddb05ffc3fa63b5225497a90c6480c8d7
Signed-off-by: Georgi Djakov <quic_c_gdjako@quicinc.com>
---
 drivers/android/vendor_hooks.c |  3 +++
 include/trace/hooks/psi.h      | 25 +++++++++++++++++++++++++
 kernel/sched/psi.c             |  5 +++++
 3 files changed, 33 insertions(+)
 create mode 100644 include/trace/hooks/psi.h

diff --git a/drivers/android/vendor_hooks.c b/drivers/android/vendor_hooks.c
index fa4670193fd1..1b9fdf8badd4 100644
--- a/drivers/android/vendor_hooks.c
+++ b/drivers/android/vendor_hooks.c
@@ -62,6 +62,7 @@
 #include <trace/hooks/typec.h>
 #include <trace/hooks/user.h>
 #include <trace/hooks/signal.h>
+#include <trace/hooks/psi.h>
 
 /*
  * Export tracepoints that act as a bare tracehook (ie: have no trace event
@@ -238,3 +239,5 @@ EXPORT_TRACEPOINT_SYMBOL_GPL(android_vh_binder_reply);
 EXPORT_TRACEPOINT_SYMBOL_GPL(android_vh_binder_trans);
 EXPORT_TRACEPOINT_SYMBOL_GPL(android_vh_mmap_region);
 EXPORT_TRACEPOINT_SYMBOL_GPL(android_vh_try_to_unmap_one);
+EXPORT_TRACEPOINT_SYMBOL_GPL(android_vh_psi_event);
+EXPORT_TRACEPOINT_SYMBOL_GPL(android_vh_psi_group);
diff --git a/include/trace/hooks/psi.h b/include/trace/hooks/psi.h
new file mode 100644
index 000000000000..91a008655031
--- /dev/null
+++ b/include/trace/hooks/psi.h
@@ -0,0 +1,25 @@
+/* SPDX-License-Identifier: GPL-2.0-only */
+#undef TRACE_SYSTEM
+#define TRACE_SYSTEM psi
+
+#define TRACE_INCLUDE_PATH trace/hooks
+
+#if !defined(_TRACE_HOOK_PSI_H) || defined(TRACE_HEADER_MULTI_READ)
+#define _TRACE_HOOK_PSI_H
+
+#include <trace/hooks/vendor_hooks.h>
+
+struct psi_trigger;
+struct psi_group;
+DECLARE_HOOK(android_vh_psi_event,
+	TP_PROTO(struct psi_trigger *t),
+	TP_ARGS(t));
+
+DECLARE_HOOK(android_vh_psi_group,
+	TP_PROTO(struct psi_group *group),
+	TP_ARGS(group));
+
+#endif /* _TRACE_HOOK_PSI_H */
+
+/* This part must be outside protection */
+#include <trace/define_trace.h>
diff --git a/kernel/sched/psi.c b/kernel/sched/psi.c
index 8fc99df5cc78..5e5a602d8bae 100644
--- a/kernel/sched/psi.c
+++ b/kernel/sched/psi.c
@@ -136,6 +136,7 @@
  * cost-wise, yet way more sensitive and accurate than periodic
  * sampling of the aggregate task states would be.
  */
+#include <trace/hooks/psi.h>
 
 static int psi_bug __read_mostly;
 
@@ -550,6 +551,8 @@ static u64 update_triggers(struct psi_group *group, u64 now)
 		if (now < t->last_event_time + t->win.size)
 			continue;
 
+		trace_android_vh_psi_event(t);
+
 		/* Generate an event */
 		if (cmpxchg(&t->event, 0, 1) == 0)
 			wake_up_interruptible(&t->event_wait);
@@ -558,6 +561,8 @@ static u64 update_triggers(struct psi_group *group, u64 now)
 		t->pending_event = false;
 	}
 
+	trace_android_vh_psi_group(group);
+
 	if (update_total)
 		memcpy(group->polling_total, total,
 				sizeof(group->polling_total));