libkernel: implement sceKernelAioWaitRequests (#1109)

kernel: implement batch AIO completion waits

Co-authored-by: nmzik <Nmzik@mail.ru>

src/libs/libKernel.cpp

@@ -2831,6 +2831,8 @@ constexpr int32_t KERNEL_AIO_STATE_ABORTED    = 4;
 constexpr int32_t KERNEL_AIO_STATE_NOTIFIED   = 0x10000;
 constexpr int32_t KERNEL_AIO_MAX_QUEUE        = 512;
 constexpr int32_t KERNEL_AIO_MAX_REQUESTS     = 128;
+constexpr uint32_t KERNEL_AIO_WAIT_AND        = 1;
+constexpr uint32_t KERNEL_AIO_WAIT_OR         = 2;
 
 struct KernelAioResult {
 	int64_t  return_value;
@@ -2933,26 +2935,28 @@ int KYTY_SYSV_ABI KernelAioSubmitWriteCommands(KernelAioRwRequest* req, int32_t
 	return OK;
 }
 
+static int32_t kernel_aio_get_state(int32_t id, bool notify) {
+	if (!kernel_aio_is_valid_id(id)) {
+		return LibKernel::KERNEL_ERROR_ESRCH;
+	}
+
+	auto current = g_kernel_aio_state[id].load(std::memory_order_acquire);
+	while (notify && (current == KERNEL_AIO_STATE_COMPLETED || current == KERNEL_AIO_STATE_ABORTED)) {
+		if (g_kernel_aio_state[id].compare_exchange_weak(current, current | KERNEL_AIO_STATE_NOTIFIED,
+		                                                std::memory_order_acq_rel)) {
+			break;
+		}
+	}
+	return current == 0 ? LibKernel::KERNEL_ERROR_ESRCH : current;
+}
+
 int KYTY_SYSV_ABI KernelAioPollRequest(int32_t id, int32_t* state) {
 	PRINT_NAME();
 
 	if (state == nullptr) {
 		return LibKernel::KERNEL_ERROR_EFAULT;
 	}
-
-	if (!kernel_aio_is_valid_id(id)) {
-		*state = LibKernel::KERNEL_ERROR_ESRCH;
-		return OK;
-	}
-
-	auto current = g_kernel_aio_state[id].load(std::memory_order_acquire);
-	while (current == KERNEL_AIO_STATE_COMPLETED || current == KERNEL_AIO_STATE_ABORTED) {
-		if (g_kernel_aio_state[id].compare_exchange_weak(current, current | KERNEL_AIO_STATE_NOTIFIED,
-		                                                std::memory_order_acq_rel)) {
-			break;
-		}
-	}
-	*state = current == 0 ? LibKernel::KERNEL_ERROR_ESRCH : current;
+	*state = kernel_aio_get_state(id, true);
 	return OK;
 }
 
@@ -2969,31 +2973,53 @@ int KYTY_SYSV_ABI KernelAioPollRequests(int32_t* ids, int32_t num, int32_t* stat
 	return OK;
 }
 
-int KYTY_SYSV_ABI KernelAioWaitRequest(int32_t id, int32_t* state, uint32_t* usec) {
+int KYTY_SYSV_ABI KernelAioWaitRequests(int32_t* ids, int32_t num, int32_t* states, uint32_t mode,
+                                        uint32_t* usec) {
 	PRINT_NAME();
 
-	if (state == nullptr) {
+	if (ids == nullptr || states == nullptr) {
 		return LibKernel::KERNEL_ERROR_EFAULT;
 	}
-
-	if (!kernel_aio_is_valid_id(id)) {
-		return KernelAioPollRequest(id, state);
+	if (num <= 0 || num > KERNEL_AIO_MAX_REQUESTS ||
+	    (num != 1 && mode != KERNEL_AIO_WAIT_AND && mode != KERNEL_AIO_WAIT_OR)) {
+		return LibKernel::KERNEL_ERROR_EINVAL;
 	}
 
-	uint32_t waited  = 0;
-	auto     current = g_kernel_aio_state[id].load(std::memory_order_acquire);
-	while (current == KERNEL_AIO_STATE_PROCESSING) {
-		if (usec != nullptr && *usec != 0 && waited >= *usec) {
-			*state = current;
-			return LibKernel::KERNEL_ERROR_ETIMEDOUT;
+	const auto collect_states = [&](bool notify) {
+		int32_t finished = 0;
+		bool newly_complete = false;
+		for (int32_t i = 0; i < num; i++) {
+			const auto current = kernel_aio_get_state(ids[i], notify);
+			states[i] = current;
+			finished += current != KERNEL_AIO_STATE_PROCESSING;
+			newly_complete |= current == KERNEL_AIO_STATE_COMPLETED || current == KERNEL_AIO_STATE_ABORTED;
+		}
+		return finished == num || (mode == KERNEL_AIO_WAIT_OR && newly_complete);
+	};
+	const auto start = std::chrono::steady_clock::now();
+	for (;;) {
+		const bool complete = collect_states(false);
+		const auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(
+		                         std::chrono::steady_clock::now() - start).count();
+		const bool timed_out = usec != nullptr && elapsed >= *usec;
+		if (complete || timed_out) {
+			// A competing waiter can claim the newly completed request after the first scan.
+			if (collect_states(true)) {
+				if (usec != nullptr) {
+					*usec = elapsed < *usec ? *usec - static_cast<uint32_t>(elapsed) : 0;
+				}
+				return OK;
+			}
+			if (timed_out) {
+				return LibKernel::KERNEL_ERROR_ETIMEDOUT;
+			}
 		}
-
 		Common::Thread::SleepMicro(10);
-		waited += 10;
-		current = g_kernel_aio_state[id].load(std::memory_order_acquire);
 	}
+}
 
-	return KernelAioPollRequest(id, state);
+int KYTY_SYSV_ABI KernelAioWaitRequest(int32_t id, int32_t* state, uint32_t* usec) {
+	return KernelAioWaitRequests(&id, 1, state, KERNEL_AIO_WAIT_AND, usec);
 }
 
 int KYTY_SYSV_ABI KernelAioDeleteRequest(int32_t id, int32_t* ret) {
@@ -3376,6 +3402,7 @@ LIB_DEFINE(InitLibKernel_1) {
 	LIB_FUNC("2pOuoWoCxdk", KernelAioPollRequest);
 	LIB_FUNC("o7O4z3jwKzo", KernelAioPollRequests);
 	LIB_FUNC("KOF-oJbQVvc", KernelAioWaitRequest);
+	LIB_FUNC("lgK+oIWkJyA", KernelAioWaitRequests);
 	LIB_FUNC("XQ8C8y+de+E", KernelAioSubmitWriteCommands);
 	LIB_FUNC("nu4a0-arQis", KernelAioInitializeParam);
 	LIB_FUNC("il03nluKfMk", LibKernel::KernelRaiseException);

tests/KernelFileSystemTests.cpp

@@ -169,12 +169,14 @@ void TestAioBatches() {
   using Batch = int (KYTY_SYSV_ABI *)(int32_t *, int32_t, int32_t *);
   using Single = int (KYTY_SYSV_ABI *)(int32_t, int32_t *);
   using Wait = int (KYTY_SYSV_ABI *)(int32_t, int32_t *, uint32_t *);
+  using WaitBatch = int (KYTY_SYSV_ABI *)(int32_t *, int32_t, int32_t *, uint32_t, uint32_t *);
   const auto submit = reinterpret_cast<Submit>(find("HgX7+AORI58"));
   const auto poll = reinterpret_cast<Batch>(find("o7O4z3jwKzo"));
   const auto erase = reinterpret_cast<Batch>(find("Ft3EtsZzAoY"));
   const auto poll_one = reinterpret_cast<Single>(find("2pOuoWoCxdk"));
   const auto erase_one = reinterpret_cast<Single>(find("5TgME6AYty4"));
   const auto wait = reinterpret_cast<Wait>(find("KOF-oJbQVvc"));
+  const auto wait_batch = reinterpret_cast<WaitBatch>(find("lgK+oIWkJyA"));
   constexpr char Payload[] = "AIO payload";
   const int fd = FileSystem::KernelOpen("/savedata0/aio.dat", 0x602, 0777);
   Check(fd >= 3 && FileSystem::KernelWrite(fd, Payload, sizeof(Payload)) == sizeof(Payload),
@@ -191,14 +193,22 @@ void TestAioBatches() {
   }
   ids[2] = -1;
   std::array<int32_t, 3> states {-1, -1, -1};
-  Check(poll(ids.data(), ids.size(), states.data()) == OK &&
+  uint32_t timeout = 0;
+  Check(wait_batch(ids.data(), ids.size(), states.data(), 2, &timeout) == OK &&
             states[0] == 3 && states[1] == 3 && states[2] == Kernel::KERNEL_ERROR_ESRCH,
-        "batch poll writes every state and per-request invalid-ID error");
+        "OR wait observes completed reads and invalid-ID errors even with a zero timeout");
   Check(poll_one(ids[0], &states[0]) == OK && states[0] == (3 | 0x10000) &&
             poll(ids.data(), 2, states.data()) == OK &&
             states[0] == (3 | 0x10000) && states[1] == (3 | 0x10000) &&
             wait(ids[0], &states[0], nullptr) == OK && states[0] == (3 | 0x10000),
         "single and batch polls and waits share completion notification state");
+  timeout = 1000000;
+  Check(wait_batch(ids.data(), ids.size(), states.data(), 1, &timeout) == OK &&
+            states[0] == (3 | 0x10000) && states[1] == (3 | 0x10000) &&
+            states[2] == Kernel::KERNEL_ERROR_ESRCH && timeout <= 1000000 &&
+            wait_batch(ids.data(), ids.size(), states.data(), 2, nullptr) == OK &&
+            wait_batch(ids.data(), 1, states.data(), 0, nullptr) == OK,
+        "AND and OR waits return when all IDs are terminal and one-ID waits ignore mode");
   Check(erase(ids.data(), ids.size(), states.data()) == OK &&
             states[0] == OK && states[1] == OK && states[2] == Kernel::KERNEL_ERROR_ESRCH,
         "batch deletion writes per-request results");
@@ -221,6 +231,22 @@ void TestAioBatches() {
               }),
           "maximum-sized batch returns all invalid-ID errors");
   }
+  states.fill(42);
+  Check(wait_batch(nullptr, 1, states.data(), 1, nullptr) == Kernel::KERNEL_ERROR_EFAULT &&
+            wait_batch(ids.data(), 1, nullptr, 1, nullptr) == Kernel::KERNEL_ERROR_EFAULT &&
+            wait_batch(ids.data(), 0, states.data(), 1, nullptr) == Kernel::KERNEL_ERROR_EINVAL &&
+            wait_batch(ids.data(), 129, states.data(), 1, nullptr) == Kernel::KERNEL_ERROR_EINVAL &&
+            wait_batch(ids.data(), 2, states.data(), 0, nullptr) == Kernel::KERNEL_ERROR_EINVAL &&
+            wait_batch(ids.data(), 2, states.data(), 3, nullptr) == Kernel::KERNEL_ERROR_EINVAL &&
+            states == std::array<int32_t, 3> {42, 42, 42},
+        "invalid wait arguments do not modify outputs");
+  std::array<int32_t, 128> invalid_ids {};
+  std::array<int32_t, 128> errors {};
+  Check(wait_batch(invalid_ids.data(), invalid_ids.size(), errors.data(), 1, nullptr) == OK &&
+            std::all_of(errors.begin(), errors.end(), [](int32_t error) {
+              return error == Kernel::KERNEL_ERROR_ESRCH;
+            }),
+        "maximum-sized wait returns every invalid-ID error");
   Check(FileSystem::KernelClose(fd) == OK, "close AIO read fixture");
 }