io_uring/net: make page accounting more consistent

Make network page accounting more consistent with how buffer
registration is working, i.e. account all memory to ctx->user.

Signed-off-by: Pavel Begunkov <asml.silence@gmail.com>
Link: https://lore.kernel.org/r/4aacfe64bbb81b27f9ecf5d5c219c69a07e5aa56.1658742118.git.asml.silence@gmail.com
Signed-off-by: Jens Axboe <axboe@kernel.dk>

1 files changed, 19 insertions, 0 deletions

diff --git a/io_uring/notif.h b/io_uring/notif.h
index 3e05d2cecb6f..d6f366b1518b 100644
--- a/io_uring/notif.h
+++ b/io_uring/notif.h
@@ -5,6 +5,8 @@
 #include <net/sock.h>
 #include <linux/nospec.h>
 
+#include "rsrc.h"
+
 #define IO_NOTIF_SPLICE_BATCH	32
 #define IORING_MAX_NOTIF_SLOTS (1U << 10)
 
@@ -23,6 +25,8 @@ struct io_notif {
 	/* hook into ctx->notif_list and ctx->notif_list_locked */
 	struct list_head	cache_node;
 
+	unsigned long		account_pages;
+
 	union {
 		struct callback_head	task_work;
 		struct work_struct	commit_work;
@@ -85,3 +89,18 @@ static inline void io_notif_slot_flush_submit(struct io_notif_slot *slot,
 	}
 	io_notif_slot_flush(slot);
 }
+
+static inline int io_notif_account_mem(struct io_notif *notif, unsigned len)
+{
+	struct io_ring_ctx *ctx = notif->ctx;
+	unsigned nr_pages = (len >> PAGE_SHIFT) + 2;
+	int ret;
+
+	if (ctx->user) {
+		ret = __io_account_mem(ctx->user, nr_pages);
+		if (ret)
+			return ret;
+		notif->account_pages += nr_pages;
+	}
+	return 0;
+}