Example scripts

"""Copy a file with linked requests: C liburing's `examples/link-cp.c`.

Every chunk is a pair of requests in the ring at once: a read into a buffer, and (`IOSQE_IO_LINK`) the write of
that buffer at the same offset, which the kernel starts only after the read is done. Up to `QD` pairs are in
flight, so one `io_uring_submit` hands the kernel many reads and writes; nothing is copied through Python.
The size comes from `io_uring_prep_statx`, the files from `io_uring_prep_open`: every step goes through the ring.

    python3 example/copy_file.py <source> <destination>
"""

import sys
from liburing import *

QD = 64  # pairs in flight at most
BS = 32 * 1024  # bytes per pair

READ, WRITE = 0, 1  # the low bit of `user_data`; the rest is the offset of the pair's chunk


def one(ring, cqe, prep, *args):
    """Submit one request on its own and give its result."""
    sqe = io_uring_get_sqe(ring)
    prep(sqe, *args)
    io_uring_submit(ring)
    io_uring_wait_cqe(ring, cqe)
    result = cqe[0].res  # raises the matching `OSError` if it failed
    io_uring_cqe_seen(ring, cqe[0])
    return result


def queue_pair(ring, src, dst, offset, size):
    """A read linked to a write of the same buffer: the two requests of one chunk."""
    buffer = bytearray(size)  # both requests hold it until their completions
    sqe = io_uring_get_sqe(ring)
    io_uring_prep_read(sqe, src, buffer, offset)
    sqe.flags = IOSQE_IO_LINK  # the next entry starts when this one has finished, and only if it succeeded
    sqe.user_data = offset << 1 | READ
    sqe = io_uring_get_sqe(ring)
    io_uring_prep_write(sqe, dst, buffer, offset)
    sqe.user_data = offset << 1 | WRITE


def copy_file(ring, cqe, src, dst, size):
    offset = inflight = 0
    while offset < size or inflight:
        queued = inflight
        while offset < size and inflight < 2 * QD:  # fill the queue: as many pairs as fit
            chunk = min(BS, size - offset)
            queue_pair(ring, src, dst, offset, chunk)
            offset += chunk
            inflight += 2
        if inflight != queued:
            io_uring_submit(ring)  # one syscall for every pair queued above

        depth = 2 * QD if offset < size else 1  # collect until a pair fits again (at the end: everything)
        while inflight >= depth:
            io_uring_wait_cqe(ring, cqe)
            entry = cqe[0]
            try:
                _ = entry.res  # the read's bytes, or the write's; a failed read also cancels its write (`ECANCELED`)
            finally:
                io_uring_cqe_seen(ring, entry)
            inflight -= 1


def main(source, destination):
    ring = Ring()
    cqe = Cqe()
    stat = Statx()
    io_uring_queue_init(2 * QD, ring)  # a read and a write per pair
    try:
        one(ring, cqe, io_uring_prep_statx, stat, source)
        src = one(ring, cqe, io_uring_prep_open, source, O_RDONLY)
        dst = one(ring, cqe, io_uring_prep_open, destination, O_WRONLY | O_CREAT | O_TRUNC, 0o644)
        copy_file(ring, cqe, src, dst, stat.size)
        one(ring, cqe, io_uring_prep_close, src)
        one(ring, cqe, io_uring_prep_close, dst)
        print(f"copied {stat.size} bytes")
    finally:
        io_uring_queue_exit(ring)


if __name__ == "__main__":
    if len(sys.argv) != 3:
        sys.exit(f"usage: {sys.argv[0]} <source> <destination>")
    main(*sys.argv[1:])

"""TCP echo server: C liburing's `examples/echo-server.c`, the modern io_uring networking pattern.

- Multishot accept: armed once, one completion per new connection.
- Multishot receive: armed once per connection, one completion per chunk of data.
- A provided buffer ring: the kernel picks a buffer from a shared pool for each receive, so no buffer is
  allocated or held per connection; the buffer goes back into the pool when its send has completed.

One thread, one ring, no object per request: `Sqe(ring)` preps straight into the next free entry, and
`Ring(auto_submit=True)` submits by itself when the queue is full. Linux 6.0+.

    python3 example/echo_server.py [port=8000]
    nc localhost 8000
"""

import errno
import os
import socket
import sys
from liburing import *

QD = 64  # submission queue entries
BUF_SIZE = 4096  # bytes per provided buffer
BUFFERS = 256  # buffers in the pool (a power of 2)
BGID = 0  # the pool's buffer group id

# `user_data`: bits 63-56 the event, 31-16 the buffer id (a send: which buffer to recycle), 15-0 the fd
ACCEPT, RECV, SEND = 1, 2, 3


def encode(event, fd, bid=0):
    return event << 56 | bid << 16 | fd


def main(port):
    ring = Ring(auto_submit=True)  # a full submission queue submits itself
    cqe = Cqe()
    listener = socket.socket()
    listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    listener.bind(("127.0.0.1", port))
    listener.listen(QD)
    listen_fd = listener.fileno()
    # `DEFER_TASKRUN`: completions are processed only when this thread waits for them (needs `SINGLE_ISSUER`).
    io_uring_queue_init(QD, ring, IORING_SETUP_SINGLE_ISSUER | IORING_SETUP_DEFER_TASKRUN)
    sqe = Sqe(ring)  # each prep takes the ring's next free entry

    # the buffer pool: `BUFFERS` buffers the kernel picks from when a receive needs one
    bufs = [bytearray(BUF_SIZE) for _ in range(BUFFERS)]
    mask = io_uring_buf_ring_mask(BUFFERS)
    br = io_uring_setup_buf_ring(ring, BUFFERS, BGID, 0)
    for bid, buf in enumerate(bufs):
        io_uring_buf_ring_add(br, buf, bid, mask, bid)
    io_uring_buf_ring_advance(br, BUFFERS)
    conns = set()  # open connections
    starved = set()  # connections whose receive stopped because the pool was empty (`ENOBUFS`)

    def add_accept():
        io_uring_prep_multishot_accept(sqe, listen_fd)
        sqe.user_data = encode(ACCEPT, listen_fd)

    def add_recv(fd):
        io_uring_prep_recv_multishot(sqe, fd)  # no buffer: the kernel picks one from group `BGID`
        sqe.flags = IOSQE_BUFFER_SELECT
        io_uring_sqe_set_buf_group(sqe, BGID)
        sqe.user_data = encode(RECV, fd)

    def add_send(fd, bid, length):
        io_uring_prep_send(sqe, fd, memoryview(bufs[bid])[:length])  # held until the completion
        sqe.user_data = encode(SEND, fd, bid)

    def recycle(bid):
        io_uring_buf_ring_add(br, bufs[bid], bid, mask, 0)
        io_uring_buf_ring_advance(br, 1)

    def close(fd):
        if fd in conns:  # once: a send may fail after its connection's receive closed it
            conns.discard(fd)
            starved.discard(fd)
            os.close(fd)

    print(f"echo server on 127.0.0.1:{port}", flush=True)
    add_accept()
    try:
        while True:
            io_uring_submit_and_wait(ring, 1)
            io_uring_peek_cqe(ring, cqe)  # `cqe[0]` is the oldest completion
            ready = io_uring_cq_ready(ring)
            for i in range(ready):
                entry = cqe[i]
                flags, data = entry.flags, entry.user_data
                event, fd = data >> 56, data & 0xFFFF
                try:
                    res = entry.res  # raises the matching `OSError` for a failed request
                except OSError as e:
                    res = -e.errno
                more = flags & IORING_CQE_F_MORE  # the multishot request is still armed

                if event == ACCEPT:
                    if res >= 0:
                        conns.add(res)  # the new connection's fd
                        add_recv(res)
                    if not more:
                        add_accept()
                elif event == RECV:
                    if res > 0:
                        add_send(fd, flags >> IORING_CQE_BUFFER_SHIFT, res)  # echo from the buffer the kernel picked
                        if not more:
                            starved.add(fd)
                    elif res == -errno.ENOBUFS:
                        starved.add(fd)  # the pool is empty: arm again when a send gives a buffer back
                    else:
                        close(fd)  # `0`: the peer closed; else an error
                else:  # SEND
                    recycle(data >> 16 & 0xFFFF)  # back into the pool, sent or not
                    if res < 0:
                        close(fd)
                    while starved:  # a buffer is free again
                        add_recv(starved.pop())
            io_uring_cq_advance(ring, ready)  # done with all of them
    finally:
        io_uring_queue_exit(ring)  # frees the buffer ring too
        listener.close()


if __name__ == "__main__":
    main(int(sys.argv[1]) if len(sys.argv) > 1 else 8000)

'''Messages between threads, each with its own ring, through a futex (Linux 6.7+).

The receiver arms an `io_uring_prep_futex_wait` on its own ring: an ordinary request, with its own
`user_data`. A sender puts the message in a shared `deque` (any Python object, not only 64 bits), bumps the
futex and wakes it from its own ring. The receiver's wait completes, it takes every message queued (several
sends may give one wake-up) and arms the wait again.

The receiver reads the futex before it looks in the `deque`: a message sent after that read changes the
futex, so the wait armed on the old value returns at once. Read it after, and a message sent between the
empty `deque` and the read is never seen (the wait sleeps on the new value).

`io_uring_prep_msg_ring` also posts on another ring, but carries only `data` (below 2**63) and `len`: use it
for a small value or to pass a registered file (`io_uring_prep_msg_ring_fd`).

    python3.15t example/futex_messages.py
'''

import collections
import threading
from liburing import *

FLAGS = FUTEX2_SIZE_U32 | FUTEX2_PRIVATE  # one process: a private futex is cheaper
WAIT = 1  # the receiver's `user_data` for its futex wait

futex = Futex(1, FLAGS)
messages = collections.deque()  # thread-safe appends / pops


def receiver(got):
    ring, cqe = Ring(), Cqe()
    io_uring_queue_init(8, ring)
    try:
        while True:
            seen = futex.state  # before the `deque`: waits only while the futex still holds this value
            while messages:  # one wake-up may stand for several messages
                message = messages.popleft()
                if message is None:
                    return
                got.append(message)
            sqe = io_uring_get_sqe(ring)
            io_uring_prep_futex_wait(sqe, futex, seen, FUTEX_BITSET_MATCH_ANY, FLAGS)
            sqe.user_data = WAIT
            io_uring_submit(ring)
            io_uring_wait_cqe(ring, cqe)
            entry = cqe[0]
            assert entry.user_data == WAIT  # our own request: an entry of this ring's table
            try:
                _ = entry.res  # 0: woken
            except BlockingIOError:
                pass  # `EAGAIN`: a message came between reading `seen` and arming the wait
            io_uring_cqe_seen(ring, entry)
    finally:
        io_uring_queue_exit(ring)


def sender(items):
    ring, cqe = Ring(), Cqe()
    io_uring_queue_init(8, ring)
    try:
        for item in items:
            messages.append(item)
            futex.state = (futex.state + 1) & 0xFFFF_FFFF  # changed: a wait armed on the old value returns
            sqe = io_uring_get_sqe(ring)
            io_uring_prep_futex_wake(sqe, futex, 1, FUTEX_BITSET_MATCH_ANY, FLAGS)
            io_uring_submit_and_wait(ring, 1)
            io_uring_wait_cqe(ring, cqe)
            io_uring_cqe_seen(ring, cqe[0])  # `res`: how many waiters it woke (0 if none was armed yet)
    finally:
        io_uring_queue_exit(ring)


def main():
    got = []
    r = threading.Thread(target=receiver, args=(got,))
    r.start()
    s = threading.Thread(target=sender, args=([{'id': i, 'text': f'message {i}'} for i in range(5)] + [None],))
    s.start()
    s.join()
    r.join()
    for message in got:
        print(message)


if __name__ == '__main__':
    main()

"""Open, write, read and close a file: every step a request on the ring, one at a time."""

import os
import tempfile
from liburing import *


def submit_and_wait(ring, cqe):
    io_uring_submit(ring)  # hand the queued request to the kernel
    io_uring_wait_cqe(ring, cqe)  # wait for its completion
    entry = cqe[0]
    result = entry.res  # raises the matching `OSError` if the operation failed
    io_uring_cqe_seen(ring, entry)  # done with this entry
    return result


def main():
    path = os.path.join(tempfile.mkdtemp(), "example.txt")
    ring = Ring()
    cqe = Cqe()  # a completion queue entry
    io_uring_queue_init(8, ring)  # a ring with room for 8 requests
    try:
        sqe = io_uring_get_sqe(ring)  # a submission queue entry
        io_uring_prep_open(sqe, path, O_CREAT | O_RDWR)
        sqe.user_data = 1  # comes back in `cqe.user_data`, to tell completions apart
        fd = submit_and_wait(ring, cqe)
        print("fd:", fd)

        data = b"Hello, io_uring!"
        sqe = io_uring_get_sqe(ring)
        io_uring_prep_write(sqe, fd, data)
        sqe.user_data = 2
        print("wrote:", submit_and_wait(ring, cqe))  # bytes written

        buffer = bytearray(len(data))  # held by the request until the read completes
        sqe = io_uring_get_sqe(ring)
        io_uring_prep_read(sqe, fd, buffer, 0)  # from the start of the file
        sqe.user_data = 3
        print("read:", submit_and_wait(ring, cqe), buffer)  # bytes read

        sqe = io_uring_get_sqe(ring)
        io_uring_prep_close(sqe, fd)
        sqe.user_data = 4
        submit_and_wait(ring, cqe)  # no error means closed
        print("closed.")
    finally:
        io_uring_queue_exit(ring)
        os.unlink(path)
        os.rmdir(os.path.dirname(path))


if __name__ == "__main__":
    main()

from liburing import *

for k, v in probe().items():
    print(k, v)

import os
import tempfile
from liburing import *


def run(ring, cqe, prep, *args):
    sqe = io_uring_get_sqe(ring)
    prep(sqe, *args)
    io_uring_submit(ring)
    io_uring_wait_cqe(ring, cqe)
    entry = cqe[0]
    result = entry.res  # raises if the request failed
    io_uring_cqe_seen(ring, entry)
    return result


def main():
    path = os.path.join(tempfile.mkdtemp(), "vector.txt")
    ring = Ring()
    cqe = Cqe()
    io_uring_queue_init(32, ring)
    try:
        fd = run(ring, cqe, io_uring_prep_open, path, O_CREAT | O_RDWR)

        # a list of buffers: the request holds them (locked) until its completion
        print("written:", run(ring, cqe, io_uring_prep_writev, fd, [b"hello ", b"vector ", b"world"], 0))

        bufs = [bytearray(6), bytearray(7), bytearray(5)]
        print("read:", run(ring, cqe, io_uring_prep_readv, fd, bufs, 0))
        print(bufs)  # [bytearray(b'hello '), bytearray(b'vector '), bytearray(b'world')]
        print(bytes(bufs[1]))  # b'vector '

        run(ring, cqe, io_uring_prep_close, fd)
    finally:
        io_uring_queue_exit(ring)
        os.remove(path)


if __name__ == "__main__":
    main()

"""Four ways to fill the submission queue: `io_uring_get_sqe(ring)`, `Sqe()`, `Sqe(ring)` and `Sqe(123)`."""

from liburing import *


def run(ring, cqe, count):
    """Submit, wait for `count` completions and return their `user_data`."""
    io_uring_submit_and_wait(ring, count)
    io_uring_peek_cqe(ring, cqe)  # `cqe[0]` is the oldest completion
    ready = io_uring_cq_ready(ring)
    data = [cqe.user_data_at(i) for i in range(ready)]
    io_uring_cq_advance(ring, ready)  # done with all of them
    return data


def main():
    ring = Ring()
    cqe = Cqe()
    try:
        io_uring_queue_init(128, ring)

        # `io_uring_get_sqe(ring)`: the C way, a new `SQE` for every request; `None` when the queue is full.
        for i in range(3):
            sqe = io_uring_get_sqe(ring)
            io_uring_prep_nop(sqe)
            sqe.user_data = i
        print("get_sqe():", run(ring, cqe, 3))

        # `Sqe()`: one entry you prepare once, then `put_sqe` copies it into the ring as often as you like.
        sqe = Sqe()
        io_uring_prep_nop(sqe)
        sqe.user_data = 7
        for _ in range(3):
            put_sqe(ring, sqe)  # `True`; `False` when the queue is full: submit, then put again
        print("Sqe():    ", run(ring, cqe, 3))

        # `Sqe(ring)`: one object for the ring's lifetime, each prep writes straight into the ring's next free entry:
        # no `put_sqe`, no `io_uring_get_sqe`, no object per request.
        sqe = Sqe(ring)
        for i in range(3):
            io_uring_prep_nop(sqe)
            sqe.user_data = i  # acts on the entry of the last prep
        print("Sqe(ring):", run(ring, cqe, 3))

        # `Sqe(123)`: 123 entries prepared as a block (`sqe[i]`), then one `put_sqe` puts all of them in the ring.
        sqe = Sqe(123)
        for i in range(len(sqe)):
            io_uring_prep_nop(sqe[i])
            sqe[i].user_data = i
        put_sqe(ring, sqe)
        done = run(ring, cqe, 123)
        print("Sqe(123):", len(done), "completions, first", done[:3], "last", done[-1])
    finally:
        io_uring_queue_exit(ring)


if __name__ == "__main__":
    main()