aio: add LinuxQueue

This is used in cl/769787293.

PiperOrigin-RevId: 787248248

Author: nixprime

nogo.yaml

@@ -188,6 +188,7 @@ analyzers:
     internal:
       exclude:
         - ".*_test.go"                                             # Exclude tests.
+        - pkg/aio/aio_linux_unsafe.go                              # Special case.
         - pkg/eventfd/eventfd_unsafe.go                            # Special case.
         - "pkg/flipcall/.*_unsafe.go"                              # Special case.
         - pkg/gohacks/noescape_unsafe.go                           # Special case.

pkg/abi/linux/aio.go

@@ -16,9 +16,26 @@ package linux
 
 import "encoding/binary"
 
+// AIORing is struct aio_ring, from fs/aio.c, without the trailing
+// variable-length array.
+type AIORing struct {
+	ID               uint32
+	Nr               uint32
+	Head             uint32
+	Tail             uint32
+	Magic            uint32
+	CompatFeatures   uint32
+	IncompatFeatures uint32
+	HeaderLength     uint32
+}
+
 // AIORingSize is sizeof(struct aio_ring).
 const AIORingSize = 32
 
+// AIO_RING_MAGIC is fs/aio.c:AIO_RING_MAGIC, the expected value of
+// AIORing.Magic.
+const AIO_RING_MAGIC = 0xa10a10a1
+
 // I/O commands.
 const (
 	IOCB_CMD_PREAD  = 0

pkg/aio/BUILD

@@ -9,10 +9,12 @@ go_library(
     name = "aio",
     srcs = [
         "aio.go",
+        "aio_linux_unsafe.go",
         "aio_unsafe.go",
     ],
     visibility = ["//pkg/sentry:internal"],
     deps = [
+        "//pkg/abi/linux",
         "//pkg/gomaxprocs",
         "//pkg/sync",
         "@org_golang_x_sys//unix:go_default_library",
@@ -22,7 +24,10 @@ go_library(
 go_test(
     name = "aio_test",
     size = "small",
-    srcs = ["aio_test.go"],
+    srcs = [
+        "aio_linux_test.go",
+        "aio_test.go",
+    ],
     library = ":aio",
     deps = [
         "//pkg/bitmap",

pkg/aio/aio.go

@@ -13,6 +13,8 @@
 // limitations under the License.
 
 // Package aio provides asynchronous I/O on host file descriptors.
+//
+// +checkalignedignore
 package aio
 
 import (

pkg/aio/aio_linux_test.go

@@ -0,0 +1,29 @@
+// Copyright 2024 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build linux
+// +build linux
+
+package aio
+
+import (
+	"testing"
+)
+
+func TestLinuxQueue(t *testing.T) {
+	testQueue(t, func(cap int) (Queue, error) {
+		q, err := NewLinuxQueue(cap)
+		return q, err
+	})
+}

pkg/aio/aio_linux_unsafe.go

@@ -0,0 +1,170 @@
+// Copyright 2025 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build linux
+// +build linux
+
+package aio
+
+import (
+	"fmt"
+	"sync/atomic"
+	"unsafe"
+
+	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/abi/linux"
+)
+
+// LinuxQueue implements Queue using Linux native AIO.
+//
+// In LinuxQueue, I/O requests (struct iocb) are submitted when Wait() =>
+// io_submit(2) is called. In Linux, io_submit(2) sequentially invokes the
+// usual *synchronous* entry point into file reads/writes
+// (file_operations::read/write_iter), but provides a completion callback in
+// the kernel I/O control block (struct kiocb); file read/write implementations
+// may recognize this as indicating support for asynchronous I/O
+// (!is_sync_kiocb()), start the request, and return -EIOCBQUEUED to indicate
+// asynchronous completion. This means that submission of a read/write
+// operation will block (without submitting later requests) until the operation
+// completes, unless the implementation of that operation specifically supports
+// asynchronous. In practice, this seems to mean:
+//
+// - O_DIRECT reads, subject to the constraints described by the man page notes
+// for open(2), will probably work.
+//
+// - O_DIRECT writes, subject to the same constraints, will probably work if
+// disk blocks are already allocated.
+type LinuxQueue struct {
+	ctxID    uintptr
+	iocbs    []linux.IOCallback
+	iocbPtrs []*linux.IOCallback
+	ioevs    []linux.IOEvent
+}
+
+// NewLinuxQueue returns a new LinuxQueue with the given capacity.
+func NewLinuxQueue(cap int) (*LinuxQueue, error) {
+	var ctxID uintptr
+	if _, _, e := unix.Syscall(unix.SYS_IO_SETUP, uintptr(cap), uintptr(unsafe.Pointer(&ctxID)), 0 /* unused */); e != 0 {
+		return nil, e
+	}
+	q := &LinuxQueue{
+		ctxID:    ctxID,
+		iocbs:    make([]linux.IOCallback, cap),
+		iocbPtrs: make([]*linux.IOCallback, cap),
+		ioevs:    make([]linux.IOEvent, cap),
+	}
+	for i := range q.iocbPtrs {
+		q.iocbPtrs[i] = &q.iocbs[i]
+	}
+	q.iocbs = q.iocbs[:0]
+	return q, nil
+}
+
+// Destroy implements Queue.Destroy.
+func (q *LinuxQueue) Destroy() {
+	unix.Syscall(unix.SYS_IO_DESTROY, q.ctxID, 0 /* unused */, 0 /* unused */)
+}
+
+// Cap implements Queue.Cap.
+func (q *LinuxQueue) Cap() int {
+	return len(q.ioevs)
+}
+
+// Add implements Queue.Add.
+func (q *LinuxQueue) Add(r Request) {
+	iocb := linux.IOCallback{
+		Data:   r.ID,
+		FD:     r.FD,
+		Buf:    uint64(uintptr(r.Buf)),
+		Bytes:  uint64(r.Len),
+		Offset: r.Off,
+	}
+	switch r.Op {
+	case OpRead:
+		iocb.OpCode = linux.IOCB_CMD_PREAD
+	case OpWrite:
+		iocb.OpCode = linux.IOCB_CMD_PWRITE
+	case OpReadv:
+		iocb.OpCode = linux.IOCB_CMD_PREADV
+	case OpWritev:
+		iocb.OpCode = linux.IOCB_CMD_PWRITEV
+	default:
+		panic(fmt.Sprintf("unknown op %v", r.Op))
+	}
+	q.iocbs = append(q.iocbs, iocb)
+}
+
+// Wait implements Queue.Wait.
+func (q *LinuxQueue) Wait(cs []Completion, minCompletions int) ([]Completion, error) {
+	if len(q.iocbs) != 0 {
+		_, _, e := unix.Syscall(unix.SYS_IO_SUBMIT, q.ctxID, uintptr(len(q.iocbs)), uintptr(unsafe.Pointer(unsafe.SliceData(q.iocbPtrs))))
+		q.iocbs = q.iocbs[:0]
+		if e != 0 {
+			return cs, e
+		}
+	}
+	for {
+		if ring := q.ring(); ring.Magic == linux.AIO_RING_MAGIC {
+			nr := ring.Nr
+			head := ring.Head
+			if head >= nr {
+				panic(fmt.Sprintf("aio_ring::head (%d) >= aio_ring::nr (%d)", head, nr))
+			}
+			origHead := head
+			tail := atomic.LoadUint32(&ring.Tail)
+			for head != tail {
+				ioev := q.ringAt(head)
+				cs = append(cs, Completion{
+					ID:     ioev.Data,
+					Result: ioev.Result,
+				})
+				head++
+				if head >= nr {
+					head = 0
+				}
+				minCompletions--
+			}
+			if head != origHead {
+				atomic.StoreUint32(&ring.Head, head)
+				if minCompletions <= 0 {
+					return cs, nil
+				}
+			}
+		}
+		n, _, e := unix.Syscall6(unix.SYS_IO_GETEVENTS, q.ctxID, uintptr(minCompletions), uintptr(len(q.ioevs)), uintptr(unsafe.Pointer(unsafe.SliceData(q.ioevs))), 0 /* timeout */, 0 /* unused */)
+		if e != 0 {
+			if e == unix.EINTR {
+				continue
+			}
+			return cs, e
+		}
+		for i := range q.ioevs[:n] {
+			ioev := &q.ioevs[i]
+			cs = append(cs, Completion{
+				ID:     ioev.Data,
+				Result: ioev.Result,
+			})
+		}
+		return cs, nil
+	}
+}
+
+func (q *LinuxQueue) ring() *linux.AIORing {
+	return (*linux.AIORing)(unsafe.Pointer(q.ctxID))
+}
+
+// Preconditions: idx < q.ring().Nr.
+func (q *LinuxQueue) ringAt(idx uint32) *linux.IOEvent {
+	return (*linux.IOEvent)(unsafe.Pointer(q.ctxID + uintptr(q.ring().HeaderLength) + uintptr(idx)*unsafe.Sizeof(linux.IOEvent{})))
+}

pkg/aio/aio_test.go

@@ -25,7 +25,7 @@ import (
 	"gvisor.dev/gvisor/pkg/rand"
 )
 
-func TestRead(t *testing.T) {
+func testRead(t *testing.T, newQueue func(cap int) (Queue, error)) {
 	// Create a temp file.
 	testFile, err := os.CreateTemp(t.TempDir(), "aio_test_read")
 	if err != nil {
@@ -46,7 +46,10 @@ func TestRead(t *testing.T) {
 	}
 
 	// Read data from the file using async reads.
-	q := NewGoQueue(8)
+	q, err := newQueue(8)
+	if err != nil {
+		t.Fatalf("failed to create Queue: %v", err)
+	}
 	defer q.Destroy()
 	qavail := q.Cap()
 	off := int64(0)
@@ -82,7 +85,7 @@ func TestRead(t *testing.T) {
 	}
 }
 
-func TestReadv(t *testing.T) {
+func testReadv(t *testing.T, newQueue func(cap int) (Queue, error)) {
 	// Create a temp file.
 	testFile, err := os.CreateTemp(t.TempDir(), "aio_test_readv")
 	if err != nil {
@@ -103,7 +106,10 @@ func TestReadv(t *testing.T) {
 	}
 
 	// Read data from the file using async vectored reads.
-	q := NewGoQueue(8)
+	q, err := newQueue(8)
+	if err != nil {
+		t.Fatalf("failed to create Queue: %v", err)
+	}
 	defer q.Destroy()
 	qavail := q.Cap()
 	iovecsData := make([][2]unix.Iovec, qavail)
@@ -152,7 +158,7 @@ func TestReadv(t *testing.T) {
 	}
 }
 
-func TestWrite(t *testing.T) {
+func testWrite(t *testing.T, newQueue func(cap int) (Queue, error)) {
 	// Create a temp file.
 	testFile, err := os.CreateTemp(t.TempDir(), "aio_test_write")
 	if err != nil {
@@ -168,7 +174,10 @@ func TestWrite(t *testing.T) {
 	_, _ = rand.Read(data)
 
 	// Write data to the file using async writes.
-	q := NewGoQueue(8)
+	q, err := newQueue(8)
+	if err != nil {
+		t.Fatalf("failed to create Queue: %v", err)
+	}
 	defer q.Destroy()
 	qavail := q.Cap()
 	off := int64(0)
@@ -209,7 +218,7 @@ func TestWrite(t *testing.T) {
 	}
 }
 
-func TestWritev(t *testing.T) {
+func testWritev(t *testing.T, newQueue func(cap int) (Queue, error)) {
 	// Create a temp file.
 	testFile, err := os.CreateTemp(t.TempDir(), "aio_test_writev")
 	if err != nil {
@@ -225,7 +234,10 @@ func TestWritev(t *testing.T) {
 	_, _ = rand.Read(data)
 
 	// Write data to the file using async vectored writes.
-	q := NewGoQueue(8)
+	q, err := newQueue(8)
+	if err != nil {
+		t.Fatalf("failed to create Queue: %v", err)
+	}
 	defer q.Destroy()
 	qavail := q.Cap()
 	iovecsData := make([][2]unix.Iovec, qavail)
@@ -278,3 +290,33 @@ func TestWritev(t *testing.T) {
 		t.Errorf("bytes differ")
 	}
 }
+
+func testQueue(t *testing.T, newQueue func(cap int) (Queue, error)) {
+	t.Run("Read", func(t *testing.T) {
+		t.Helper()
+		t.Parallel()
+		testRead(t, newQueue)
+	})
+	t.Run("Readv", func(t *testing.T) {
+		t.Helper()
+		t.Parallel()
+		testReadv(t, newQueue)
+	})
+	t.Run("Write", func(t *testing.T) {
+		t.Helper()
+		t.Parallel()
+		testWrite(t, newQueue)
+	})
+	t.Run("Writev", func(t *testing.T) {
+		t.Helper()
+		t.Parallel()
+		testWritev(t, newQueue)
+	})
+}
+
+func TestGoQueue(t *testing.T) {
+	testQueue(t, func(cap int) (Queue, error) {
+		q := NewGoQueue(cap)
+		return q, nil
+	})
+}