WIP: feat: io_uring for netpoll I/O poller

go.mod

@@ -1,5 +1,8 @@
 module github.com/cloudwego/netpoll
 
-go 1.15
+go 1.17
 
-require github.com/bytedance/gopkg v0.0.0-20220413063733-65bf48ffb3a7
+require (
+	github.com/bytedance/gopkg v0.0.0-20220413063733-65bf48ffb3a7
+	golang.org/x/sys v0.0.0-20220825204002-c680a09ffe64
+)

go.sum

@@ -9,6 +9,8 @@ github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5Cc
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20220110181412-a018aaa089fe/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220825204002-c680a09ffe64 h1:UiNENfZ8gDvpiWw7IpOMQ27spWmThO1RwwdQVbJahJM=
+golang.org/x/sys v0.0.0-20220825204002-c680a09ffe64/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=

poll_default_bsd.go

@@ -25,11 +25,7 @@ import (
 	"unsafe"
 )
 
-func openPoll() Poll {
-	return openDefaultPoll()
-}
-
-func openDefaultPoll() *defaultPoll {
+func openDefaultPoll() Poll {
 	l := new(defaultPoll)
 	p, err := syscall.Kqueue()
 	if err != nil {

poll_default_linux.go

@@ -25,12 +25,7 @@ import (
 	"unsafe"
 )
 
-// Includes defaultPoll/multiPoll/uringPoll...
-func openPoll() Poll {
-	return openDefaultPoll()
-}
-
-func openDefaultPoll() *defaultPoll {
+func openDefaultPoll() Poll {
 	var poll = defaultPoll{}
 	poll.buf = make([]byte, 8)
 	var p, err = syscall.EpollCreate1(0)

poll_manager.go

@@ -23,6 +23,11 @@ import (
 	"runtime"
 )
 
+// Includes defaultPoll/multiPoll/uringPoll...
+func openPoll() Poll {
+	return pollRegister()
+}
+
 func setNumLoops(numLoops int) error {
 	return pollmanager.SetNumLoops(numLoops)
 }

poll_race_bsd.go

@@ -25,12 +25,7 @@ import (
 	"syscall"
 )
 
-// mock no race poll
-func openPoll() Poll {
-	return openDefaultPoll()
-}
-
-func openDefaultPoll() *defaultPoll {
+func openDefaultPoll() Poll {
 	l := new(defaultPoll)
 	p, err := syscall.Kqueue()
 	if err != nil {

poll_race_linux.go

@@ -25,12 +25,7 @@ import (
 	"syscall"
 )
 
-// mock no race poll
-func openPoll() Poll {
-	return openDefaultPoll()
-}
-
-func openDefaultPoll() *defaultPoll {
+func openDefaultPoll() Poll {
 	var poll = defaultPoll{}
 	poll.buf = make([]byte, 8)
 	var p, err = syscall.EpollCreate1(0)

poll_register.go

@@ -0,0 +1,31 @@
+// Copyright 2022 CloudWeGo 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 !windows
+// +build !windows
+
+package netpoll
+
+// registerPoll is the func of openning Poller
+var pollRegister = openDefaultPoll
+
+// RegisterEpoll implement Epoll
+func RegisterEpoll() {
+	pollRegister = openDefaultPoll
+}
+
+// RegisterURingPoll implement URing Poller
+func RegisterURingPoll() {
+	pollRegister = openURingPoll
+}

poll_uring.go

@@ -0,0 +1,250 @@
+// Copyright 2022 CloudWeGo 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 !race
+// +build !race
+
+package netpoll
+
+import (
+	"log"
+	"sync/atomic"
+	"syscall"
+	"unsafe"
+
+	. "github.com/cloudwego/netpoll/uring"
+)
+
+func openURingPoll() Poll {
+	poll := &uringPoll{}
+	uring, err := IOURing(128)
+	if err != nil {
+		panic(err)
+	}
+	poll.uring = uring
+	return poll
+}
+
+type uringPoll struct {
+	size    int
+	caps    int
+	trigger uint32
+
+	uring    *URing
+	cqes     []*URingCQE
+	barriers []barrier
+	hups     []func(p Poll) error
+}
+
+// Wait implements Poll.
+func (p *uringPoll) Wait() error {
+	// init
+	var caps, n = barriercap, 0
+	p.reset(128, caps)
+	// wait
+	for {
+		if n == p.size && p.size < 128*1024 {
+			p.reset(p.size<<1, caps)
+		}
+		n = p.uring.PeekBatchCQE(p.cqes)
+		if n == 0 {
+			continue
+		}
+		p.uring.Advance(uint32(n))
+		if p.handler(p.cqes[:n]) {
+			return nil
+		}
+	}
+}
+
+// Close implements Poll.
+func (p *uringPoll) Close() error {
+	var userData uint64
+	*(**FDOperator)(unsafe.Pointer(&userData)) = &FDOperator{FD: p.uring.Fd(), state: -1}
+	err := p.trig(userData)
+	return err
+}
+
+// Trigger implements Poll.
+func (p *uringPoll) Trigger() error {
+	if atomic.AddUint32(&p.trigger, 1) > 1 {
+		return nil
+	}
+	var userData uint64
+	*(**FDOperator)(unsafe.Pointer(&userData)) = &FDOperator{FD: p.uring.Fd()}
+	err := p.trig(userData)
+	return err
+}
+
+// Control implements Poll.
+func (p *uringPoll) Control(operator *FDOperator, event PollEvent) (err error) {
+	var pollOp Op
+	var mask uint32
+	switch event {
+	case PollReadable:
+		operator.inuse()
+		mask = syscall.EPOLLIN | syscall.EPOLLRDHUP | syscall.EPOLLERR
+		pollOp = PollAdd(uintptr(operator.FD), mask)
+	case PollModReadable:
+		operator.inuse()
+		mask = syscall.EPOLLIN | syscall.EPOLLRDHUP | syscall.EPOLLERR
+		pollOp = PollAdd(uintptr(operator.FD), mask)
+	case PollDetach:
+		pollOp = PollRemove(uint64(uintptr(unsafe.Pointer(operator))))
+	case PollWritable:
+		operator.inuse()
+		mask = EPOLLET | syscall.EPOLLOUT | syscall.EPOLLRDHUP | syscall.EPOLLERR
+		pollOp = PollAdd(uintptr(operator.FD), mask)
+	case PollR2RW:
+		mask = syscall.EPOLLIN | syscall.EPOLLOUT | syscall.EPOLLRDHUP | syscall.EPOLLERR
+		pollOp = PollAdd(uintptr(operator.FD), mask)
+	case PollRW2R:
+		mask = syscall.EPOLLIN | syscall.EPOLLRDHUP | syscall.EPOLLERR
+		pollOp = PollAdd(uintptr(operator.FD), mask)
+	}
+
+	var userData uint64
+	*(**FDOperator)(unsafe.Pointer(&userData)) = operator
+
+	err = p.uring.Queue(pollOp, 0, userData)
+	if err != nil {
+		panic(err)
+	}
+
+	_, err = p.uring.Submit()
+	return err
+}
+
+func (p *uringPoll) reset(size, caps int) {
+	p.size, p.caps = size, caps
+	p.cqes, p.barriers = make([]*URingCQE, size), make([]barrier, size)
+	for i := range p.barriers {
+		p.barriers[i].bs = make([][]byte, caps)
+		p.barriers[i].ivs = make([]syscall.Iovec, caps)
+	}
+}
+
+func (p *uringPoll) handler(cqes []*URingCQE) (closed bool) {
+	for i := range cqes {
+		var operator = *(**FDOperator)(unsafe.Pointer(&cqes[i].UserData))
+		// trigger or exit gracefully
+		if operator.FD == p.uring.Fd() {
+			// must clean trigger first
+			atomic.StoreUint32(&p.trigger, 0)
+			// if closed & exit
+			if operator.state == -1 {
+				p.uring.Close()
+				return true
+			}
+			operator.done()
+			continue
+		}
+
+		if !operator.do() {
+			continue
+		}
+
+		var events = cqes[i].Res
+		// check poll in
+		if events&syscall.EPOLLIN != 0 {
+			if operator.OnRead != nil {
+				// for non-connection
+				operator.OnRead(p)
+			} else {
+				// only for connection
+				var bs = operator.Inputs(p.barriers[i].bs)
+				if len(bs) > 0 {
+					var n, err = readv(operator.FD, bs, p.barriers[i].ivs)
+					operator.InputAck(n)
+					if err != nil && err != syscall.EAGAIN && err != syscall.EINTR {
+						log.Printf("readv(fd=%d) failed: %s", operator.FD, err.Error())
+						p.appendHup(operator)
+						continue
+					}
+				}
+			}
+		}
+
+		// check hup
+		if events&(syscall.EPOLLHUP|syscall.EPOLLRDHUP) != 0 {
+			p.appendHup(operator)
+			continue
+		}
+		if events&syscall.EPOLLERR != 0 {
+			// Under block-zerocopy, the kernel may give an error callback, which is not a real error, just an EAGAIN.
+			// So here we need to check this error, if it is EAGAIN then do nothing, otherwise still mark as hup.
+			if _, _, _, _, err := syscall.Recvmsg(operator.FD, nil, nil, syscall.MSG_ERRQUEUE); err != syscall.EAGAIN {
+				p.appendHup(operator)
+			} else {
+				operator.done()
+			}
+			continue
+		}
+
+		// check poll out
+		if events&syscall.EPOLLOUT != 0 {
+			if operator.OnWrite != nil {
+				// for non-connection
+				operator.OnWrite(p)
+			} else {
+				// only for connection
+				var bs, supportZeroCopy = operator.Outputs(p.barriers[i].bs)
+				if len(bs) > 0 {
+					// TODO: Let the upper layer pass in whether to use ZeroCopy.
+					var n, err = sendmsg(operator.FD, bs, p.barriers[i].ivs, false && supportZeroCopy)
+					operator.OutputAck(n)
+					if err != nil && err != syscall.EAGAIN {
+						log.Printf("sendmsg(fd=%d) failed: %s", operator.FD, err.Error())
+						p.appendHup(operator)
+						continue
+					}
+				}
+			}
+		}
+		operator.done()
+	}
+	// hup conns together to avoid blocking the poll.
+	p.detaches()
+	return false
+}
+
+func (p *uringPoll) trig(userData uint64) error {
+	err := p.uring.Queue(Nop(), 0, userData)
+	if err != nil {
+		return err
+	}
+	_, err = p.uring.Submit()
+	return err
+}
+
+func (p *uringPoll) appendHup(operator *FDOperator) {
+	p.hups = append(p.hups, operator.OnHup)
+	operator.Control(PollDetach)
+	operator.done()
+}
+
+func (p *uringPoll) detaches() {
+	if len(p.hups) == 0 {
+		return
+	}
+	hups := p.hups
+	p.hups = nil
+	go func(onhups []func(p Poll) error) {
+		for i := range onhups {
+			if onhups[i] != nil {
+				onhups[i](p)
+			}
+		}
+	}(hups)
+}