feat: add WithReadThreshold API

Author: joway

connection_errors.go

@@ -35,6 +35,8 @@ const (
 	ErrEOF = syscall.Errno(0x106)
 	// Write I/O buffer timeout, calling by Connection.Writer
 	ErrWriteTimeout = syscall.Errno(0x107)
+	// The wait read size large than read threshold
+	ErrReadOutOfThreshold = syscall.Errno(0x108)
 )
 
 const ErrnoMask = 0xFF
@@ -97,4 +99,5 @@ var errnos = [...]string{
 	ErrnoMask & ErrUnsupported:    "netpoll dose not support",
 	ErrnoMask & ErrEOF:            "EOF",
 	ErrnoMask & ErrWriteTimeout:   "connection write timeout",
+	ErrnoMask & ErrReadOutOfThreshold: "connection read size is out of threshold",
 }

connection_impl.go

@@ -45,8 +45,9 @@ type connection struct {
 	outputBuffer    *LinkBuffer
 	outputBarrier   *barrier
 	supportZeroCopy bool
-	maxSize         int // The maximum size of data between two Release().
-	bookSize        int // The size of data that can be read at once.
+	maxSize         int   // The maximum size of data between two Release().
+	bookSize        int   // The size of data that can be read at once.
+	readThreshold   int64 // The readThreshold of connection max read.
 }
 
 var (
@@ -94,6 +95,12 @@ func (c *connection) SetWriteTimeout(timeout time.Duration) error {
 	return nil
 }
 
+// SetReadThreshold implements Connection.
+func (c *connection) SetReadThreshold(readThreshold int64) error {
+	c.readThreshold = readThreshold
+	return nil
+}
+
 // ------------------------------------------ implement zero-copy reader ------------------------------------------
 
 // Next implements Connection.
@@ -394,28 +401,44 @@ func (c *connection) triggerWrite(err error) {
 // waitRead will wait full n bytes.
 func (c *connection) waitRead(n int) (err error) {
 	if n <= c.inputBuffer.Len() {
-		return nil
+		goto CLEANUP
 	}
+	// cannot wait read with an out of threshold size
+	if c.readThreshold > 0 && int64(n) > c.readThreshold {
+		// just return error and dont do cleanup
+		return Exception(ErrReadOutOfThreshold, "wait read")
+	}
+
 	atomic.StoreInt64(&c.waitReadSize, int64(n))
-	defer atomic.StoreInt64(&c.waitReadSize, 0)
 	if c.readTimeout > 0 {
-		return c.waitReadWithTimeout(n)
+		err = c.waitReadWithTimeout(n)
+		goto CLEANUP
 	}
 	// wait full n
 	for c.inputBuffer.Len() < n {
 		switch c.status(closing) {
 		case poller:
-			return Exception(ErrEOF, "wait read")
+			err = Exception(ErrEOF, "wait read")
 		case user:
-			return Exception(ErrConnClosed, "wait read")
+			err = Exception(ErrConnClosed, "wait read")
 		default:
 			err = <-c.readTrigger
-			if err != nil {
-				return err
-			}
+		}
+		if err != nil {
+			goto CLEANUP
 		}
 	}
-	return nil
+CLEANUP:
+	atomic.StoreInt64(&c.waitReadSize, 0)
+	if c.readThreshold > 0 && err == nil {
+		// only resume read when current read size could make newBufferSize < readThreshold
+		bufferSize := int64(c.inputBuffer.Len())
+		newBufferSize := bufferSize - int64(n)
+		if bufferSize >= c.readThreshold && newBufferSize < c.readThreshold {
+			c.resumeRead()
+		}
+	}
+	return err
 }
 
 // waitReadWithTimeout will wait full n bytes or until timeout.

connection_onevent.go

@@ -103,6 +103,7 @@ func (c *connection) onPrepare(opts *options) (err error) {
 		c.SetReadTimeout(opts.readTimeout)
 		c.SetWriteTimeout(opts.writeTimeout)
 		c.SetIdleTimeout(opts.idleTimeout)
+		c.SetReadThreshold(opts.readThreshold)
 
 		// calling prepare first and then register.
 		if opts.onPrepare != nil {

connection_reactor.go

@@ -104,6 +104,11 @@ func (c *connection) inputAck(n int) (err error) {
 		c.maxSize = mallocMax
 	}
 
+	// trigger throttle
+	if c.readThreshold > 0 && int64(length) >= c.readThreshold {
+		c.pauseRead()
+	}
+
 	var needTrigger = true
 	if length == n { // first start onRequest
 		needTrigger = c.onRequest()
@@ -138,6 +143,29 @@ func (c *connection) outputAck(n int) (err error) {
 
 // rw2r removed the monitoring of write events.
 func (c *connection) rw2r() {
-	c.operator.Control(PollRW2R)
+	switch c.operator.getMode() {
+	case opreadwrite:
+		c.operator.Control(PollRW2R)
+	case opwrite:
+		c.operator.Control(PollW2RW)
+	}
 	c.triggerWrite(nil)
 }
+
+func (c *connection) pauseRead() {
+	switch c.operator.getMode() {
+	case opread:
+		c.operator.Control(PollR2Hup)
+	case opreadwrite:
+		c.operator.Control(PollRW2W)
+	}
+}
+
+func (c *connection) resumeRead() {
+	switch c.operator.getMode() {
+	case ophup:
+		c.operator.Control(PollHup2R)
+	case opwrite:
+		c.operator.Control(PollW2RW)
+	}
+}

connection_test.go

@@ -675,3 +675,95 @@ func TestConnectionDailTimeoutAndClose(t *testing.T) {
 		wg.Wait()
 	}
 }
+
+func TestConnectionReadOutOfThreshold(t *testing.T) {
+	var readThreshold = 1024 * 100
+	var readSize = readThreshold + 1
+	var opts = &options{}
+	var wg sync.WaitGroup
+	wg.Add(1)
+	opts.onRequest = func(ctx context.Context, connection Connection) error {
+		if connection.Reader().Len() < readThreshold {
+			return nil
+		}
+		defer wg.Done()
+		// read throttled data
+		_, err := connection.Reader().Next(readSize)
+		Assert(t, errors.Is(err, ErrReadOutOfThreshold), err)
+		connection.Close()
+		return nil
+	}
+
+	WithReadThreshold(int64(readThreshold)).f(opts)
+	r, w := GetSysFdPairs()
+	rconn, wconn := &connection{}, &connection{}
+	rconn.init(&netFD{fd: r}, opts)
+	wconn.init(&netFD{fd: w}, opts)
+
+	msg := make([]byte, readThreshold)
+	_, err := wconn.Writer().WriteBinary(msg)
+	MustNil(t, err)
+	err = wconn.Writer().Flush()
+	MustNil(t, err)
+	wg.Wait()
+}
+
+func TestConnectionReadThreshold(t *testing.T) {
+	var readThreshold int64 = 1024 * 100
+	var opts = &options{}
+	var wg sync.WaitGroup
+	var throttled int32
+	wg.Add(1)
+	opts.onRequest = func(ctx context.Context, connection Connection) error {
+		if int64(connection.Reader().Len()) < readThreshold {
+			return nil
+		}
+		defer wg.Done()
+
+		atomic.StoreInt32(&throttled, 1)
+		// check if no more read data when throttled
+		inbuffered := connection.Reader().Len()
+		t.Logf("Inbuffered: %d", inbuffered)
+		time.Sleep(time.Millisecond * 100)
+		Equal(t, inbuffered, connection.Reader().Len())
+
+		// read non-throttled data
+		buf, err := connection.Reader().Next(int(readThreshold))
+		Equal(t, int64(len(buf)), readThreshold)
+		MustNil(t, err)
+		err = connection.Reader().Release()
+		MustNil(t, err)
+		t.Logf("read non-throttled data")
+
+		// continue read throttled data
+		buf, err = connection.Reader().Next(5)
+		MustNil(t, err)
+		t.Logf("read throttled data: [%s]", buf)
+		Equal(t, len(buf), 5)
+		MustNil(t, err)
+		err = connection.Reader().Release()
+		MustNil(t, err)
+		Equal(t, connection.Reader().Len(), 0)
+		return nil
+	}
+
+	WithReadThreshold(readThreshold).f(opts)
+	r, w := GetSysFdPairs()
+	rconn, wconn := &connection{}, &connection{}
+	rconn.init(&netFD{fd: r}, opts)
+	wconn.init(&netFD{fd: w}, opts)
+	Assert(t, rconn.readThreshold == readThreshold)
+
+	msg := make([]byte, readThreshold)
+	_, err := wconn.Writer().WriteBinary(msg)
+	MustNil(t, err)
+	err = wconn.Writer().Flush()
+	MustNil(t, err)
+	_, err = wconn.Writer().WriteString("hello")
+	MustNil(t, err)
+	err = wconn.Writer().Flush()
+	MustNil(t, err)
+	t.Logf("flush final msg")
+
+	wg.Wait()
+}

docs/guide/guide_cn.md

@@ -519,6 +519,26 @@ func callback(connection netpoll.Connection) error {
 }
 ```
 
+## 8. 如何配置连接的读取阈值大小 ？
+
+Netpoll 默认不会对端发送数据的读取速度有任何限制，每当连接有数据时，Netpoll 会尽可能快地将数据存放在自己的 buffer 中。但有时候可能用户不希望数据过快发送，或者是希望控制服务内存使用量，又或者业务 OnRequest 回调处理速度很慢需要限制发送方速度，此时可以使用 `WithReadThreshold` 来控制读取的最大阈值。
+
+### Client 侧使用
+
+```
+dialer := netpoll.NewDialer(netpoll.WithReadThreshold(1024 * 1024 * 1024 * 1)) // 1GB
+conn, _ = dialer.DialConnection(network, address, timeout)
+```
+
+### Server 侧使用
+
+```
+eventLoop, _ := netpoll.NewEventLoop(
+	handle,
+	netpoll.WithReadThreshold(1024 * 1024 * 1024 * 1), // 1GB
+)
+```
+
 # 注意事项
 
 ## 1. 错误设置 NumLoops

docs/guide/guide_en.md

@@ -558,6 +558,30 @@ func callback(connection netpoll.Connection) error {
 }
 ```
 
+## 8. How to configure the read threshold of the connection?
+
+By default, Netpoll does not place any limit on the reading speed of data sent by the end.
+Whenever there have more data on the connection, Netpoll will read the data into its own buffer as quickly as possible. 
+
+But sometimes users may not want data to be read too quickly, or they want to control the service memory usage, or the user's OnRequest callback processing data very slowly and need to control the peer's send speed.
+In this case, you can use `WithReadThreshold` to control the maximum reading threshold.
+
+### Client side use
+
+```
+dialer := netpoll.NewDialer(netpoll.WithReadThreshold(1024 * 1024 * 1024 * 1)) // 1GB
+conn, _ = dialer.DialConnection(network, address, timeout)
+```
+
+### Server side use
+
+```
+eventLoop, _ := netpoll.NewEventLoop(
+   handle,
+   netpoll.WithReadThreshold(1024 * 1024 * 1024 * 1), // 1GB
+)
+```
+
 # Attention
 
 ## 1. Wrong setting of NumLoops

fd_operator.go

@@ -19,6 +19,15 @@ import (
 	"sync/atomic"
 )
 
+const (
+	opdetach    int32 = -1
+	_           int32 = 0 // default op mode, means nothing
+	opread      int32 = 1
+	opwrite     int32 = 2
+	opreadwrite int32 = 3
+	ophup       int32 = 4
+)
+
 // FDOperator is a collection of operations on file descriptors.
 type FDOperator struct {
 	// FD is file descriptor, poll will bind when register.
@@ -42,8 +51,7 @@ type FDOperator struct {
 	// poll is the registered location of the file descriptor.
 	poll Poll
 
-	// protect only detach once
-	detached int32
+	mode int32
 
 	// private, used by operatorCache
 	next  *FDOperator
@@ -52,16 +60,21 @@ type FDOperator struct {
 }
 
 func (op *FDOperator) Control(event PollEvent) error {
-	if event == PollDetach && atomic.AddInt32(&op.detached, 1) > 1 {
-		return nil
-	}
 	return op.poll.Control(op, event)
 }
 
 func (op *FDOperator) Free() {
 	op.poll.Free(op)
 }
 
+func (op *FDOperator) getMode() int32 {
+	return atomic.LoadInt32(&op.mode)
+}
+
+func (op *FDOperator) setMode(mode int32) {
+	atomic.StoreInt32(&op.mode, mode)
+}
+
 func (op *FDOperator) do() (can bool) {
 	return atomic.CompareAndSwapInt32(&op.state, 1, 2)
 }
@@ -98,5 +111,5 @@ func (op *FDOperator) reset() {
 	op.Inputs, op.InputAck = nil, nil
 	op.Outputs, op.OutputAck = nil, nil
 	op.poll = nil
-	op.detached = 0
+	op.mode = 0
 }