[RFC PATCH] mm/slub: barn sheaf counts can exceed their limit

[RFC PATCH] mm/slub: barn sheaf counts can exceed their limit

[Junseo Lim]

Hello,

Following the suggestion to discuss this publicly, I would like to
report a repeatable overshoot of the SLUB barn limits. Although
MAX_FULL_SHEAVES and MAX_EMPTY_SHEAVES are both 10, a userspace workload
can repeatedly trigger the refill recovery path. In one run, I observed
states like nr_full=91.

I would appreciate feedback on the appropriate way to enforce these
limits.

Description
-----------

While the sheaf layer provides a simple array structure for
amortization, the edge cases that involve refills and allocating
structures are inherently complex when combined with multiple threads.
One problematic case occurs in the allocation path. Consider a state
where every object is drained and TH1 on CPU1 is running:

  CPU1
  =======================================================================
  TH1                    TH2                    TH3
  ---------------------  ---------------------  --------------------- ...
  pcs->main->size=0      not entered            not entered
  pcs->spare->size=0
  nr_full=0 nr_empty=10

    <observed states>      <observed states>      <observed states>

If an allocation request hits this state, __pcs_replace_empty_main is
called to retrieve a full sheaf and swap it with main. However, no full
sheaf is available in the barn, so TH1 removes pcs->spare, sets it to
NULL, and refills the removed sheaf through the slow path. Here, the
per-CPU sheaf lock is released during refill and reacquired afterward.
If preemption occurs during this period, a subsequent thread (TH2) sees
pcs->spare=NULL and tries to get a remaining sheaf from the barn's empty
list and refill it. Likewise, this part is not protected by the lock.
Chaining this results in a state where new threads keep draining the
sheaf and ultimately consume all the empty sheaves in the barn and start
to allocate new sheaves.

  CPU1
  =======================================================================
  TH1 (preempted)        TH2 (preempted)        TH3
  ---------------------  ---------------------  --------------------- ...
  pcs->main->size=0      pcs->main->size=0      pcs->main->size=0
  pcs->spare->size=0     pcs->spare=NULL        pcs->spare=NULL
  nr_full=0 nr_empty=10  nr_full=0 nr_empty=10  nr_full=0 nr_empty=9

    <observed states>      <observed states>      <observed states>

The construction of the sheaf layer allows such racy scenarios, as
mentioned in the comments. Nevertheless, this overshoot is much larger
than what we would expect, and the fact that this state can be recreated
with moderate effort implies that a fix would be appropriate.

Impact
------

This itself does not pose a direct threat or system failure. But, in a
security sense, this behavior can be leveraged for heap grooming that
targets the underlying slab_sheaf.

Design considerations and Question
----------------------------------

One way to limit barn overshoot is to constrain the racy allocations.
Simply using the current nr_full and nr_empty values does not help in
detecting excessive allocation. What I came up with is adding an
nr_alloc counter that keeps track of allocations during the unlocked
path. From what I've tested, this restricted the targeted workload that
induces the excessive allocation. Before this patch, nr_full reached 60
to 90. After the patch, it remained at 22 under the adversarial
workload.

I haven't checked every possible interleaving yet, but I would like to
know whether you think this is the right direction. The possible
regressions are that it requires an additional field just to track the
adversarial input and that, in a highly contended scenario, every
allocation after hitting the nr_alloc limit will go to the slow path
until the preempted task returns and continues the allocation.

Is tracking in-flight sheaf allocations per barn, and falling back when
too many are in progress, the right way to limit this refill race?

Note: This is only an implementation on the allocation side and is still
a prototype.

---
 mm/slub.c | 34 +++++++++++++++++++++++++++++-----
 1 file changed, 29 insertions(+), 5 deletions(-)

diff --git a/mm/slub.c b/mm/slub.c
index 45e077c7526c..472fcae4ac91 100644
--- a/mm/slub.c
+++ b/mm/slub.c
@@ -456,6 +456,7 @@ struct node_barn {
 	struct slab_sheaf *sheaf_partial;
 	unsigned int nr_full;
 	unsigned int nr_empty;
+	unsigned int nr_alloc;
 };
 
 struct slab_sheaf {
@@ -3333,6 +3334,7 @@ static void barn_init(struct node_barn *barn)
 	barn->sheaf_partial = NULL;
 	barn->nr_full = 0;
 	barn->nr_empty = 0;
+	barn->nr_alloc = 0;
 }
 
 static void barn_shrink(struct kmem_cache *s, struct node_barn *barn)
@@ -4754,7 +4756,9 @@ __pcs_replace_empty_main(struct kmem_cache *s, struct slub_percpu_sheaves *pcs,
 	struct slab_sheaf *empty = NULL;
 	struct slab_sheaf *full;
 	struct node_barn *barn;
+	unsigned long flags;
 	bool allow_spin;
+	bool alloc_new = false;
 
 	slab_lockdep_assert_held(this_cpu_ptr(&s->cpu_sheaves->lock));
 
@@ -4796,6 +4800,19 @@ __pcs_replace_empty_main(struct kmem_cache *s, struct slub_percpu_sheaves *pcs,
 		}
 	}
 
+	if (allow_spin && !empty) {
+		/* The limit is advisory; concurrent reservations may exceed it. */
+		if (data_race(barn->nr_alloc) > MAX_EMPTY_SHEAVES) {
+			local_unlock(&s->cpu_sheaves->lock);
+			return NULL;
+		}
+
+		spin_lock_irqsave(&barn->lock, flags);
+		barn->nr_alloc++;
+		spin_unlock_irqrestore(&barn->lock, flags);
+		alloc_new = true;
+	}
+
 	local_unlock(&s->cpu_sheaves->lock);
 	pcs = NULL;
 
@@ -4805,7 +4822,7 @@ __pcs_replace_empty_main(struct kmem_cache *s, struct slub_percpu_sheaves *pcs,
 	if (!empty) {
 		empty = alloc_empty_sheaf(s, gfp, alloc_flags);
 		if (!empty)
-			return NULL;
+			goto out;
 	}
 
 	if (refill_sheaf(s, empty, gfp | __GFP_NOMEMALLOC | __GFP_NOWARN)) {
@@ -4816,7 +4833,7 @@ __pcs_replace_empty_main(struct kmem_cache *s, struct slub_percpu_sheaves *pcs,
 		sheaf_flush_unused(s, empty);
 		free_empty_sheaf(s, empty);
 
-		return NULL;
+		goto out;
 	}
 
 	full = empty;
@@ -4839,24 +4856,31 @@ __pcs_replace_empty_main(struct kmem_cache *s, struct slub_percpu_sheaves *pcs,
 		else
 			barn_put_empty_sheaf(barn, pcs->main);
 		pcs->main = full;
-		return pcs;
+		goto out;
 	}
 
 	if (!pcs->spare) {
 		pcs->spare = full;
-		return pcs;
+		goto out;
 	}
 
 	if (pcs->spare->size == 0) {
 		barn_put_empty_sheaf(barn, pcs->spare);
 		pcs->spare = full;
-		return pcs;
+		goto out;
 	}
 
 barn_put:
 	barn_put_full_sheaf(barn, full);
 	stat(s, BARN_PUT);
 
+out:
+	if (alloc_new) {
+		spin_lock_irqsave(&barn->lock, flags);
+		barn->nr_alloc--;
+		spin_unlock_irqrestore(&barn->lock, flags);
+	}
+
 	return pcs;
 }
 
-- 
2.56.0