perf(transfer): batch io_uring l2 reads

daser/transfer/iouring/layer.py

@@ -19,6 +19,7 @@
 logger = init_logger(__name__)
 
 _DIRECT_IO_ALIGNMENT = 4096
+_IO_URING_ENTRIES = 64
 
 
 class TieredIOUringTransferLayer(TransferLayer):
@@ -65,7 +66,10 @@ def __init__(
 
         self._path = path
         self._fd = os.open(path, os.O_RDWR | os.O_DIRECT)
-        self._urings = [NativeIOUring(entries=64) for _ in range(io_workers)]
+        self._urings = [
+            NativeIOUring(entries=_IO_URING_ENTRIES) for _ in range(io_workers)
+        ]
+        self._l2_read_batch_spans = io_workers * _IO_URING_ENTRIES
         self._uring_lock = threading.Lock()
         self._next_uring_index = 0
         self._io_executor = concurrent.futures.ThreadPoolExecutor(
@@ -405,6 +409,17 @@ def _read_l2_into(
         """Blocking io_uring L2 read into pinned memory."""
         return uring.read_into(self._fd, file_offset, dst.view())
 
+    def _readv_l2_into(
+        self,
+        reads: list[tuple[int, PinnedMemorySlice]],
+        uring: NativeIOUring,
+    ) -> int:
+        """Blocking batched io_uring L2 reads into pinned memory."""
+        return uring.readv_into(
+            self._fd,
+            [(file_offset, pinned.view()) for file_offset, pinned in reads],
+        )
+
     def _write_l2(
         self,
         file_offset: int,
@@ -537,16 +552,20 @@ async def _load_l2_miss_batch(
                 pinned = await self._reserve_l1_buffer(key, nbytes)
                 reads.append((span, pinned))
 
+            read_pairs = [(int(span["file_offset"]), pinned) for span, pinned in reads]
             await asyncio.gather(
                 *(
                     loop.run_in_executor(
                         self._io_executor,
-                        self._read_l2_into,
-                        int(span["file_offset"]),
-                        pinned,
+                        self._readv_l2_into,
+                        read_pairs[start : start + _IO_URING_ENTRIES],
                         self._next_uring(),
                     )
-                    for span, pinned in reads
+                    for start in range(
+                        0,
+                        len(read_pairs),
+                        _IO_URING_ENTRIES,
+                    )
                 )
             )
 
@@ -589,7 +608,8 @@ def _next_l2_miss_batch(
         for span in misses[start:]:
             nbytes = int(span["nbytes"])
             if batch and (
-                batch_bytes + nbytes > self._l1_bytes or len(batch) >= len(self._urings)
+                batch_bytes + nbytes > self._l1_bytes
+                or len(batch) >= self._l2_read_batch_spans
             ):
                 break
             batch.append(span)

daser/transfer/iouring/native.py

@@ -239,6 +239,74 @@ def read_into(self, fd: int, file_offset: int, dst: memoryview) -> int:
             cursor += chunk
         return total
 
+    def readv_into(
+        self,
+        fd: int,
+        reads: list[tuple[int, memoryview]],
+    ) -> int:
+        """Read multiple positioned ranges into writable buffers.
+
+        Args:
+            fd: Open file descriptor.
+            reads: ``(file_offset, dst)`` pairs. Each destination is read fully.
+
+        Returns:
+            Total number of bytes read.
+
+        Thread-safety:
+            Serialized by the wrapper lock. Destination buffers must not be
+            mutated by another thread until this call returns.
+        """
+        total = 0
+        pending: list[tuple[int, memoryview]] = []
+        for file_offset, dst in reads:
+            view = memoryview(dst).cast("B")
+            if view.readonly:
+                raise ValueError("dst must be writable")
+            cursor = 0
+            while cursor < len(view):
+                chunk = min(_MAX_RW_COUNT, len(view) - cursor)
+                pending.append((file_offset + cursor, view[cursor : cursor + chunk]))
+                total += chunk
+                cursor += chunk
+        if not pending:
+            return 0
+
+        results: list[int] = []
+        batch_limit = int(self._params.sq_entries)
+        with self._lock:
+            for batch_start in range(0, len(pending), batch_limit):
+                batch = pending[batch_start : batch_start + batch_limit]
+                start_user_data = self._next_user_data
+                self._next_user_data += len(batch)
+                for idx, (file_offset, view) in enumerate(batch):
+                    self._submit_locked(
+                        _IORING_OP_READ,
+                        fd,
+                        file_offset,
+                        view,
+                        len(view),
+                        start_user_data + idx,
+                    )
+                self._enter(
+                    len(batch),
+                    len(batch),
+                    _IORING_ENTER_GETEVENTS,
+                )
+                results.extend(
+                    self._wait_completions_locked(
+                        start_user_data,
+                        [len(view) for _offset, view in batch],
+                    )
+                )
+        for idx, res in enumerate(results):
+            if res < 0:
+                raise OSError(-res, os.strerror(-res))
+            expected = len(pending[idx][1])
+            if res != expected:
+                raise IOError(f"short io_uring result: {res} != {expected}")
+        return total
+
     def write(self, fd: int, file_offset: int, data: bytes | memoryview) -> int:
         """Write bytes at a file offset through io_uring.
 
@@ -305,15 +373,16 @@ def _submit_and_wait(
         with self._lock:
             user_data = self._next_user_data
             self._next_user_data += 1
-            self._submit(opcode, fd, file_offset, buf, nbytes, user_data)
-            res = self._wait_completion(user_data)
+            self._submit_locked(opcode, fd, file_offset, buf, nbytes, user_data)
+            self._enter(1, 1, _IORING_ENTER_GETEVENTS)
+            res = self._wait_completion_locked(user_data)
         if res < 0:
             raise OSError(-res, os.strerror(-res))
         if res != nbytes:
             raise IOError(f"short io_uring result: {res} != {nbytes}")
         return res
 
-    def _submit(
+    def _submit_locked(
         self,
         opcode: int,
         fd: int,
@@ -350,21 +419,33 @@ def _submit(
         )
         self._write_u32(self._sq_ring, self._params.sq_off.tail, sq_tail + 1)
 
+    def _enter(self, to_submit: int, min_complete: int, flags: int) -> None:
+        """Enter the kernel to submit SQEs and optionally wait for CQEs."""
         ret = self._libc.syscall(
             _SYS_IO_URING_ENTER,
             ctypes.c_int(self._ring_fd),
-            ctypes.c_uint32(1),
-            ctypes.c_uint32(0),
-            ctypes.c_uint32(0),
+            ctypes.c_uint32(to_submit),
+            ctypes.c_uint32(min_complete),
+            ctypes.c_uint32(flags),
             ctypes.c_void_p(0),
             ctypes.c_size_t(0),
         )
         if ret < 0:
             errno = ctypes.get_errno()
             raise OSError(errno, os.strerror(errno))
 
-    def _wait_completion(self, user_data: int) -> int:
+    def _wait_completion_locked(self, user_data: int) -> int:
         """Wait for a CQE matching ``user_data`` and return its result."""
+        return self._wait_completions_locked(user_data, [0])[0]
+
+    def _wait_completions_locked(
+        self,
+        start_user_data: int,
+        expected_sizes: list[int],
+    ) -> list[int]:
+        """Wait for a contiguous user-data range and return ordered results."""
+        remaining = len(expected_sizes)
+        results: list[int | None] = [None] * remaining
         while True:
             cq_head = self._read_u32(self._cq_ring, self._params.cq_off.head)
             cq_tail = self._read_u32(self._cq_ring, self._params.cq_off.tail)
@@ -379,25 +460,20 @@ def _wait_completion(self, user_data: int) -> int:
                 res = int(cqe.res)
                 seen_user_data = int(cqe.user_data)
                 self._write_u32(self._cq_ring, self._params.cq_off.head, cq_head + 1)
-                if seen_user_data != user_data:
+                index = seen_user_data - start_user_data
+                if index < 0 or index >= len(expected_sizes):
                     raise RuntimeError(
                         "unexpected io_uring completion "
-                        f"{seen_user_data}, expected {user_data}"
+                        f"{seen_user_data}, expected range "
+                        f"[{start_user_data}, {start_user_data + len(expected_sizes)})"
                     )
-                return res
-
-            ret = self._libc.syscall(
-                _SYS_IO_URING_ENTER,
-                ctypes.c_int(self._ring_fd),
-                ctypes.c_uint32(0),
-                ctypes.c_uint32(1),
-                ctypes.c_uint32(_IORING_ENTER_GETEVENTS),
-                ctypes.c_void_p(0),
-                ctypes.c_size_t(0),
-            )
-            if ret < 0:
-                errno = ctypes.get_errno()
-                raise OSError(errno, os.strerror(errno))
+                results[index] = res
+                remaining -= 1
+                if remaining == 0:
+                    return [int(result) for result in results]
+                continue
+
+            self._enter(0, 1, _IORING_ENTER_GETEVENTS)
 
     def _read_u32(self, buf: mmap.mmap, offset: int) -> int:
         """Read a uint32 from a ring mapping without exporting the buffer."""

tests/transfer/test_tiered_iouring_transfer.py

@@ -72,6 +72,27 @@ def _copy_grouped_to_dst(
         super()._copy_grouped_to_dst(dst, chunks)
 
 
+class MissBatchProbe(TieredIOUringTransferLayer):
+    """Test transfer layer that records grouped L2 miss batch sizes."""
+
+    def __init__(self, path: str, l1_bytes: int, l2_bytes: int) -> None:
+        super().__init__(
+            path=path,
+            l1_bytes=l1_bytes,
+            l2_bytes=l2_bytes,
+        )
+        self.miss_batches: list[int] = []
+
+    async def _load_l2_miss_batch(
+        self,
+        dst: object,
+        misses: list[dict[str, int]],
+    ) -> None:
+        """Record miss batch size before delegating to production logic."""
+        self.miss_batches.append(len(misses))
+        await super()._load_l2_miss_batch(dst, misses)
+
+
 def _run(coro: object) -> object:
     """Run a coroutine on the current test event loop."""
     return asyncio.get_event_loop().run_until_complete(coro)
@@ -192,6 +213,57 @@ def test_native_iouring_read_into(tmp_path) -> None:
         os.close(fd)
 
 
+def test_native_iouring_readv_into_batches_submit(tmp_path, monkeypatch) -> None:
+    """Native io_uring submits multiple positioned reads in one enter call."""
+    path = tmp_path / "readv.store"
+    path.write_bytes(b"abcdefghijklmnop")
+    fd = os.open(path, os.O_RDWR)
+    calls: list[tuple[int, int, int]] = []
+    original_cdll = native_iouring.ctypes.CDLL
+
+    def as_int(value: object) -> int:
+        """Return an int from a ctypes scalar or plain Python value."""
+        return int(getattr(value, "value", value))
+
+    class TrackingCDLL:
+        """Proxy libc calls while recording io_uring_enter arguments."""
+
+        def __init__(self, *args: object, **kwargs: object) -> None:
+            self._real = original_cdll(*args, **kwargs)
+
+        def syscall(self, number: int, *args: object) -> int:
+            """Record io_uring_enter calls and delegate to libc."""
+            if number == 426:
+                calls.append((as_int(args[1]), as_int(args[2]), as_int(args[3])))
+            return self._real.syscall(number, *args)
+
+    monkeypatch.setattr(native_iouring.ctypes, "CDLL", TrackingCDLL)
+
+    uring = NativeIOUring(entries=8)
+    one = bytearray(4)
+    two = bytearray(4)
+    three = bytearray(4)
+    try:
+        assert (
+            uring.readv_into(
+                fd,
+                [
+                    (0, memoryview(one)),
+                    (4, memoryview(two)),
+                    (8, memoryview(three)),
+                ],
+            )
+            == 12
+        )
+        assert bytes(one) == b"abcd"
+        assert bytes(two) == b"efgh"
+        assert bytes(three) == b"ijkl"
+        assert calls == [(3, 3, 1)]
+    finally:
+        uring.close()
+        os.close(fd)
+
+
 def test_iouring_direct_io_aligned_roundtrip(tmp_path) -> None:
     """Production io_uring mode uses O_DIRECT for aligned L2 ranges."""
     path = str(tmp_path / "direct.store")
@@ -424,6 +496,55 @@ async def scenario() -> None:
     _run(scenario())
 
 
+def test_iouring_grouped_l2_misses_batch_beyond_ring_count(tmp_path) -> None:
+    """Grouped L2 misses are batched by queue capacity, not just ring count."""
+
+    async def scenario() -> None:
+        path = str(tmp_path / "daser.store")
+        block_count = 16
+        layer = TieredIOUringTransferLayer(
+            path=path,
+            l1_bytes=ALIGNMENT * block_count,
+            l2_bytes=ALIGNMENT * (block_count + 1),
+        )
+        try:
+            for idx in range(block_count):
+                await layer.store_bytes(
+                    _block(bytes([idx])),
+                    file_offset=idx * ALIGNMENT,
+                    nbytes=ALIGNMENT,
+                )
+            await layer.drain()
+        finally:
+            layer.close()
+
+        layer = MissBatchProbe(
+            path=path,
+            l1_bytes=ALIGNMENT * block_count,
+            l2_bytes=ALIGNMENT * (block_count + 1),
+        )
+        try:
+            dst = bytearray(ALIGNMENT * block_count)
+            loaded = await layer.load_bytes_grouped(
+                dst,
+                [
+                    {
+                        "target_offset": idx * ALIGNMENT,
+                        "file_offset": idx * ALIGNMENT,
+                        "nbytes": ALIGNMENT,
+                    }
+                    for idx in range(block_count)
+                ],
+            )
+
+            assert loaded == ALIGNMENT * block_count
+            assert layer.miss_batches == [block_count]
+        finally:
+            layer.close()
+
+    _run(scenario())
+
+
 def test_iouring_parallel_l2_loads_use_independent_offsets(tmp_path) -> None:
     """Concurrent L2 loads read their requested byte ranges exactly."""