[RFC 00/10] Reclaimable kernel stacks

[RFC 00/10] Reclaimable kernel stacks

[David Stevens]

This RFC is a different approach to reducing kernel stack usage from the
earlier dynamic kernel stack RFC [1]. This patch series aims to reduce
the cost of kernel stacks by partially reclaiming stacks of blocked
tasks when it is safe to do so.

On Android, system processes typically have 2000-3000 threads. App
processes add 1000s more threads on top of this. The number of app
processes varies based on device RAM size, but the end result is that
1-2% of system RAM is consumed by kernel stacks. However, most of these
threads spend extended periods of time blocked. As such, reclaiming
blocked kernel stacks can reduce total kernel stack memory usage by
upwards of 50% in various multi-tasking test cases.

When a task is blocked, we know exactly where the top of its stack is
and can reclaim any pages past that point. Since any accesses to that
portion of the stack are bugs like use-after-return or buffer overflow,
turning those invalid accesses into hard crashes could even be
considered a positive.

Tracking blocked state and when it is safe to reclaim a stack is done
via a series of hooks in the scheduler. The actual reclaim of stacks is
done asynchronously in a shrinker.

Once a task's stack has been reclaimed, it cannot be rescheduled until
its stack is repopulated. Although there can be a repopulation fast path
within the scheduler, reliably allocating memory to repopulate the stack
requires a fallback path that defers the repopulation and wakeup to a
workqueue context that can use GFP_KERNEL.

The primary challenge is avoiding reclaim deadlocks. If a task blocks
while holding a lock used by direct reclaim and then has its stack
reclaimed, using GFP_KERNEL to reallocate its stack risks deadlock. To
avoid this, only tasks which are known not to hold any locks upon which
reclaim depends are considered eligible for stack reclaim. Automatically
inferring this property is not feasible, so instead a new
PF_RECLAIMABLE_STACK task flag is used to annotate blocking locations
that are known safe. While annotating all safe blocking locations is not
feasible, the vast majority of userspace threads block using a fairly
small number of syscalls - futex, epoll, nanosleep, etc. The 10
annotations added in this series cover >95% of userspace threads on
Android based on my testing. Since missing annotations are leaving an
optimization on the table rather than an actual bug, other annotations
can be added later as needed.

Although reclaimable stacks will not cause reclaim to deadlock, it does
introduce a dependency on needing to allocate memory before an OOM
victim can exit, as reclaimed stacks need to be repopulated before their
tasks can run. While the OOM reaper will still be able to immediately
free the victim's mm, the freeing of non-mm memory may be delayed. This
can lead to more OOM kills. While this is not a significant concern on
Android due to the reliance on lmkd over the kernel OOM killer, it may
be a concern on other systems.

This RFC was developed primarily on 6.18 and 7.1 based kernels. I have
done fairly heavy stress testing, but it has not yet been deployed to
any production systems. If initial feedback on the RFC is somewhat
positive, I will work on deploying it to production systems for further
stability and performance testing as well as resolving the handful of
TODOs left in the RFC.

[1] https://lore.kernel.org/linux-mm/20260424191456.2679717-1-stevensd@google.com/

David Stevens (10):
  Add !MEMCG memcg_list_lru_alloc implementation
  mm/vmalloc: Skip vmallocinfo NUMA stats for VM_SPARSE
  fork: refactor vmap stack alloc/free into helpers
  mm: vmalloc: support creating aligned vm areas
  fork: allocate reclaimable stacks with VM_SPARSE
  Reclaim memory from blocked kernel stacks
  Reclaim stacks via a shrinker
  Set PF_RECLAIMABLE_STACK in various places
  x86: Enable reclaimable stacks
  arm64: Enable reclaimable stacks

 arch/Kconfig                       |  18 +
 arch/arm64/Kconfig                 |   1 +
 arch/arm64/include/asm/processor.h |   5 +
 arch/x86/Kconfig                   |   1 +
 arch/x86/include/asm/processor.h   |   5 +
 drivers/android/binder/thread.rs   |  14 +
 fs/eventpoll.c                     |   3 +
 fs/pipe.c                          |  28 +-
 fs/select.c                        |   3 +
 include/linux/list_lru.h           |   7 +-
 include/linux/sched.h              |  44 +-
 include/linux/sched/task_stack.h   |  23 +
 include/linux/vmalloc.h            |   2 +
 kernel/Makefile                    |   2 +
 kernel/fork.c                      | 140 +++++-
 kernel/futex/waitwake.c            |   3 +
 kernel/sched/core.c                |  18 +-
 kernel/sched/sched.h               |   3 +
 kernel/signal.c                    |  36 +-
 kernel/stack_shrinker.c            | 760 +++++++++++++++++++++++++++++
 kernel/stack_shrinker.h            |  58 +++
 kernel/time/hrtimer.c              |   3 +
 mm/vmalloc.c                       |  27 +-
 rust/kernel/task.rs                |  16 +
 24 files changed, 1175 insertions(+), 45 deletions(-)
 create mode 100644 kernel/stack_shrinker.c
 create mode 100644 kernel/stack_shrinker.h


base-commit: 8d3ae59288f1e7d58d76558a6ee96d533bc5019f
-- 
2.55.0.897.gb25b4bd76c-goog

[RFC 01/10] Add !MEMCG memcg_list_lru_alloc implementation

[David Stevens]

Add a stub implementation of memcg_list_lru_alloc() when CONFIG_MEMCG
isn't enabled.

Signed-off-by: David Stevens <stevensd@google.com>
---
 include/linux/list_lru.h | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/include/linux/list_lru.h b/include/linux/list_lru.h
index a450fffe1550..40eaeaa3424c 100644
--- a/include/linux/list_lru.h
+++ b/include/linux/list_lru.h
@@ -79,10 +79,10 @@ static inline int list_lru_init_memcg_key(struct list_lru *lru, struct shrinker
 	return list_lru_init_memcg(lru, shrinker);
 }
 
+#ifdef CONFIG_MEMCG
 int memcg_list_lru_alloc(struct mem_cgroup *memcg, struct list_lru *lru,
 			 gfp_t gfp);
 
-#ifdef CONFIG_MEMCG
 /**
  * folio_memcg_list_lru_alloc - allocate list_lru heads for shrinkable folio
  * @folio: the newly allocated & charged folio
@@ -106,6 +106,11 @@ static inline int folio_memcg_list_lru_alloc(struct folio *folio,
 {
 	return 0;
 }
+
+static inline int memcg_list_lru_alloc(struct mem_cgroup *memcg, struct list_lru *lru, gfp_t gfp)
+{
+	return 0;
+}
 #endif
 
 void memcg_reparent_list_lrus(struct mem_cgroup *memcg, struct mem_cgroup *parent);
-- 
2.55.0.897.gb25b4bd76c-goog

[RFC 02/10] mm/vmalloc: Skip vmallocinfo NUMA stats for VM_SPARSE

[David Stevens]

Bail from show_numa_info() in vmalloc_info_show() for areas with
VM_SPARSE set. This lets clients that allocate VM_SPARSE areas directly
use vm_struct's nr_pages and pages fields.

Signed-off-by: David Stevens <stevensd@google.com>
---
 mm/vmalloc.c | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/mm/vmalloc.c b/mm/vmalloc.c
index f4fa227a8d7f..d39897300f29 100644
--- a/mm/vmalloc.c
+++ b/mm/vmalloc.c
@@ -5294,6 +5294,9 @@ static void show_numa_info(struct seq_file *m, struct vm_struct *v,
 	if (!counters)
 		return;
 
+	if (v->flags & VM_SPARSE)
+		return;
+
 	memset(counters, 0, nr_node_ids * sizeof(unsigned int));
 
 	for (nr = 0; nr < v->nr_pages; nr += step)
-- 
2.55.0.897.gb25b4bd76c-goog

[RFC 03/10] fork: refactor vmap stack alloc/free into helpers

[David Stevens]

Prepare for alternate implementations of some functions for reclaimable
stacks by creating helpers for allocating and freeing vmap stacks. Also
reorder some functions to consolidate the functions that will have
alternate implementations.

Signed-off-by: David Stevens <stevensd@google.com>
---
 kernel/fork.c | 41 +++++++++++++++++++++++++++--------------
 1 file changed, 27 insertions(+), 14 deletions(-)

diff --git a/kernel/fork.c b/kernel/fork.c
index f0e2e131a9a5..10bd76f6dd20 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -207,11 +207,6 @@ static DEFINE_PER_CPU(struct vm_struct *, cached_stacks[NR_CACHED_STACKS]);
  */
 #define GFP_VMAP_STACK (GFP_KERNEL | __GFP_ZERO | __GFP_SKIP_KASAN)
 
-struct vm_stack {
-	struct rcu_head rcu;
-	struct vm_struct *stack_vm_area;
-};
-
 static struct vm_struct *alloc_thread_stack_node_from_cache(struct task_struct *tsk, int node)
 {
 	struct vm_struct *vm_area;
@@ -272,6 +267,26 @@ static bool try_release_thread_stack_to_cache(struct vm_struct *vm_area)
 	return false;
 }
 
+static void free_vmap_stack(struct vm_struct *vm_area)
+{
+	vfree(vm_area->addr);
+}
+
+static struct vm_struct *alloc_vmap_stack(int node)
+{
+	void *stack = __vmalloc_node(THREAD_SIZE, THREAD_ALIGN,
+				     GFP_VMAP_STACK,
+				     node, __builtin_return_address(0));
+	if (!stack)
+		return NULL;
+	return find_vm_area(stack);
+}
+
+struct vm_stack {
+	struct rcu_head rcu;
+	struct vm_struct *stack_vm_area;
+};
+
 static void thread_stack_free_rcu(struct rcu_head *rh)
 {
 	struct vm_stack *vm_stack = container_of(rh, struct vm_stack, rcu);
@@ -280,7 +295,7 @@ static void thread_stack_free_rcu(struct rcu_head *rh)
 	if (try_release_thread_stack_to_cache(vm_stack->stack_vm_area))
 		return;
 
-	vfree(vm_area->addr);
+	free_vmap_stack(vm_area);
 }
 
 static void thread_stack_delayed_free(struct task_struct *tsk)
@@ -302,7 +317,7 @@ static int free_vm_stack_cache(unsigned int cpu)
 		if (!vm_area)
 			continue;
 
-		vfree(vm_area->addr);
+		free_vmap_stack(vm_area);
 		cached_vm_stack_areas[i] = NULL;
 	}
 
@@ -338,7 +353,7 @@ static int alloc_thread_stack_node(struct task_struct *tsk, int node)
 	vm_area = alloc_thread_stack_node_from_cache(tsk, node);
 	if (vm_area) {
 		if (memcg_charge_kernel_stack(vm_area)) {
-			vfree(vm_area->addr);
+			free_vmap_stack(vm_area);
 			return -ENOMEM;
 		}
 
@@ -356,15 +371,13 @@ static int alloc_thread_stack_node(struct task_struct *tsk, int node)
 		return 0;
 	}
 
-	stack = __vmalloc_node(THREAD_SIZE, THREAD_ALIGN,
-				     GFP_VMAP_STACK,
-				     node, __builtin_return_address(0));
-	if (!stack)
+	vm_area = alloc_vmap_stack(node);
+	if (!vm_area)
 		return -ENOMEM;
+	stack = vm_area->addr;
 
-	vm_area = find_vm_area(stack);
 	if (memcg_charge_kernel_stack(vm_area)) {
-		vfree(stack);
+		free_vmap_stack(vm_area);
 		return -ENOMEM;
 	}
 	/*
-- 
2.55.0.897.gb25b4bd76c-goog

[RFC 04/10] mm: vmalloc: support creating aligned vm areas

[David Stevens]

Expose a get_vm_area_node() function that supports passing an alignment
when allocating a vm area.

Signed-off-by: David Stevens <stevensd@google.com>
---
 include/linux/vmalloc.h |  2 ++
 mm/vmalloc.c            | 24 +++++++++++++++++++++++-
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/include/linux/vmalloc.h b/include/linux/vmalloc.h
index d87dc7f77f4e..5ed9c5b25026 100644
--- a/include/linux/vmalloc.h
+++ b/include/linux/vmalloc.h
@@ -246,6 +246,8 @@ static inline size_t get_vm_area_size(const struct vm_struct *area)
 extern struct vm_struct *get_vm_area(unsigned long size, unsigned long flags);
 extern struct vm_struct *get_vm_area_caller(unsigned long size,
 					unsigned long flags, const void *caller);
+extern struct vm_struct *get_vm_area_node(unsigned long size, unsigned long align,
+					  unsigned long flags, int node);
 extern struct vm_struct *__get_vm_area_caller(unsigned long size,
 					unsigned long flags,
 					unsigned long start, unsigned long end,
diff --git a/mm/vmalloc.c b/mm/vmalloc.c
index d39897300f29..059cd0748514 100644
--- a/mm/vmalloc.c
+++ b/mm/vmalloc.c
@@ -3280,7 +3280,7 @@ struct vm_struct *__get_vm_area_caller(unsigned long size, unsigned long flags,
  * @flags:	 %VM_IOREMAP for I/O mappings or VM_ALLOC
  *
  * Search an area of @size in the kernel virtual mapping area,
- * and reserved it for out purposes.  Returns the area descriptor
+ * and reserved it for our purposes.  Returns the area descriptor
  * on success or %NULL on failure.
  *
  * Return: the area descriptor on success or %NULL on failure.
@@ -3301,6 +3301,28 @@ struct vm_struct *get_vm_area_caller(unsigned long size, unsigned long flags,
 				  NUMA_NO_NODE, GFP_KERNEL, caller);
 }
 
+/**
+ * get_vm_area_node - reserve a contiguous kernel virtual area
+ * @size:	 size of the area
+ * @align:	 alignment of the area
+ * @flags:	 %VM_IOREMAP for I/O mappings or VM_ALLOC
+ * @node:	 node from which to allocate data structures
+ *
+ * Search an area of @size in the kernel virtual mapping area,
+ * and reserve it for our purposes.  Returns the area descriptor
+ * on success or %NULL on failure.
+ *
+ * Return: the area descriptor on success or %NULL on failure.
+ */
+struct vm_struct *get_vm_area_node(unsigned long size, unsigned long align,
+				   unsigned long flags, int node)
+{
+	return __get_vm_area_node(size, align, PAGE_SHIFT, flags,
+				  VMALLOC_START, VMALLOC_END,
+				  node, GFP_KERNEL,
+				  __builtin_return_address(0));
+}
+
 /**
  * find_vm_area - find a continuous kernel virtual area
  * @addr:	  base address
-- 
2.55.0.897.gb25b4bd76c-goog

[RFC 05/10] fork: allocate reclaimable stacks with VM_SPARSE

[David Stevens]

Since the pages of reclaimable stacks will not always be populated, we
need to set VM_SPARSE on their vm_structs to avoid crashes when
vread_iter sees a partially reclaimed stack. Note that stacks will only
be partially reclaimed when they are associated with a live task, so the
stack management and caching code in fork.c won't actually ever see a
partially reclaimed stack.

Directly managing the vm_structs instead of going through vmalloc also
requires doing the freeing of the stack on a work queue instead of in an
RCU callback. Previously, we were relying on deferred vfree work to move
much of the cleanup work from softirq to process context.

The fact that vmap stack pages are included in vmalloc's vmstat is well
known. We should continue that to avoid changing procfs.

Signed-off-by: David Stevens <stevensd@google.com>
---
 kernel/fork.c | 92 +++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 92 insertions(+)

diff --git a/kernel/fork.c b/kernel/fork.c
index 10bd76f6dd20..6acad0038b78 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -267,6 +267,97 @@ static bool try_release_thread_stack_to_cache(struct vm_struct *vm_area)
 	return false;
 }
 
+#ifdef CONFIG_RECLAIMABLE_STACK
+static void free_vmap_stack(struct vm_struct *vm_area)
+{
+	int i;
+
+	remove_vm_area(vm_area->addr);
+
+	for (i = 0; i < vm_area->nr_pages; i++) {
+		mod_node_page_state(page_pgdat(vm_area->pages[i]), NR_VMALLOC, -1);
+		__free_page(vm_area->pages[i]);
+	}
+
+	kfree(vm_area->pages);
+	kfree(vm_area);
+}
+
+static struct vm_struct *alloc_vmap_stack(int node)
+{
+	struct vm_struct *vm_area;
+	int ret;
+
+	vm_area = get_vm_area_node(THREAD_SIZE, THREAD_ALIGN, VM_MAP | VM_SPARSE, node);
+	if (!vm_area)
+		return NULL;
+
+	vm_area->pages = kcalloc_node(THREAD_SIZE >> PAGE_SHIFT, sizeof(*vm_area->pages),
+				      GFP_KERNEL | __GFP_ZERO, node);
+	if (!vm_area->pages)
+		goto alloc_failure;
+
+	while (vm_area->nr_pages < THREAD_SIZE >> PAGE_SHIFT) {
+		struct page *page;
+		gfp_t gfp = GFP_VMAP_STACK | __GFP_HIGHMEM;
+
+		if (node == NUMA_NO_NODE)
+			page = alloc_pages(gfp, 0);
+		else
+			page = alloc_pages_node(node, gfp, 0);
+
+		if (!page)
+			goto alloc_failure;
+
+		/*
+		 * Non-reclaimable vmap stacks pages aren't charged against an
+		 * memcg until account_kernel_stack(), but they are added to
+		 * the node's NR_VMALLOC counter. Copy that behavior to avoid
+		 * confusing userspace.
+		 */
+		mod_node_page_state(page_pgdat(page), NR_VMALLOC, 1);
+		vm_area->pages[vm_area->nr_pages++] = page;
+	}
+
+	ret = vmap_pages_range((unsigned long)vm_area->addr,
+			       (unsigned long)vm_area->addr + THREAD_SIZE,
+			       PAGE_KERNEL, vm_area->pages, PAGE_SHIFT);
+	if (ret)
+		goto alloc_failure;
+
+	return vm_area;
+
+alloc_failure:
+	free_vmap_stack(vm_area);
+	return NULL;
+}
+
+struct vm_stack {
+	struct rcu_work work;
+	struct vm_struct *stack_vm_area;
+};
+
+static void thread_stack_free_work(struct work_struct *work)
+{
+	struct vm_stack *vm_stack = container_of(to_rcu_work(work), struct vm_stack, work);
+	struct vm_struct *vm_area = vm_stack->stack_vm_area;
+
+	if (try_release_thread_stack_to_cache(vm_stack->stack_vm_area))
+		return;
+
+	free_vmap_stack(vm_area);
+}
+
+static void thread_stack_delayed_free(struct task_struct *tsk)
+{
+	struct vm_stack *vm_stack = tsk->stack;
+
+	vm_stack->stack_vm_area = tsk->stack_vm_area;
+	INIT_RCU_WORK(&vm_stack->work, thread_stack_free_work);
+	queue_rcu_work(system_wq, &vm_stack->work);
+}
+
+#else /* !CONFIG_RECLAIMABLE_STACK */
 static void free_vmap_stack(struct vm_struct *vm_area)
 {
 	vfree(vm_area->addr);
@@ -305,6 +396,7 @@ static void thread_stack_delayed_free(struct task_struct *tsk)
 	vm_stack->stack_vm_area = tsk->stack_vm_area;
 	call_rcu(&vm_stack->rcu, thread_stack_free_rcu);
 }
+#endif /* CONFIG_RECLAIMABLE_STACK */
 
 static int free_vm_stack_cache(unsigned int cpu)
 {
-- 
2.55.0.897.gb25b4bd76c-goog

[RFC 06/10] Reclaim memory from blocked kernel stacks

[David Stevens]

Although an individual kernel stack is cheap, the cost can add up on
thread heavy systems. On such systems, many threads are often blocked
for extended periods of time, waiting for userspace or external events.
When a task is blocked, its stack is in a stable, well-defined state.
Since we can tell exactly which part of those stacks are unused, we can
free the unused portions and reduce the global memory footprint of
kernel stacks.

That said, not all blocked tasks can safely have their stacks reclaimed.
Since a running task may need all of its stack, a blocked task whose
stack has been partially freed cannot be rescheduled until its stack is
fully repopulated. Since allocating that memory may require direct
reclaim, if a task holding a lock that direct reclaim depends on has its
stack reclaimed, we would run the risk of deadlock.

Knowing when tasks hold locks that direct reclaim may depend on is not
viable in production systems, especially without lockdep. But noting
that most userspace threads will be blocked on events external to
the kernel (e.g. waiting on a futex for a userspace controlled wakeup,
waiting in epoll for a network packet, waiting on a timer for time to
pass, etc), it is possible to enumerate a relatively small number of
call sites where most userspace threads will be blocked while holding
zero kernel locks.

This change adds a new component that reclaims task stacks. A set of
scheduler hooks are used to determine when it is safe to reclaim stacks
and to ensure that a task's stack is fully repopulated before being
rescheduled. A new PF_RECLAIMABLE_STACK task flag is added that will be
used to annotate such safe-to-reclaim call sites. Users of the flag must
ensure the task doesn't block while holding a lock while the flag is
set. A new TASK_STACK_RECLAIM state is introduced for tasks that are
blocked waiting for the stacks to be repopulated.

While reclaimable stacks will not cause direct reclaim to deadlock, it
does introduce a dependency on allocating memory before an OOM victim
can exit, since any reclaimed stacks need to be repopulated before their
threads can run again. This dependency can lead to further OOM kills or
potentially an OOM panic if no additional victims are available.

Reclaimable stacks is built on VMAP_STACK. There is no fundamental
dependency on !STACK_GROWSUP, but since the only STACK_GROWSUP
architecture doesn't support VMAP_STACK, support for STACK_GROWSUP is
not implemented.

This feature does work on PREEMPT_RT, but is likely undesirable due to
the extra uncertainty. The fact that alloc_pages_nolock_noprof() cannot
be called from under the scheduler lock also makes repopulating stacks
more expensive.

Enabling reclaimable stacks is mutually exclusive with enabling
DEBUG_STACK_USAGE. Making stack_not_used() safe (including from NMIs)
may be technically feasible. Even then, the number that it would report
would be subtly different: max since last reclaim vs max ever used. This
difference seems likely to confuse anything consuming the data, so
disabling the feature seems prudent.

Since kernel stacks are not visible in PROC_KCORE when reclaimable
stacks is enabled, it is disabled by default when that config is
enabled. If a user wants to enable the manually enable the feature and
give up visibility, they can make that choice.

In the followup patch, reclaiming task stacks will be done in a
shrinker. However, this patch does the reclaim using per-task work
structs for simplicity.

Signed-off-by: David Stevens <stevensd@google.com>
---
 arch/Kconfig                     |  18 ++
 include/linux/sched.h            |  39 ++-
 include/linux/sched/task_stack.h |  23 ++
 kernel/Makefile                  |   2 +
 kernel/fork.c                    |   7 +
 kernel/sched/core.c              |  18 +-
 kernel/sched/sched.h             |   3 +
 kernel/stack_shrinker.c          | 464 +++++++++++++++++++++++++++++++
 kernel/stack_shrinker.h          |  58 ++++
 9 files changed, 627 insertions(+), 5 deletions(-)
 create mode 100644 kernel/stack_shrinker.c
 create mode 100644 kernel/stack_shrinker.h

diff --git a/arch/Kconfig b/arch/Kconfig
index fa7507ac8e13..adb4a5957996 100644
--- a/arch/Kconfig
+++ b/arch/Kconfig
@@ -1534,6 +1534,24 @@ config VMAP_STACK
 	  backing virtual mappings with real shadow memory, and KASAN_VMALLOC
 	  must be enabled.
 
+config HAVE_ARCH_RECLAIMABLE_STACK
+	def_bool n
+
+config RECLAIMABLE_STACK
+	default !PREEMPT_RT && !PROC_KCORE
+	bool "Allow stacks of some blocked threads to be reclaimed"
+	depends on VMAP_STACK && !STACK_GROWSUP
+	depends on HAVE_ARCH_RECLAIMABLE_STACK
+	depends on !DEBUG_STACK_USAGE
+	depends on !KASAN_VMALLOC # TODO: add support for this
+	depends on !DEBUG_KMEMLEAK # TODO: add support for this
+	help
+	  Enable this to allow the unused portion of kernel stacks of most
+	  blocked tasks to be reclaimed.
+
+	  The wakeup latency of tasks with reclaimed stacks may increase,
+	  especially while the system is under memory pressure.
+
 config HAVE_ARCH_RANDOMIZE_KSTACK_OFFSET
 	def_bool n
 	help
diff --git a/include/linux/sched.h b/include/linux/sched.h
index 373bcc0598d1..c93a234fac96 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -83,6 +83,7 @@ struct sched_dl_entity;
 struct seq_file;
 struct sighand_struct;
 struct signal_struct;
+struct stack_reclaim_work;
 struct task_delay_info;
 struct task_exec_state;
 struct task_group;
@@ -124,7 +125,8 @@ struct user_event_mm;
 #define TASK_FREEZABLE			0x00002000
 #define __TASK_FREEZABLE_UNSAFE	       (0x00004000 * IS_ENABLED(CONFIG_LOCKDEP))
 #define TASK_FROZEN			0x00008000
-#define TASK_STATE_MAX			0x00010000
+#define TASK_STACK_RECLAIM		0x00010000
+#define TASK_STATE_MAX			0x00020000
 
 #define TASK_ANY			(TASK_STATE_MAX-1)
 
@@ -823,6 +825,29 @@ struct kmap_ctrl {
 #endif
 };
 
+#ifdef CONFIG_RECLAIMABLE_STACK
+enum stack_reclaim_enum {
+	STACK_IN_USE		= 0,
+	STACK_PREPARE_RECLAIM	= 1,
+	STACK_RECLAIMABLE	= 2,
+	STACK_RECLAIMING	= 3,
+	STACK_RECLAIMING_IN_USE = 4,
+	STACK_RECLAIMED		= 5,
+} __packed;
+
+union stack_reclaim_state {
+	struct {
+		enum stack_reclaim_enum stack_state;
+		u16 node;
+	};
+	u32 val;
+};
+
+union stack_reclaim_list {
+	struct llist_node refill_entry;
+};
+#endif
+
 struct task_struct {
 #ifdef CONFIG_THREAD_INFO_IN_TASK
 	/*
@@ -1587,6 +1612,16 @@ struct task_struct {
 #endif
 #ifdef CONFIG_VMAP_STACK
 	struct vm_struct		*stack_vm_area;
+#ifdef CONFIG_RECLAIMABLE_STACK
+	union stack_reclaim_state	stack_reclaim_state;
+	union stack_reclaim_list	stack_reclaim_list;
+#ifdef CONFIG_MEMCG
+	struct obj_cgroup		*stack_obj_cgroup;
+#endif
+
+	// TODO: Replace these with a shrinker
+	struct stack_reclaim_work	*stack_reclaim_work;
+#endif
 #endif
 #ifdef CONFIG_THREAD_INFO_IN_TASK
 	/* A live task holds one reference: */
@@ -1796,7 +1831,7 @@ extern struct pid *cad_pid;
 						 * I am cleaning dirty pages from some other bdi. */
 #define PF_KTHREAD		0x00200000	/* I am a kernel thread */
 #define PF_RANDOMIZE		0x00400000	/* Randomize virtual address space */
-#define PF__HOLE__00800000	0x00800000
+#define PF_RECLAIMABLE_STACK	0x00800000	/* This task's stack is reclaimable */
 #define PF__HOLE__01000000	0x01000000
 #define PF__HOLE__02000000	0x02000000
 #define PF_NO_SETAFFINITY	0x04000000	/* Userland is not allowed to meddle with cpus_mask */
diff --git a/include/linux/sched/task_stack.h b/include/linux/sched/task_stack.h
index 1fab7e9043a3..5d5567899d51 100644
--- a/include/linux/sched/task_stack.h
+++ b/include/linux/sched/task_stack.h
@@ -6,6 +6,7 @@
  * task->stack (kernel stack) handling interfaces:
  */
 
+#include <linux/cleanup.h>
 #include <linux/sched.h>
 #include <linux/magic.h>
 #include <linux/refcount.h>
@@ -83,6 +84,10 @@ static inline void put_task_stack(struct task_struct *tsk) {}
 
 void exit_task_stack_account(struct task_struct *tsk);
 
+/*
+ * Must only be called on current or from inside __schedule() on
+ * prev, to avoid crashing when CONFIG_RECLAIM_STACK is enabled.
+ */
 #define task_stack_end_corrupted(task) \
 		(*(end_of_stack(task)) != STACK_END_MAGIC)
 
@@ -114,4 +119,22 @@ static inline int kstack_end(void *addr)
 	return !(((unsigned long)addr+sizeof(void*)-1) & (THREAD_SIZE-sizeof(void*)));
 }
 
+#ifdef CONFIG_RECLAIMABLE_STACK
+
+DEFINE_CLASS(allow_stack_reclaim, bool,
+	     ({
+		if (!_T)
+			current->flags &= ~PF_RECLAIMABLE_STACK;
+	      }),
+	     ({
+		bool was_set = current->flags & PF_RECLAIMABLE_STACK;
+
+		current->flags |= PF_RECLAIMABLE_STACK;
+		was_set;
+	      }),
+	     void)
+#else /* !CONFIG_RECLAIMABLE_STACK */
+DEFINE_CLASS(allow_stack_reclaim, bool, ({ (void)_T; }), ({ false; }), void)
+#endif /* !CONFIG_VMAP_STACK */
+
 #endif /* _LINUX_SCHED_TASK_STACK_H */
diff --git a/kernel/Makefile b/kernel/Makefile
index 1e1a31673577..01547102e13a 100644
--- a/kernel/Makefile
+++ b/kernel/Makefile
@@ -12,6 +12,8 @@ obj-y     = fork.o exec_domain.o exec_state.o panic.o \
 	    notifier.o ksysfs.o cred.o reboot.o \
 	    async.o range.o smpboot.o ucount.o regset.o ksyms_common.o
 
+obj-$(CONFIG_RECLAIMABLE_STACK) += stack_shrinker.o
+
 obj-$(CONFIG_MULTIUSER) += groups.o
 obj-$(CONFIG_VHOST_TASK) += vhost_task.o
 
diff --git a/kernel/fork.c b/kernel/fork.c
index 6acad0038b78..9b2cc3d01dd1 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -120,6 +120,8 @@
 /* For dup_mmap(). */
 #include "../mm/internal.h"
 
+#include "stack_shrinker.h"
+
 #include <trace/events/sched.h>
 
 #define CREATE_TRACE_POINTS
@@ -460,6 +462,7 @@ static int alloc_thread_stack_node(struct task_struct *tsk, int node)
 
 		tsk->stack_vm_area = vm_area;
 		tsk->stack = stack;
+		add_to_stack_shrinker(tsk, node);
 		return 0;
 	}
 
@@ -472,6 +475,7 @@ static int alloc_thread_stack_node(struct task_struct *tsk, int node)
 		free_vmap_stack(vm_area);
 		return -ENOMEM;
 	}
+
 	/*
 	 * We can't call find_vm_area() in interrupt context, and
 	 * free_thread_stack() can be called in interrupt context,
@@ -480,6 +484,7 @@ static int alloc_thread_stack_node(struct task_struct *tsk, int node)
 	tsk->stack_vm_area = vm_area;
 	stack = kasan_reset_tag(stack);
 	tsk->stack = stack;
+	add_to_stack_shrinker(tsk, node);
 	return 0;
 }
 
@@ -898,6 +903,7 @@ void __put_task_struct(struct task_struct *tsk)
 	delayacct_tsk_free(tsk);
 	put_signal_struct(tsk->signal);
 	sched_core_free(tsk);
+	remove_from_stack_shrinker(tsk);
 	free_task(tsk);
 }
 EXPORT_SYMBOL_GPL(__put_task_struct);
@@ -1124,6 +1130,7 @@ static struct task_struct *dup_task_struct(struct task_struct *orig, int node)
 free_stack:
 	exit_task_stack_account(tsk);
 	free_thread_stack(tsk);
+	remove_from_stack_shrinker(tsk);
 free_tsk:
 	free_task_struct(tsk);
 	return NULL;
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 96226707c2f6..ab7db60d5dc2 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -4252,6 +4252,7 @@ int try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags)
 {
 	guard(preempt)();
 	int cpu, success = 0;
+	bool need_deferred_repopulate, do_deferred_repopulate_wake = false;
 
 	wake_flags |= WF_TTWU;
 
@@ -4295,8 +4296,6 @@ int try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags)
 		if (!ttwu_state_match(p, state, &success))
 			break;
 
-		trace_sched_waking(p);
-
 		/*
 		 * Ensure we load p->on_rq _after_ p->state, otherwise it would
 		 * be possible to, falsely, observe p->on_rq == 0 and get stuck
@@ -4320,8 +4319,18 @@ int try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags)
 		 * A similar smp_rmb() lives in __task_needs_rq_lock().
 		 */
 		smp_rmb();
-		if (READ_ONCE(p->on_rq) && ttwu_runnable(p, wake_flags))
+		if (READ_ONCE(p->on_rq) && ttwu_runnable(p, wake_flags)) {
+			trace_sched_waking(p);
+			break;
+		}
+
+		if (!ensure_stack_is_present(p, &need_deferred_repopulate)) {
+			WRITE_ONCE(p->__state, TASK_STACK_RECLAIM);
+			do_deferred_repopulate_wake = need_deferred_repopulate;
 			break;
+		}
+
+		trace_sched_waking(p);
 
 		/*
 		 * Ensure we load p->on_cpu _after_ p->on_rq, otherwise it would be
@@ -4418,6 +4427,8 @@ int try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags)
 	if (success)
 		ttwu_stat(p, task_cpu(p), wake_flags);
 
+	if (unlikely(do_deferred_repopulate_wake))
+		wake_stack_repopulate();
 	return success;
 }
 
@@ -5354,6 +5365,7 @@ static struct rq *finish_task_switch(struct task_struct *prev)
 	prev_state = READ_ONCE(prev->__state);
 	vtime_task_switch(prev);
 	perf_event_task_sched_in(prev, current);
+	allow_stack_reclaim(prev);
 	finish_task(prev);
 	tick_nohz_task_switch();
 	finish_lock_switch(rq);
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 56acf502ba26..8a81fb83fe8f 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -79,6 +79,7 @@
 #include <trace/events/sched.h>
 
 #include "../workqueue_internal.h"
+#include "../stack_shrinker.h"
 
 struct rq;
 struct cfs_rq;
@@ -3039,6 +3040,8 @@ static inline void __block_task(struct rq *rq, struct task_struct *p)
 		delayacct_blkio_start();
 	}
 
+	prepare_stack_for_reclaim(p);
+
 	ASSERT_EXCLUSIVE_WRITER(p->on_rq);
 
 	/*
diff --git a/kernel/stack_shrinker.c b/kernel/stack_shrinker.c
new file mode 100644
index 000000000000..d7b1a7dfa716
--- /dev/null
+++ b/kernel/stack_shrinker.c
@@ -0,0 +1,464 @@
+// SPDX-License-Identifier: GPL-2.0-only
+
+#include <linux/cpuhotplug.h>
+#include <linux/highmem.h>
+#include <linux/irq_work.h>
+#include <linux/list.h>
+#include <linux/list_lru.h>
+#include <linux/llist.h>
+#include <linux/memcontrol.h>
+#include <linux/oom.h>
+#include <linux/percpu_counter.h>
+#include <linux/pgtable.h>
+#include <linux/sched/task.h>
+#include <linux/spinlock.h>
+#include <linux/swait.h>
+#include <linux/vmalloc.h>
+
+#include "stack_shrinker.h"
+
+#include "../mm/internal.h"
+
+struct repopulate_work {
+	struct work_struct work;
+	struct llist_head stacks;
+};
+
+static DEFINE_PER_CPU(struct repopulate_work, repopulate_work);
+
+// TODO: replace with shrinker
+struct stack_reclaim_work {
+	struct task_struct *tsk;
+	struct irq_work irq_work;
+	struct work_struct work;
+};
+
+static void schedule_stack_reclaim_work(struct irq_work *w)
+{
+	struct stack_reclaim_work *work = container_of(w, typeof(*work), irq_work);
+
+	if (!queue_work(system_wq, &work->work))
+		put_task_struct(work->tsk);
+}
+
+static void do_reclaim_stack(struct task_struct *tsk);
+
+static void do_stack_reclaim_work(struct work_struct *w)
+{
+	struct task_struct *tsk = container_of(w, struct stack_reclaim_work, work)->tsk;
+
+	do_reclaim_stack(tsk);
+}
+
+#ifdef CONFIG_MEMCG
+static void set_stack_obj_cgroup(struct task_struct *tsk)
+{
+	tsk->stack_obj_cgroup = get_obj_cgroup_from_current();
+}
+
+static void put_stack_obj_cgroup(struct task_struct *tsk)
+{
+	if (tsk->stack_obj_cgroup)
+		obj_cgroup_put(tsk->stack_obj_cgroup);
+}
+
+static inline struct mem_cgroup *get_stack_memcg(struct task_struct *tsk)
+{
+	return tsk->stack_obj_cgroup ? get_mem_cgroup_from_objcg(tsk->stack_obj_cgroup)
+				     : root_mem_cgroup;
+}
+#else /* !CONFIG_MEMCG */
+static void set_stack_obj_cgroup(struct task_struct *tsk) { }
+
+static void put_stack_obj_cgroup(struct task_struct *tsk) { }
+
+static inline struct mem_cgroup *get_stack_memcg(struct task_struct *tsk)
+{
+	return NULL;
+}
+#endif /* CONFIG_MEMCG */
+
+void add_to_stack_shrinker(struct task_struct *tsk, int node)
+{
+	BUILD_BUG_ON(sizeof(tsk->stack_reclaim_state) != sizeof(tsk->stack_reclaim_state.val));
+	BUILD_BUG_ON(NODES_SHIFT > 15);
+
+	tsk->stack_reclaim_state.val = 0;
+	tsk->stack_reclaim_state.stack_state = STACK_IN_USE;
+	tsk->stack_reclaim_state.node = node;
+	set_stack_obj_cgroup(tsk);
+	init_llist_node(&tsk->stack_reclaim_list.refill_entry);
+
+	// TODO: replace with shrinker
+	tsk->stack_reclaim_work = kmalloc_obj(*tsk->stack_reclaim_work, GFP_KERNEL);
+	BUG_ON(!tsk->stack_reclaim_work);
+
+	tsk->stack_reclaim_work->tsk = tsk;
+	init_irq_work(&tsk->stack_reclaim_work->irq_work, schedule_stack_reclaim_work);
+	INIT_WORK(&tsk->stack_reclaim_work->work, do_stack_reclaim_work);
+}
+
+static inline int calculate_num_unused_pages(struct task_struct *tsk)
+{
+	unsigned long top_of_stack = (unsigned long)end_of_stack(tsk);
+	/*
+	 * Since tsk is !on_rq and !on_cpu, top_of_blocked_task_stack() safely
+	 * tells us the end of the stack frame of the inner most context switch
+	 * function. Any parts of the stack above that are stale stack frames
+	 * or never used, and thus can be unmapped and discarded.
+	 */
+	return (top_of_blocked_task_stack(&tsk->thread) - top_of_stack) >> PAGE_SHIFT;
+}
+
+static bool repopulate_stack(struct task_struct *tsk, bool is_deferred,
+			     struct llist_head *fail_list)
+{
+	int num_missing_pages = 0, nr_allocated = 0, ret;
+	struct page *pages[THREAD_SIZE >> PAGE_SHIFT] = {};
+	struct vm_struct *vm_area = tsk->stack_vm_area;
+	unsigned long addr = (unsigned long)vm_area->addr;
+	struct mem_cgroup *tsk_memcg, *old_active_memcg;
+	int node = tsk->stack_reclaim_state.node == U16_MAX ? NUMA_NO_NODE
+							    : tsk->stack_reclaim_state.node;
+
+	num_missing_pages = (THREAD_SIZE >> PAGE_SHIFT) - vm_area->nr_pages;
+	if (num_missing_pages == 0)
+		return true;
+
+	tsk_memcg = get_stack_memcg(tsk);
+	old_active_memcg = set_active_memcg(tsk_memcg);
+
+	if (is_deferred) {
+		gfp_t gfp = GFP_KERNEL_ACCOUNT | __GFP_ZERO;
+		/*
+		 * If the oom killer wants to free memory from this process,
+		 * allow access to reserves so the task can hopefully run
+		 * sooner to die and thus make progress towards freeing the
+		 * process's non-mm memory.
+		 */
+		if (tsk_is_oom_victim(tsk))
+			gfp |= __GFP_MEMALLOC;
+
+		for (; nr_allocated < num_missing_pages; nr_allocated++) {
+			pages[nr_allocated] = alloc_pages_node_noprof(node, gfp, 0);
+			if (!pages[nr_allocated])
+				goto repopulate_fail;
+		}
+	} else {
+		/*
+		 * PREEMPT_RT turns spin_trylock() from an atomic cmpxchg into
+		 * an operation that takes a rt_mutex's internal raw spin lock.
+		 * Doing that from inside the scheduler would result in a
+		 * circular locking dependency.
+		 */
+		if (IS_ENABLED(CONFIG_PREEMPT_RT))
+			goto repopulate_fail;
+
+		for (; nr_allocated < num_missing_pages; nr_allocated++) {
+			pages[nr_allocated] = alloc_pages_nolock_noprof(__GFP_ACCOUNT,
+									node, 0);
+			if (!pages[nr_allocated])
+				goto repopulate_fail;
+		}
+	}
+
+	set_active_memcg(old_active_memcg);
+	mem_cgroup_put(tsk_memcg);
+
+	for (int i = 0; i < num_missing_pages; i++) {
+		vm_area->pages[i] = pages[i];
+		mod_lruvec_page_state(pages[i], NR_KERNEL_STACK_KB, PAGE_SIZE / 1024);
+		mod_node_page_state(page_pgdat(pages[i]), NR_VMALLOC, 1);
+	}
+	vm_area->nr_pages = THREAD_SIZE >> PAGE_SHIFT;
+
+	/*
+	 * The page tables for the stack were allocated when the stack was
+	 * originally created, so we're guaranteed not to need to allocate new
+	 * ones. As such, vmap_pages_range() won't acquire any locks and can be
+	 * called under the scheduler's raw spinlocks.
+	 *
+	 * The only way it can fail is if we're trying to colbber an existing
+	 * mapping or if the page allocator gave us an invalid page. Neither
+	 * case is recoverable.
+	 */
+	ret = vmap_pages_range(addr, addr + num_missing_pages * PAGE_SIZE,
+			       PAGE_KERNEL, vm_area->pages, PAGE_SHIFT);
+	BUG_ON(ret != 0);
+
+	// TODO: Clearing pages under the scheduler lock is probably too much
+	// work under a raw spinlock. We could try maintaining our own small
+	// pool of pre-zero'ed pages instead of using alloc_pages_nolock.
+	if (!is_deferred)
+		clear_pages((void *)addr, num_missing_pages);
+
+	set_task_stack_end_magic(tsk);
+	return true;
+
+repopulate_fail:
+	set_active_memcg(old_active_memcg);
+	mem_cgroup_put(tsk_memcg);
+
+	while (nr_allocated--)
+		free_pages_nolock(pages[nr_allocated], 0);
+
+	if (!fail_list) {
+		/*
+		 * Preemption is left disabled until wake_stack_repopulate(), to
+		 * guarantee that we queue the correct work.
+		 */
+		preempt_disable();
+		fail_list = &this_cpu_ptr(&repopulate_work)->stacks;
+	}
+	llist_add(&tsk->stack_reclaim_list.refill_entry, fail_list);
+	return false;
+}
+
+static void release_stack(struct task_struct *tsk)
+{
+	struct vm_struct *vm_area = tsk->stack_vm_area;
+	unsigned long addr = (unsigned long)vm_area->addr;
+	int nr_to_free;
+
+	nr_to_free = calculate_num_unused_pages(tsk);
+
+	if (unlikely(nr_to_free == 0))
+		return;
+
+	vm_area_unmap_pages(vm_area, addr, addr + nr_to_free * PAGE_SIZE);
+	for (int i = 0; i < nr_to_free; i++) {
+		mod_lruvec_page_state(vm_area->pages[i], NR_KERNEL_STACK_KB,
+				      -(int)(PAGE_SIZE / 1024));
+		mod_node_page_state(page_pgdat(vm_area->pages[i]), NR_VMALLOC, -1);
+		__free_pages(vm_area->pages[i], 0);
+	}
+	vm_area->nr_pages -= nr_to_free;
+}
+
+static void do_reclaim_stack(struct task_struct *tsk)
+{
+	union stack_reclaim_state prev_state, target_state;
+
+	prev_state.val = READ_ONCE(tsk->stack_reclaim_state.val);
+	do {
+		target_state.val = prev_state.val;
+		if (prev_state.stack_state == STACK_RECLAIMABLE)
+			target_state.stack_state = STACK_RECLAIMING;
+	} while (!try_cmpxchg(&tsk->stack_reclaim_state.val, &prev_state.val, target_state.val));
+
+	/*
+	 * If target_state.stack_state == STACK_RECLAIMING, we know tsk is still
+	 * alive and can't run until we're done, so putting the ref here is safe.
+	 */
+	put_task_struct(tsk);
+	if (target_state.stack_state != STACK_RECLAIMING)
+		return;
+
+	release_stack(tsk);
+
+	prev_state.val = READ_ONCE(tsk->stack_reclaim_state.val);
+	do {
+		target_state.val = prev_state.val;
+		target_state.stack_state = STACK_RECLAIMED;
+	} while (!try_cmpxchg(&tsk->stack_reclaim_state.val, &prev_state.val, target_state.val));
+
+	if (prev_state.stack_state == STACK_RECLAIMING_IN_USE) {
+		struct repopulate_work *work = get_cpu_ptr(&repopulate_work);
+
+		llist_add(&tsk->stack_reclaim_list.refill_entry, &work->stacks);
+		queue_work(system_highpri_wq, &work->work);
+		put_cpu_ptr(work);
+	}
+}
+
+static void do_repopulate_stacks(struct work_struct *w)
+{
+	struct repopulate_work *work = container_of(w, struct repopulate_work, work);
+	struct llist_node *head;
+
+	while ((head = llist_del_all(&work->stacks))) {
+		struct task_struct *tsk, *tmp;
+
+		llist_for_each_entry_safe(tsk, tmp, head, stack_reclaim_list.refill_entry) {
+			init_llist_node(&tsk->stack_reclaim_list.refill_entry);
+			if (repopulate_stack(tsk, true, &work->stacks)) {
+				wake_up_state(tsk, TASK_STACK_RECLAIM);
+			} else {
+				/*
+				 * Repopulate only failes due to low memory. If
+				 * that happens, give the rest of the system a
+				 * chance to free some memory.
+				 */
+				cond_resched();
+			}
+		}
+	}
+}
+
+bool __ensure_stack_is_present(struct task_struct *tsk, bool *need_deferred_repopulate)
+{
+	union stack_reclaim_state prev_state, target_state;
+
+	*need_deferred_repopulate = false;
+	prev_state.val = READ_ONCE(tsk->stack_reclaim_state.val);
+	do {
+		target_state.val = prev_state.val;
+
+		switch (prev_state.stack_state) {
+		case STACK_IN_USE:
+			return true;
+		case STACK_PREPARE_RECLAIM:
+		case STACK_RECLAIMABLE:
+		case STACK_RECLAIMED:
+			/*
+			 * Transitioning STACK_RECLAIM -> STACK_IN_USE won't
+			 * combine with the prior STACK_IN_USE case to lead to
+			 * tasks with unpopulated stacks running. If immediate
+			 * repopulation fails, then ttwu() puts the task in the
+			 * TASK_STACK_RECLAIM state, so the only ttwu() that can
+			 * wake up the task is after we repopulate the stack.
+			 */
+			target_state.stack_state = STACK_IN_USE;
+			break;
+		case STACK_RECLAIMING:
+			target_state.stack_state = STACK_RECLAIMING_IN_USE;
+			break;
+		case STACK_RECLAIMING_IN_USE:
+			/*
+			 * Tasks with stacks in state STACK_RECLAIMING_IN_USE
+			 * should have __state == TASK_STACK_RECLAIM state, so
+			 * ttwu_state_match() should reject any wakeups other
+			 * than the one after the stack gets repopulated.
+			 */
+			WARN(1, "TASK_STACK_RECLAIM violation");
+			return false;
+		}
+	} while (!try_cmpxchg(&tsk->stack_reclaim_state.val, &prev_state.val, target_state.val));
+
+	switch (prev_state.stack_state) {
+	case STACK_RECLAIMABLE:
+	case STACK_PREPARE_RECLAIM:
+		return true;
+	case STACK_RECLAIMING:
+		return false;
+	case STACK_RECLAIMED:
+		if (repopulate_stack(tsk, false, NULL))
+			return true;
+		*need_deferred_repopulate = true;
+		return false;
+	default:
+		// Unreachable due to return statements in cmpxchg loop
+		unreachable();
+	}
+}
+
+void __prepare_stack_for_reclaim(struct task_struct *tsk)
+{
+	union stack_reclaim_state prev_state, target_state;
+
+	// TODO: Skip rt threads
+
+	prev_state.val = READ_ONCE(tsk->stack_reclaim_state.val);
+	do {
+		target_state.val = prev_state.val;
+
+		if (prev_state.stack_state == STACK_IN_USE) {
+			target_state.stack_state = STACK_PREPARE_RECLAIM;
+		} else {
+			WARN(1, "Runnable thread with reclaimable stack state=%x",
+			     prev_state.stack_state);
+			return;
+		}
+	} while (!try_cmpxchg(&tsk->stack_reclaim_state.val, &prev_state.val, target_state.val));
+
+	/*
+	 * With delayed dequeue, __allow_stack_reclaim() can be called by
+	 * finish_task() before __block_task() calls this function. When
+	 * that happens, __allow_stack_reclaim() sees STACK_IN_USE and is
+	 * thus a no-op. We need a call here to progress the state machine.
+	 *
+	 * Note that __block_task() is called under the rq lock, so we don't
+	 * need to worry about concurrent calls.
+	 */
+	if (!tsk->on_cpu)
+		__allow_stack_reclaim(tsk);
+}
+
+void __allow_stack_reclaim(struct task_struct *tsk)
+{
+	union stack_reclaim_state prev_state, target_state;
+
+	if (WARN_ON_ONCE(tsk->__state == TASK_DEAD))
+		return;
+
+	prev_state.val = READ_ONCE(tsk->stack_reclaim_state.val);
+	do {
+		target_state.val = prev_state.val;
+
+		if (prev_state.stack_state != STACK_PREPARE_RECLAIM) {
+			WARN(prev_state.stack_state != STACK_IN_USE,
+			     "Reclaimable state %x for previously running task", prev_state.val);
+			return;
+		}
+		target_state.stack_state = STACK_RECLAIMABLE;
+	} while (!try_cmpxchg(&tsk->stack_reclaim_state.val, &prev_state.val, target_state.val));
+
+	if (irq_work_queue(&tsk->stack_reclaim_work->irq_work)) {
+		/*
+		 * Take a ref that gets released by do_reclaim_stack() so we don't
+		 * have to worry about races with remove_from_stack_shrinker().
+		 */
+		get_task_struct(tsk);
+	}
+}
+
+/*
+ * This function is called when the task is deleted, which can happen well
+ * after the task releases its stack. However, a dead task will never have
+ * TASK_STACK_RECLAIM set, so its last context switch will leave the stack
+ * state as STACK_IN_USE. As such, if the shrinker processes a task after its
+ * death, it will remove the task from the lru without accessing the stack.
+ */
+void remove_from_stack_shrinker(struct task_struct *tsk)
+{
+	put_stack_obj_cgroup(tsk);
+	kfree(tsk->stack_reclaim_work);
+}
+
+void wake_stack_repopulate(void)
+{
+	queue_work(system_highpri_wq, &this_cpu_ptr(&repopulate_work)->work);
+	preempt_enable();
+}
+
+static int stack_shrinker_cpuhp_setup(unsigned int cpu)
+{
+	struct repopulate_work *work = per_cpu_ptr(&repopulate_work, cpu);
+
+	init_llist_head(&work->stacks);
+	INIT_WORK(&work->work, do_repopulate_stacks);
+	return 0;
+}
+
+static int stack_shrinker_cpuhp_teardown(unsigned int cpu)
+{
+	flush_work(&per_cpu_ptr(&repopulate_work, cpu)->work);
+	return 0;
+}
+
+static int __init fork_late_init(void)
+{
+	int ret;
+
+	ret = cpuhp_setup_state(CPUHP_BP_PREPARE_DYN, "stack_shrinker",
+				stack_shrinker_cpuhp_setup,
+				stack_shrinker_cpuhp_teardown);
+	if (ret < 0) {
+		WARN(1, "Failed to initialize stack_shrinker cpuhp %d\n", ret);
+		return 0;
+	}
+
+	return 0;
+}
+
+module_init(fork_late_init);
diff --git a/kernel/stack_shrinker.h b/kernel/stack_shrinker.h
new file mode 100644
index 000000000000..242a35bddf3e
--- /dev/null
+++ b/kernel/stack_shrinker.h
@@ -0,0 +1,58 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+#ifndef _LINUX_STACK_SHRINKER_H
+#define _LINUX_STACK_SHRINKER_H
+
+#include <linux/cleanup.h>
+#include <linux/list.h>
+#include <linux/llist.h>
+#include <linux/sched.h>
+#include <linux/seq_file.h>
+
+#ifdef CONFIG_RECLAIMABLE_STACK
+
+void add_to_stack_shrinker(struct task_struct *tsk, int node);
+void remove_from_stack_shrinker(struct task_struct *tsk);
+
+bool __ensure_stack_is_present(struct task_struct *tsk, bool *need_deferred_repopulate);
+void __prepare_stack_for_reclaim(struct task_struct *tsk);
+void __allow_stack_reclaim(struct task_struct *tsk);
+void wake_stack_repopulate(void);
+
+static inline bool ensure_stack_is_present(struct task_struct *tsk, bool *need_deferred_repopulate)
+{
+	if (unlikely(tsk->flags & PF_RECLAIMABLE_STACK))
+		return __ensure_stack_is_present(tsk, need_deferred_repopulate);
+	return true;
+}
+
+static inline void prepare_stack_for_reclaim(struct task_struct *tsk)
+{
+	if (unlikely(tsk->flags & PF_RECLAIMABLE_STACK))
+		__prepare_stack_for_reclaim(tsk);
+}
+
+static inline void allow_stack_reclaim(struct task_struct *tsk)
+{
+	if (unlikely(tsk->flags & PF_RECLAIMABLE_STACK))
+		__allow_stack_reclaim(tsk);
+}
+
+#else /* !CONFIG_RECLAIMABLE_STACK */
+
+static inline void add_to_stack_shrinker(struct task_struct *tsk, int node) {}
+static inline void remove_from_stack_shrinker(struct task_struct *tsk) {}
+
+static inline bool ensure_stack_is_present(struct task_struct *tsk, bool *need_deferred_repopulate)
+{
+	return true;
+}
+
+static inline void prepare_stack_for_reclaim(struct task_struct *tsk) {}
+
+static inline void allow_stack_reclaim(struct task_struct *tsk) {}
+
+static inline void wake_stack_repopulate(void) {}
+
+#endif /* CONFIG_RECLAIMABLE_STACK */
+
+#endif /* _LINUX_STACK_SHRINKER_H */
-- 
2.55.0.897.gb25b4bd76c-goog

[RFC 07/10] Reclaim stacks via a shrinker

[David Stevens]

Use a shrinker to reclaim unused portions of blocked task stacks. The
core of the shrinker is a list_lru. A referenced bit is used to require
that a stack remains unused for long enough to be seen by the shrinker
twice.

To minimize how much work is done in the scheduler path and under the
pi_lock, tasks are initially put onto a regular list before being moved
onto the list_lru by a shrinker. Additionally, tasks are not proactively
removed from the lru when they become unblocked - again, to reduce what
happens in the scheduler path.

Signed-off-by: David Stevens <stevensd@google.com>
---
 include/linux/sched.h   |  13 +-
 kernel/stack_shrinker.c | 400 ++++++++++++++++++++++++++++++++++------
 2 files changed, 357 insertions(+), 56 deletions(-)

diff --git a/include/linux/sched.h b/include/linux/sched.h
index c93a234fac96..dc241d1c058f 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -83,7 +83,6 @@ struct sched_dl_entity;
 struct seq_file;
 struct sighand_struct;
 struct signal_struct;
-struct stack_reclaim_work;
 struct task_delay_info;
 struct task_exec_state;
 struct task_group;
@@ -835,9 +834,17 @@ enum stack_reclaim_enum {
 	STACK_RECLAIMED		= 5,
 } __packed;
 
+enum stack_reclaim_lru_state_enum {
+	STACK_LRU_NOT_PRESENT	= 0,
+	STACK_LRU_PENDING	= 1,
+	STACK_LRU_NEW		= 2,
+	STACK_LRU_OLD		= 3,
+} __packed;
+
 union stack_reclaim_state {
 	struct {
 		enum stack_reclaim_enum stack_state;
+		enum stack_reclaim_lru_state_enum lru_state;
 		u16 node;
 	};
 	u32 val;
@@ -845,6 +852,7 @@ union stack_reclaim_state {
 
 union stack_reclaim_list {
 	struct llist_node refill_entry;
+	struct list_head reclaim_entry;
 };
 #endif
 
@@ -1618,9 +1626,6 @@ struct task_struct {
 #ifdef CONFIG_MEMCG
 	struct obj_cgroup		*stack_obj_cgroup;
 #endif
-
-	// TODO: Replace these with a shrinker
-	struct stack_reclaim_work	*stack_reclaim_work;
 #endif
 #endif
 #ifdef CONFIG_THREAD_INFO_IN_TASK
diff --git a/kernel/stack_shrinker.c b/kernel/stack_shrinker.c
index d7b1a7dfa716..62cc5303d69b 100644
--- a/kernel/stack_shrinker.c
+++ b/kernel/stack_shrinker.c
@@ -26,28 +26,41 @@ struct repopulate_work {
 
 static DEFINE_PER_CPU(struct repopulate_work, repopulate_work);
 
-// TODO: replace with shrinker
-struct stack_reclaim_work {
-	struct task_struct *tsk;
-	struct irq_work irq_work;
-	struct work_struct work;
-};
+static DEFINE_RAW_SPINLOCK(new_reclaimable_stacks_lock);
+static LIST_HEAD(new_reclaimable_stacks);
 
-static void schedule_stack_reclaim_work(struct irq_work *w)
-{
-	struct stack_reclaim_work *work = container_of(w, typeof(*work), irq_work);
+static struct list_lru reclaimable_stacks_lru;
 
-	if (!queue_work(system_wq, &work->work))
-		put_task_struct(work->tsk);
-}
+/*
+ * do_shrink_slab() wants us to scan (nr_obj / (1 << priority)), but unless
+ * reclaim is really struggling, that can round down to 0 for a lot of
+ * memcgs. Compensate for that by scaling count and batch size.
+ */
+#define STACK_COUNT_SHIFT DEF_PRIORITY
+/*
+ * Large batch sizes (like the 128 default) can result in spiky behavior, where
+ * deferred work acculmulates and then all of a memcg's stacks get scanned all
+ * at once.
+ */
+#define STACK_BATCH_SIZE 4
 
-static void do_reclaim_stack(struct task_struct *tsk);
+/*
+ * For purposes of shrinker iteration, a stack allocated with NUMA_NO_NODE is
+ * associated with the node that created it. This extra bit allows us to
+ * determine if NUMA_NO_NODE or the saved node should be used for repopulation.
+ */
+#define TASK_NUMA_NO_NODE_FLAG BIT(15)
 
-static void do_stack_reclaim_work(struct work_struct *w)
+static int task_shrinker_node(struct task_struct *tsk)
 {
-	struct task_struct *tsk = container_of(w, struct stack_reclaim_work, work)->tsk;
+	return tsk->stack_reclaim_state.node & ~TASK_NUMA_NO_NODE_FLAG;
+}
 
-	do_reclaim_stack(tsk);
+static int task_alloc_node(struct task_struct *tsk)
+{
+	if (tsk->stack_reclaim_state.node & TASK_NUMA_NO_NODE_FLAG)
+		return NUMA_NO_NODE;
+	return tsk->stack_reclaim_state.node;
 }
 
 #ifdef CONFIG_MEMCG
@@ -85,17 +98,13 @@ void add_to_stack_shrinker(struct task_struct *tsk, int node)
 
 	tsk->stack_reclaim_state.val = 0;
 	tsk->stack_reclaim_state.stack_state = STACK_IN_USE;
-	tsk->stack_reclaim_state.node = node;
+	tsk->stack_reclaim_state.lru_state = STACK_LRU_NOT_PRESENT;
 	set_stack_obj_cgroup(tsk);
-	init_llist_node(&tsk->stack_reclaim_list.refill_entry);
+	INIT_LIST_HEAD(&tsk->stack_reclaim_list.reclaim_entry);
 
-	// TODO: replace with shrinker
-	tsk->stack_reclaim_work = kmalloc_obj(*tsk->stack_reclaim_work, GFP_KERNEL);
-	BUG_ON(!tsk->stack_reclaim_work);
-
-	tsk->stack_reclaim_work->tsk = tsk;
-	init_irq_work(&tsk->stack_reclaim_work->irq_work, schedule_stack_reclaim_work);
-	INIT_WORK(&tsk->stack_reclaim_work->work, do_stack_reclaim_work);
+	if (node == NUMA_NO_NODE)
+		node = numa_node_id() | TASK_NUMA_NO_NODE_FLAG;
+	tsk->stack_reclaim_state.node = node;
 }
 
 static inline int calculate_num_unused_pages(struct task_struct *tsk)
@@ -118,8 +127,7 @@ static bool repopulate_stack(struct task_struct *tsk, bool is_deferred,
 	struct vm_struct *vm_area = tsk->stack_vm_area;
 	unsigned long addr = (unsigned long)vm_area->addr;
 	struct mem_cgroup *tsk_memcg, *old_active_memcg;
-	int node = tsk->stack_reclaim_state.node == U16_MAX ? NUMA_NO_NODE
-							    : tsk->stack_reclaim_state.node;
+	int node = task_alloc_node(tsk);
 
 	num_missing_pages = (THREAD_SIZE >> PAGE_SHIFT) - vm_area->nr_pages;
 	if (num_missing_pages == 0)
@@ -239,21 +247,6 @@ static void do_reclaim_stack(struct task_struct *tsk)
 {
 	union stack_reclaim_state prev_state, target_state;
 
-	prev_state.val = READ_ONCE(tsk->stack_reclaim_state.val);
-	do {
-		target_state.val = prev_state.val;
-		if (prev_state.stack_state == STACK_RECLAIMABLE)
-			target_state.stack_state = STACK_RECLAIMING;
-	} while (!try_cmpxchg(&tsk->stack_reclaim_state.val, &prev_state.val, target_state.val));
-
-	/*
-	 * If target_state.stack_state == STACK_RECLAIMING, we know tsk is still
-	 * alive and can't run until we're done, so putting the ref here is safe.
-	 */
-	put_task_struct(tsk);
-	if (target_state.stack_state != STACK_RECLAIMING)
-		return;
-
 	release_stack(tsk);
 
 	prev_state.val = READ_ONCE(tsk->stack_reclaim_state.val);
@@ -271,6 +264,245 @@ static void do_reclaim_stack(struct task_struct *tsk)
 	}
 }
 
+static void process_one_new_reclaimable_stack(struct task_struct *tsk, struct list_head *new_stacks)
+{
+	union stack_reclaim_state prev_state, target_state;
+
+	prev_state.val = READ_ONCE(tsk->stack_reclaim_state.val);
+	do {
+		target_state.val = prev_state.val;
+
+		switch (prev_state.stack_state) {
+		case STACK_IN_USE:
+		case STACK_PREPARE_RECLAIM:
+			target_state.lru_state = STACK_LRU_NOT_PRESENT;
+			break;
+		case STACK_RECLAIMABLE:
+			target_state.lru_state = STACK_LRU_NEW;
+			break;
+		case STACK_RECLAIMING:
+		case STACK_RECLAIMING_IN_USE:
+		case STACK_RECLAIMED:
+			/*
+			 * These states only happen after a shrinker has started
+			 * processing a stack, so seeing one of these means
+			 * multiple shrinkers are somehow targeting one stack.
+			 */
+			WARN(1, "task with stack state %x on list\n", prev_state.val);
+			put_task_struct(tsk);
+			return;
+		}
+	} while (!try_cmpxchg(&tsk->stack_reclaim_state.val, &prev_state.val, target_state.val));
+
+	if (target_state.stack_state == STACK_RECLAIMABLE)
+		list_add(&tsk->stack_reclaim_list.reclaim_entry, new_stacks);
+	else
+		put_task_struct(tsk);
+}
+
+static void process_new_reclaimable_stacks(void)
+{
+	LIST_HEAD(new_stacks);
+	struct task_struct *tsk, *tmp;
+
+	scoped_guard(raw_spinlock_irq, &new_reclaimable_stacks_lock) {
+		while ((tsk = list_first_entry_or_null(&new_reclaimable_stacks,
+						       typeof(*tsk),
+						       stack_reclaim_list.reclaim_entry))) {
+			list_del_init(&tsk->stack_reclaim_list.reclaim_entry);
+
+			/*
+			 * We hold a ref on the task during the interval when a
+			 * stack is being moved from new_reclaimable_stacks
+			 * onto the lru, to avoid needing to deal with races
+			 * against remove_from_stack_shrinker(). If tryget
+			 * fails here, tsk is about to be deleted, so we can
+			 * just skip it.
+			 */
+			if (!tryget_task_struct(tsk))
+				continue;
+
+			raw_spin_unlock_irq(&new_reclaimable_stacks_lock);
+
+			process_one_new_reclaimable_stack(tsk, &new_stacks);
+
+			raw_spin_lock_irq(&new_reclaimable_stacks_lock);
+		}
+	}
+
+	list_for_each_entry_safe(tsk, tmp, &new_stacks, stack_reclaim_list.reclaim_entry) {
+		struct mem_cgroup *memcg = get_stack_memcg(tsk);
+
+		list_del_init(&tsk->stack_reclaim_list.reclaim_entry);
+
+		if (memcg_list_lru_alloc(memcg, &reclaimable_stacks_lru, GFP_ATOMIC) == 0) {
+			local_bh_disable();
+			list_lru_add(&reclaimable_stacks_lru,
+				     &tsk->stack_reclaim_list.reclaim_entry,
+				     task_shrinker_node(tsk), memcg);
+			local_bh_enable();
+		} else {
+			union stack_reclaim_state prev_state, target_state;
+
+			pr_warn_ratelimited("failed to allocate stack reclaim metadata\n");
+
+			/*
+			 * Just give up if memcg_list_lru_alloc() fails. We
+			 * could try immediately reclaiming the stack, but
+			 * reclaiming a single stack isn't going to help if
+			 * things are so bad a GFP_ATOMIC slab allocation fails.
+			 */
+			prev_state.val = READ_ONCE(tsk->stack_reclaim_state.val);
+			do {
+				target_state.val = prev_state.val;
+				if (prev_state.stack_state == STACK_RECLAIMABLE)
+					target_state.stack_state = STACK_RECLAIMED;
+				target_state.lru_state = STACK_LRU_NOT_PRESENT;
+			} while (!try_cmpxchg(&tsk->stack_reclaim_state.val,
+					      &prev_state.val, target_state.val));
+		}
+
+		mem_cgroup_put(memcg);
+		put_task_struct(tsk);
+	}
+}
+
+static enum lru_status isolate_lru_stack(struct list_head *item,
+					 struct list_lru_one *lru, void *arg)
+{
+	struct list_head *to_reclaim = arg;
+	struct task_struct *tsk = container_of(item, struct task_struct,
+					       stack_reclaim_list.reclaim_entry);
+	union stack_reclaim_state prev_state, target_state;
+	enum lru_status ret;
+	bool is_isolated = false;
+
+	prev_state.val = READ_ONCE(tsk->stack_reclaim_state.val);
+	do {
+		target_state.val = prev_state.val;
+
+		switch (prev_state.stack_state) {
+		case STACK_IN_USE:
+		case STACK_PREPARE_RECLAIM:
+			target_state.lru_state = STACK_LRU_NOT_PRESENT;
+			ret = LRU_REMOVED;
+			break;
+		case STACK_RECLAIMABLE:
+			switch (prev_state.lru_state) {
+			case STACK_LRU_NEW:
+				target_state.lru_state = STACK_LRU_OLD;
+				ret = LRU_ROTATE;
+				break;
+			case STACK_LRU_OLD:
+				target_state.stack_state = STACK_RECLAIMING;
+				target_state.lru_state = STACK_LRU_NOT_PRESENT;
+				ret = LRU_REMOVED;
+				break;
+			case STACK_LRU_NOT_PRESENT:
+			case STACK_LRU_PENDING:
+				WARN(1, "item on stack reclaim lru with bad state\n");
+				list_lru_isolate(lru, item);
+				return LRU_REMOVED;
+			}
+			break;
+		case STACK_RECLAIMING:
+		case STACK_RECLAIMING_IN_USE:
+		case STACK_RECLAIMED:
+			/*
+			 * These states only happen after a shrinker has started
+			 * processing a stack, so seeing one of these means
+			 * multiple shrinkers are somehow targeting one stack.
+			 */
+			WARN(1, "stack with state %x on list\n", prev_state.val);
+			list_lru_isolate(lru, item);
+			return LRU_REMOVED;
+		}
+
+		/*
+		 * We need to isolate the item before the cmpxchg to prevent
+		 * races with process_one_new_reclaimable_stack() adding the
+		 * entry to its new_stacks list.
+		 */
+		if (ret == LRU_REMOVED && !is_isolated) {
+			is_isolated = true;
+			list_lru_isolate(lru, item);
+		}
+	} while (!try_cmpxchg(&tsk->stack_reclaim_state.val, &prev_state.val, target_state.val));
+
+	if (target_state.stack_state == STACK_RECLAIMING) {
+		/*
+		 * The task can't be freed since STACK_RECLAIMING prevents
+		 * it from running and exiting, so no need to hold a ref.
+		 */
+		list_add_tail(item, to_reclaim);
+	} else if (target_state.stack_state == STACK_RECLAIMABLE) {
+		/*
+		 * If we isolated but then lost a cmpxchg race against
+		 * __allow_stack_reclaim(), we need to undo the isolation.
+		 *
+		 * STACK_RECLAIMABLE means we observed the task blocked (i.e.
+		 * not dead). Even if the task dies and its refcount hits zero,
+		 * RCU cleanup of the task will be delayed because the lru walk
+		 * is guarded by local_bh_disable(), so unlocking won't cause
+		 * races with remove_from_stack_shrinker().
+		 */
+		if (is_isolated) {
+			struct mem_cgroup *lru_memcg;
+
+			spin_unlock(&lru->lock);
+
+			lru_memcg = get_stack_memcg(tsk);
+			list_lru_add(&reclaimable_stacks_lru, item,
+				     task_shrinker_node(tsk), lru_memcg);
+			mem_cgroup_put(lru_memcg);
+
+			ret = LRU_REMOVED_RETRY;
+
+		} else {
+			WARN(target_state.lru_state != STACK_LRU_OLD,
+			     "Reclaimable stack with bad state %x\n", target_state.val);
+		}
+	} else {
+		WARN(target_state.lru_state != STACK_LRU_NOT_PRESENT,
+		     "Isolate stack with bad state %x\n", target_state.val);
+	}
+
+	return ret;
+}
+
+static unsigned long scan_reclaimable_stacks(struct shrinker *shrinker,
+					     struct shrink_control *sc)
+{
+	unsigned long freed = 0;
+	LIST_HEAD(to_reclaim);
+	struct task_struct *tsk, *tmp;
+
+	sc->nr_to_scan >>= STACK_COUNT_SHIFT;
+	local_bh_disable();
+	list_lru_shrink_walk(&reclaimable_stacks_lru, sc,
+			     isolate_lru_stack, &to_reclaim);
+	local_bh_enable();
+
+	list_for_each_entry_safe(tsk, tmp, &to_reclaim, stack_reclaim_list.reclaim_entry) {
+		freed++;
+		list_del_init(&tsk->stack_reclaim_list.reclaim_entry);
+		do_reclaim_stack(tsk);
+	}
+
+	return freed << STACK_COUNT_SHIFT;
+}
+
+static unsigned long get_reclaimable_stack_count(struct shrinker *shrinker,
+						 struct shrink_control *sc)
+{
+	unsigned long count;
+
+	process_new_reclaimable_stacks();
+	count = list_lru_shrink_count(&reclaimable_stacks_lru, sc);
+
+	return count ? (unsigned long)(count << STACK_COUNT_SHIFT) : SHRINK_EMPTY;
+}
+
 static void do_repopulate_stacks(struct work_struct *w)
 {
 	struct repopulate_work *work = container_of(w, struct repopulate_work, work);
@@ -285,7 +517,7 @@ static void do_repopulate_stacks(struct work_struct *w)
 				wake_up_state(tsk, TASK_STACK_RECLAIM);
 			} else {
 				/*
-				 * Repopulate only failes due to low memory. If
+				 * Repopulate only fails due to low memory. If
 				 * that happens, give the rest of the system a
 				 * chance to free some memory.
 				 */
@@ -387,6 +619,7 @@ void __prepare_stack_for_reclaim(struct task_struct *tsk)
 void __allow_stack_reclaim(struct task_struct *tsk)
 {
 	union stack_reclaim_state prev_state, target_state;
+	bool add_to_list = false;
 
 	if (WARN_ON_ONCE(tsk->__state == TASK_DEAD))
 		return;
@@ -401,14 +634,17 @@ void __allow_stack_reclaim(struct task_struct *tsk)
 			return;
 		}
 		target_state.stack_state = STACK_RECLAIMABLE;
+		if (prev_state.lru_state == STACK_LRU_NOT_PRESENT) {
+			target_state.lru_state = STACK_LRU_PENDING;
+			add_to_list = true;
+		} else if (prev_state.lru_state == STACK_LRU_OLD) {
+			target_state.lru_state = STACK_LRU_NEW;
+		}
 	} while (!try_cmpxchg(&tsk->stack_reclaim_state.val, &prev_state.val, target_state.val));
 
-	if (irq_work_queue(&tsk->stack_reclaim_work->irq_work)) {
-		/*
-		 * Take a ref that gets released by do_reclaim_stack() so we don't
-		 * have to worry about races with remove_from_stack_shrinker().
-		 */
-		get_task_struct(tsk);
+	if (add_to_list) {
+		guard(raw_spinlock_irqsave)(&new_reclaimable_stacks_lock);
+		list_add(&tsk->stack_reclaim_list.reclaim_entry, &new_reclaimable_stacks);
 	}
 }
 
@@ -421,8 +657,34 @@ void __allow_stack_reclaim(struct task_struct *tsk)
  */
 void remove_from_stack_shrinker(struct task_struct *tsk)
 {
+	struct mem_cgroup *lru_memcg;
+
+	/*
+	 * Since tsk is being deleted, the tryget_task_struct() call in
+	 * process_new_reclaimable_stacks() excludes racing with a shrinker
+	 * moving the task from STACK_LRU_PENDING -> STACK_LRU_NEW. A race can
+	 * just result in the delete operation being a no-op.
+	 */
+	switch (READ_ONCE(tsk->stack_reclaim_state.lru_state)) {
+	case STACK_LRU_NOT_PRESENT:
+		break;
+	case STACK_LRU_PENDING:
+		scoped_guard(raw_spinlock_irqsave, &new_reclaimable_stacks_lock) {
+			if (!list_empty(&tsk->stack_reclaim_list.reclaim_entry))
+				list_del_init(&tsk->stack_reclaim_list.reclaim_entry);
+		}
+		break;
+	case STACK_LRU_NEW:
+	case STACK_LRU_OLD:
+		lru_memcg = get_stack_memcg(tsk);
+		local_bh_disable();
+		list_lru_del(&reclaimable_stacks_lru, &tsk->stack_reclaim_list.reclaim_entry,
+			     task_shrinker_node(tsk), lru_memcg);
+		local_bh_enable();
+		mem_cgroup_put(lru_memcg);
+		break;
+	}
 	put_stack_obj_cgroup(tsk);
-	kfree(tsk->stack_reclaim_work);
 }
 
 void wake_stack_repopulate(void)
@@ -431,6 +693,8 @@ void wake_stack_repopulate(void)
 	preempt_enable();
 }
 
+static struct lock_class_key stack_shrinker_key;
+
 static int stack_shrinker_cpuhp_setup(unsigned int cpu)
 {
 	struct repopulate_work *work = per_cpu_ptr(&repopulate_work, cpu);
@@ -448,16 +712,48 @@ static int stack_shrinker_cpuhp_teardown(unsigned int cpu)
 
 static int __init fork_late_init(void)
 {
+	struct shrinker *shrinker;
+	const char *msg;
 	int ret;
+	enum cpuhp_state cpuhp_val;
 
 	ret = cpuhp_setup_state(CPUHP_BP_PREPARE_DYN, "stack_shrinker",
 				stack_shrinker_cpuhp_setup,
 				stack_shrinker_cpuhp_teardown);
 	if (ret < 0) {
-		WARN(1, "Failed to initialize stack_shrinker cpuhp %d\n", ret);
-		return 0;
+		msg = "cpuhp failure";
+		goto cpuhp_setup_fail;
+	}
+	cpuhp_val = ret;
+
+	shrinker = shrinker_alloc(SHRINKER_NUMA_AWARE | SHRINKER_MEMCG_AWARE,
+				  "stack_shrinker");
+	if (!shrinker) {
+		msg = "shrinker alloc failure";
+		ret = -ENOMEM;
+		goto shrinker_alloc_fail;
 	}
 
+	ret = list_lru_init_memcg_key(&reclaimable_stacks_lru, shrinker,
+				      &stack_shrinker_key);
+	if (ret != 0) {
+		msg = "list_lru_init failure";
+		goto list_lru_init_fail;
+	}
+
+	shrinker->count_objects = get_reclaimable_stack_count;
+	shrinker->scan_objects = scan_reclaimable_stacks;
+	shrinker->seeks = 4;
+	shrinker->batch = STACK_BATCH_SIZE << STACK_COUNT_SHIFT;
+	shrinker_register(shrinker);
+	return 0;
+
+list_lru_init_fail:
+	shrinker_free(shrinker);
+shrinker_alloc_fail:
+	cpuhp_remove_state(cpuhp_val);
+cpuhp_setup_fail:
+	WARN(1, "Failed to initialize stack_shrinker %s: %d\n", msg, ret);
 	return 0;
 }
 
-- 
2.55.0.897.gb25b4bd76c-goog

[RFC 08/10] Set PF_RECLAIMABLE_STACK in various places

[David Stevens]

Annotate various blocking locations with the PF_RECLAIMABLE_STACK.

Signed-off-by: David Stevens <stevensd@google.com>
---
 drivers/android/binder/thread.rs | 14 +++++++++++++
 fs/eventpoll.c                   |  3 +++
 fs/pipe.c                        | 28 ++++++++++++++++---------
 fs/select.c                      |  3 +++
 kernel/futex/waitwake.c          |  3 +++
 kernel/signal.c                  | 36 +++++++++++++++++++-------------
 kernel/time/hrtimer.c            |  3 +++
 rust/kernel/task.rs              | 16 ++++++++++++++
 8 files changed, 82 insertions(+), 24 deletions(-)

diff --git a/drivers/android/binder/thread.rs b/drivers/android/binder/thread.rs
index bc0ef8927905..be62f7fd43ca 100644
--- a/drivers/android/binder/thread.rs
+++ b/drivers/android/binder/thread.rs
@@ -538,13 +538,21 @@ fn get_work_local(self: &Arc<Self>, wait: bool) -> Result<Option<DLArc<dyn Deliv
 
         // Loop waiting only on the local queue (i.e., not registering with the process queue).
         let mut inner = self.inner.lock();
+        // SAFETY: Only accessed locally in this function
+        let current = unsafe { Task::current() };
         loop {
             if let Some(work) = inner.pop_work() {
                 return Ok(Some(work));
             }
 
             inner.looper_flags |= LOOPER_WAITING;
+
+            current.set_flag_bits(bindings::PF_RECLAIMABLE_STACK);
+
             let signal_pending = self.work_condvar.wait_interruptible_freezable(&mut inner);
+
+            current.clear_flag_bits(bindings::PF_RECLAIMABLE_STACK);
+
             inner.looper_flags &= !LOOPER_WAITING;
 
             if signal_pending {
@@ -592,15 +600,21 @@ fn get_work(self: &Arc<Self>, wait: bool) -> Result<Option<DLArc<dyn DeliverToRe
         };
 
         let mut inner = self.inner.lock();
+        // SAFETY: Only accessed locally in this function
+        let current = unsafe { Task::current() };
         loop {
             if let Some(work) = inner.pop_work() {
                 return Ok(Some(work));
             }
 
+            current.set_flag_bits(bindings::PF_RECLAIMABLE_STACK);
+
             inner.looper_flags |= LOOPER_WAITING | LOOPER_WAITING_PROC;
             let signal_pending = self.work_condvar.wait_interruptible_freezable(&mut inner);
             inner.looper_flags &= !(LOOPER_WAITING | LOOPER_WAITING_PROC);
 
+            current.clear_flag_bits(bindings::PF_RECLAIMABLE_STACK);
+
             if signal_pending || inner.looper_need_return {
                 // We need to return now. We need to pull the thread off the list of ready threads
                 // (by dropping `reg`), then check the state again after it's off the list to
diff --git a/fs/eventpoll.c b/fs/eventpoll.c
index eed8cecd94e3..4e9c09446b48 100644
--- a/fs/eventpoll.c
+++ b/fs/eventpoll.c
@@ -9,6 +9,7 @@
 #include <linux/init.h>
 #include <linux/kernel.h>
 #include <linux/sched/signal.h>
+#include <linux/sched/task_stack.h>
 #include <linux/fs.h>
 #include <linux/file.h>
 #include <linux/signal.h>
@@ -2302,6 +2303,8 @@ static int ep_poll(struct eventpoll *ep, struct epoll_event __user *events,
 		if (signal_pending(current))
 			return -EINTR;
 
+		guard(allow_stack_reclaim)();
+
 		/*
 		 * Internally init_wait() uses autoremove_wake_function(),
 		 * thus wait entry is removed from the wait queue on each
diff --git a/fs/pipe.c b/fs/pipe.c
index 429b0714ec57..15503b12fe8f 100644
--- a/fs/pipe.c
+++ b/fs/pipe.c
@@ -27,6 +27,7 @@
 #include <linux/watch_queue.h>
 #include <linux/sysctl.h>
 #include <linux/sort.h>
+#include <linux/sched/task_stack.h>
 
 #include <linux/uaccess.h>
 #include <asm/ioctls.h>
@@ -469,16 +470,23 @@ anon_pipe_read(struct kiocb *iocb, struct iov_iter *to)
 			break;
 		}
 		mutex_unlock(&pipe->mutex);
-		/*
-		 * We only get here if we didn't actually read anything.
-		 *
-		 * But because we didn't read anything, at this point we can
-		 * just return directly with -ERESTARTSYS if we're interrupted,
-		 * since we've done any required wakeups and there's no need
-		 * to mark anything accessed. And we've dropped the lock.
-		 */
-		if (wait_event_interruptible_exclusive(pipe->rd_wait, pipe_readable(pipe)) < 0)
-			return -ERESTARTSYS;
+
+		{
+			guard(allow_stack_reclaim)();
+			/*
+			 * We only get here if we didn't actually read
+			 * anything.
+			 *
+			 * But because we didn't read anything, at this point
+			 * we can just return directly with -ERESTARTSYS if
+			 * we're interrupted, since we've done any required
+			 * wakeups and there's no need to mark anything
+			 * accessed. And we've dropped the lock.
+			 */
+			if (wait_event_interruptible_exclusive(pipe->rd_wait,
+							       pipe_readable(pipe)) < 0)
+				return -ERESTARTSYS;
+		}
 
 		wake_next_reader = true;
 		mutex_lock(&pipe->mutex);
diff --git a/fs/select.c b/fs/select.c
index 95d76531015a..3af1bf4d74a9 100644
--- a/fs/select.c
+++ b/fs/select.c
@@ -19,6 +19,7 @@
 #include <linux/kernel.h>
 #include <linux/sched/signal.h>
 #include <linux/sched/rt.h>
+#include <linux/sched/task_stack.h>
 #include <linux/syscalls.h>
 #include <linux/export.h>
 #include <linux/slab.h>
@@ -236,6 +237,8 @@ static int poll_schedule_timeout(struct poll_wqueues *pwq, int state,
 {
 	int rc = -EINTR;
 
+	guard(allow_stack_reclaim)();
+
 	set_current_state(state);
 	if (!READ_ONCE(pwq->triggered))
 		rc = schedule_hrtimeout_range(expires, slack, HRTIMER_MODE_ABS);
diff --git a/kernel/futex/waitwake.c b/kernel/futex/waitwake.c
index d4483d15d30a..0fbf7bfcd904 100644
--- a/kernel/futex/waitwake.c
+++ b/kernel/futex/waitwake.c
@@ -2,6 +2,7 @@
 
 #include <linux/plist.h>
 #include <linux/sched/task.h>
+#include <linux/sched/task_stack.h>
 #include <linux/sched/signal.h>
 #include <linux/freezer.h>
 
@@ -378,6 +379,7 @@ void futex_do_wait(struct futex_q *q, struct hrtimer_sleeper *timeout)
 	 * has tried to wake us, and we can skip the call to schedule().
 	 */
 	if (likely(!plist_node_empty(&q->list))) {
+		guard(allow_stack_reclaim)();
 		/*
 		 * If the timer has already expired, current will already be
 		 * flagged for rescheduling. Only call schedule if there
@@ -547,6 +549,7 @@ static void futex_sleep_multiple(struct futex_vector *vs, unsigned int count,
 			return;
 	}
 
+	guard(allow_stack_reclaim)();
 	schedule();
 }
 
diff --git a/kernel/signal.c b/kernel/signal.c
index bbc0fd4cc4d7..ca0d79ee013f 100644
--- a/kernel/signal.c
+++ b/kernel/signal.c
@@ -2718,17 +2718,21 @@ static void do_freezer_trap(void)
 		return;
 	}
 
-	/*
-	 * Now we're sure that there is no pending fatal signal and no
-	 * pending traps. Clear TIF_SIGPENDING to not get out of schedule()
-	 * immediately (if there is a non-fatal signal pending), and
-	 * put the task into sleep.
-	 */
-	__set_current_state(TASK_INTERRUPTIBLE|TASK_FREEZABLE);
-	clear_thread_flag(TIF_SIGPENDING);
-	spin_unlock_irq(&current->sighand->siglock);
-	cgroup_enter_frozen();
-	schedule();
+	{
+		guard(allow_stack_reclaim)();
+
+		/*
+		 * Now we're sure that there is no pending fatal signal and no
+		 * pending traps. Clear TIF_SIGPENDING to not get out of schedule()
+		 * immediately (if there is a non-fatal signal pending), and
+		 * put the task into sleep.
+		 */
+		__set_current_state(TASK_INTERRUPTIBLE | TASK_FREEZABLE);
+		clear_thread_flag(TIF_SIGPENDING);
+		spin_unlock_irq(&current->sighand->siglock);
+		cgroup_enter_frozen();
+		schedule();
+	}
 
 	/*
 	 * We could've been woken by task_work, run it to clear
@@ -3788,9 +3792,13 @@ static int do_sigtimedwait(const sigset_t *which, kernel_siginfo_t *info,
 		recalc_sigpending();
 		spin_unlock_irq(&tsk->sighand->siglock);
 
-		__set_current_state(TASK_INTERRUPTIBLE|TASK_FREEZABLE);
-		ret = schedule_hrtimeout_range(to, tsk->timer_slack_ns,
-					       HRTIMER_MODE_REL);
+		{
+			guard(allow_stack_reclaim)();
+			__set_current_state(TASK_INTERRUPTIBLE | TASK_FREEZABLE);
+			ret = schedule_hrtimeout_range(to, tsk->timer_slack_ns,
+						       HRTIMER_MODE_REL);
+		}
+
 		spin_lock_irq(&tsk->sighand->siglock);
 		__set_task_blocked(tsk, &tsk->real_blocked);
 		sigemptyset(&tsk->real_blocked);
diff --git a/kernel/time/hrtimer.c b/kernel/time/hrtimer.c
index 313dcea127fe..30c85482dfd5 100644
--- a/kernel/time/hrtimer.c
+++ b/kernel/time/hrtimer.c
@@ -39,6 +39,7 @@
 #include <linux/sched/nohz.h>
 #include <linux/sched/debug.h>
 #include <linux/sched/isolation.h>
+#include <linux/sched/task_stack.h>
 #include <linux/timer.h>
 #include <linux/freezer.h>
 #include <linux/compat.h>
@@ -2392,6 +2393,8 @@ static int __sched do_nanosleep(struct hrtimer_sleeper *t, enum hrtimer_mode mod
 	struct restart_block *restart;
 
 	do {
+		guard(allow_stack_reclaim)();
+
 		set_current_state(TASK_INTERRUPTIBLE|TASK_FREEZABLE);
 		hrtimer_sleeper_start_expires(t, mode);
 
diff --git a/rust/kernel/task.rs b/rust/kernel/task.rs
index 38273f4eedb5..6168f058343a 100644
--- a/rust/kernel/task.rs
+++ b/rust/kernel/task.rs
@@ -344,6 +344,22 @@ pub fn group_leader(&self) -> &Task {
         // only be used while `current` is still valid, thus still running.
         unsafe { &*ptr.cast() }
     }
+
+    /// Sets the given task flag bits on the current task.
+    #[inline]
+    pub fn set_flag_bits(&self, set: u32) {
+            // SAFETY: The `flags` field of `current` is not modified from other threads, so
+	    // the non-atomic update isn't a race.
+	    unsafe { (*self.as_ptr()).flags |= set }
+    }
+
+    /// Clears the given task flag bits on the current task.
+    #[inline]
+    pub fn clear_flag_bits(&self, clear: u32) {
+            // SAFETY: The `flags` field of `current` is not modified from other threads, so
+	    // the non-atomic update isn't a race.
+	    unsafe { (*self.as_ptr()).flags &= !clear }
+    }
 }
 
 // SAFETY: The type invariants guarantee that `Task` is always refcounted.
-- 
2.55.0.897.gb25b4bd76c-goog

[RFC 09/10] x86: Enable reclaimable stacks

[David Stevens]

Implement top_of_blocked_task_stack() to get the stack pointer out of
a thread_struct and enable HAVE_ARCH_RECLAIMABLE_STACK.

Signed-off-by: David Stevens <stevensd@google.com>
---
 arch/x86/Kconfig                 | 1 +
 arch/x86/include/asm/processor.h | 5 +++++
 2 files changed, 6 insertions(+)

diff --git a/arch/x86/Kconfig b/arch/x86/Kconfig
index bdad90f210e4..652d5a9c0830 100644
--- a/arch/x86/Kconfig
+++ b/arch/x86/Kconfig
@@ -211,6 +211,7 @@ config X86
 	select HAVE_ARCH_USERFAULTFD_WP         if X86_64 && USERFAULTFD
 	select HAVE_ARCH_USERFAULTFD_MINOR	if X86_64 && USERFAULTFD
 	select HAVE_ARCH_VMAP_STACK		if X86_64
+	select HAVE_ARCH_RECLAIMABLE_STACK	if X86_64
 	select HAVE_ARCH_RANDOMIZE_KSTACK_OFFSET
 	select HAVE_ARCH_WITHIN_STACK_FRAMES
 	select HAVE_ASM_MODVERSIONS
diff --git a/arch/x86/include/asm/processor.h b/arch/x86/include/asm/processor.h
index 87b1d4c0727e..207e7adced64 100644
--- a/arch/x86/include/asm/processor.h
+++ b/arch/x86/include/asm/processor.h
@@ -690,6 +690,11 @@ extern void start_thread(struct pt_regs *regs, unsigned long new_ip,
 #define GET_TSC_CTL(adr)	get_tsc_mode((adr))
 #define SET_TSC_CTL(val)	set_tsc_mode((val))
 
+static inline unsigned long top_of_blocked_task_stack(struct thread_struct *thread)
+{
+	return thread->sp;
+}
+
 extern int get_tsc_mode(unsigned long adr);
 extern int set_tsc_mode(unsigned int val);
 
-- 
2.55.0.897.gb25b4bd76c-goog

[RFC 10/10] arm64: Enable reclaimable stacks

[David Stevens]

Implement top_of_blocked_task_stack() to get the stack pointer out of
a thread_struct and enable HAVE_ARCH_RECLAIMABLE_STACK.

Signed-off-by: David Stevens <stevensd@google.com>
---
 arch/arm64/Kconfig                 | 1 +
 arch/arm64/include/asm/processor.h | 5 +++++
 2 files changed, 6 insertions(+)

diff --git a/arch/arm64/Kconfig b/arch/arm64/Kconfig
index b3afe0688919..262a4fddce19 100644
--- a/arch/arm64/Kconfig
+++ b/arch/arm64/Kconfig
@@ -173,6 +173,7 @@ config ARM64
 	select HAVE_ARCH_TRACEHOOK
 	select HAVE_ARCH_TRANSPARENT_HUGEPAGE
 	select HAVE_ARCH_VMAP_STACK
+	select HAVE_ARCH_RECLAIMABLE_STACK
 	select HAVE_ARM_SMCCC
 	select HAVE_ASM_MODVERSIONS
 	select HAVE_EBPF_JIT
diff --git a/arch/arm64/include/asm/processor.h b/arch/arm64/include/asm/processor.h
index c2a627f39314..f5a550835c94 100644
--- a/arch/arm64/include/asm/processor.h
+++ b/arch/arm64/include/asm/processor.h
@@ -204,6 +204,11 @@ struct thread_struct {
 #endif
 };
 
+static inline unsigned long top_of_blocked_task_stack(struct thread_struct *thread)
+{
+	return thread->cpu_context.sp;
+}
+
 static inline unsigned int thread_get_vl(struct thread_struct *thread,
 					 enum vec_type type)
 {
-- 
2.55.0.897.gb25b4bd76c-goog

Re: [RFC 08/10] Set PF_RECLAIMABLE_STACK in various places

[sashiko-bot]

Thank you for your contribution! Sashiko AI review found 2 potential issue(s) to consider:
- [High] Safe Rust methods `set_flag_bits` and `clear_flag_bits` on `Task` allow data races when called on non-current tasks.
- [Medium] Unconditional clearing of `PF_RECLAIMABLE_STACK` in Rust binder breaks flag preservation logic.
--

commit 6f8ac02f62f897200c71ad824d4019095dc4428c
Author: David Stevens <stevensd@google.com>

Set PF_RECLAIMABLE_STACK in various places

Annotate various blocking locations with the PF_RECLAIMABLE_STACK.

> diff --git a/drivers/android/binder/thread.rs b/drivers/android/binder/thread.rs
> index bc0ef8927905c..be62f7fd43ca2 100644
> --- a/drivers/android/binder/thread.rs
> +++ b/drivers/android/binder/thread.rs
> @@ -538,13 +538,21 @@ fn get_work_local(self: &Arc<Self>, wait: bool) -> Result<Option<DLArc<dyn Deliv
[ ... ]
>              inner.looper_flags |= LOOPER_WAITING;
> +
> +            current.set_flag_bits(bindings::PF_RECLAIMABLE_STACK);
> +
>              let signal_pending = self.work_condvar.wait_interruptible_freezable(&mut inner);
> +
> +            current.clear_flag_bits(bindings::PF_RECLAIMABLE_STACK);

[Severity: Medium]
Will unconditionally clearing the flag break flag preservation logic if stack
reclaim states are nested?

The C allow_stack_reclaim guard explicitly saves and restores the flag to
support nesting. If get_work_local() is called from a context where
PF_RECLAIMABLE_STACK is already set, this will unexpectedly disable stack
reclamation for the outer scope upon exit.

> +
>              inner.looper_flags &= !LOOPER_WAITING;
>  
>              if signal_pending {
> @@ -592,15 +600,21 @@ fn get_work(self: &Arc<Self>, wait: bool) -> Result<Option<DLArc<dyn DeliverToRe
[ ... ]
> +            current.set_flag_bits(bindings::PF_RECLAIMABLE_STACK);
> +
>              inner.looper_flags |= LOOPER_WAITING | LOOPER_WAITING_PROC;
>              let signal_pending = self.work_condvar.wait_interruptible_freezable(&mut inner);
>              inner.looper_flags &= !(LOOPER_WAITING | LOOPER_WAITING_PROC);
>  
> +            current.clear_flag_bits(bindings::PF_RECLAIMABLE_STACK);

[Severity: Medium]
Similar to get_work_local(), is it safe to unconditionally clear the flag here
in get_work() instead of restoring its previous state?

> +
>              if signal_pending || inner.looper_need_return {
>                  // We need to return now. We need to pull the thread off the list of ready threads
>                  // (by dropping `reg`), then check the state again after it's off the list to
> diff --git a/rust/kernel/task.rs b/rust/kernel/task.rs
> index 38273f4eedb51..6168f058343a2 100644
> --- a/rust/kernel/task.rs
> +++ b/rust/kernel/task.rs
> @@ -344,6 +344,22 @@ pub fn group_leader(&self) -> &Task {
>          // only be used while `current` is still valid, thus still running.
>          unsafe { &*ptr.cast() }
>      }
> +
> +    /// Sets the given task flag bits on the current task.
> +    #[inline]
> +    pub fn set_flag_bits(&self, set: u32) {
> +            // SAFETY: The `flags` field of `current` is not modified from other threads, so
> +	    // the non-atomic update isn't a race.
> +	    unsafe { (*self.as_ptr()).flags |= set }
> +    }

[Severity: High]
Does this expose a soundness hole by allowing data races on the flags field
when called on a non-current task?

The safety comment assumes self is the current task, but the method signature
accepts any Task reference. Because Task implements Send + Sync, another thread
could obtain a reference to a non-current task (for example, via group_leader())
and safely call this method. This allows an unsafe non-atomic read-modify-write
and undefined behavior in safe Rust.

> +
> +    /// Clears the given task flag bits on the current task.
> +    #[inline]
> +    pub fn clear_flag_bits(&self, clear: u32) {
> +            // SAFETY: The `flags` field of `current` is not modified from other threads, so
> +	    // the non-atomic update isn't a race.
> +	    unsafe { (*self.as_ptr()).flags &= !clear }
> +    }
>  }

[Severity: High]
Does this method suffer from the same soundness hole as set_flag_bits() since it
is a safe method performing a non-atomic update without verifying self is the
current task?

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20260827232948.2520558-1-stevensd@google.com?part=8

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[sashiko-bot]

Thank you for your contribution! Sashiko AI review found 2 potential issue(s) to consider:
- [High] Memory leak of `obj_cgroup` in the `copy_process()` error path.
- [Medium] Repopulating kernel stack pages with `NUMA_NO_NODE` silently bypasses the task's NUMA mempolicy.
--

commit 19da8486785bcd64133ab34a423835e83ba21942
Author: David Stevens <stevensd@google.com>
Subject: Reclaim memory from blocked kernel stacks

This commit introduces a mechanism to reclaim unused portions of kernel
stacks for blocked tasks. It adds scheduler hooks and a new task flag
PF_RECLAIMABLE_STACK to manage safe stack reclaim and repopulation.

> diff --git a/kernel/fork.c b/kernel/fork.c
> index 6acad0038b78e..9b2cc3d01dd10 100644
> --- a/kernel/fork.c
> +++ b/kernel/fork.c
[ ... ]
> @@ -460,6 +462,7 @@ static int alloc_thread_stack_node(struct task_struct *tsk, int node)
>  
>  		tsk->stack_vm_area = vm_area;
>  		tsk->stack = stack;
> +		add_to_stack_shrinker(tsk, node);
>  		return 0;
>  	}

[Severity: High]
Does this leak the obj_cgroup reference in the copy_process() error path?

When dup_task_struct() completes successfully, we acquire an obj_cgroup
reference inside add_to_stack_shrinker(). However, if a subsequent step in
copy_process() fails (such as copy_mm()), the kernel jumps to the 
bad_fork_free label:

kernel/fork.c:copy_process() {
    ...
    bad_fork_free:
        WRITE_ONCE(p->__state, TASK_DEAD);
        exit_task_stack_account(p);
        put_task_stack(p);
        delayed_free_task(p);
    ...
}

Because delayed_free_task() directly frees the task without calling
__put_task_struct() or remove_from_stack_shrinker(), it seems the
obj_cgroup reference is permanently leaked, preventing the associated
mem_cgroup from being destroyed.

> diff --git a/kernel/stack_shrinker.c b/kernel/stack_shrinker.c
> new file mode 100644
> index 0000000000000..d7b1a7dfa7165
> --- /dev/null
> +++ b/kernel/stack_shrinker.c
[ ... ]
> +static bool repopulate_stack(struct task_struct *tsk, bool is_deferred,
> +			     struct llist_head *fail_list)
> +{
[ ... ]
> +	if (is_deferred) {
> +		gfp_t gfp = GFP_KERNEL_ACCOUNT | __GFP_ZERO;
[ ... ]
> +		for (; nr_allocated < num_missing_pages; nr_allocated++) {
> +			pages[nr_allocated] = alloc_pages_node_noprof(node, gfp, 0);

[Severity: Medium]
Can this silently bypass the task's NUMA mempolicy?

If a task has no specific node forced, the target node is set to NUMA_NO_NODE. 

Calling alloc_pages_node_noprof() with NUMA_NO_NODE unconditionally ignores
the task's explicit NUMA mempolicy (like MPOL_BIND or MPOL_INTERLEAVE) and
allocates from the current CPU's local node instead.

Should this branch and use alloc_pages_noprof() when the target node is 
NUMA_NO_NODE to properly respect the policy?

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20260827232948.2520558-1-stevensd@google.com?part=6

Re: [RFC 08/10] Set PF_RECLAIMABLE_STACK in various places

[K Prateek Nayak]

Hello David,

On 8/28/2026 4:59 AM, David Stevens wrote:
> @@ -469,16 +470,23 @@ anon_pipe_read(struct kiocb *iocb, struct iov_iter *to)
>                         break;
>                 }
>                 mutex_unlock(&pipe->mutex);
> -               /*
> -                * We only get here if we didn't actually read anything.
> -                *
> -                * But because we didn't read anything, at this point we can
> -                * just return directly with -ERESTARTSYS if we're interrupted,
> -                * since we've done any required wakeups and there's no need
> -                * to mark anything accessed. And we've dropped the lock.
> -                */
> -               if (wait_event_interruptible_exclusive(pipe->rd_wait, pipe_readable(pipe)) < 0)
> -                       return -ERESTARTSYS;
> +
> +               {
> +                       guard(allow_stack_reclaim)();

nit. You can just use a

    scoped_guard(allow_stack_reclaim) {
        if (wait_event_interruptible_exclusive(...))
            return -ERESTARTSYS;
    }

here.

Same comment for rest of the thread where a scoped_guard() can be used
instead of this pattern.

> +                       /*
> +                        * We only get here if we didn't actually read
> +                        * anything.
> +                        *
> +                        * But because we didn't read anything, at this point
> +                        * we can just return directly with -ERESTARTSYS if
> +                        * we're interrupted, since we've done any required
> +                        * wakeups and there's no need to mark anything
> +                        * accessed. And we've dropped the lock.
> +                        */
> +                       if (wait_event_interruptible_exclusive(pipe->rd_wait,
> +                                                              pipe_readable(pipe)) < 0)
> +                               return -ERESTARTSYS;
> +               }

-- 
Thanks and Regards,
Prateek

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[Peter Zijlstra]

On Thu, Aug 27, 2026 at 04:29:44PM -0700, David Stevens wrote:
> A new TASK_STACK_RECLAIM state is introduced for tasks that are
> blocked waiting for the stacks to be repopulated.

Not yet read the full patch, but *why*.

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[Peter Zijlstra]

On Thu, Aug 27, 2026 at 04:29:44PM -0700, David Stevens wrote:
> @@ -4295,8 +4296,6 @@ int try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags)
>  		if (!ttwu_state_match(p, state, &success))
>  			break;
>  
> -		trace_sched_waking(p);
> -
>  		/*
>  		 * Ensure we load p->on_rq _after_ p->state, otherwise it would
>  		 * be possible to, falsely, observe p->on_rq == 0 and get stuck
> @@ -4320,8 +4319,18 @@ int try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags)
>  		 * A similar smp_rmb() lives in __task_needs_rq_lock().
>  		 */
>  		smp_rmb();
> -		if (READ_ONCE(p->on_rq) && ttwu_runnable(p, wake_flags))
> +		if (READ_ONCE(p->on_rq) && ttwu_runnable(p, wake_flags)) {
> +			trace_sched_waking(p);
> +			break;
> +		}
> +
> +		if (!ensure_stack_is_present(p, &need_deferred_repopulate)) {
> +			WRITE_ONCE(p->__state, TASK_STACK_RECLAIM);
> +			do_deferred_repopulate_wake = need_deferred_repopulate;
>  			break;
> +		}
> +
> +		trace_sched_waking(p);
>  
>  		/*
>  		 * Ensure we load p->on_cpu _after_ p->on_rq, otherwise it would be

Don't move that tracepoint. That's simply lying about how long it takes
to wake up the task.

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[Peter Zijlstra]

On Thu, Aug 27, 2026 at 04:29:44PM -0700, David Stevens wrote:
> @@ -4320,8 +4319,18 @@ int try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags)
>  		 * A similar smp_rmb() lives in __task_needs_rq_lock().
>  		 */
>  		smp_rmb();
> -		if (READ_ONCE(p->on_rq) && ttwu_runnable(p, wake_flags))
> +		if (READ_ONCE(p->on_rq) && ttwu_runnable(p, wake_flags)) {
> +			trace_sched_waking(p);
> +			break;
> +		}
> +
> +		if (!ensure_stack_is_present(p, &need_deferred_repopulate)) {
> +			WRITE_ONCE(p->__state, TASK_STACK_RECLAIM);
> +			do_deferred_repopulate_wake = need_deferred_repopulate;
>  			break;
> +		}
> +
> +		trace_sched_waking(p);

Absolutely not; ensure_stack_is_present() must not call
repopulate_stack() while holding ->pi_lock. Not happening.

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[Peter Zijlstra]

On Thu, Aug 27, 2026 at 04:29:44PM -0700, David Stevens wrote:
> This feature does work on PREEMPT_RT, but is likely undesirable due to
> the extra uncertainty. The fact that alloc_pages_nolock_noprof() cannot
> be called from under the scheduler lock also makes repopulating stacks
> more expensive.

Aaah, so you knew it was hot garbage :-(

Please don't ever post anything like this again.

Re: [RFC 00/10] Reclaimable kernel stacks

[Peter Zijlstra]

On Thu, Aug 27, 2026 at 04:29:38PM -0700, David Stevens wrote:

> On Android, system processes typically have 2000-3000 threads. App
> processes add 1000s more threads on top of this.

WTF ?!? Why does that all spawn *that* many threads? Perhaps work on
reducing that some?

> Tracking blocked state and when it is safe to reclaim a stack is done
> via a series of hooks in the scheduler. The actual reclaim of stacks is
> done asynchronously in a shrinker.
> 
> Once a task's stack has been reclaimed, it cannot be rescheduled until
> its stack is repopulated. Although there can be a repopulation fast path
> within the scheduler, reliably allocating memory to repopulate the stack
> requires a fallback path that defers the repopulation and wakeup to a
> workqueue context that can use GFP_KERNEL.

I am really confused. On the one hand you have John working on proxy
execution, with the aim on reducing latencies, and then here you are,
posting something that will introduce basically unbound latencies.

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[Peter Zijlstra]

On Thu, Aug 27, 2026 at 04:29:44PM -0700, David Stevens wrote:
> +DEFINE_CLASS(allow_stack_reclaim, bool,
> +	     ({
> +		if (!_T)
> +			current->flags &= ~PF_RECLAIMABLE_STACK;
> +	      }),
> +	     ({
> +		bool was_set = current->flags & PF_RECLAIMABLE_STACK;
> +
> +		current->flags |= PF_RECLAIMABLE_STACK;
> +		was_set;
> +	      }),
> +	     void)

> +	allow_stack_reclaim(prev);

> +void __allow_stack_reclaim(struct task_struct *tsk)
> +{
> +	union stack_reclaim_state prev_state, target_state;
> +
> +	if (WARN_ON_ONCE(tsk->__state == TASK_DEAD))
> +		return;
> +
> +	prev_state.val = READ_ONCE(tsk->stack_reclaim_state.val);
> +	do {
> +		target_state.val = prev_state.val;
> +
> +		if (prev_state.stack_state != STACK_PREPARE_RECLAIM) {
> +			WARN(prev_state.stack_state != STACK_IN_USE,
> +			     "Reclaimable state %x for previously running task", prev_state.val);
> +			return;
> +		}
> +		target_state.stack_state = STACK_RECLAIMABLE;
> +	} while (!try_cmpxchg(&tsk->stack_reclaim_state.val, &prev_state.val, target_state.val));
> +
> +	if (irq_work_queue(&tsk->stack_reclaim_work->irq_work)) {
> +		/*
> +		 * Take a ref that gets released by do_reclaim_stack() so we don't
> +		 * have to worry about races with remove_from_stack_shrinker().
> +		 */
> +		get_task_struct(tsk);
> +	}
> +}

> +static inline void allow_stack_reclaim(struct task_struct *tsk)
> +{
> +	if (unlikely(tsk->flags & PF_RECLAIMABLE_STACK))
> +		__allow_stack_reclaim(tsk);
> +}

So you have a guard with the same name as a function, but the function
only functions when inside the guard of the same name. WTF ?!

Anyway, it looks like you're sprinkling this guard around a few specific
block sites. Which seems to suggest your PF_ flag *should* have been a
TASK_ flag, no?

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[Sebastian Andrzej Siewior]

On 2026-08-27 16:29:44 [-0700], David Stevens wrote:
> diff --git a/arch/Kconfig b/arch/Kconfig
> index fa7507ac8e13..adb4a5957996 100644
> --- a/arch/Kconfig
> +++ b/arch/Kconfig
> @@ -1534,6 +1534,24 @@ config VMAP_STACK
>  	  backing virtual mappings with real shadow memory, and KASAN_VMALLOC
>  	  must be enabled.
>  
> +config HAVE_ARCH_RECLAIMABLE_STACK
> +	def_bool n
> +
> +config RECLAIMABLE_STACK
> +	default !PREEMPT_RT && !PROC_KCORE

This shouldn't default like this for RT. It either is useable or it is
not.

> +	bool "Allow stacks of some blocked threads to be reclaimed"
> +	depends on VMAP_STACK && !STACK_GROWSUP
> +	depends on HAVE_ARCH_RECLAIMABLE_STACK
> +	depends on !DEBUG_STACK_USAGE
> +	depends on !KASAN_VMALLOC # TODO: add support for this
> +	depends on !DEBUG_KMEMLEAK # TODO: add support for this
> +	help
> +	  Enable this to allow the unused portion of kernel stacks of most
> +	  blocked tasks to be reclaimed.
> +
> +	  The wakeup latency of tasks with reclaimed stacks may increase,
> +	  especially while the system is under memory pressure.

It says *may* increase and on RT it _definitely_ will increase since
there is a kworker involved not to mention the memory allocation itself.
Anyway. This either needs to stay away from PREEMPT_RT or find a way to
exclude at the very least mlock()ed tasks.
Did lockdep see this?

If I understood the whole exercise correct then you have a kernel stack
of two pages and in best case you can unmap and release the second page
while the task is napping.

What might be a tad simpler is to memset(,0,) the remaining part of the
stack. Since the stack is vmap-ed it should be swapped out on its own
without additional tricks. That memset() would help zram to compress
better so it uses less memory. ta-da.

What also should be simpler (and I am not saying just to move you away
from the scheduler) is to have a shrinker which iterates over all tasks
which are marked for reclaim and then similar to swap just unmap both
stack pages and release the second page which is not used.
Upon wake up the task should create a page_fault which would be used to
allocate the second stack page and map the whole stack again.

This sounds simpler.

Sebastian

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[Peter Zijlstra]

On Fri, Aug 28, 2026 at 03:36:20PM +0200, Sebastian Andrzej Siewior wrote:
> On 2026-08-27 16:29:44 [-0700], David Stevens wrote:
> > diff --git a/arch/Kconfig b/arch/Kconfig
> > index fa7507ac8e13..adb4a5957996 100644
> > --- a/arch/Kconfig
> > +++ b/arch/Kconfig
> > @@ -1534,6 +1534,24 @@ config VMAP_STACK
> >  	  backing virtual mappings with real shadow memory, and KASAN_VMALLOC
> >  	  must be enabled.
> >  
> > +config HAVE_ARCH_RECLAIMABLE_STACK
> > +	def_bool n
> > +
> > +config RECLAIMABLE_STACK
> > +	default !PREEMPT_RT && !PROC_KCORE
> 
> This shouldn't default like this for RT. It either is useable or it is
> not.
> 
> > +	bool "Allow stacks of some blocked threads to be reclaimed"
> > +	depends on VMAP_STACK && !STACK_GROWSUP
> > +	depends on HAVE_ARCH_RECLAIMABLE_STACK
> > +	depends on !DEBUG_STACK_USAGE
> > +	depends on !KASAN_VMALLOC # TODO: add support for this
> > +	depends on !DEBUG_KMEMLEAK # TODO: add support for this
> > +	help
> > +	  Enable this to allow the unused portion of kernel stacks of most
> > +	  blocked tasks to be reclaimed.
> > +
> > +	  The wakeup latency of tasks with reclaimed stacks may increase,
> > +	  especially while the system is under memory pressure.
> 
> It says *may* increase and on RT it _definitely_ will increase since
> there is a kworker involved not to mention the memory allocation itself.
> Anyway. This either needs to stay away from PREEMPT_RT or find a way to
> exclude at the very least mlock()ed tasks.
> Did lockdep see this?

It should have. They're taking spinlock inside raw_spinlock and lockdep
should very much warn about that by default.

> If I understood the whole exercise correct then you have a kernel stack
> of two pages and in best case you can unmap and release the second page
> while the task is napping.

THREAD_SIZE_ORDER	2
THREAD_SIZE		(PAGE_SIZE << THREAD_SIZE_ORDER)

that makes for 4 pages.

> What might be a tad simpler is to memset(,0,) the remaining part of the
> stack. Since the stack is vmap-ed it should be swapped out on its own
> without additional tricks. That memset() would help zram to compress
> better so it uses less memory. ta-da.

That would still be a 12k memset with IRQs-disabled and rq->lock held.

> What also should be simpler (and I am not saying just to move you away
> from the scheduler) is to have a shrinker which iterates over all tasks
> which are marked for reclaim and then similar to swap just unmap both
> stack pages and release the second page which is not used.
> Upon wake up the task should create a page_fault which would be used to
> allocate the second stack page and map the whole stack again.

Right, so you can FREEZE the task, unmap its stack and then thaw it or
something. But there should be a definite opt-out on all this, because
taking faults on your stack will be horrible.

Not to mention you'll suffer wakeup latencies while frozen.

This all really sounds like what should be addressed is this insane
number of tasks rather than trying to cope with the consequences of
that.

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[Peter Zijlstra]

On Fri, Aug 28, 2026 at 03:59:47PM +0200, Peter Zijlstra wrote:

> > What also should be simpler (and I am not saying just to move you away
> > from the scheduler) is to have a shrinker which iterates over all tasks
> > which are marked for reclaim and then similar to swap just unmap both
> > stack pages and release the second page which is not used.
> > Upon wake up the task should create a page_fault which would be used to
> > allocate the second stack page and map the whole stack again.
> 
> Right, so you can FREEZE the task, unmap its stack and then thaw it or
> something. But there should be a definite opt-out on all this, because
> taking faults on your stack will be horrible.
> 
> Not to mention you'll suffer wakeup latencies while frozen.
> 
> This all really sounds like what should be addressed is this insane
> number of tasks rather than trying to cope with the consequences of
> that.

Taking faults on the task-stack is not going to work. That fault will
happen while you have locks held and IRQs disabled. Ideally, it'll
happen when you're knee deep in the allocator.

That fault will then again call the allocator to allocator your stack
page, and...

Re: [RFC 00/10] Reclaimable kernel stacks

[Steven Rostedt]

On Fri, 28 Aug 2026 14:47:58 +0200
Peter Zijlstra <peterz@infradead.org> wrote:

> > On Android, system processes typically have 2000-3000 threads. App
> > processes add 1000s more threads on top of this.  
> 
> WTF ?!? Why does that all spawn *that* many threads? Perhaps work on
> reducing that some?

I guess you never ran Java ;-)

-- Steve

Re: [RFC 00/10] Reclaimable kernel stacks

[Peter Zijlstra]

On Fri, Aug 28, 2026 at 10:33:28AM -0400, Steven Rostedt wrote:
> On Fri, 28 Aug 2026 14:47:58 +0200
> Peter Zijlstra <peterz@infradead.org> wrote:
> 
> > > On Android, system processes typically have 2000-3000 threads. App
> > > processes add 1000s more threads on top of this.  
> > 
> > WTF ?!? Why does that all spawn *that* many threads? Perhaps work on
> > reducing that some?
> 
> I guess you never ran Java ;-)

Of course not; why would you want to do something that silly ;-)

Re: [RFC 00/10] Reclaimable kernel stacks

[Peter Zijlstra]

On Fri, Aug 28, 2026 at 04:35:40PM +0200, Peter Zijlstra wrote:
> On Fri, Aug 28, 2026 at 10:33:28AM -0400, Steven Rostedt wrote:
> > On Fri, 28 Aug 2026 14:47:58 +0200
> > Peter Zijlstra <peterz@infradead.org> wrote:
> > 
> > > > On Android, system processes typically have 2000-3000 threads. App
> > > > processes add 1000s more threads on top of this.  
> > > 
> > > WTF ?!? Why does that all spawn *that* many threads? Perhaps work on
> > > reducing that some?
> > 
> > I guess you never ran Java ;-)
> 
> Of course not; why would you want to do something that silly ;-)

I realized I actually have a Java thing, the kids have Minecraft, so I
booted that up and I found all of 108 threads. Which is still kinda
insane, but nowhere near the 1000s per app as claimed.

Reducing the thread count from O(1e3) to O(1e2) would win far more
memory than any thread stack shenanigans would.

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[Sebastian Andrzej Siewior]

On 2026-08-28 15:59:47 [+0200], Peter Zijlstra wrote:
> > > +
> > > +	  The wakeup latency of tasks with reclaimed stacks may increase,
> > > +	  especially while the system is under memory pressure.
> > 
> > It says *may* increase and on RT it _definitely_ will increase since
> > there is a kworker involved not to mention the memory allocation itself.
> > Anyway. This either needs to stay away from PREEMPT_RT or find a way to
> > exclude at the very least mlock()ed tasks.
> > Did lockdep see this?
> 
> It should have. They're taking spinlock inside raw_spinlock and lockdep
> should very much warn about that by default.

Yes. My point was that this was hidden from lockdep.

> > If I understood the whole exercise correct then you have a kernel stack
> > of two pages and in best case you can unmap and release the second page
> > while the task is napping.
> 
> THREAD_SIZE_ORDER	2
> THREAD_SIZE		(PAGE_SIZE << THREAD_SIZE_ORDER)
> 
> that makes for 4 pages.

Oh. I wasn't aware that we have 16kib stacks these days. But looking
at it we have it now for over 10 years… Judging from 6538b8ea886e4
("x86_64: expand kernel stack to 16K") it might be temporary and things
are better now? Arm64 has a different story according to 845ad05ec31e0
("arm64: Change kernel stack size to 16K"). Risc-V also mentions "for
now" in 0cac21b02ba5f ("riscv: use 16KB kernel stack on 64-bit").

I just booted my XFS kvm box and did things and 8KiB works so far.

> > What might be a tad simpler is to memset(,0,) the remaining part of the
> > stack. Since the stack is vmap-ed it should be swapped out on its own
> > without additional tricks. That memset() would help zram to compress
> > better so it uses less memory. ta-da.
> 
> That would still be a 12k memset with IRQs-disabled and rq->lock held.

Right, because the stack grew a bit. Probably still cheaper compared to
the other things done here ;)

> > What also should be simpler (and I am not saying just to move you away
> > from the scheduler) is to have a shrinker which iterates over all tasks
> > which are marked for reclaim and then similar to swap just unmap both
> > stack pages and release the second page which is not used.
> > Upon wake up the task should create a page_fault which would be used to
> > allocate the second stack page and map the whole stack again.
> 
> Right, so you can FREEZE the task, unmap its stack and then thaw it or
> something. But there should be a definite opt-out on all this, because
> taking faults on your stack will be horrible.

Definitely. Not something for the currently visible app.

> Not to mention you'll suffer wakeup latencies while frozen.

Right but you would use it under memory pressure and steal the stack
from the most idle tasks rather from everyone. 

> This all really sounds like what should be addressed is this insane
> number of tasks rather than trying to cope with the consequences of
> that.

Sebastian

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[Sebastian Andrzej Siewior]

On 2026-08-28 16:25:52 [+0200], Peter Zijlstra wrote:
> Taking faults on the task-stack is not going to work. That fault will
> happen while you have locks held and IRQs disabled. Ideally, it'll
> happen when you're knee deep in the allocator.
> 
> That fault will then again call the allocator to allocator your stack
> page, and...

indeed. So I do have a few ideas floating around but I leave it at it.

Sebastian

Re: [RFC 00/10] Reclaimable kernel stacks

[Steven Rostedt]

On Fri, 28 Aug 2026 16:45:00 +0200
Peter Zijlstra <peterz@infradead.org> wrote:

> I realized I actually have a Java thing, the kids have Minecraft, so I
> booted that up and I found all of 108 threads. Which is still kinda
> insane, but nowhere near the 1000s per app as claimed.

Have you looked at how many threads Chome uses?

 $ ps -eLf|grep chrome | wc -l
 930

Of course I probably have 900 tabs open :-p

-- Steve

Re: [RFC 00/10] Reclaimable kernel stacks

[David Stevens]

On Fri, Aug 28, 2026 at 5:48 AM Peter Zijlstra <peterz@infradead.org> wrote:
>
> On Thu, Aug 27, 2026 at 04:29:38PM -0700, David Stevens wrote:
>
> > On Android, system processes typically have 2000-3000 threads. App
> > processes add 1000s more threads on top of this.
>
> WTF ?!? Why does that all spawn *that* many threads? Perhaps work on
> reducing that some?

I agree that the number of threads is rather excessive, and reducing
it in userspace is an ongoing effort. However, given Android's use of
Java and its heavily multi-process architecture, especially once apps
are involved, there are limits to how much the thread count can
realistically be reduced. And even if a 50% reduction in thread count
were somehow achieved, kernel stacks would still consume upwards of 1%
of system RAM on lower spec devices with 4GB of memory.

> > Tracking blocked state and when it is safe to reclaim a stack is done
> > via a series of hooks in the scheduler. The actual reclaim of stacks is
> > done asynchronously in a shrinker.
> >
> > Once a task's stack has been reclaimed, it cannot be rescheduled until
> > its stack is repopulated. Although there can be a repopulation fast path
> > within the scheduler, reliably allocating memory to repopulate the stack
> > requires a fallback path that defers the repopulation and wakeup to a
> > workqueue context that can use GFP_KERNEL.
>
> I am really confused. On the one hand you have John working on proxy
> execution, with the aim on reducing latencies, and then here you are,
> posting something that will introduce basically unbound latencies.

The two projects aren't contradictory because they deal with different
types of latency.

Proxy execution aims to solve latency introduced by priority inversion
between background tasks and foreground tasks. This can happen at
basically any point if you get unlucky, resulting in unpredictable
latency spikes in high priority tasks.

The latency introduced by reclaimable stacks isn't substantially
different from the latency from refaulting evicted anon or file pages.
It's a tradeoff between the reduced cost of cold memory and the
increased latency when re-accessing that cold memory - clearly a
tradeoff that Linux already makes. In situations where reclaimed
stacks introduce latency, there will almost certainly already be extra
latency incurred from refaulting userspace pages. I don't yet have
production data, but the testing I've done so far doesn't show any
measurable negative impact on user visible latency metrics.

-David

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[Steven Rostedt]

On Fri, 28 Aug 2026 17:10:18 +0200
Sebastian Andrzej Siewior <bigeasy@linutronix.de> wrote:

> > THREAD_SIZE_ORDER	2
> > THREAD_SIZE		(PAGE_SIZE << THREAD_SIZE_ORDER)
> > 
> > that makes for 4 pages.  
> 
> Oh. I wasn't aware that we have 16kib stacks these days. But looking
> at it we have it now for over 10 years… Judging from 6538b8ea886e4
> ("x86_64: expand kernel stack to 16K") it might be temporary and things
> are better now? Arm64 has a different story according to 845ad05ec31e0
> ("arm64: Change kernel stack size to 16K"). Risc-V also mentions "for
> now" in 0cac21b02ba5f ("riscv: use 16KB kernel stack on 64-bit").
> 
> I just booted my XFS kvm box and did things and 8KiB works so far.

If you want to see how big your stack is and where its used, you can run
the stack tracer:

 # trace-cmd stack --start

[wait]

 # trace-cmd stack
(stack tracer running)
        Depth    Size   Location    (41 entries)
        -----    ----   --------
  0)     5032      48   choose_idle_cpu+0x5/0x100
  1)     4984     296   select_task_rq_fair+0xa84/0x2b40
  2)     4688     104   try_to_wake_up+0x168/0x7e0
  3)     4584      16   task_work_add+0xd7/0xf0
  4)     4568      16   io_req_normal_work_add+0x75/0xb0
  5)     4552      40   io_poll_wake+0x10f/0x160
  6)     4512      64   __wake_up_common+0x72/0xa0
  7)     4448      40   __wake_up_sync_key+0x43/0x60
  8)     4408      24   sock_def_readable+0x46/0xe0
  9)     4384      80   tun_net_xmit+0x240/0x550 [tun]
 10)     4304      80   dev_hard_start_xmit+0x63/0x1e0
 11)     4224     256   __dev_queue_xmit+0x880/0x10b0
 12)     3968      24   br_dev_queue_push_xmit+0x62/0xf0 [bridge]
 13)     3944      80   br_dev_xmit+0x15e/0x4a0 [bridge]
 14)     3864      80   dev_hard_start_xmit+0x63/0x1e0
 15)     3784     256   __dev_queue_xmit+0x880/0x10b0
 16)     3528      72   ip_finish_output2+0x2ce/0x650
 17)     3456     104   ip_output+0x63/0x110
 18)     3352      72   __ip_queue_xmit+0x16f/0x470
 19)     3280     232   __tcp_transmit_skb+0xcb5/0x1140
 20)     3048     120   tcp_write_xmit+0x5aa/0x1760
 21)     2928      24   __tcp_push_pending_frames+0x39/0x110
 22)     2904     104   tcp_rcv_established+0x37f/0xeb0
 23)     2800      32   tcp_v4_do_rcv+0x13f/0x300
 24)     2768     168   tcp_v4_rcv+0xc01/0x1300
 25)     2600      48   ip_protocol_deliver_rcu+0x35/0x1b0
 26)     2552      24   ip_local_deliver_finish+0x85/0x100
 27)     2528      40   __netif_receive_skb_one_core+0x85/0xa0
 28)     2488      56   netif_receive_skb+0x127/0x180
 29)     2432     104   br_handle_frame_finish+0x428/0x680 [bridge]
 30)     2328      40   br_handle_frame+0x132/0x2a0 [bridge]
 31)     2288     280   __netif_receive_skb_core.constprop.0+0x16e/0xdf0
 32)     2008      40   __netif_receive_skb_one_core+0x39/0xa0
 33)     1968      56   netif_receive_skb+0x127/0x180
 34)     1912     256   tun_get_user+0xbd0/0x1260 [tun]
 35)     1656      56   tun_chr_write_iter+0x77/0xba [tun]
 36)     1600      88   do_iter_readv_writev+0x161/0x260
 37)     1512     256   vfs_writev+0x168/0x3c0
 38)     1256      80   do_writev+0x7f/0x110
 39)     1176     984   do_syscall_64+0xaa/0x670
 40)      192     192   entry_SYSCALL_64_after_hwframe+0x76/0x7e

Note, it only monitors task context (not interrupt).

-- Steve

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[Steven Rostedt]

On Fri, 28 Aug 2026 15:08:36 -0400
Steven Rostedt <rostedt@goodmis.org> wrote:

> If you want to see how big your stack is and where its used, you can run
> the stack tracer:
> 
>  # trace-cmd stack --start
> 
> [wait]
> 
>  # trace-cmd stack
> (stack tracer running)
>         Depth    Size   Location    (41 entries)
>         -----    ----   --------
>   0)     5032      48   choose_idle_cpu+0x5/0x100
>   1)     4984     296   select_task_rq_fair+0xa84/0x2b40
>   2)     4688     104   try_to_wake_up+0x168/0x7e0
>   3)     4584      16   task_work_add+0xd7/0xf0
>   4)     4568      16   io_req_normal_work_add+0x75/0xb0
>   5)     4552      40   io_poll_wake+0x10f/0x160
>   6)     4512      64   __wake_up_common+0x72/0xa0
>   7)     4448      40   __wake_up_sync_key+0x43/0x60
>   8)     4408      24   sock_def_readable+0x46/0xe0
>   9)     4384      80   tun_net_xmit+0x240/0x550 [tun]
>  10)     4304      80   dev_hard_start_xmit+0x63/0x1e0
>  11)     4224     256   __dev_queue_xmit+0x880/0x10b0
>  12)     3968      24   br_dev_queue_push_xmit+0x62/0xf0 [bridge]
>  13)     3944      80   br_dev_xmit+0x15e/0x4a0 [bridge]
>  14)     3864      80   dev_hard_start_xmit+0x63/0x1e0
>  15)     3784     256   __dev_queue_xmit+0x880/0x10b0
>  16)     3528      72   ip_finish_output2+0x2ce/0x650
>  17)     3456     104   ip_output+0x63/0x110
>  18)     3352      72   __ip_queue_xmit+0x16f/0x470
>  19)     3280     232   __tcp_transmit_skb+0xcb5/0x1140
>  20)     3048     120   tcp_write_xmit+0x5aa/0x1760
>  21)     2928      24   __tcp_push_pending_frames+0x39/0x110
>  22)     2904     104   tcp_rcv_established+0x37f/0xeb0
>  23)     2800      32   tcp_v4_do_rcv+0x13f/0x300
>  24)     2768     168   tcp_v4_rcv+0xc01/0x1300
>  25)     2600      48   ip_protocol_deliver_rcu+0x35/0x1b0
>  26)     2552      24   ip_local_deliver_finish+0x85/0x100
>  27)     2528      40   __netif_receive_skb_one_core+0x85/0xa0
>  28)     2488      56   netif_receive_skb+0x127/0x180
>  29)     2432     104   br_handle_frame_finish+0x428/0x680 [bridge]
>  30)     2328      40   br_handle_frame+0x132/0x2a0 [bridge]
>  31)     2288     280   __netif_receive_skb_core.constprop.0+0x16e/0xdf0
>  32)     2008      40   __netif_receive_skb_one_core+0x39/0xa0
>  33)     1968      56   netif_receive_skb+0x127/0x180
>  34)     1912     256   tun_get_user+0xbd0/0x1260 [tun]
>  35)     1656      56   tun_chr_write_iter+0x77/0xba [tun]
>  36)     1600      88   do_iter_readv_writev+0x161/0x260
>  37)     1512     256   vfs_writev+0x168/0x3c0
>  38)     1256      80   do_writev+0x7f/0x110

>  39)     1176     984   do_syscall_64+0xaa/0x670

If one worries about stack size, they may want to turn off
CONFIG_RANDOMIZE_KSTACK_OFFSET as I see some big numbers for do_syscall_64()

-- Steve


>  40)      192     192   entry_SYSCALL_64_after_hwframe+0x76/0x7e

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[Steven Rostedt]

On Fri, 28 Aug 2026 15:13:13 -0400
Steven Rostedt <rostedt@goodmis.org> wrote:

> If one worries about stack size, they may want to turn off
> CONFIG_RANDOMIZE_KSTACK_OFFSET as I see some big numbers for do_syscall_64()

And the scheduler has some heavy stack usage here!

  # trace-cmd stack
(stack tracer running)
        Depth    Size   Location    (31 entries)
        -----    ----   --------
  0)     8120      48   __msecs_to_jiffies+0x9/0x30
  1)     8072     104   update_group_capacity+0x94/0x960
  2)     7968     528   update_sd_lb_stats.constprop.0+0x426/0x39b0
  3)     7440     424   sched_balance_find_src_group+0x8f/0x1150
  4)     7016     552   sched_balance_rq+0x934/0x4130
  5)     6464     408   pick_task_fair+0xa71/0x1fe0
  6)     6056     496   __schedule+0x628/0x76f0
  7)     5560      32   schedule+0xde/0x2c0
  8)     5528      32   io_schedule+0x8c/0x100
  9)     5496     256   rq_qos_wait+0x12b/0x230
 10)     5240     136   wbt_wait+0x150/0x260
 11)     5104      40   __rq_qos_throttle+0x51/0xa0
 12)     5064     304   blk_mq_submit_bio+0xc40/0x28e0
 13)     4760     240   submit_bio_noacct_nocheck+0x410/0xb30
 14)     4520      40   ext4_io_submit+0xca/0x1a0
 15)     4480     144   ext4_bio_write_folio+0x5e2/0x1470
 16)     4336      96   mpage_process_page_bufs+0x392/0x700
 17)     4240     616   mpage_prepare_extent_to_map+0xaae/0x1120
 18)     3624     480   ext4_do_writepages+0x971/0x37f0
 19)     3144     312   ext4_writepages+0x2d2/0x610
 20)     2832     152   do_writepages+0x21e/0x560
 21)     2680     184   __writeback_single_inode+0x119/0x1240
 22)     2496     752   writeback_sb_inodes+0x67d/0x17b0
 23)     1744     168   __writeback_inodes_wb+0xf2/0x270
 24)     1576     304   wb_writeback+0x63e/0x890
 25)     1272     320   wb_workfn+0x844/0xca0
 26)      952     360   process_one_work+0x8b4/0x15d0
 27)      592     176   worker_thread+0x5d3/0xfb0
 28)      416      64   kthread+0x33c/0x420
 29)      352     160   ret_from_fork+0x65c/0x9e0
 30)      192     192   ret_from_fork_asm+0x1a/0x30

Note, I'm running this on my server. That is, this is a production machine.

-- Steve

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[David Stevens]

On Fri, Aug 28, 2026 at 6:59 AM Peter Zijlstra <peterz@infradead.org> wrote:
>
> On Fri, Aug 28, 2026 at 03:36:20PM +0200, Sebastian Andrzej Siewior wrote:
> > On 2026-08-27 16:29:44 [-0700], David Stevens wrote:
> > > diff --git a/arch/Kconfig b/arch/Kconfig
> > > index fa7507ac8e13..adb4a5957996 100644
> > > --- a/arch/Kconfig
> > > +++ b/arch/Kconfig
> > > @@ -1534,6 +1534,24 @@ config VMAP_STACK
> > >       backing virtual mappings with real shadow memory, and KASAN_VMALLOC
> > >       must be enabled.
> > >
> > > +config HAVE_ARCH_RECLAIMABLE_STACK
> > > +   def_bool n
> > > +
> > > +config RECLAIMABLE_STACK
> > > +   default !PREEMPT_RT && !PROC_KCORE
> >
> > This shouldn't default like this for RT. It either is useable or it is
> > not.
> >
> > > +   bool "Allow stacks of some blocked threads to be reclaimed"
> > > +   depends on VMAP_STACK && !STACK_GROWSUP
> > > +   depends on HAVE_ARCH_RECLAIMABLE_STACK
> > > +   depends on !DEBUG_STACK_USAGE
> > > +   depends on !KASAN_VMALLOC # TODO: add support for this
> > > +   depends on !DEBUG_KMEMLEAK # TODO: add support for this
> > > +   help
> > > +     Enable this to allow the unused portion of kernel stacks of most
> > > +     blocked tasks to be reclaimed.
> > > +
> > > +     The wakeup latency of tasks with reclaimed stacks may increase,
> > > +     especially while the system is under memory pressure.
> >
> > It says *may* increase and on RT it _definitely_ will increase since
> > there is a kworker involved not to mention the memory allocation itself.
> > Anyway. This either needs to stay away from PREEMPT_RT or find a way to
> > exclude at the very least mlock()ed tasks.
> > Did lockdep see this?
>
> It should have. They're taking spinlock inside raw_spinlock and lockdep
> should very much warn about that by default.

It does trylocks on spinlocks and local locks, inside of
alloc_pages_nolock_noprof(). Lockdep considers that safe.

On PREEMPT_RT, it is necessary to skip the call to
alloc_pages_nolock_noprof() from under pi_lock, since with that
configuration spin_trylock can end up needing to take pi_lock. But at
least on !PREEMPT_RT, there is no risk of deadlock.

-David

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[David Stevens]

On Fri, Aug 28, 2026 at 6:36 AM Sebastian Andrzej Siewior
<bigeasy@linutronix.de> wrote:
>
> On 2026-08-27 16:29:44 [-0700], David Stevens wrote:
> > diff --git a/arch/Kconfig b/arch/Kconfig
> > index fa7507ac8e13..adb4a5957996 100644
> > --- a/arch/Kconfig
> > +++ b/arch/Kconfig
> > @@ -1534,6 +1534,24 @@ config VMAP_STACK
> >         backing virtual mappings with real shadow memory, and KASAN_VMALLOC
> >         must be enabled.
> >
> > +config HAVE_ARCH_RECLAIMABLE_STACK
> > +     def_bool n
> > +
> > +config RECLAIMABLE_STACK
> > +     default !PREEMPT_RT && !PROC_KCORE
>
> This shouldn't default like this for RT. It either is useable or it is
> not.
>
> > +     bool "Allow stacks of some blocked threads to be reclaimed"
> > +     depends on VMAP_STACK && !STACK_GROWSUP
> > +     depends on HAVE_ARCH_RECLAIMABLE_STACK
> > +     depends on !DEBUG_STACK_USAGE
> > +     depends on !KASAN_VMALLOC # TODO: add support for this
> > +     depends on !DEBUG_KMEMLEAK # TODO: add support for this
> > +     help
> > +       Enable this to allow the unused portion of kernel stacks of most
> > +       blocked tasks to be reclaimed.
> > +
> > +       The wakeup latency of tasks with reclaimed stacks may increase,
> > +       especially while the system is under memory pressure.
>
> It says *may* increase and on RT it _definitely_ will increase since
> there is a kworker involved not to mention the memory allocation itself.

You're correct that this phrasing is wrong. A more accurate way to say
it might be: "Under memory pressure, tasks that are not scheduled for
extended periods may have their stacks reclaimed. Such tasks will
experience increased wakeup latency."

> Anyway. This either needs to stay away from PREEMPT_RT or find a way to
> exclude at the very least mlock()ed tasks.

Disabling kernel stack reclaim for mlockall() tasks definitely makes sense.

> Did lockdep see this?

What are you referring to by "this"? I did stress testing with lockdep
enabled and didn't see any errors, although I will admit the bulk of
that stress testing was !PREEMPT_RT.

> If I understood the whole exercise correct then you have a kernel stack
> of two pages and in best case you can unmap and release the second page
> while the task is napping.
>
> What might be a tad simpler is to memset(,0,) the remaining part of the
> stack. Since the stack is vmap-ed it should be swapped out on its own
> without additional tricks. That memset() would help zram to compress
> better so it uses less memory. ta-da.
>
> What also should be simpler (and I am not saying just to move you away
> from the scheduler) is to have a shrinker which iterates over all tasks
> which are marked for reclaim and then similar to swap just unmap both
> stack pages and release the second page which is not used.
> Upon wake up the task should create a page_fault which would be used to
> allocate the second stack page and map the whole stack again.
>
> This sounds simpler.

There is no swap for kernel memory, only user memory. I did contribute
to some previous work that aimed to only only prepopulate one page of
each kernel stack and then to dynamically fault in further pages as
needed [1], but handling that architecturally and guaranteeing that
memory is available when needed is difficult. Full-on swap to zram
within the kernel would complicate pre-allocating pages for kernel
stack faults further, since you could start getting inter-thread
kernel stack faults due to blocked tasks putting pointers to their
stacks into waitqueues and similar structures.

[1] https://lore.kernel.org/linux-mm/20260424191456.2679717-1-stevensd@google.com/#r

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[David Stevens]

On Fri, Aug 28, 2026 at 5:57 AM Peter Zijlstra <peterz@infradead.org> wrote:
>
> On Thu, Aug 27, 2026 at 04:29:44PM -0700, David Stevens wrote:
> > +DEFINE_CLASS(allow_stack_reclaim, bool,
> > +          ({
> > +             if (!_T)
> > +                     current->flags &= ~PF_RECLAIMABLE_STACK;
> > +           }),
> > +          ({
> > +             bool was_set = current->flags & PF_RECLAIMABLE_STACK;
> > +
> > +             current->flags |= PF_RECLAIMABLE_STACK;
> > +             was_set;
> > +           }),
> > +          void)
>
> > +     allow_stack_reclaim(prev);
>
> > +void __allow_stack_reclaim(struct task_struct *tsk)
> > +{
> > +     union stack_reclaim_state prev_state, target_state;
> > +
> > +     if (WARN_ON_ONCE(tsk->__state == TASK_DEAD))
> > +             return;
> > +
> > +     prev_state.val = READ_ONCE(tsk->stack_reclaim_state.val);
> > +     do {
> > +             target_state.val = prev_state.val;
> > +
> > +             if (prev_state.stack_state != STACK_PREPARE_RECLAIM) {
> > +                     WARN(prev_state.stack_state != STACK_IN_USE,
> > +                          "Reclaimable state %x for previously running task", prev_state.val);
> > +                     return;
> > +             }
> > +             target_state.stack_state = STACK_RECLAIMABLE;
> > +     } while (!try_cmpxchg(&tsk->stack_reclaim_state.val, &prev_state.val, target_state.val));
> > +
> > +     if (irq_work_queue(&tsk->stack_reclaim_work->irq_work)) {
> > +             /*
> > +              * Take a ref that gets released by do_reclaim_stack() so we don't
> > +              * have to worry about races with remove_from_stack_shrinker().
> > +              */
> > +             get_task_struct(tsk);
> > +     }
> > +}
>
> > +static inline void allow_stack_reclaim(struct task_struct *tsk)
> > +{
> > +     if (unlikely(tsk->flags & PF_RECLAIMABLE_STACK))
> > +             __allow_stack_reclaim(tsk);
> > +}
>
> So you have a guard with the same name as a function, but the function
> only functions when inside the guard of the same name. WTF ?!
>
> Anyway, it looks like you're sprinkling this guard around a few specific
> block sites. Which seems to suggest your PF_ flag *should* have been a
> TASK_ flag, no?

A TASK_ flag is definitely better. I was thinking in terms of defining
safe scopes, but that just complicates things for no real benefit.

-David

Re: [RFC 06/10] Reclaim memory from blocked kernel stacks

[David Stevens]

On Fri, Aug 28, 2026 at 5:04 AM Peter Zijlstra <peterz@infradead.org> wrote:
>
> On Thu, Aug 27, 2026 at 04:29:44PM -0700, David Stevens wrote:
> > @@ -4320,8 +4319,18 @@ int try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags)
> >                * A similar smp_rmb() lives in __task_needs_rq_lock().
> >                */
> >               smp_rmb();
> > -             if (READ_ONCE(p->on_rq) && ttwu_runnable(p, wake_flags))
> > +             if (READ_ONCE(p->on_rq) && ttwu_runnable(p, wake_flags)) {
> > +                     trace_sched_waking(p);
> > +                     break;
> > +             }
> > +
> > +             if (!ensure_stack_is_present(p, &need_deferred_repopulate)) {
> > +                     WRITE_ONCE(p->__state, TASK_STACK_RECLAIM);
> > +                     do_deferred_repopulate_wake = need_deferred_repopulate;
> >                       break;
> > +             }
> > +
> > +             trace_sched_waking(p);
>
> Absolutely not; ensure_stack_is_present() must not call
> repopulate_stack() while holding ->pi_lock. Not happening.

The optimistic fast path for repopulate_stack() could be modified to
try pulling from a pre-allocated pool of zero'ed pages. That would
reduce the function to a couple of memcg_kmem_charge_page() calls and
then vmap_pages_range() to repopulate the stack's page tables. That
wouldn't require touching any locks except a raw_spinlock protecting
the pre-allocated pool (or just make it per_cpu). In terms of cost,
this would involve a couple of atomic operations for the page pool
lock and the memcg charging plus non-atomic operations on 5-10 other
cache lines.

Is that within the scope of what can be done under the pi_lock? If
that's still not happening, I can see how things look if we always
defer wakeup to a workqueue.

-David