[PATCH v7 2/4] mm,page_owner: Implement the tracking of the stacks count

[Oscar Salvador <osalvador@xxxxxxx>]

page_owner needs to increment a stack_record refcount when a new allocation
occurs, and decrement it on a free operation.
In order to do that, we need to have a way to get a stack_record from a
handle.
Implement stack_depot_get_stack() which just does that, and make it public
so page_owner can use it.

Also implement {inc,dec}_stack_record_count() which increments
or decrements on respective allocation and free operations, via
__reset_page_owner() (free operation) and __set_page_owner() (alloc
operation).

Signed-off-by: Oscar Salvador <osalvador@xxxxxxx>
---
 include/linux/stackdepot.h |  8 ++++++++
 lib/stackdepot.c           |  8 ++++++++
 mm/page_owner.c            | 30 ++++++++++++++++++++++++++++++
 3 files changed, 46 insertions(+)

diff --git a/include/linux/stackdepot.h b/include/linux/stackdepot.h
index d0dcf4aebfb4..ac62de4d4999 100644
--- a/include/linux/stackdepot.h
+++ b/include/linux/stackdepot.h
@@ -175,6 +175,14 @@ depot_stack_handle_t stack_depot_save_flags(unsigned long *entries,
 depot_stack_handle_t stack_depot_save(unsigned long *entries,
 				      unsigned int nr_entries, gfp_t gfp_flags);
 
+/**
+ * stack_depo_get_stack - Get a pointer to a stack struct
+ * @handle: Stack depot handle
+ *
+ * Return: Returns a pointer to a stack struct
+ */
+struct stack_record *stack_depot_get_stack(depot_stack_handle_t handle);
+
 /**
  * stack_depot_fetch - Fetch a stack trace from stack depot
  *
diff --git a/lib/stackdepot.c b/lib/stackdepot.c
index 16c8a1bf0008..197c355601f9 100644
--- a/lib/stackdepot.c
+++ b/lib/stackdepot.c
@@ -681,6 +681,14 @@ depot_stack_handle_t stack_depot_save(unsigned long *entries,
 }
 EXPORT_SYMBOL_GPL(stack_depot_save);
 
+struct stack_record *stack_depot_get_stack(depot_stack_handle_t handle)
+{
+	if (!handle)
+		return NULL;
+
+	return depot_fetch_stack(handle);
+}
+
 unsigned int stack_depot_fetch(depot_stack_handle_t handle,
 			       unsigned long **entries)
 {
diff --git a/mm/page_owner.c b/mm/page_owner.c
index 5634e5d890f8..0adf41702b9d 100644
--- a/mm/page_owner.c
+++ b/mm/page_owner.c
@@ -61,6 +61,22 @@ static __init bool need_page_owner(void)
 	return page_owner_enabled;
 }
 
+static void inc_stack_record_count(depot_stack_handle_t handle)
+{
+	struct stack_record *stack = stack_depot_get_stack(handle);
+
+	if (stack)
+		refcount_inc(&stack->count);
+}
+
+static void dec_stack_record_count(depot_stack_handle_t handle)
+{
+	struct stack_record *stack = stack_depot_get_stack(handle);
+
+	if (stack)
+		refcount_dec(&stack->count);
+}
+
 static __always_inline depot_stack_handle_t create_dummy_stack(void)
 {
 	unsigned long entries[4];
@@ -140,6 +156,7 @@ void __reset_page_owner(struct page *page, unsigned short order)
 	int i;
 	struct page_ext *page_ext;
 	depot_stack_handle_t handle;
+	depot_stack_handle_t alloc_handle;
 	struct page_owner *page_owner;
 	u64 free_ts_nsec = local_clock();
 
@@ -147,6 +164,9 @@ void __reset_page_owner(struct page *page, unsigned short order)
 	if (unlikely(!page_ext))
 		return;
 
+	page_owner = get_page_owner(page_ext);
+	alloc_handle = page_owner->handle;
+
 	handle = save_stack(GFP_NOWAIT | __GFP_NOWARN);
 	for (i = 0; i < (1 << order); i++) {
 		__clear_bit(PAGE_EXT_OWNER_ALLOCATED, &page_ext->flags);
@@ -158,6 +178,15 @@ void __reset_page_owner(struct page *page, unsigned short order)
 		page_ext = page_ext_next(page_ext);
 	}
 	page_ext_put(page_ext);
+	if (alloc_handle != early_handle)
+		/*
+		 * early_handle is being set as a handle for all those
+		 * early allocated pages. See init_pages_in_zone().
+		 * Since their refcount is not being incremented because
+		 * the machinery is not ready yet, we cannot decrement
+		 * their refcount either.
+		 */
+		dec_stack_record_count(alloc_handle);
 }
 
 static inline void __set_page_owner_handle(struct page_ext *page_ext,
@@ -199,6 +228,7 @@ noinline void __set_page_owner(struct page *page, unsigned short order,
 		return;
 	__set_page_owner_handle(page_ext, handle, order, gfp_mask);
 	page_ext_put(page_ext);
+	inc_stack_record_count(handle);
 }
 
 void __set_page_owner_migrate_reason(struct page *page, int reason)
-- 
2.43.0