Re: [PATCH v6 5/6] fuse: add zero-copy over io-uring

diff --git a/fs/fuse/dev_uring.c b/fs/fuse/dev_uring.c
index d7d517497a1e..e1b55f7c0c09 100644
--- a/fs/fuse/dev_uring.c
+++ b/fs/fuse/dev_uring.c
@@ -34,11 +34,6 @@ struct fuse_uring_pdu {
 	struct fuse_ring_ent *ent;
 };
 
-struct fuse_zero_copy_bvs {
-	unsigned int nr_bvs;
-	struct bio_vec bvs[];
-};
-
 static const struct fuse_iqueue_ops fuse_io_uring_ops;
 
 enum fuse_uring_header_type {
@@ -786,24 +781,12 @@ static int fuse_uring_copy_from_ring(struct fuse_req *req,
 	return err;
 }
 
-static void fuse_zero_copy_release(void *priv)
-{
-	struct fuse_zero_copy_bvs *zc_bvs = priv;
-	unsigned int i;
-
-	for (i = 0; i < zc_bvs->nr_bvs; i++)
-		folio_put(page_folio(zc_bvs->bvs[i].bv_page));
-
-	kfree(zc_bvs);
-}
-
 static int fuse_uring_set_up_zero_copy(struct fuse_ring_ent *ent,
 				       struct fuse_req *req,
 				       unsigned int issue_flags)
 {
 	struct fuse_args_pages *ap;
 	int err, i, ddir = 0;
-	struct fuse_zero_copy_bvs *zc_bvs;
 	struct bio_vec *bvs;
 
 	/* out_pages indicates a read, in_pages indicates a write */
@@ -814,13 +797,11 @@ static int fuse_uring_set_up_zero_copy(struct fuse_ring_ent *ent,
 
 	ap = container_of(req->args, typeof(*ap), args);
 
-	zc_bvs = kmalloc(struct_size(zc_bvs, bvs, ap->num_folios),
-			 GFP_KERNEL_ACCOUNT);
-	if (!zc_bvs)
+	/* only needed until registration copies it into the ring */
+	bvs = kmalloc_array(ap->num_folios, sizeof(*bvs), GFP_KERNEL_ACCOUNT);
+	if (!bvs)
 		return -ENOMEM;
 
-	zc_bvs->nr_bvs = ap->num_folios;
-	bvs = zc_bvs->bvs;
 	for (i = 0; i < ap->num_folios; i++) {
 		bvs[i].bv_page = folio_page(ap->folios[i], 0);
 		bvs[i].bv_offset = ap->descs[i].offset;
@@ -828,12 +809,13 @@ static int fuse_uring_set_up_zero_copy(struct fuse_ring_ent *ent,
 		folio_get(ap->folios[i]);
 	}
 
-	err = io_buffer_register_bvec(ent->cmd, bvs, ap->num_folios,
-				      fuse_zero_copy_release, zc_bvs,
-				      ddir, ent->zero_copy_index,
-				      issue_flags);
+	err = io_buffer_register_bvec(ent->cmd, bvs, ap->num_folios, NULL, NULL,
+				      ddir, IO_BUF_F_PUT_FOLIOS,
+				      ent->zero_copy_index, issue_flags);
+	kfree(bvs);
 	if (err) {
-		fuse_zero_copy_release(zc_bvs);
+		for (i = 0; i < ap->num_folios; i++)
+			folio_put(ap->folios[i]);
 		return err;
 	}
 
diff --git a/include/linux/io_uring/cmd.h b/include/linux/io_uring/cmd.h
index 42801f0b6456..980948c86c62 100644
--- a/include/linux/io_uring/cmd.h
+++ b/include/linux/io_uring/cmd.h
@@ -41,6 +41,15 @@ static inline void io_uring_cmd_private_sz_check(size_t cmd_sz)
 	((pdu_type *)&(cmd)->pdu) \
 )
 
+/* Flags for io_buffer_register_bvec() */
+enum io_buffer_register_flags {
+	/*
+	 * The caller holds a folio reference per bvec, to be dropped when the
+	 * buffer is released. Registers without release/priv.
+	 */
+	IO_BUF_F_PUT_FOLIOS	= (1u << 0),
+};
+
 #if defined(CONFIG_IO_URING)
 int io_uring_cmd_import_fixed(u64 ubuf, unsigned long len, int rw,
 			      struct iov_iter *iter,
@@ -96,8 +105,8 @@ int io_buffer_register_request(struct io_uring_cmd *cmd, struct request *rq,
 			       unsigned int issue_flags);
 int io_buffer_register_bvec(struct io_uring_cmd *cmd, const struct bio_vec *bvs,
 			    unsigned int nr_bvecs, void (*release)(void *),
-			    void *priv, u8 dir, unsigned int index,
-			    unsigned int issue_flags);
+			    void *priv, u8 dir, unsigned int flags,
+			    unsigned int index, unsigned int issue_flags);
 int io_buffer_unregister(struct io_uring_cmd *cmd, unsigned int index,
 			 unsigned int issue_flags);
 #else
@@ -154,7 +163,8 @@ static inline int io_buffer_register_bvec(struct io_uring_cmd *cmd,
 					  const struct bio_vec *bvs,
 					  unsigned int nr_bvecs,
 					  void (*release)(void *), void *priv,
-					  u8 dir, unsigned int index,
+					  u8 dir, unsigned int flags,
+					  unsigned int index,
 					  unsigned int issue_flags)
 {
 	return -EOPNOTSUPP;
diff --git a/io_uring/rsrc.c b/io_uring/rsrc.c
index f3f01e0c8102..bd73a7ee98cc 100644
--- a/io_uring/rsrc.c
+++ b/io_uring/rsrc.c
@@ -158,6 +158,15 @@ static void io_release_ubuf(void *priv)
 	}
 }
 
+static void io_release_kbuf_folios(void *priv)
+{
+	struct io_mapped_ubuf *imu = priv;
+	unsigned int i;
+
+	for (i = 0; i < imu->nr_bvecs; i++)
+		folio_put(bvec_folio(&imu->bvec[i]));
+}
+
 static struct io_mapped_ubuf *io_alloc_imu(struct io_ring_ctx *ctx,
 					   int nr_bvecs)
 {
@@ -1101,8 +1110,8 @@ EXPORT_SYMBOL_GPL(io_buffer_register_request);
  */
 int io_buffer_register_bvec(struct io_uring_cmd *cmd, const struct bio_vec *bvs,
 			    unsigned int nr_bvecs, void (*release)(void *),
-			    void *priv, u8 dir, unsigned int index,
-			    unsigned int issue_flags)
+			    void *priv, u8 dir, unsigned int flags,
+			    unsigned int index, unsigned int issue_flags)
 {
 	struct io_ring_ctx *ctx = cmd_to_io_kiocb(cmd)->ctx;
 	struct io_mapped_ubuf *imu;
@@ -1125,6 +1134,12 @@ int io_buffer_register_bvec(struct io_uring_cmd *cmd, const struct bio_vec *bvs,
 	for (i = 0; i < nr_bvecs; i++)
 		bvec[i] = bvs[i];
 
+	/* imu is only known here, and needs the bvecs in place */
+	if (flags & IO_BUF_F_PUT_FOLIOS) {
+		imu->release = io_release_kbuf_folios;
+		imu->priv = imu;
+	}
+
 unlock:
 	io_ring_submit_unlock(ctx, issue_flags);
 	return ret;