From 1f9021e2873e0b4950be4f2c90df4f268e97ff6b Mon Sep 17 00:00:00 2001
From: kadmin <julianknodt@gmail.com>
Date: Wed, 7 Jul 2021 18:15:52 +0000
Subject: [PATCH] add capability check

---
 net/uring/all.go            | 39 +++++++++++++++++++++++++++++++++++++
 net/uring/io_uring.c        |  9 +++++++++
 net/uring/io_uring_linux.go | 18 +++++++++++++++++
 net/uring/udp_test.go       | 11 ++++++-----
 4 files changed, 72 insertions(+), 5 deletions(-)

diff --git a/net/uring/all.go b/net/uring/all.go
index 596c98565..4a936563e 100644
--- a/net/uring/all.go
+++ b/net/uring/all.go
@@ -21,3 +21,42 @@ func Available() bool {
 
 // DisabledError indicates that io_uring was explicitly disabled.
 var DisabledError = errors.New("io_uring disabled")
+
+type IORingOp = int
+
+//https://unixism.net/loti/tutorial/probe_liburing.html
+const (
+	IORING_OP_NOP IORingOp = iota
+	IORING_OP_READV
+	IORING_OP_WRITEV
+	IORING_OP_FSYNC
+	IORING_OP_READ_FIXED
+	IORING_OP_WRITE_FIXED
+	IORING_OP_POLL_ADD
+	IORING_OP_POLL_REMOVE
+	IORING_OP_SYNC_FILE_RANGE
+	IORING_OP_SENDMSG
+	IORING_OP_RECVMSG
+	IORING_OP_TIMEOUT
+	IORING_OP_TIMEOUT_REMOVE
+	IORING_OP_ACCEPT
+	IORING_OP_ASYNC_CANCEL
+	IORING_OP_LINK_TIMEOUT
+	IORING_OP_CONNECT
+	IORING_OP_FALLOCATE
+	IORING_OP_OPENAT
+	IORING_OP_CLOSE
+	IORING_OP_FILES_UPDATE
+	IORING_OP_STATX
+	IORING_OP_READ
+	IORING_OP_WRITE
+	IORING_OP_FADVISE
+	IORING_OP_MADVISE
+	IORING_OP_SEND
+	IORING_OP_RECV
+	IORING_OP_OPENAT2
+	IORING_OP_EPOLL_CTL
+	IORING_OP_SPLICE
+	IORING_OP_PROVIDE_BUFFERS
+	IORING_OP_REMOVE_BUFFERS
+)
diff --git a/net/uring/io_uring.c b/net/uring/io_uring.c
index afa25b5d3..9ade6e560 100644
--- a/net/uring/io_uring.c
+++ b/net/uring/io_uring.c
@@ -182,6 +182,15 @@ static int set_deadline(struct io_uring *ring, int64_t sec, long long ns) {
   return 0;
 }
 
+// index of io uring capability
+static int has_capability(int i) {
+    int supported;
+    struct io_uring_probe *probe = io_uring_get_probe();
+    supported = io_uring_opcode_supported(probe, i);
+    free(probe);
+    return supported;
+}
+
 #endif
 
 static int has_io_uring(void) {
diff --git a/net/uring/io_uring_linux.go b/net/uring/io_uring_linux.go
index 1e72ad6c4..2c34f4a00 100644
--- a/net/uring/io_uring_linux.go
+++ b/net/uring/io_uring_linux.go
@@ -71,6 +71,24 @@ type UDPConn struct {
 	reads int32
 }
 
+var (
+	mu sync.Mutex
+	// checks capabilities available on this system
+	capabilities map[IORingOp]bool
+)
+
+func checkCapability(op IORingOp) bool {
+	mu.Lock()
+	defer mu.Unlock()
+	if v, ok := capabilities[op]; ok {
+		return v
+	}
+
+	has_op := C.has_capability(C.int(op)) == 1
+	capabilities[op] = has_op
+	return has_op
+}
+
 func NewUDPConn(pconn net.PacketConn) (*UDPConn, error) {
 	if !*useIOURing {
 		return nil, DisabledError
diff --git a/net/uring/udp_test.go b/net/uring/udp_test.go
index 54808a854..b5bb0bcca 100644
--- a/net/uring/udp_test.go
+++ b/net/uring/udp_test.go
@@ -52,10 +52,10 @@ func TestUDPSendRecv(t *testing.T) {
 	Port: TestPort,
 }
 
-func NewUDPTestServer(t *testing.T) error {
+func NewUDPTestServer(t *testing.T) (closer func() error, err error) {
 	conn, err := net.ListenUDP("udp", serverAddr)
 	if err != nil {
-		return err
+		return nil, err
 	}
 	go func() {
 		for {
@@ -67,7 +67,7 @@ func NewUDPTestServer(t *testing.T) error {
 			}
 		}
 	}()
-	return nil
+	return conn.Close, nil
 }
 
 func TestUDPConn(t *testing.T) {
@@ -76,8 +76,9 @@ func TestUDPConn(t *testing.T) {
 	}
 	c := qt.New(t)
 	// TODO add a closer here
-	err := NewUDPTestServer(t)
+	closer, err := NewUDPTestServer(t)
 	c.Assert(err, qt.IsNil)
+	t.Cleanup(func() { closer() })
 	udpConn, err := net.DialUDP("udp", nil, serverAddr)
 	c.Assert(err, qt.IsNil)
 	defer udpConn.Close()
@@ -94,7 +95,7 @@ func TestUDPConn(t *testing.T) {
 	}
 
 	// Test many writes at once
-	for i := 0; i < 1000; i++ {
+	for i := 0; i < 256; i++ {
 		n, err := conn.WriteTo(content, serverAddr)
 		c.Assert(err, qt.IsNil)
 		if n != len(content) {