4 files changed, 38 insertions, 6 deletions

diff --git a/tools/testing/selftests/ublk/file_backed.c b/tools/testing/selftests/ublk/file_backed.c
index a2e8793390a8..e2287eedaac8 100644
--- a/tools/testing/selftests/ublk/file_backed.c
+++ b/tools/testing/selftests/ublk/file_backed.c
@@ -123,7 +123,7 @@ exit:
 	q->io_inflight--;
 }
 
-static int ublk_loop_tgt_init(struct ublk_dev *dev)
+static int ublk_loop_tgt_init(const struct dev_ctx *ctx, struct ublk_dev *dev)
 {
 	unsigned long long bytes;
 	int ret;
diff --git a/tools/testing/selftests/ublk/kublk.c b/tools/testing/selftests/ublk/kublk.c
index 0080cad1f3ae..2dd17663ef30 100644
--- a/tools/testing/selftests/ublk/kublk.c
+++ b/tools/testing/selftests/ublk/kublk.c
@@ -381,7 +381,7 @@ static int ublk_queue_init(struct ublk_queue *q)
 
 #define WAIT_USEC 	100000
 #define MAX_WAIT_USEC 	(3 * 1000000)
-static int ublk_dev_prep(struct ublk_dev *dev)
+static int ublk_dev_prep(const struct dev_ctx *ctx, struct ublk_dev *dev)
 {
 	int dev_id = dev->dev_info.dev_id;
 	unsigned int wait_usec = 0;
@@ -404,7 +404,7 @@ static int ublk_dev_prep(struct ublk_dev *dev)
 
 	dev->fds[0] = fd;
 	if (dev->tgt.ops->init_tgt)
-		ret = dev->tgt.ops->init_tgt(dev);
+		ret = dev->tgt.ops->init_tgt(ctx, dev);
 	if (ret)
 		close(dev->fds[0]);
 	return ret;
@@ -666,7 +666,7 @@ static int ublk_start_daemon(const struct dev_ctx *ctx, struct ublk_dev *dev)
 
 	ublk_dbg(UBLK_DBG_DEV, "%s enter\n", __func__);
 
-	ret = ublk_dev_prep(dev);
+	ret = ublk_dev_prep(ctx, dev);
 	if (ret)
 		return ret;
 
diff --git a/tools/testing/selftests/ublk/kublk.h b/tools/testing/selftests/ublk/kublk.h
index eaadd7364e25..4eee9ad2bead 100644
--- a/tools/testing/selftests/ublk/kublk.h
+++ b/tools/testing/selftests/ublk/kublk.h
@@ -94,11 +94,14 @@ struct ublk_io {
 	unsigned short refs;		/* used by target code only */
 
 	int result;
+
+	unsigned short tgt_ios;
+	void *private_data;
 };
 
 struct ublk_tgt_ops {
 	const char *name;
-	int (*init_tgt)(struct ublk_dev *);
+	int (*init_tgt)(const struct dev_ctx *ctx, struct ublk_dev *);
 	void (*deinit_tgt)(struct ublk_dev *);
 
 	int (*queue_io)(struct ublk_queue *, int tag);
@@ -146,6 +149,8 @@ struct ublk_dev {
 	int nr_fds;
 	int ctrl_fd;
 	struct io_uring ring;
+
+	void *private_data;
 };
 
 #ifndef offsetof
@@ -303,6 +308,11 @@ static inline void ublk_set_sqe_cmd_op(struct io_uring_sqe *sqe, __u32 cmd_op)
 	addr[1] = 0;
 }
 
+static inline struct ublk_io *ublk_get_io(struct ublk_queue *q, unsigned tag)
+{
+	return &q->ios[tag];
+}
+
 static inline int ublk_complete_io(struct ublk_queue *q, unsigned tag, int res)
 {
 	struct ublk_io *io = &q->ios[tag];
@@ -312,6 +322,28 @@ static inline int ublk_complete_io(struct ublk_queue *q, unsigned tag, int res)
 	return ublk_queue_io_cmd(q, io, tag);
 }
 
+static inline void ublk_queued_tgt_io(struct ublk_queue *q, unsigned tag, int queued)
+{
+	if (queued < 0)
+		ublk_complete_io(q, tag, queued);
+	else {
+		struct ublk_io *io = ublk_get_io(q, tag);
+
+		q->io_inflight += queued;
+		io->tgt_ios = queued;
+		io->result = 0;
+	}
+}
+
+static inline int ublk_completed_tgt_io(struct ublk_queue *q, unsigned tag)
+{
+	struct ublk_io *io = ublk_get_io(q, tag);
+
+	q->io_inflight--;
+
+	return --io->tgt_ios == 0;
+}
+
 static inline int ublk_queue_use_zc(const struct ublk_queue *q)
 {
 	return q->state & UBLKSRV_ZC;
diff --git a/tools/testing/selftests/ublk/null.c b/tools/testing/selftests/ublk/null.c
index b6ef16a8f514..975a11db22fd 100644
--- a/tools/testing/selftests/ublk/null.c
+++ b/tools/testing/selftests/ublk/null.c
@@ -2,7 +2,7 @@
 
 #include "kublk.h"
 
-static int ublk_null_tgt_init(struct ublk_dev *dev)
+static int ublk_null_tgt_init(const struct dev_ctx *ctx, struct ublk_dev *dev)
 {
 	const struct ublksrv_ctrl_dev_info *info = &dev->dev_info;
 	unsigned long dev_size = 250UL << 30;