[PATCH v3 1/8] memcg: Add MEMCG_DMABUF accounting infrastructure

[Albert Esteve]

From: "T.J. Mercier" <tjmercier@xxxxxxxxxx>

Introduce a MEMCG_DMABUF stat counter and helpers for
reserve/commit/cancel and uncharging dma-buf memory against a
memory cgroup:

The charge is tracked on struct dma_buf via a new memcg pointer. The
page count is derived internally from PAGE_ALIGN(dmabuf->size), and
the memcg reference is stored into (or cleared from) dmabuf->memcg by
the helpers themselves, removing the need for callers to manage those
details. The release path in dma_buf_release() is reduced to a single
mem_cgroup_uncharge_dmabuf() call.

This decouples accounting from the backing page allocator, avoiding the
need for every exporter to pass __GFP_ACCOUNT and enabling cross-cgroup
charge transfer, which is impractical with page-level accounting.

The next patch wires up the charging policy in the dma-heap ioctl.

Signed-off-by: T.J. Mercier <tjmercier@xxxxxxxxxx>
Signed-off-by: Albert Esteve <aesteve@xxxxxxxxxx>
---
 Documentation/admin-guide/cgroup-v2.rst |  4 ++
 drivers/dma-buf/dma-buf.c               |  3 ++
 include/linux/dma-buf.h                 |  4 ++
 include/linux/memcontrol.h              | 36 ++++++++++++++
 mm/memcontrol.c                         | 87 +++++++++++++++++++++++++++++++++
 5 files changed, 134 insertions(+)

diff --git a/Documentation/admin-guide/cgroup-v2.rst b/Documentation/admin-guide/cgroup-v2.rst
index 86a2a0099178..a3265a4a953e 100644
--- a/Documentation/admin-guide/cgroup-v2.rst
+++ b/Documentation/admin-guide/cgroup-v2.rst
@@ -1662,6 +1662,10 @@ The following nested keys are defined.
 		Amount of memory used for storing in-kernel data
 		structures.
 
+	  dmabuf (npn)
+		Amount of memory used for exported DMA buffers allocated by the cgroup.
+		Stays with the allocating cgroup regardless of how the buffer is shared.
+
 	  workingset_refault_anon
 		Number of refaults of previously evicted anonymous pages.
 
diff --git a/drivers/dma-buf/dma-buf.c b/drivers/dma-buf/dma-buf.c
index 4c9add51f9ef..f31f49ebd0af 100644
--- a/drivers/dma-buf/dma-buf.c
+++ b/drivers/dma-buf/dma-buf.c
@@ -14,6 +14,7 @@
 #include <linux/fs.h>
 #include <linux/slab.h>
 #include <linux/dma-buf.h>
+#include <linux/memcontrol.h>
 #include <linux/dma-fence.h>
 #include <linux/dma-fence-unwrap.h>
 #include <linux/anon_inodes.h>
@@ -180,6 +181,8 @@ static void dma_buf_release(struct dentry *dentry)
 	 */
 	BUG_ON(dmabuf->cb_in.active || dmabuf->cb_out.active);
 
+	mem_cgroup_uncharge_dmabuf(dmabuf);
+
 	dmabuf->ops->release(dmabuf);
 
 	if (dmabuf->resv == (struct dma_resv *)&dmabuf[1])
diff --git a/include/linux/dma-buf.h b/include/linux/dma-buf.h
index d15b2b31d3c9..23ec0ff23fba 100644
--- a/include/linux/dma-buf.h
+++ b/include/linux/dma-buf.h
@@ -27,6 +27,7 @@
 struct device;
 struct dma_buf;
 struct dma_buf_attachment;
+struct mem_cgroup;
 
 /**
  * struct dma_buf_ops - operations possible on struct dma_buf
@@ -429,6 +430,9 @@ struct dma_buf {
 
 		__poll_t active;
 	} cb_in, cb_out;
+
+	/** @memcg: the cgroup to which this buffer is currently attributed */
+	struct mem_cgroup *memcg;
 };
 
 /**
diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index 7d1c0ce189a8..4c8711fcf4fd 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -10,6 +10,9 @@
 
 #ifndef _LINUX_MEMCONTROL_H
 #define _LINUX_MEMCONTROL_H
+
+struct dma_buf;
+
 #include <linux/cgroup.h>
 #include <linux/vm_event_item.h>
 #include <linux/hardirq.h>
@@ -40,6 +43,7 @@ enum memcg_stat_item {
 	MEMCG_ZSWAP_B,
 	MEMCG_ZSWAPPED,
 	MEMCG_ZSWAP_INCOMP,
+	MEMCG_DMABUF,
 	MEMCG_NR_STAT,
 };
 
@@ -667,6 +671,12 @@ int mem_cgroup_charge_hugetlb(struct folio* folio, gfp_t gfp);
 int mem_cgroup_swapin_charge_folio(struct folio *folio, unsigned short id,
 				   struct mm_struct *mm, gfp_t gfp);
 
+bool mem_cgroup_reserve_dmabuf(struct mem_cgroup *memcg, gfp_t gfp,
+			       unsigned int nr_pages);
+
+void mem_cgroup_commit_dmabuf(struct dma_buf *dmabuf, struct mem_cgroup *memcg,
+			      unsigned int nr_pages);
+
 void __mem_cgroup_uncharge(struct folio *folio);
 
 /**
@@ -682,6 +692,10 @@ static inline void mem_cgroup_uncharge(struct folio *folio)
 	__mem_cgroup_uncharge(folio);
 }
 
+void mem_cgroup_cancel_dmabuf(struct mem_cgroup *memcg, unsigned int nr_pages);
+
+void mem_cgroup_uncharge_dmabuf(struct dma_buf *dmabuf);
+
 void __mem_cgroup_uncharge_folios(struct folio_batch *folios);
 static inline void mem_cgroup_uncharge_folios(struct folio_batch *folios)
 {
@@ -1167,10 +1181,32 @@ static inline int mem_cgroup_swapin_charge_folio(struct folio *folio,
 	return 0;
 }
 
+static inline bool mem_cgroup_reserve_dmabuf(struct mem_cgroup *memcg,
+					     gfp_t gfp,
+					     unsigned int nr_pages)
+{
+	return true;
+}
+
+static inline void mem_cgroup_commit_dmabuf(struct dma_buf *dmabuf,
+					    struct mem_cgroup *memcg,
+					    unsigned int nr_pages)
+{
+}
+
+static inline void mem_cgroup_cancel_dmabuf(struct mem_cgroup *memcg,
+					    unsigned int nr_pages)
+{
+}
+
 static inline void mem_cgroup_uncharge(struct folio *folio)
 {
 }
 
+static inline void mem_cgroup_uncharge_dmabuf(struct dma_buf *dmabuf)
+{
+}
+
 static inline void mem_cgroup_uncharge_folios(struct folio_batch *folios)
 {
 }
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 856a7d07586c..bbd46ba40350 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -69,6 +69,7 @@
 #include <net/ip.h>
 #include "slab.h"
 #include "memcontrol-v1.h"
+#include <linux/dma-buf.h>
 
 #include <linux/uaccess.h>
 
@@ -442,6 +443,7 @@ static const unsigned int memcg_stat_items[] = {
 	MEMCG_ZSWAP_B,
 	MEMCG_ZSWAPPED,
 	MEMCG_ZSWAP_INCOMP,
+	MEMCG_DMABUF,
 };
 
 #define NR_MEMCG_NODE_STAT_ITEMS ARRAY_SIZE(memcg_node_stat_items)
@@ -1625,6 +1627,7 @@ static const struct memory_stat memory_stats[] = {
 #ifdef CONFIG_HUGETLB_PAGE
 	{ "hugetlb",			NR_HUGETLB			},
 #endif
+	{ "dmabuf",			MEMCG_DMABUF			},
 
 	/* The memory events */
 	{ "workingset_refault_anon",	WORKINGSET_REFAULT_ANON		},
@@ -5658,6 +5661,90 @@ void mem_cgroup_flush_workqueue(void)
 	flush_workqueue(memcg_wq);
 }
 
+/**
+ * mem_cgroup_reserve_dmabuf - charge a memcg for a dma-buf allocation
+ * @memcg: memcg to charge (may be %NULL)
+ * @gfp_mask: reclaim mode
+ * @nr_pages: number of pages to charge
+ *
+ * Charges @nr_pages against @memcg. Call this before allocating pages so
+ * the cgroup limit is enforced up front. On success the caller must either
+ * call mem_cgroup_commit_dmabuf() to finalize the charge, or
+ * mem_cgroup_cancel_dmabuf() to undo it.
+ *
+ * Returns %true on success, %false if the charge would exceed the limit.
+ */
+bool mem_cgroup_reserve_dmabuf(struct mem_cgroup *memcg, gfp_t gfp_mask,
+			       unsigned int nr_pages)
+{
+	if (!memcg || mem_cgroup_disabled())
+		return true;
+
+	if (try_charge(memcg, gfp_mask, nr_pages))
+		return false;
+
+	return true;
+}
+
+/**
+ * mem_cgroup_commit_dmabuf - finalize a dma-buf charge after export
+ * @dmabuf: the exported dma-buf
+ * @memcg: memcg previously charged via mem_cgroup_reserve_dmabuf()
+ * @nr_pages: same value passed to mem_cgroup_reserve_dmabuf()
+ *
+ * Updates the MEMCG_DMABUF counter and stores @memcg in @dmabuf->memcg,
+ * transferring the caller's reference.
+ */
+void mem_cgroup_commit_dmabuf(struct dma_buf *dmabuf, struct mem_cgroup *memcg,
+			      unsigned int nr_pages)
+{
+	if (!memcg || mem_cgroup_disabled())
+		return;
+
+	mod_memcg_state(memcg, MEMCG_DMABUF, nr_pages);
+	dmabuf->memcg = memcg;
+}
+
+/**
+ * mem_cgroup_cancel_dmabuf - undo a dma-buf charge
+ * @memcg: memcg to uncharge (may be %NULL)
+ * @nr_pages: same value passed to mem_cgroup_reserve_dmabuf()
+ *
+ * Returns @nr_pages of charge to @memcg. Use when the allocation fails
+ * after mem_cgroup_reserve_dmabuf() but before mem_cgroup_commit_dmabuf().
+ */
+void mem_cgroup_cancel_dmabuf(struct mem_cgroup *memcg, unsigned int nr_pages)
+{
+	if (!memcg)
+		return;
+
+	if (!mem_cgroup_is_root(memcg))
+		refill_stock(memcg, nr_pages);
+}
+
+/**
+ * mem_cgroup_uncharge_dmabuf - uncharge a dma-buf from its memcg
+ * @dmabuf: buffer to uncharge
+ *
+ * Reverses a successful mem_cgroup_commit_dmabuf(). Drops the memcg
+ * reference stored in @dmabuf->memcg and clears the pointer.
+ * Safe to call when @dmabuf->memcg is %NULL.
+ */
+void mem_cgroup_uncharge_dmabuf(struct dma_buf *dmabuf)
+{
+	unsigned int nr_pages;
+	struct mem_cgroup *memcg = dmabuf->memcg;
+
+	if (!memcg)
+		return;
+
+	nr_pages = PAGE_ALIGN(dmabuf->size) >> PAGE_SHIFT;
+	mod_memcg_state(memcg, MEMCG_DMABUF, -nr_pages);
+	mem_cgroup_cancel_dmabuf(memcg, nr_pages);
+	mem_cgroup_put(memcg);
+	dmabuf->memcg = NULL;
+}
+
 static int __init cgroup_memory(char *s)
 {
 	char *token;

-- 
2.55.0