From b79d1815c400c3929bb78a209b6e51a38d6ffcb2 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: 178721511
Change-Id: Ibe39263ddb05ffc3fa63b5225497a90c6480c8d7
Signed-off-by: Georgi Djakov <quic_c_gdjako@quicinc.com>
---
 drivers/android/vendor_hooks.c |  3 +++
 include/trace/hooks/psi.h      | 29 +++++++++++++++++++++++++++++
 kernel/sched/psi.c             |  6 ++++++
 3 files changed, 38 insertions(+)
 create mode 100644 include/trace/hooks/psi.h

diff --git a/drivers/android/vendor_hooks.c b/drivers/android/vendor_hooks.c
index c953faf0b674..43dbf97dacf0 100644
--- a/drivers/android/vendor_hooks.c
+++ b/drivers/android/vendor_hooks.c
@@ -52,6 +52,7 @@
 #include <trace/hooks/signal.h>
 #include <trace/hooks/logbuf.h>
 #include <trace/hooks/vmscan.h>
+#include <trace/hooks/psi.h>
 
 /*
  * Export tracepoints that act as a bare tracehook (ie: have no trace event
@@ -243,3 +244,5 @@ EXPORT_TRACEPOINT_SYMBOL_GPL(android_vh_typec_tcpm_adj_current_limit);
 EXPORT_TRACEPOINT_SYMBOL_GPL(android_vh_logbuf);
 EXPORT_TRACEPOINT_SYMBOL_GPL(android_vh_tune_scan_type);
 EXPORT_TRACEPOINT_SYMBOL_GPL(android_vh_tune_swappiness);
+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..1cbdaef6dfa7
--- /dev/null
+++ b/include/trace/hooks/psi.h
@@ -0,0 +1,29 @@
+/* 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 <linux/psi_types.h>
+#include <linux/tracepoint.h>
+#include <trace/hooks/vendor_hooks.h>
+
+#if defined(CONFIG_TRACEPOINTS) && defined(CONFIG_ANDROID_VENDOR_HOOKS)
+
+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
+
+#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 fcc46032ddde..29cb3245c67e 100644
--- a/kernel/sched/psi.c
+++ b/kernel/sched/psi.c
@@ -142,6 +142,8 @@
 #include <linux/psi.h>
 #include "sched.h"
 
+#include <trace/hooks/psi.h>
+
 static int psi_bug __read_mostly;
 
 DEFINE_STATIC_KEY_FALSE(psi_disabled);
@@ -533,12 +535,16 @@ 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);
 		t->last_event_time = now;
 	}
 
+	trace_android_vh_psi_group(group);
+
 	if (new_stall)
 		memcpy(group->polling_total, total,
 				sizeof(group->polling_total));