rtio: Add an sqe_acquire_n helper

Makes it safer/easier to obtain n many submissions for use in one go.

Signed-off-by: Tom Burdick <thomas.burdick@intel.com>

Author: Tom Burdick

include/zephyr/rtio/rtio.h

@@ -1121,6 +1121,51 @@ static inline struct rtio_sqe *rtio_sqe_acquire(struct rtio *r)
 	return &iodev_sqe->sqe;
 }
 
+/**
+ * @brief Acquire a number of submission queue events if available
+ *
+ * @warning The sqe array is in an undefined state if the return value is -ENOMEM
+ *
+ * @param r RTIO context
+ * @param n Number of submission queue events to acquire
+ * @param sqes A pointer to an array of rtio_sqe struct pointers of size n
+ *
+ * @retval 0 success
+ * @retval -ENOMEM out of memory
+ */
+static inline int rtio_sqe_acquire_array(struct rtio *r, size_t n, struct rtio_sqe **sqes)
+{
+	struct rtio_iodev_sqe *iodev_sqe;
+	size_t i;
+
+	for (i = 0; i < n; i++) {
+		iodev_sqe = rtio_sqe_pool_alloc(r->sqe_pool);
+		if (iodev_sqe == NULL) {
+			break;
+		}
+		sqes[i] = &iodev_sqe->sqe;
+	}
+
+	/* Not enough SQEs in the pool */
+	if (i < n) {
+		while (i > 0) {
+			i--;
+			iodev_sqe = CONTAINER_OF(sqes[i], struct rtio_iodev_sqe, sqe);
+			rtio_sqe_pool_free(r->sqe_pool, iodev_sqe);
+			sqes[i] = NULL;
+		}
+
+		return -ENOMEM;
+	}
+
+	for (i = 0; i < n; i++) {
+		iodev_sqe = CONTAINER_OF(sqes[i], struct rtio_iodev_sqe, sqe);
+		mpsc_push(&r->sq, &iodev_sqe->q);
+	}
+
+	return 0;
+}
+
 /**
  * @brief Drop all previously acquired sqe
  *

tests/subsys/rtio/rtio_api/src/test_rtio_api.c

@@ -1079,7 +1079,6 @@ ZTEST(rtio_api, test_rtio_await)
 }
 
 
-
 RTIO_DEFINE(r_callback_result, SQE_POOL_SIZE, CQE_POOL_SIZE);
 RTIO_IODEV_TEST_DEFINE(iodev_test_callback_result);
 static int callback_count;
@@ -1140,6 +1139,37 @@ ZTEST(rtio_api, test_rtio_callbacks)
 }
 
 
+RTIO_DEFINE(r_acquire_array, SQE_POOL_SIZE, CQE_POOL_SIZE);
+
+ZTEST(rtio_api, test_rtio_acquire_array)
+{
+	TC_PRINT("rtio acquire array\n");
+
+	struct rtio_sqe *sqes[SQE_POOL_SIZE];
+
+	int res = rtio_sqe_acquire_array(&r_acquire_array, SQE_POOL_SIZE, sqes);
+
+	zassert_ok(res, "Expected to acquire sqes");
+
+	struct rtio_sqe *last_sqe;
+
+	res = rtio_sqe_acquire_array(&r_acquire_array, 1, &last_sqe);
+	zassert_equal(res, -ENOMEM, "Expected to have no more sqes available");
+
+	rtio_sqe_drop_all(&r_acquire_array);
+
+	res = rtio_sqe_acquire_array(&r_acquire_array, SQE_POOL_SIZE - 1, sqes);
+	zassert_ok(res, "Expected to acquire sqes");
+	res = rtio_sqe_acquire_array(&r_acquire_array, 2, &last_sqe);
+	zassert_equal(res, -ENOMEM, "Expected to have only have a single sqe available");
+	res = rtio_sqe_acquire_array(&r_acquire_array, 1, &last_sqe);
+	zassert_equal(res, 0, "Expected a single sqe available");
+	res = rtio_sqe_acquire_array(&r_acquire_array, 1, &last_sqe);
+	zassert_equal(res, -ENOMEM, "Expected to have no more sqes available");
+
+	rtio_sqe_drop_all(&r_acquire_array);
+}
+
 static void *rtio_api_setup(void)
 {
 #ifdef CONFIG_USERSPACE