memcg: Use trylock to access memcg stock_lock.

Teach memcg to operate under trylock conditions when spinning locks
cannot be used.

localtry_trylock might fail and this would lead to charge cache bypass
if the calling context doesn't allow spinning (gfpflags_allow_spinning).
In those cases charge the memcg counter directly and fail early if
that is not possible. This might cause a pre-mature charge failing
but it will allow an opportunistic charging that is safe from
try_alloc_pages path.

Acked-by: Michal Hocko <mhocko@suse.com>
Acked-by: Vlastimil Babka <vbabka@suse.cz>
Acked-by: Shakeel Butt <shakeel.butt@linux.dev>
Signed-off-by: Alexei Starovoitov <ast@kernel.org>
Link: https://lore.kernel.org/r/20250222024427.30294-5-alexei.starovoitov@gmail.com
Signed-off-by: Alexei Starovoitov <ast@kernel.org>

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 46f8b372d212b13e541a7cf4fc13a5a3674205ff..092cab99dec70b6ff9145087a859045867b5b9c4 100644 (file)
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -1739,7 +1739,7 @@ void mem_cgroup_print_oom_group(struct mem_cgroup *memcg)
 }
 
 struct memcg_stock_pcp {
-       local_lock_t stock_lock;
+       localtry_lock_t stock_lock;
        struct mem_cgroup *cached; /* this never be root cgroup */
        unsigned int nr_pages;
 
@@ -1754,7 +1754,7 @@ struct memcg_stock_pcp {
 #define FLUSHING_CACHED_CHARGE 0
 };
 static DEFINE_PER_CPU(struct memcg_stock_pcp, memcg_stock) = {
-       .stock_lock = INIT_LOCAL_LOCK(stock_lock),
+       .stock_lock = INIT_LOCALTRY_LOCK(stock_lock),
 };
 static DEFINE_MUTEX(percpu_charge_mutex);
 
@@ -1766,6 +1766,7 @@ static bool obj_stock_flush_required(struct memcg_stock_pcp *stock,
  * consume_stock: Try to consume stocked charge on this cpu.
  * @memcg: memcg to consume from.
  * @nr_pages: how many pages to charge.
+ * @gfp_mask: allocation mask.
  *
  * The charges will only happen if @memcg matches the current cpu's memcg
  * stock, and at least @nr_pages are available in that stock.  Failure to
@@ -1773,7 +1774,8 @@ static bool obj_stock_flush_required(struct memcg_stock_pcp *stock,
  *
  * returns true if successful, false otherwise.
  */
-static bool consume_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
+static bool consume_stock(struct mem_cgroup *memcg, unsigned int nr_pages,
+                         gfp_t gfp_mask)
 {
        struct memcg_stock_pcp *stock;
        unsigned int stock_pages;
@@ -1783,7 +1785,11 @@ static bool consume_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
        if (nr_pages > MEMCG_CHARGE_BATCH)
                return ret;
 
-       local_lock_irqsave(&memcg_stock.stock_lock, flags);
+       if (!localtry_trylock_irqsave(&memcg_stock.stock_lock, flags)) {
+               if (!gfpflags_allow_spinning(gfp_mask))
+                       return ret;
+               localtry_lock_irqsave(&memcg_stock.stock_lock, flags);
+       }
 
        stock = this_cpu_ptr(&memcg_stock);
        stock_pages = READ_ONCE(stock->nr_pages);
@@ -1792,7 +1798,7 @@ static bool consume_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
                ret = true;
        }
 
-       local_unlock_irqrestore(&memcg_stock.stock_lock, flags);
+       localtry_unlock_irqrestore(&memcg_stock.stock_lock, flags);
 
        return ret;
 }
@@ -1831,14 +1837,14 @@ static void drain_local_stock(struct work_struct *dummy)
         * drain_stock races is that we always operate on local CPU stock
         * here with IRQ disabled
         */
-       local_lock_irqsave(&memcg_stock.stock_lock, flags);
+       localtry_lock_irqsave(&memcg_stock.stock_lock, flags);
 
        stock = this_cpu_ptr(&memcg_stock);
        old = drain_obj_stock(stock);
        drain_stock(stock);
        clear_bit(FLUSHING_CACHED_CHARGE, &stock->flags);
 
-       local_unlock_irqrestore(&memcg_stock.stock_lock, flags);
+       localtry_unlock_irqrestore(&memcg_stock.stock_lock, flags);
        obj_cgroup_put(old);
 }
 
@@ -1868,9 +1874,20 @@ static void refill_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
 {
        unsigned long flags;
 
-       local_lock_irqsave(&memcg_stock.stock_lock, flags);
+       if (!localtry_trylock_irqsave(&memcg_stock.stock_lock, flags)) {
+               /*
+                * In case of unlikely failure to lock percpu stock_lock
+                * uncharge memcg directly.
+                */
+               if (mem_cgroup_is_root(memcg))
+                       return;
+               page_counter_uncharge(&memcg->memory, nr_pages);
+               if (do_memsw_account())
+                       page_counter_uncharge(&memcg->memsw, nr_pages);
+               return;
+       }
        __refill_stock(memcg, nr_pages);
-       local_unlock_irqrestore(&memcg_stock.stock_lock, flags);
+       localtry_unlock_irqrestore(&memcg_stock.stock_lock, flags);
 }
 
 /*
@@ -2213,9 +2230,13 @@ int try_charge_memcg(struct mem_cgroup *memcg, gfp_t gfp_mask,
        unsigned long pflags;
 
 retry:
-       if (consume_stock(memcg, nr_pages))
+       if (consume_stock(memcg, nr_pages, gfp_mask))
                return 0;
 
+       if (!gfpflags_allow_spinning(gfp_mask))
+               /* Avoid the refill and flush of the older stock */
+               batch = nr_pages;
+
        if (!do_memsw_account() ||
            page_counter_try_charge(&memcg->memsw, batch, &counter)) {
                if (page_counter_try_charge(&memcg->memory, batch, &counter))
@@ -2699,7 +2720,7 @@ static void mod_objcg_state(struct obj_cgroup *objcg, struct pglist_data *pgdat,
        unsigned long flags;
        int *bytes;
 
-       local_lock_irqsave(&memcg_stock.stock_lock, flags);
+       localtry_lock_irqsave(&memcg_stock.stock_lock, flags);
        stock = this_cpu_ptr(&memcg_stock);
 
        /*
@@ -2752,7 +2773,7 @@ static void mod_objcg_state(struct obj_cgroup *objcg, struct pglist_data *pgdat,
        if (nr)
                __mod_objcg_mlstate(objcg, pgdat, idx, nr);
 
-       local_unlock_irqrestore(&memcg_stock.stock_lock, flags);
+       localtry_unlock_irqrestore(&memcg_stock.stock_lock, flags);
        obj_cgroup_put(old);
 }
 
@@ -2762,7 +2783,7 @@ static bool consume_obj_stock(struct obj_cgroup *objcg, unsigned int nr_bytes)
        unsigned long flags;
        bool ret = false;
 
-       local_lock_irqsave(&memcg_stock.stock_lock, flags);
+       localtry_lock_irqsave(&memcg_stock.stock_lock, flags);
 
        stock = this_cpu_ptr(&memcg_stock);
        if (objcg == READ_ONCE(stock->cached_objcg) && stock->nr_bytes >= nr_bytes) {
@@ -2770,7 +2791,7 @@ static bool consume_obj_stock(struct obj_cgroup *objcg, unsigned int nr_bytes)
                ret = true;
        }
 
-       local_unlock_irqrestore(&memcg_stock.stock_lock, flags);
+       localtry_unlock_irqrestore(&memcg_stock.stock_lock, flags);
 
        return ret;
 }
@@ -2862,7 +2883,7 @@ static void refill_obj_stock(struct obj_cgroup *objcg, unsigned int nr_bytes,
        unsigned long flags;
        unsigned int nr_pages = 0;
 
-       local_lock_irqsave(&memcg_stock.stock_lock, flags);
+       localtry_lock_irqsave(&memcg_stock.stock_lock, flags);
 
        stock = this_cpu_ptr(&memcg_stock);
        if (READ_ONCE(stock->cached_objcg) != objcg) { /* reset if necessary */
@@ -2880,7 +2901,7 @@ static void refill_obj_stock(struct obj_cgroup *objcg, unsigned int nr_bytes,
                stock->nr_bytes &= (PAGE_SIZE - 1);
        }
 
-       local_unlock_irqrestore(&memcg_stock.stock_lock, flags);
+       localtry_unlock_irqrestore(&memcg_stock.stock_lock, flags);
        obj_cgroup_put(old);
 
        if (nr_pages)