feat: 加上流量背压机制的选项

Author: guonaihong

bytes_pool.go

@@ -49,7 +49,7 @@ func init() {
 	// go debug()
 }
 
-func debug() {
+func debugPool() {
 	for {
 		time.Sleep(time.Second * 1)
 		slog.Info("debug", "index7AllocCount", atomic.LoadInt64(&index7AllocCount),

conn.go

@@ -14,16 +14,19 @@ import (
 )
 
 type Conn struct {
-	fd             int64
-	wbufList       []*[]byte // write buffer, 为了理精细控制内存使用量
-	mu             sync.Mutex
-	safeConns      *safeConns[Conn]
-	task           driver.TaskExecutor
-	eventLoop      core.PollingApi
-	readTimer      *time.Timer
-	writeTimer     *time.Timer
-	session        any // 会话数据
-	readBufferSize int // 读缓冲区大小
+	fd         int64
+	wbufList   []*[]byte // write buffer, 为了理精细控制内存使用量
+	mu         sync.Mutex
+	safeConns  *safeConns[Conn]
+	task       driver.TaskExecutor
+	eventLoop  core.PollingApi
+	readTimer  *time.Timer
+	writeTimer *time.Timer
+	session    any // 会话数据
+
+	// 如果再加字段，可以改成对options的指针的访问, 目前只是浪费了8个字节
+	readBufferSize             int  // 读缓冲区大小
+	flowBackPressureRemoveRead bool // 流量背压机制，当连接的写缓冲区满了，会移除读事件
 }
 
 func (c *Conn) SetNoDelay(nodelay bool) error {
@@ -36,7 +39,7 @@ func (c *Conn) getFd() int {
 
 func newConn(fd int, safeConns *safeConns[Conn],
 	task selectTasks, taskType TaskType,
-	eventLoop core.PollingApi, readBufferSize int) *Conn {
+	eventLoop core.PollingApi, readBufferSize int, flowBackPressureRemoveRead bool) *Conn {
 	var taskExecutor driver.TaskExecutor
 	switch taskType {
 	case TaskTypeInConnectionGoroutine:
@@ -50,11 +53,12 @@ func newConn(fd int, safeConns *safeConns[Conn],
 	}
 
 	return &Conn{
-		fd:             int64(fd),
-		safeConns:      safeConns,
-		task:           taskExecutor,
-		eventLoop:      eventLoop,
-		readBufferSize: readBufferSize,
+		fd:                         int64(fd),
+		safeConns:                  safeConns,
+		task:                       taskExecutor,
+		eventLoop:                  eventLoop,
+		readBufferSize:             readBufferSize,
+		flowBackPressureRemoveRead: flowBackPressureRemoveRead,
 	}
 }
 
@@ -109,7 +113,7 @@ func (c *Conn) writeToSocket(data []byte) (int, error) {
 	if err == syscall.EAGAIN {
 		return 0, err // 资源暂时不可用
 	}
-	return n, err // 其他错误直接返回
+	return 0, err // 其他错误直接返回
 
 }
 
@@ -179,10 +183,19 @@ func (c *Conn) handlePartialWrite(data *[]byte, n int, needAppend bool) error {
 		*data = (*data)[:len(*data)-n]
 	}
 
-	if err := c.eventLoop.AddWrite(c.getFd()); err != nil {
-		slog.Error("failed to add write event", "error", err)
-		return err
+	// 部分写入成功，或者全部失败
+	// 如果启用了流量背压机制且有部分写入，先删除读事件
+	if c.flowBackPressureRemoveRead {
+		if delErr := c.eventLoop.DelRead(c.getFd()); delErr != nil {
+			slog.Error("failed to delete read event", "error", delErr)
+		}
+	} else {
+		if err := c.eventLoop.AddWrite(c.getFd()); err != nil {
+			slog.Error("failed to add write event", "error", err)
+			return err
+		}
 	}
+
 	return nil
 }
 
@@ -201,7 +214,9 @@ func (c *Conn) Write(data []byte) (int, error) {
 	if len(c.wbufList) == 0 {
 		n, err := c.writeToSocket(data)
 		if errors.Is(err, core.EAGAIN) || errors.Is(err, core.EINTR) || err == nil {
-			// 部分写入成功，或者全部失败
+			if n == len(data) {
+				return n, nil
+			}
 			// 把剩余数据放到缓冲区
 			if err := c.handlePartialWrite(&data, n, true); err != nil {
 				c.close()
@@ -223,8 +238,14 @@ func (c *Conn) Write(data []byte) (int, error) {
 	for i < len(c.wbufList) {
 		wbuf := c.wbufList[i]
 		n, err := c.writeToSocket(*wbuf)
-		if errors.Is(err, core.EAGAIN) || errors.Is(err, core.EINTR) {
-			// 部分写入，移动剩余数据到缓冲区开始位置
+		if errors.Is(err, core.EAGAIN) || errors.Is(err, core.EINTR) || err == nil /*写入成功，也有n != len(*wbuf)的情况*/ {
+			if n == len(*wbuf) {
+				putBytes(wbuf)
+				c.wbufList[i] = nil
+				i++
+				continue
+			}
+			// 移动剩余数据到缓冲区开始位置
 			if err := c.handlePartialWrite(wbuf, n, false); err != nil {
 				c.close()
 				return 0, err
@@ -236,18 +257,13 @@ func (c *Conn) Write(data []byte) (int, error) {
 			return len(data), nil
 		}
 
-		if err != nil {
-			c.close()
-			return 0, err
-		}
-
-		putBytes(wbuf)
-		c.wbufList[i] = nil
-		i++
+		c.close()
+		return n, err
 	}
 
 	// 所有数据都已写入
 	c.wbufList = c.wbufList[:0]
+	// 如果启用了流量背压机制，重新添加读事件
 	if err := c.eventLoop.ResetRead(c.getFd()); err != nil {
 		slog.Error("failed to reset read event", "error", err)
 	}

core/api.go

@@ -47,6 +47,7 @@ type PollingApi interface {
 	AddRead(fd int) error
 	AddWrite(fd int) error
 	ResetRead(fd int) error
+	DelRead(fd int) error
 	Del(fd int) error
 	Poll(tv time.Duration, cb func(int, State, error)) (retVal int, err error)
 	Free()

core/api_epoll.go

@@ -81,6 +81,9 @@ func Create(triggerType TriggerType) (la PollingApi, err error) {
 	slog.Info("create epoll", "triggerType", triggerType)
 	e.events = make([]syscall.EpollEvent, 1024)
 	e.rev, e.wev, e.dwEv, e.resetEv = getReadWriteDeleteReset(triggerType == TriggerTypeEdge)
+	if e.dwEv == 0 {
+		panic("dwEv is 0")
+	}
 	return &e, nil
 }
 
@@ -136,6 +139,18 @@ func (e *eventPollState) DelWrite(fd int) error {
 	return nil
 }
 
+// 删除读事件
+func (e *eventPollState) DelRead(fd int) error {
+	if fd > 0 {
+		// 移除读事件，只保留写事件
+		return syscall.EpollCtl(e.epfd, syscall.EPOLL_CTL_MOD, fd, &syscall.EpollEvent{
+			Fd:     int32(fd),
+			Events: uint32(syscall.EPOLLOUT),
+		})
+	}
+	return nil
+}
+
 // 删除事件
 func (e *eventPollState) Del(fd int) error {
 	return syscall.EpollCtl(e.epfd, syscall.EPOLL_CTL_DEL, fd, &syscall.EpollEvent{Fd: int32(fd)})

core/api_iocp.go

@@ -56,6 +56,11 @@ func (i *iocp) ResetRead(fd int) error {
 	return nil
 }
 
+func (i *iocp) DelRead(fd int) error {
+	i.events[fd] &^= READ
+	return nil
+}
+
 func (i *iocp) Del(fd int) error {
 	delete(i.events, fd)
 	return nil

core/api_kqueue.go

@@ -79,6 +79,17 @@ func (as *eventPollState) ResetRead(fd int) error {
 	return err
 }
 
+func (as *eventPollState) DelRead(fd int) error {
+	if fd == -1 {
+		return nil
+	}
+
+	_, err := unix.Kevent(as.kqfd, []unix.Kevent_t{
+		{Ident: uint64(fd), Flags: unix.EV_DELETE, Filter: unix.EVFILT_READ},
+	}, nil, nil)
+	return err
+}
+
 func (as *eventPollState) Del(fd int) error {
 	// _, err := unix.Kevent(as.kqfd, []unix.Kevent_t{
 	// 	{Ident: uint64(fd), Flags: unix.EV_DELETE, Filter: unix.EVFILT_READ},

example/flow_backpressure/main.go

@@ -0,0 +1,66 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"time"
+
+	"github.com/antlabs/pulse"
+)
+
+type handler struct {
+	name string
+}
+
+func (h *handler) OnOpen(c *pulse.Conn) {
+	slog.Info("connection opened", "name", h.name)
+}
+
+func (h *handler) OnData(c *pulse.Conn, data []byte) {
+	slog.Info("received data", "name", h.name, "size", len(data))
+
+	// 模拟回显大量数据，可能触发流量背压
+	largeResponse := make([]byte, 10240) // 10KB 响应
+	copy(largeResponse, data)
+
+	n, err := c.Write(largeResponse)
+	if err != nil {
+		slog.Error("failed to write response", "error", err)
+		return
+	}
+	slog.Info("sent response", "bytes", n)
+}
+
+func (h *handler) OnClose(c *pulse.Conn, err error) {
+	if err != nil {
+		slog.Info("connection closed with error", "name", h.name, "error", err)
+	} else {
+		slog.Info("connection closed normally", "name", h.name)
+	}
+}
+
+func main() {
+	// 创建启用流量背压删除读事件机制的服务器
+	server, err := pulse.NewMultiEventLoop(
+		context.Background(),
+		pulse.WithCallback(&handler{name: "flow-backpressure-server"}),
+		pulse.WithTaskType(pulse.TaskTypeInEventLoop),
+		pulse.WithTriggerType(pulse.TriggerTypeEdge),
+		pulse.WithLogLevel(slog.LevelInfo),
+		pulse.WithFlowBackPressureRemoveRead(true), // 启用流量背压删除读事件机制
+	)
+	if err != nil {
+		panic(err)
+	}
+
+	go func() {
+		fmt.Println("Server starting on :8080 with flow backpressure (remove read events) enabled")
+		if err := server.ListenAndServe(":8080"); err != nil {
+			panic(err)
+		}
+	}()
+
+	// 保持服务器运行
+	time.Sleep(time.Hour)
+}

multi_event_loops.go

@@ -137,7 +137,8 @@ func (e *MultiEventLoop) ListenAndServe(addr string) error {
 			c2 := newConn(fd, &safeConns, e.localTask,
 				e.options.taskType,
 				e.eventLoops[index],
-				e.options.eventLoopReadBufferSize)
+				e.options.eventLoopReadBufferSize,
+				e.options.flowBackPressureRemoveRead)
 			safeConns.Add(fd, c2)
 			e.options.callback.OnOpen(c2)
 			err = e.eventLoops[index].AddRead(fd)

multi_event_loops_options.go

@@ -100,6 +100,7 @@ func WithFlowBackPressure(enable bool) func(*Options) {
 }
 
 // 设置流量背压机制，当连接的写缓冲区满了，会移除读事件，直到写缓冲区有空闲空间
+// 第二种背压机制会比第一种背压机制更高效, 7945hx cpu上，第二种是3.4GB/s的读写 第一种是3.0GB/s的读写
 func WithFlowBackPressureRemoveRead(enable bool) func(*Options) {
 	return func(o *Options) {
 		o.flowBackPressureRemoveRead = enable