From 5bc4a832f197f57eba9e3f23f72997575c744a32 Mon Sep 17 00:00:00 2001
From: Abdur-Rahman Mansoor <armansoor@missingno.dev>
Date: Thu, 27 Jun 2024 15:19:32 -0400
Subject: [PATCH] feat: add support for consuming completion queues entries
 from multiple queues

---
 drivers/nvme/nvme.asm | 20 +++++++++++++++-----
 drivers/nvme/nvme.inc |  1 +
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/drivers/nvme/nvme.asm b/drivers/nvme/nvme.asm
index 908aabe..8f57d0e 100644
--- a/drivers/nvme/nvme.asm
+++ b/drivers/nvme/nvme.asm
@@ -1023,6 +1023,8 @@ proc is_queue_full stdcall, tail:word, head:word
 		push 	bx
 		mov 	ax, [tail]
 		mov 	bx, [head]
+		cmp 	ax, bx
+		je 		.not_full
 		test 	bx, bx
 		jnz 	@f
 		cmp 	ax, NVM_ASQS
@@ -1033,15 +1035,15 @@ proc is_queue_full stdcall, tail:word, head:word
 
 @@:
 		cmp 	ax, bx
-		jge 	@f
+		jge 	.not_full
 		sub 	ax, bx
 		cmp 	ax, 1
-		jne 	@f
+		jne 	.not_full
 		xor 	eax, eax
 		inc 	eax
 		ret
 		
-@@:
+.not_full:
 		pop 	bx
 		xor 	eax, eax
 		ret
@@ -1093,10 +1095,18 @@ proc irq_handler
 		mov 	edi, esi
 		mov 	edi, dword [edi + pcidev.io_addr]
 		mov 	dword [edi + NVME_MMIO.INTMS], 0x3
-		stdcall consume_cq_entries, [p_nvme_devices], 0
-		mov 	dword [edi + NVME_MMIO.INTMC], 0x3
+		xor 	ecx, ecx
+
+@@:
+		push 	ecx
+		stdcall consume_cq_entries, [p_nvme_devices], ecx
+		pop 	ecx
+		inc 	ecx
+		cmp 	ecx, LAST_QUEUE_ID
+		jne 	@b
 
 		; Interrupt handled by driver, return 1
+		mov 	dword [edi + NVME_MMIO.INTMC], 0x3
 		pop 	edi esi
 		xor 	eax, eax
 		inc 	eax
diff --git a/drivers/nvme/nvme.inc b/drivers/nvme/nvme.inc
index b37c22a..ac12bf1 100644
--- a/drivers/nvme/nvme.inc
+++ b/drivers/nvme/nvme.inc
@@ -20,6 +20,7 @@ NVM_CMDS = 64       ; Number of Commands
 NVM_MPS = 0         ; Memory Page Size (2 ^ (12 + MPS))
 NVM_ASQS = 63        ; Admin Submission Queue Size
 NVM_ACQS = NVM_ASQS ; Admin Completion Queue Size
+LAST_QUEUE_ID = 1 ; Index of the last queue
 
 ADMIN_QUEUE = 0 ; Admin Queue ID