Finally, call the page_owner swap hooks from the generic swap hooks.
This plugs the new functionality into the kernel.

Add a missing include in 'page_owner.h' to fix build errors in files
that include 'pgtable.h' (thus now 'page_owner.h') without an earlier
include that provides 'pg_data_t':

    In file included from .../include/linux/pgtable.h:19,
                     from .../arch/x86/boot/startup/map_kernel.c:7:
    .../include/linux/page_owner.h:20:41: error: unknown type name ‘pg_data_t’

Signed-off-by: Mauricio Faria de Oliveira <mfo@igalia.com>
---
 include/linux/page_owner.h |  1 +
 include/linux/pgtable.h    | 12 +++++++++++-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/include/linux/page_owner.h b/include/linux/page_owner.h
index cd95aacceba7..d7029d1e41de 100644
--- a/include/linux/page_owner.h
+++ b/include/linux/page_owner.h
@@ -3,6 +3,7 @@
 #define __LINUX_PAGE_OWNER_H
 
 #include <linux/jump_label.h>
+#include <linux/mmzone.h>
 
 #ifdef CONFIG_PAGE_OWNER
 extern struct static_key_false page_owner_inited;
diff --git a/include/linux/pgtable.h b/include/linux/pgtable.h
index 23d30afe439e..ae397f5884d5 100644
--- a/include/linux/pgtable.h
+++ b/include/linux/pgtable.h
@@ -16,6 +16,7 @@
 #include <linux/errno.h>
 #include <asm-generic/pgtable_uffd.h>
 #include <linux/page_table_check.h>
+#include <linux/page_owner.h>
 
 #if 5 - defined(__PAGETABLE_P4D_FOLDED) - defined(__PAGETABLE_PUD_FOLDED) - \
 	defined(__PAGETABLE_PMD_FOLDED) != CONFIG_PGTABLE_LEVELS
@@ -1222,22 +1223,31 @@ static inline void arch_swap_restore(swp_entry_t entry, struct folio *folio)
 
 static inline int hook_prepare_to_swap(struct folio *folio)
 {
-	return arch_prepare_to_swap(folio);
+	int ret;
+
+	ret = arch_prepare_to_swap(folio);
+	if (ret)
+		return ret;
+
+	return page_owner_prepare_to_swap(folio);
 }
 
 static inline void hook_swap_invalidate_page(int type, pgoff_t offset)
 {
 	arch_swap_invalidate_page(type, offset);
+	page_owner_swap_invalidate_page(type, offset);
 }
 
 static inline void hook_swap_invalidate_area(int type)
 {
 	arch_swap_invalidate_area(type);
+	page_owner_swap_invalidate_area(type);
 }
 
 static inline void hook_swap_restore(swp_entry_t entry, struct folio *folio)
 {
 	arch_swap_restore(entry, folio);
+	page_owner_swap_restore(entry, folio);
 }
 
 #ifndef __HAVE_ARCH_MOVE_PTE
-- 
2.51.0