[PATCH 0/4] io_uring: improve alloc_cache branch prediction

[PATCH 0/4] io_uring: improve alloc_cache branch prediction

[Caleb Sander Mateos]

A call to io_uring_alloc_async_data() that provides a struct io_alloc_cache
currently has 2 unnecessary NULL checks in the fast path:
- Of the struct io_alloc_cache *, even though it points to a struct io_ring_ctx
  field
- Of the void * entry popped from the cache, even though all cache entries are
  non-NULL
The compiler also incorrectly assumes the success path in io_alloc_cache_get()
is the unlikely one because it's written as an early return.

Avoid the first NULL check by splitting io_uring_alloc_async_data()'s cache
logic to a separate function io_uring_alloc_cached_async_data() that uses the
cache unconditionally. Avoid the second by returning a separate success bool
from io_cache_alloc() instead of checking the NULL-ness of the returned entry.
Also mark io_alloc_cache_get() and io_alloc_cache_put()'s success paths as
likely so the compiler can avoid jumps in the fast path.

The final commit is a small code simplification in io_futex_wait() noticed while
touching the neighboring io_cache_alloc() call.

Caleb Sander Mateos (4):
  io_uring: split io_uring_alloc_async_data()
  io_uring/alloc_cache: mark cache success paths as likely
  io_uring/alloc_cache: assume io_alloc_cache_get() doesn't return NULL
  io_uring/futex: avoid ifd NULL check in io_futex_wait()

 io_uring/alloc_cache.c |  2 +-
 io_uring/alloc_cache.h | 25 ++++++++++++------------
 io_uring/filetable.c   |  3 +--
 io_uring/futex.c       | 13 +++++--------
 io_uring/io_uring.h    | 24 ++++++++++++++---------
 io_uring/net.c         |  8 ++++----
 io_uring/poll.c        | 13 ++++++++-----
 io_uring/rsrc.c        | 43 +++++++++++++++++++++---------------------
 io_uring/rsrc.h        |  3 ++-
 io_uring/rw.c          |  4 ++--
 io_uring/timeout.c     |  2 +-
 io_uring/uring_cmd.c   |  4 ++--
 io_uring/waitid.c      |  2 +-
 13 files changed, 76 insertions(+), 70 deletions(-)

-- 
2.56.0

[PATCH 1/4] io_uring: split io_uring_alloc_async_data()

[Caleb Sander Mateos]

io_uring_alloc_async_data() accepts an optional struct io_alloc_cache *
argument to use for the allocation in place of kmalloc(). All callers
either always pass NULL or a non-NULL io_alloc_cache. However, the
compiler still emits a NULL check for callers passing a non-NULL
io_alloc_cache. Split out a separate io_uring_alloc_cached_async_data()
for the non-NULL io_alloc_cache cache. This avoids the NULL check at the
cost of duplicating the code to set REQ_F_ASYNC_DATA.

Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
---
 io_uring/io_uring.h  | 23 ++++++++++++++---------
 io_uring/net.c       |  6 +++---
 io_uring/rw.c        |  2 +-
 io_uring/timeout.c   |  2 +-
 io_uring/uring_cmd.c |  2 +-
 io_uring/waitid.c    |  2 +-
 6 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/io_uring/io_uring.h b/io_uring/io_uring.h
index 32247529076c..a36a177559f4 100644
--- a/io_uring/io_uring.h
+++ b/io_uring/io_uring.h
@@ -335,21 +335,26 @@ static inline ssize_t io_fixup_restart_res(ssize_t ret)
 	default:
 		return ret;
 	}
 }
 
-static inline void *io_uring_alloc_async_data(struct io_alloc_cache *cache,
-					      struct io_kiocb *req)
+static inline void *
+io_uring_alloc_cached_async_data(struct io_alloc_cache *cache,
+				 struct io_kiocb *req)
 {
-	if (cache) {
-		req->async_data = io_cache_alloc(cache, GFP_KERNEL);
-	} else {
-		const struct io_issue_def *def = &io_issue_defs[req->opcode];
+	req->async_data = io_cache_alloc(cache, GFP_KERNEL);
+	if (req->async_data)
+		req->flags |= REQ_F_ASYNC_DATA;
+	return req->async_data;
+}
 
-		WARN_ON_ONCE(!def->async_size);
-		req->async_data = kmalloc(def->async_size, GFP_KERNEL);
-	}
+static inline void *io_uring_alloc_async_data(struct io_kiocb *req)
+{
+	const struct io_issue_def *def = &io_issue_defs[req->opcode];
+
+	WARN_ON_ONCE(!def->async_size);
+	req->async_data = kmalloc(def->async_size, GFP_KERNEL);
 	if (req->async_data)
 		req->flags |= REQ_F_ASYNC_DATA;
 	return req->async_data;
 }
 
diff --git a/io_uring/net.c b/io_uring/net.c
index 533afd335627..e5f7c13519fb 100644
--- a/io_uring/net.c
+++ b/io_uring/net.c
@@ -187,11 +187,11 @@ static void io_netmsg_recycle(struct io_kiocb *req, unsigned int issue_flags)
 static struct io_async_msghdr *io_msg_alloc_async(struct io_kiocb *req)
 {
 	struct io_ring_ctx *ctx = req->ctx;
 	struct io_async_msghdr *hdr;
 
-	hdr = io_uring_alloc_async_data(&ctx->netmsg_cache, req);
+	hdr = io_uring_alloc_cached_async_data(&ctx->netmsg_cache, req);
 	if (!hdr)
 		return NULL;
 
 	/* If the async data was cached, we might have an iov cached inside. */
 	if (hdr->vec.iovec)
@@ -1846,11 +1846,11 @@ int io_connect_prep(struct io_kiocb *req, const struct io_uring_sqe *sqe)
 
 	conn->addr = u64_to_user_ptr(READ_ONCE(sqe->addr));
 	conn->addr_len =  READ_ONCE(sqe->addr2);
 	conn->in_progress = conn->seen_econnaborted = false;
 
-	addr = io_uring_alloc_async_data(NULL, req);
+	addr = io_uring_alloc_async_data(req);
 	if (unlikely(!addr))
 		return -ENOMEM;
 
 	return move_addr_to_kernel(conn->addr, conn->addr_len, addr);
 }
@@ -1944,11 +1944,11 @@ int io_bind_prep(struct io_kiocb *req, const struct io_uring_sqe *sqe)
 		return -EINVAL;
 
 	uaddr = u64_to_user_ptr(READ_ONCE(sqe->addr));
 	bind->addr_len =  READ_ONCE(sqe->addr2);
 
-	addr = io_uring_alloc_async_data(NULL, req);
+	addr = io_uring_alloc_async_data(req);
 	if (unlikely(!addr))
 		return -ENOMEM;
 	ret = move_addr_to_kernel(uaddr, bind->addr_len, addr);
 	if (unlikely(ret))
 		return ret;
diff --git a/io_uring/rw.c b/io_uring/rw.c
index 61bd0641dd77..4a9538157d94 100644
--- a/io_uring/rw.c
+++ b/io_uring/rw.c
@@ -191,11 +191,11 @@ static void io_req_rw_cleanup(struct io_kiocb *req, unsigned int issue_flags)
 static int io_rw_alloc_async(struct io_kiocb *req)
 {
 	struct io_ring_ctx *ctx = req->ctx;
 	struct io_async_rw *rw;
 
-	rw = io_uring_alloc_async_data(&ctx->rw_cache, req);
+	rw = io_uring_alloc_cached_async_data(&ctx->rw_cache, req);
 	if (!rw)
 		return -ENOMEM;
 	if (rw->vec.iovec)
 		req->flags |= REQ_F_NEED_CLEANUP;
 	rw->bytes_done = 0;
diff --git a/io_uring/timeout.c b/io_uring/timeout.c
index fa7b0d9a81d6..fa54846ed4ce 100644
--- a/io_uring/timeout.c
+++ b/io_uring/timeout.c
@@ -638,11 +638,11 @@ static int __io_timeout_prep(struct io_kiocb *req,
 	if ((flags & IORING_TIMEOUT_MULTISHOT) && off > 0)
 		timeout->repeats = off;
 
 	if (WARN_ON_ONCE(req_has_async_data(req)))
 		return -EFAULT;
-	data = io_uring_alloc_async_data(NULL, req);
+	data = io_uring_alloc_async_data(req);
 	if (!data)
 		return -ENOMEM;
 	data->req = req;
 	data->flags = flags;
 
diff --git a/io_uring/uring_cmd.c b/io_uring/uring_cmd.c
index b5486b56ee0c..6f417bda52ea 100644
--- a/io_uring/uring_cmd.c
+++ b/io_uring/uring_cmd.c
@@ -188,11 +188,11 @@ int io_uring_cmd_prep(struct io_kiocb *req, const struct io_uring_sqe *sqe)
 	    !!(req->flags & REQ_F_BUFFER_SELECT))
 		return -EINVAL;
 
 	ioucmd->cmd_op = READ_ONCE(sqe->cmd_op);
 
-	ac = io_uring_alloc_async_data(&req->ctx->cmd_cache, req);
+	ac = io_uring_alloc_cached_async_data(&req->ctx->cmd_cache, req);
 	if (!ac)
 		return -ENOMEM;
 	if (ac->vec.iovec)
 		req->flags |= REQ_F_NEED_CLEANUP;
 	ioucmd->sqe = sqe;
diff --git a/io_uring/waitid.c b/io_uring/waitid.c
index 986fb46639e8..0d3c248f375c 100644
--- a/io_uring/waitid.c
+++ b/io_uring/waitid.c
@@ -292,11 +292,11 @@ int io_waitid_prep(struct io_kiocb *req, const struct io_uring_sqe *sqe)
 	struct io_waitid_async *iwa;
 
 	if (sqe->addr || sqe->buf_index || sqe->addr3 || sqe->waitid_flags)
 		return -EINVAL;
 
-	iwa = io_uring_alloc_async_data(NULL, req);
+	iwa = io_uring_alloc_async_data(req);
 	if (unlikely(!iwa))
 		return -ENOMEM;
 	iwa->req = req;
 
 	iw->which = READ_ONCE(sqe->len);
-- 
2.56.0

[PATCH 2/4] io_uring/alloc_cache: mark cache success paths as likely

[Caleb Sander Mateos]

The point of a struct io_alloc_cache is to serve allocations from the
cache to avoid a call to kmalloc(). Only very high queue depth workloads
require more concurrent allocations than the fixed-size io_alloc_caches
can provide. So mark the success paths in io_alloc_cache_get() and
io_alloc_cache_put() as likely.

Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
---
 io_uring/alloc_cache.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/io_uring/alloc_cache.h b/io_uring/alloc_cache.h
index 962b6e2d04cc..aa007b380f1b 100644
--- a/io_uring/alloc_cache.h
+++ b/io_uring/alloc_cache.h
@@ -19,22 +19,22 @@ bool io_alloc_cache_init(struct io_alloc_cache *cache,
 void *io_cache_alloc_new(struct io_alloc_cache *cache, gfp_t gfp);
 
 static inline bool io_alloc_cache_put(struct io_alloc_cache *cache,
 				      void *entry)
 {
-	if (cache->nr_cached < cache->max_cached) {
+	if (likely(cache->nr_cached < cache->max_cached)) {
 		if (!kasan_mempool_poison_object(entry))
 			return false;
 		cache->entries[cache->nr_cached++] = entry;
 		return true;
 	}
 	return false;
 }
 
 static inline void *io_alloc_cache_get(struct io_alloc_cache *cache)
 {
-	if (cache->nr_cached) {
+	if (likely(cache->nr_cached)) {
 		void *entry = cache->entries[--cache->nr_cached];
 
 		/*
 		 * If KASAN is enabled, always clear the initial bytes that
 		 * must be zeroed post alloc, in case any of them overlap
-- 
2.56.0

[PATCH 3/4] io_uring/alloc_cache: assume io_alloc_cache_get() doesn't return NULL

[Caleb Sander Mateos]

All entries in a struct io_alloc_cache are non-NULL, but the compiler
doesn't know this. That results in an unnecessary NULL check on the
return value from io_alloc_cache_get() in the successful case. Return
the success/failure as a separate bool value to avoid the extra NULL
check. Propagate this change to all io_alloc_cache_get()'s transitive
callers.

Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
---
 io_uring/alloc_cache.c |  2 +-
 io_uring/alloc_cache.h | 21 +++++++++++----------
 io_uring/filetable.c   |  3 +--
 io_uring/futex.c       |  3 +--
 io_uring/io_uring.h    | 11 ++++++-----
 io_uring/net.c         |  4 ++--
 io_uring/poll.c        | 13 ++++++++-----
 io_uring/rsrc.c        | 43 +++++++++++++++++++++---------------------
 io_uring/rsrc.h        |  3 ++-
 io_uring/rw.c          |  4 ++--
 io_uring/uring_cmd.c   |  4 ++--
 11 files changed, 57 insertions(+), 54 deletions(-)

diff --git a/io_uring/alloc_cache.c b/io_uring/alloc_cache.c
index 58423888b736..75cd0b07b97e 100644
--- a/io_uring/alloc_cache.c
+++ b/io_uring/alloc_cache.c
@@ -8,11 +8,11 @@ void io_alloc_cache_free(struct io_alloc_cache *cache,
 	void *entry;
 
 	if (!cache->entries)
 		return;
 
-	while ((entry = io_alloc_cache_get(cache)) != NULL)
+	while (io_alloc_cache_get(cache, &entry))
 		free(entry);
 
 	kvfree(cache->entries);
 	cache->entries = NULL;
 }
diff --git a/io_uring/alloc_cache.h b/io_uring/alloc_cache.h
index aa007b380f1b..6aacbff1fa3a 100644
--- a/io_uring/alloc_cache.h
+++ b/io_uring/alloc_cache.h
@@ -28,11 +28,12 @@ static inline bool io_alloc_cache_put(struct io_alloc_cache *cache,
 		return true;
 	}
 	return false;
 }
 
-static inline void *io_alloc_cache_get(struct io_alloc_cache *cache)
+static inline bool io_alloc_cache_get(struct io_alloc_cache *cache,
+				      void **entry_out)
 {
 	if (likely(cache->nr_cached)) {
 		void *entry = cache->entries[--cache->nr_cached];
 
 		/*
@@ -43,24 +44,24 @@ static inline void *io_alloc_cache_get(struct io_alloc_cache *cache)
 #if defined(CONFIG_KASAN)
 		kasan_mempool_unpoison_object(entry, cache->elem_size);
 		if (cache->init_clear)
 			memset(entry, 0, cache->init_clear);
 #endif
-		return entry;
+		*entry_out = entry;
+		return true;
 	}
 
-	return NULL;
+	return false;
 }
 
-static inline void *io_cache_alloc(struct io_alloc_cache *cache, gfp_t gfp)
+static inline bool io_cache_alloc(struct io_alloc_cache *cache, gfp_t gfp,
+				  void **obj)
 {
-	void *obj;
-
-	obj = io_alloc_cache_get(cache);
-	if (obj)
-		return obj;
-	return io_cache_alloc_new(cache, gfp);
+	if (io_alloc_cache_get(cache, obj))
+		return true;
+	*obj = io_cache_alloc_new(cache, gfp);
+	return *obj;
 }
 
 static inline void io_cache_free(struct io_alloc_cache *cache, void *obj)
 {
 	if (!io_alloc_cache_put(cache, obj))
diff --git a/io_uring/filetable.c b/io_uring/filetable.c
index cb1838c9fc37..3c847c478235 100644
--- a/io_uring/filetable.c
+++ b/io_uring/filetable.c
@@ -70,12 +70,11 @@ static int io_install_fixed_file(struct io_ring_ctx *ctx, struct file *file,
 	if (!ctx->file_table.data.nr)
 		return -ENXIO;
 	if (slot_index >= ctx->file_table.data.nr)
 		return -EINVAL;
 
-	node = io_rsrc_node_alloc(ctx, IORING_RSRC_FILE);
-	if (!node)
+	if (!io_rsrc_node_alloc(ctx, IORING_RSRC_FILE, &node))
 		return -ENOMEM;
 
 	if (!io_reset_rsrc_node(ctx, &ctx->file_table.data, slot_index))
 		io_file_bitmap_set(&ctx->file_table, slot_index);
 
diff --git a/io_uring/futex.c b/io_uring/futex.c
index 9e4e63ee7bc1..9c7d384bab9b 100644
--- a/io_uring/futex.c
+++ b/io_uring/futex.c
@@ -326,12 +326,11 @@ int io_futex_wait(struct io_kiocb *req, unsigned int issue_flags)
 		ret = -EINVAL;
 		goto done;
 	}
 
 	io_ring_submit_lock(ctx, issue_flags);
-	ifd = io_cache_alloc(&ctx->futex_cache, GFP_NOWAIT);
-	if (!ifd) {
+	if (!io_cache_alloc(&ctx->futex_cache, GFP_NOWAIT, (void **)&ifd)) {
 		ret = -ENOMEM;
 		goto done_unlock;
 	}
 
 	req->flags |= REQ_F_ASYNC_DATA;
diff --git a/io_uring/io_uring.h b/io_uring/io_uring.h
index a36a177559f4..4a563bd32c77 100644
--- a/io_uring/io_uring.h
+++ b/io_uring/io_uring.h
@@ -335,18 +335,19 @@ static inline ssize_t io_fixup_restart_res(ssize_t ret)
 	default:
 		return ret;
 	}
 }
 
-static inline void *
+static inline bool
 io_uring_alloc_cached_async_data(struct io_alloc_cache *cache,
 				 struct io_kiocb *req)
 {
-	req->async_data = io_cache_alloc(cache, GFP_KERNEL);
-	if (req->async_data)
-		req->flags |= REQ_F_ASYNC_DATA;
-	return req->async_data;
+	if (!io_cache_alloc(cache, GFP_KERNEL, &req->async_data))
+		return false;
+
+	req->flags |= REQ_F_ASYNC_DATA;
+	return true;
 }
 
 static inline void *io_uring_alloc_async_data(struct io_kiocb *req)
 {
 	const struct io_issue_def *def = &io_issue_defs[req->opcode];
diff --git a/io_uring/net.c b/io_uring/net.c
index e5f7c13519fb..183425f28b0c 100644
--- a/io_uring/net.c
+++ b/io_uring/net.c
@@ -187,14 +187,14 @@ static void io_netmsg_recycle(struct io_kiocb *req, unsigned int issue_flags)
 static struct io_async_msghdr *io_msg_alloc_async(struct io_kiocb *req)
 {
 	struct io_ring_ctx *ctx = req->ctx;
 	struct io_async_msghdr *hdr;
 
-	hdr = io_uring_alloc_cached_async_data(&ctx->netmsg_cache, req);
-	if (!hdr)
+	if (!io_uring_alloc_cached_async_data(&ctx->netmsg_cache, req))
 		return NULL;
 
+	hdr = req->async_data;
 	/* If the async data was cached, we might have an iov cached inside. */
 	if (hdr->vec.iovec)
 		req->flags |= REQ_F_NEED_CLEANUP;
 	return hdr;
 }
diff --git a/io_uring/poll.c b/io_uring/poll.c
index 4f04bda42f7f..8aca992b72c4 100644
--- a/io_uring/poll.c
+++ b/io_uring/poll.c
@@ -681,16 +681,19 @@ static struct async_poll *io_req_alloc_apoll(struct io_kiocb *req,
 
 	if (req->flags & REQ_F_POLLED) {
 		apoll = req->apoll;
 		kfree(apoll->double_poll);
 	} else {
-		if (!(issue_flags & IO_URING_F_UNLOCKED))
-			apoll = io_cache_alloc(&ctx->apoll_cache, GFP_ATOMIC);
-		else
+		if (!(issue_flags & IO_URING_F_UNLOCKED)) {
+			if (!io_cache_alloc(&ctx->apoll_cache, GFP_ATOMIC,
+					    (void **)&apoll))
+				return NULL;
+		} else {
 			apoll = kmalloc_obj(*apoll, GFP_ATOMIC);
-		if (!apoll)
-			return NULL;
+			if (!apoll)
+				return NULL;
+		}
 		apoll->poll.retries = APOLL_MAX_RETRY;
 	}
 	apoll->double_poll = NULL;
 	req->apoll = apoll;
 	if (unlikely(!--apoll->poll.retries))
diff --git a/io_uring/rsrc.c b/io_uring/rsrc.c
index d569cea12df9..11e52db52373 100644
--- a/io_uring/rsrc.c
+++ b/io_uring/rsrc.c
@@ -156,17 +156,19 @@ static void io_release_ubuf(void *priv)
 
 		unpin_user_folio(folio, 1);
 	}
 }
 
-static struct io_mapped_ubuf *io_alloc_imu(struct io_ring_ctx *ctx,
-					   int nr_bvecs)
+static bool io_alloc_imu(struct io_ring_ctx *ctx, int nr_bvecs,
+			 struct io_mapped_ubuf **imu)
 {
 	if (nr_bvecs <= IO_CACHED_BVECS_SEGS)
-		return io_cache_alloc(&ctx->imu_cache, GFP_KERNEL_ACCOUNT);
-	return kvmalloc_flex(struct io_mapped_ubuf, bvec, nr_bvecs,
+		return io_cache_alloc(&ctx->imu_cache, GFP_KERNEL_ACCOUNT,
+				      (void **)imu);
+	*imu = kvmalloc_flex(struct io_mapped_ubuf, bvec, nr_bvecs,
 			     GFP_KERNEL_ACCOUNT);
+	return *imu;
 }
 
 static void io_free_imu(struct io_ring_ctx *ctx, struct io_mapped_ubuf *imu)
 {
 	if (imu->nr_bvecs <= IO_CACHED_BVECS_SEGS)
@@ -241,22 +243,25 @@ static void io_buffer_unmap(struct io_ring_ctx *ctx, struct io_mapped_ubuf *imu)
 		io_unaccount_mem(ctx->user, ctx->mm_account, acct_pages);
 	imu->release(imu->priv);
 	io_free_imu(ctx, imu);
 }
 
-struct io_rsrc_node *io_rsrc_node_alloc(struct io_ring_ctx *ctx, int type)
+bool io_rsrc_node_alloc(struct io_ring_ctx *ctx, int type,
+			struct io_rsrc_node **node_out)
 {
 	struct io_rsrc_node *node;
 
-	node = io_cache_alloc(&ctx->node_cache, GFP_KERNEL_ACCOUNT);
-	if (node) {
+	if (io_cache_alloc(&ctx->node_cache, GFP_KERNEL_ACCOUNT,
+			   (void **)&node)) {
 		node->type = type;
 		node->refs = 1;
 		node->tag = 0;
 		node->file_ptr = 0;
+		*node_out = node;
+		return true;
 	}
-	return node;
+	return false;
 }
 
 bool io_rsrc_cache_init(struct io_ring_ctx *ctx)
 {
 	const int imu_cache_size = struct_size_t(struct io_mapped_ubuf, bvec,
@@ -364,12 +369,11 @@ static int __io_sqe_files_update(struct io_ring_ctx *ctx,
 			if (io_is_uring_fops(file)) {
 				fput(file);
 				err = -EBADF;
 				break;
 			}
-			node = io_rsrc_node_alloc(ctx, IORING_RSRC_FILE);
-			if (!node) {
+			if (!io_rsrc_node_alloc(ctx, IORING_RSRC_FILE, &node)) {
 				err = -ENOMEM;
 				fput(file);
 				break;
 			}
 			ctx->file_table.data.nodes[i] = node;
@@ -678,12 +682,11 @@ int io_sqe_files_register(struct io_ring_ctx *ctx, void __user *arg,
 		if (io_is_uring_fops(file)) {
 			fput(file);
 			goto fail;
 		}
 		ret = -ENOMEM;
-		node = io_rsrc_node_alloc(ctx, IORING_RSRC_FILE);
-		if (!node) {
+		if (!io_rsrc_node_alloc(ctx, IORING_RSRC_FILE, &node)) {
 			fput(file);
 			goto fail;
 		}
 		if (tag)
 			node->tag = tag;
@@ -892,12 +895,11 @@ static struct io_rsrc_node *io_sqe_buffer_register(struct io_ring_ctx *ctx,
 	ret = io_validate_user_buf_range((unsigned long)iov->iov_base,
 					 iov->iov_len);
 	if (ret)
 		return ERR_PTR(ret);
 
-	node = io_rsrc_node_alloc(ctx, IORING_RSRC_BUFFER);
-	if (!node)
+	if (!io_rsrc_node_alloc(ctx, IORING_RSRC_BUFFER, &node))
 		return ERR_PTR(-ENOMEM);
 
 	ret = -ENOMEM;
 	pages = io_pin_pages((unsigned long) iov->iov_base, iov->iov_len,
 				&nr_pages);
@@ -911,12 +913,11 @@ static struct io_rsrc_node *io_sqe_buffer_register(struct io_ring_ctx *ctx,
 	if (nr_pages > 1 && io_check_coalesce_buffer(pages, nr_pages, &data)) {
 		if (data.nr_pages_mid != 1)
 			coalesced = io_coalesce_buffer(&pages, &nr_pages, &data);
 	}
 
-	imu = io_alloc_imu(ctx, nr_pages);
-	if (!imu)
+	if (!io_alloc_imu(ctx, nr_pages, &imu))
 		goto done;
 
 	imu->nr_bvecs = nr_pages;
 	ret = io_buffer_account_pin(ctx, pages, nr_pages);
 	if (ret)
@@ -1050,16 +1051,14 @@ static struct io_mapped_ubuf *io_kernel_buffer_init(struct io_ring_ctx *ctx,
 	index = array_index_nospec(index, data->nr);
 
 	if (data->nodes[index])
 		return ERR_PTR(-EBUSY);
 
-	node = io_rsrc_node_alloc(ctx, IORING_RSRC_BUFFER);
-	if (!node)
+	if (!io_rsrc_node_alloc(ctx, IORING_RSRC_BUFFER, &node))
 		return ERR_PTR(-ENOMEM);
 
-	imu = io_alloc_imu(ctx, nr_bvecs);
-	if (!imu) {
+	if (!io_alloc_imu(ctx, nr_bvecs, &imu)) {
 		io_cache_free(&ctx->node_cache, node);
 		return ERR_PTR(-ENOMEM);
 	}
 
 	imu->ubuf = 0;
@@ -1421,12 +1420,12 @@ static int io_clone_buffers(struct io_ring_ctx *ctx, struct io_ring_ctx *src_ctx
 
 		src_node = io_rsrc_node_lookup(&src_ctx->buf_table, i);
 		if (!src_node) {
 			dst_node = NULL;
 		} else {
-			dst_node = io_rsrc_node_alloc(ctx, IORING_RSRC_BUFFER);
-			if (!dst_node) {
+			if (!io_rsrc_node_alloc(ctx, IORING_RSRC_BUFFER,
+						&dst_node)) {
 				io_rsrc_data_free(ctx, &data);
 				return -ENOMEM;
 			}
 
 			refcount_inc(&src_node->buf->refs);
diff --git a/io_uring/rsrc.h b/io_uring/rsrc.h
index 9ef88383b363..2632aa577d62 100644
--- a/io_uring/rsrc.h
+++ b/io_uring/rsrc.h
@@ -52,11 +52,12 @@ struct io_imu_folio_data {
 	unsigned long	first_folio_page_idx;
 };
 
 bool io_rsrc_cache_init(struct io_ring_ctx *ctx);
 void io_rsrc_cache_free(struct io_ring_ctx *ctx);
-struct io_rsrc_node *io_rsrc_node_alloc(struct io_ring_ctx *ctx, int type);
+bool io_rsrc_node_alloc(struct io_ring_ctx *ctx, int type,
+			struct io_rsrc_node **node);
 void io_free_rsrc_node(struct io_ring_ctx *ctx, struct io_rsrc_node *node);
 void io_rsrc_data_free(struct io_ring_ctx *ctx, struct io_rsrc_data *data);
 int io_rsrc_data_alloc(struct io_rsrc_data *data, unsigned nr);
 
 struct io_rsrc_node *io_find_buf_node(struct io_kiocb *req,
diff --git a/io_uring/rw.c b/io_uring/rw.c
index 4a9538157d94..29fb96f76191 100644
--- a/io_uring/rw.c
+++ b/io_uring/rw.c
@@ -191,13 +191,13 @@ static void io_req_rw_cleanup(struct io_kiocb *req, unsigned int issue_flags)
 static int io_rw_alloc_async(struct io_kiocb *req)
 {
 	struct io_ring_ctx *ctx = req->ctx;
 	struct io_async_rw *rw;
 
-	rw = io_uring_alloc_cached_async_data(&ctx->rw_cache, req);
-	if (!rw)
+	if (!io_uring_alloc_cached_async_data(&ctx->rw_cache, req))
 		return -ENOMEM;
+	rw = req->async_data;
 	if (rw->vec.iovec)
 		req->flags |= REQ_F_NEED_CLEANUP;
 	rw->bytes_done = 0;
 	return 0;
 }
diff --git a/io_uring/uring_cmd.c b/io_uring/uring_cmd.c
index 6f417bda52ea..ba2d8c9028eb 100644
--- a/io_uring/uring_cmd.c
+++ b/io_uring/uring_cmd.c
@@ -188,13 +188,13 @@ int io_uring_cmd_prep(struct io_kiocb *req, const struct io_uring_sqe *sqe)
 	    !!(req->flags & REQ_F_BUFFER_SELECT))
 		return -EINVAL;
 
 	ioucmd->cmd_op = READ_ONCE(sqe->cmd_op);
 
-	ac = io_uring_alloc_cached_async_data(&req->ctx->cmd_cache, req);
-	if (!ac)
+	if (!io_uring_alloc_cached_async_data(&req->ctx->cmd_cache, req))
 		return -ENOMEM;
+	ac = req->async_data;
 	if (ac->vec.iovec)
 		req->flags |= REQ_F_NEED_CLEANUP;
 	ioucmd->sqe = sqe;
 	return 0;
 }
-- 
2.56.0

[PATCH 4/4] io_uring/futex: avoid ifd NULL check in io_futex_wait()

[Caleb Sander Mateos]

There are two paths that hit the done_unlock label in io_futex_wait():
the ENOMEM error path allocating ifd (via goto done_unlock) and the
futex_wait_setup() error path (via a fallthrough). In the ENOMEM path,
ifd is always NULL; in the futex_wait_setup() error path, it's always
non-NULL. Move the done_unlock label so the goto skips the ifd freeing
and the success path can free it unconditionally. This also avoids the
need to initialize ifd before the io_cache_alloc() call, as it's only
used if io_cache_alloc() succeeds.

Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
---
 io_uring/futex.c | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/io_uring/futex.c b/io_uring/futex.c
index 9c7d384bab9b..e0cb815fd37a 100644
--- a/io_uring/futex.c
+++ b/io_uring/futex.c
@@ -317,11 +317,11 @@ int io_futexv_wait(struct io_kiocb *req, unsigned int issue_flags)
 
 int io_futex_wait(struct io_kiocb *req, unsigned int issue_flags)
 {
 	struct io_futex *iof = io_kiocb_to_cmd(req, struct io_futex);
 	struct io_ring_ctx *ctx = req->ctx;
-	struct io_futex_data *ifd = NULL;
+	struct io_futex_data *ifd;
 	int ret;
 
 	if (!iof->futex_mask) {
 		ret = -EINVAL;
 		goto done;
@@ -349,16 +349,14 @@ int io_futex_wait(struct io_kiocb *req, unsigned int issue_flags)
 		io_ring_submit_unlock(ctx, issue_flags);
 
 		return IOU_ISSUE_SKIP_COMPLETE;
 	}
 
+	/* As in io_futex_complete(), or failed arms would drain it. */
+	io_cache_free(&ctx->futex_cache, ifd);
+	io_req_async_data_clear(req, 0);
 done_unlock:
-	if (ifd) {
-		/* As in io_futex_complete(), or failed arms would drain it. */
-		io_cache_free(&ctx->futex_cache, ifd);
-		io_req_async_data_clear(req, 0);
-	}
 	io_ring_submit_unlock(ctx, issue_flags);
 done:
 	if (ret < 0)
 		req_set_fail(req);
 	io_req_set_res(req, ret, 0);
-- 
2.56.0