IO_URING. Часть 2. Let's GO

package main

import (
	"errors"
	"flag"
	"fmt"
	"github.com/godzie44/go-uring/uring"
	"log"
	"strconv"
	"syscall"
)

const MaxConns = 4096
const Backlog = 512
const MaxMsgLen = 2048

type connType int

const (
	_ connType = iota
	ACCEPT
	READ
	WRITE
)

// Для каждого активного соединения будем держать в памяти connInfo структуру.
// fd файловый дескриптор сокета.
// typ - состояние в котором находится сокет (ожидает чтения/записи/accept'а).
// sendOp, recvOp - закешированные операции чтения из сокета/записи в сокет.
type connInfo struct {
	fd  int
	typ connType

	sendOp *uring.SendOp
	recvOp *uring.RecvOp
}

// Для каждого соединения предалоцируем операции чтения и записи.
// Переиспользование операция чтения и записи снизит нагрузку на GC.
func makeConns() [MaxConns]connInfo {
	var conns [MaxConns]connInfo
	for fd := range conns {
		conns[fd].recvOp = uring.Recv(uintptr(fd), nil, 0)
		conns[fd].sendOp = uring.Send(uintptr(fd), nil, 0)
	}
	return conns
}

// Буфер для соединений.
var conns = makeConns()

// Для каждого соединения инициализируем буфера для записи/чтения.
func makeBuffers() [][]byte {
	buffs := make([][]byte, MaxConns)
	for i := range buffs {
		buffs[i] = make([]byte, MaxMsgLen)
	}
	return buffs
}

var buffs = makeBuffers()

func main() {
	flag.Parse()
	port, _ := strconv.Atoi(flag.Arg(0))

  // Создаем серверный сокет и слушаем порт.
 	// Отметим, что при создании сокета флаг O_NON_BLOCK не устанавливается,
	// при этом все операции read/write и тд преобразуются в неблокирующие системные вызовы внутри io_uring
	serverSockFd, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_STREAM, 0)
	checkErr(err)
	defer syscall.Close(serverSockFd)

	checkErr(syscall.SetsockoptInt(serverSockFd, syscall.SOL_SOCKET, syscall.SO_REUSEADDR, 1))

	checkErr(syscall.Bind(serverSockFd, &syscall.SockaddrInet4{
		Port: port,
	}))
	checkErr(syscall.Listen(serverSockFd, Backlog))

	fmt.Printf("io_uring echo server listening for connections on port: %d\n", port)

	// Создаем экземпляр io_uring, не используем никакие кастомные настройки.
	// Вместимость SQ/SQ буферов устанавливаем в 4096 элементов.
	ring, err := uring.New(4096)
	checkErr(err)
	defer ring.Close()

	// Проверяем наличие фичи IORING_FEAT_FAST_POLL.
  // Для нас это наиболее "перформящая" фича в данном приложении,
  // фактически это встроенный в io_uring движок для поллинга I/O.
	if !ring.Params.FastPollFeature() {
		checkErr(errors.New("IORING_FEAT_FAST_POLL not available"))
	}

	// Добавляем первую операцию в SQ - слушаем серверный сокет на предмет новых входящих соединений.
	acceptOp := uring.Accept(uintptr(serverSockFd), 0)
	addAccept(ring, acceptOp)

	cqes := make([]*uring.CQEvent, Backlog)

	var cqe *uring.CQEvent
	for {
		// Сабмитим все SQE которые были добавлены на предыдущей итерации..
		_, err = ring.Submit()
		checkErr(err)

		// Ждем когда в CQ буфере появится хотя бы одно CQE.
		_, err = ring.WaitCQEvents(1)
		if errors.Is(err, syscall.EAGAIN) || errors.Is(err, syscall.EINTR) {
			continue
		}
		checkErr(err)

		// Помещаем все "готовые" CQE в буфер cqes.
		n := ring.PeekCQEventBatch(cqes)

		for i := 0; i < n; i++ {
			cqe = cqes[i]

      // В поле user_data находится индекс соответствующего connInfo
      // в которой находится служебная информация по сокету.
			ud := conns[cqe.UserData]
			typ := ud.typ
			res := cqe.Res

			ring.SeenCQE(cqe)

			// Проверяя тип мы идентифицируем операцию результат которой находится в CQE (accept / recv / send).
			switch typ {
			case ACCEPT:
				addRead(ring, int(res))
				addAccept(ring, acceptOp)
			case READ:
				if res <= 0 {
					_ = syscall.Shutdown(ud.fd, syscall.SHUT_RDWR)
				} else {
					addWrite(ring, ud.fd, res)
				}
			case WRITE:
				addRead(ring, ud.fd)
			}
		}
	}
}

// addAccept - добавляет в SQ accept операцию, fd - файловый дескриптор сервер сокета.
func addAccept(ring *uring.Ring, acceptOp *uring.AcceptOp) {
	conns[acceptOp.Fd()].fd = acceptOp.Fd()
	conns[acceptOp.Fd()].typ = ACCEPT

	err := ring.QueueSQE(acceptOp, 0, uint64(acceptOp.Fd()))
	checkErr(err)
}

// addRead - добавляет в SQ read операцию, fd - файловый дескриптор клиентского сокета.
func addRead(ring *uring.Ring, fd int) {
	buff := buffs[fd]

	ci := &conns[fd]
	ci.fd = fd
	ci.typ = READ
	ci.recvOp.SetBuffer(buff)

	err := ring.QueueSQE(ci.recvOp, 0, uint64(fd))
	checkErr(err)
}

// addWrite - добавляет в SQ write операцию, fd - файловый дескриптор клиентского сокета.
func addWrite(ring *uring.Ring, fd int, bytesRead int32) {
	buff := buffs[fd]

	ci := &conns[fd]
	ci.fd = fd
	ci.typ = WRITE
	ci.sendOp.SetBuffer(buff[:bytesRead])

	err := ring.QueueSQE(ci.sendOp, 0, uint64(fd))
	checkErr(err)
}

func checkErr(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

package main

import (
	"io"
	"log"
	"net"
)

const MaxConns = 4096
const MaxMsgLen = 2048

func initBuffs() [][]byte {
	buffs := make([][]byte, MaxConns)
	for i := range buffs {
		buffs[i] = make([]byte, MaxMsgLen)
	}
	return buffs
}

var buffs = initBuffs()

func main() {
	listener, err := net.ListenTCP("tcp", &net.TCPAddr{
		Port: 8080,
	})
	checkErr(err)

	for {
		conn, err := listener.Accept()
		checkErr(err)

		go handleConn(conn)
	}
}

func handleConn(conn net.Conn) {
	f, _ := conn.(*net.TCPConn).File()
	buff := buffs[int(f.Fd())]
	for {
		n, err := conn.Read(buff)
		if err == io.EOF || n == 0 {
			checkErr(conn.Close())
			return
		}
		checkErr(err)

		_, err = conn.Write(buff[:n])
		checkErr(err)
	}
}

func checkErr(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

type Callback func(event uring.CQEvent)

type Reactor struct {
	ring *uring.Ring
	callbacks     map[uint64]Callback
	callbacksLock sync.Mutex
	queueSQELock sync.Mutex
	currentNonce uint64
	submitSignal chan struct{}
}

func New(ring *uring.Ring) *Reactor {
	return &Reactor{
		ring:         ring,
		submitSignal: make(chan struct{}),
		callbacks:    map[uint64]Callback{},
	}
}

func (r *Reactor) Run(ctx context.Context) {
	wg := &sync.WaitGroup{}

	//запускаем два компонента реактора
	//consumer - обрабатывает новые CQE
	//publisher - отправляет новые SQE на обработку
	wg.Add(2)
	go func() {
		defer wg.Done()
		r.runConsumer(ctx)
	}()
	go func() {
		defer wg.Done()
		r.runPublisher(ctx)
	}()
	wg.Wait()
}

// сабмитим новые SQE в отдельной горутине
func (r *Reactor) runPublisher(ctx context.Context) {
	defer close(r.submitSignal)
	for {
		select {
		//при поступлении сигнала обрабатываем новые SQE
		case <-r.submitSignal:
			r.queueSQELock.Lock()
			_, err := r.ring.Submit()
			r.queueSQELock.Unlock()
			if err != nil {
				log.Println("io_uring submit", err)
			}
		case <-ctx.Done():
			return
		}
	}
}

// получаем результаты выполнения команд и дергаем соответствующие колбеки
func (r *Reactor) runConsumer(ctx context.Context) {
	cqeBuff := make([]*uring.CQEvent, 512)
	// цикл обработки результатов операций (CQE)
	for {
		// командуем Publisher'у обработать новые SQE
		r.submitSignal <- struct{}{}

		// ждем хотябы одну завершенную операцию
		_, err := r.ring.WaitCQEventsWithTimeout(1, time.Millisecond)
		if errors.Is(err, syscall.EAGAIN) || errors.Is(err, syscall.EINTR) || errors.Is(err, syscall.ETIME) {
			runtime.Gosched()
			goto CheckCtxAndContinue
		}

		if err != nil {
			log.Println("io_uring wait", err)
			goto CheckCtxAndContinue
		}

		for n := r.ring.PeekCQEventBatch(cqeBuff); n > 0; n = r.ring.PeekCQEventBatch(cqeBuff) {
			for i := 0; i < n; i++ {
				cqe := cqeBuff[i]

				nonce := cqe.UserData

				//находим соответствующий callback
				r.callbacksLock.Lock()
				cb := r.callbacks[nonce]
				delete(r.callbacks, nonce)
				r.callbacksLock.Unlock()

				cb(uring.CQEvent{
					UserData: cqe.UserData,
					Res:      cqe.Res,
					Flags:    cqe.Flags,
				})
			}

			//сообщаем io_uring о N просмотренных CQE
			r.ring.AdvanceCQ(uint32(n))
		}

	CheckCtxAndContinue:
		select {
		case <-ctx.Done():
			return
		default:
			continue
		}
	}
}

//Queue добавляет операцию в SQ.
func (r *Reactor) Queue(op uring.Operation, cb Callback) (uint64, error) {
	//генерируем уникальное значение UserData
	nonce := r.nextNonce()

	r.queueSQELock.Lock()
	defer r.queueSQELock.Unlock()

	//помещаем операцию в SQ
	err := r.ring.QueueSQE(op, 0, nonce)
	if err == nil {
		r.callbacksLock.Lock()
		r.callbacks[nonce] = cb
		r.callbacksLock.Unlock()
	}

	return nonce, err
}

func (r *Reactor) nextNonce() uint64 {
	return atomic.AddUint64(&r.currentNonce, 1)
}

type Listener struct {
	fd         int
	reactor    *Reactor
	acceptChan chan uring.CQEvent
	addr       net.Addr
}

func NewListener(addr string, reactor *Reactor) (*Listener, error) {
	tcpAddr, err := net.ResolveTCPAddr("tcp", addr)
	if err != nil {
		return nil, err
	}
	
	// serverSocket - создает и биндит серверный сокет
	sockFd, err := serverSocket(tcpAddr)
	if err != nil {
		return nil, err
	}

	return &Listener{
		fd:         sockFd,
		addr:       tcpAddr,
		reactor:    reactor,
		acceptChan: make(chan uring.CQEvent),
	}, nil
}

func (l *Listener) Accept() (net.Conn, error) {
	//помещаем accept операцию в реактор
	op := uring.Accept(uintptr(l.fd), 0)
	l.reactor.Queue(op, func(event uring.CQEvent) {
		l.acceptChan <- event
	})
	// ждем появление подключения
	cqe := <-l.acceptChan

	if err := cqe.Error(); err != nil {
		return nil, err
	}

	fd := int(cqe.Res)
	rAddr, _ := op.Addr()
	tc := newConn(fd, l.addr, rAddr, l.reactor)
	return tc, nil
}

type Conn struct {
	Fd                  int
	lAddr, rAddr        net.Addr
	reactor             *Reactor
	readChan, writeChan chan uring.CQEvent
	readLock, writeLock sync.Mutex
}

func newConn(fd int, lAddr, rAddr net.Addr, r *Reactor) *Conn {
	return &Conn{
		lAddr:     lAddr,
		rAddr:     rAddr,
		Fd:        fd,
		reactor:   r,
		readChan:  make(chan uring.CQEvent),
		writeChan: make(chan uring.CQEvent),
	}
}

func (c *Conn) Read(b []byte) (n int, err error) {
	c.readLock.Lock()
	defer c.readLock.Unlock()

	//помещаем Recv операцию в реактор
	op := uring.Recv(uintptr(c.Fd), b, 0)
	c.reactor.Queue(op, func(event uring.CQEvent) {
		c.readChan <- event
	})
	//ждем завершения чтения из сокета
	cqe := <-c.readChan

	if err = cqe.Error(); err != nil {
		return 0, &net.OpError{Op: "read", Net: "tcp", Source: c.lAddr, Addr: c.rAddr, Err: err}
	}
	if cqe.Res == 0 {
		err = &net.OpError{Op: "read", Net: "tcp", Source: c.lAddr, Addr: c.rAddr, Err: io.EOF}
	}

	runtime.KeepAlive(b)
	return int(cqe.Res), err
}

func (c *Conn) Write(b []byte) (n int, err error) {
	c.writeLock.Lock()
	defer c.writeLock.Unlock()

	//помещаем Send операцию в реактор
	op := uring.Send(uintptr(c.Fd), b, 0)
	c.reactor.Queue(op, func(event uring.CQEvent) {
		c.writeChan <- event
	})
	//ждем завершения записи в сокет
	cqe := <-c.writeChan

	if err = cqe.Error(); err != nil {
		return 0, &net.OpError{Op: "write", Net: "tcp", Source: c.lAddr, Addr: c.rAddr, Err: err}
	}
	if cqe.Res == 0 {
		err = &net.OpError{Op: "write", Net: "tcp", Source: c.lAddr, Addr: c.rAddr, Err: io.ErrUnexpectedEOF}
	}

	runtime.KeepAlive(b)
	return int(cqe.Res), err
}

func (c *Conn) Close() error {
	return syscall.Close(c.Fd)
}