perf: 优化下流量暴力压测

Author: guonaihong

bytes_pool.go

@@ -98,7 +98,6 @@ func putBytes(bytes *[]byte) {
 	if cap(*bytes)%page != 0 {
 		index-- // 向前挪一格, 可以保证空间是够的
 	}
-	// fmt.Printf("putBytes index: %d\n", index)
 	smallPools[index].Put(bytes)
 }
 
@@ -107,7 +106,6 @@ func putBytes(bytes *[]byte) {
 var bigPools = make([]sync.Pool, 0, 4)
 var bigPoolsSize = []int{
 	512 * 1024,
-	1024 * 1024,
 }
 
 func init() {

conn.go

@@ -2,7 +2,7 @@ package pulse
 
 import (
 	"errors"
-	"log"
+	"log/slog"
 	"net"
 	"sync"
 	"sync/atomic"
@@ -79,11 +79,13 @@ func (c *Conn) close() {
 		c.safeConns.Del(int(oldFd))
 		if err := core.Close(int(oldFd)); err != nil {
 			// Log the error but don't panic as this is a cleanup function
-			log.Printf("failed to close fd %d: %v", oldFd, err)
+			slog.Error("failed to close fd", "fd", oldFd, "error", err)
 		}
 	}
 	for _, wbuf := range c.wbufList {
-		putBytes(wbuf)
+		if wbuf != nil {
+			putBytes(wbuf)
+		}
 	}
 	c.wbufList = c.wbufList[:0]
 }
@@ -129,13 +131,19 @@ func (c *Conn) Write(data []byte) (int, error) {
 		if errors.Is(err, core.EAGAIN) || errors.Is(err, core.EINTR) || err == nil {
 			// 部分写入成功，或者全部失败
 			// 把剩余数据放到缓冲区
+			if n < 0 {
+				n = 0
+			}
 			if n < len(data) {
 				newBuf := getBytes(len(data) - n)
 				copy(*newBuf, data[n:])
-				c.wbufList = append(c.wbufList, newBuf)
 				if err := c.eventLoop.AddWrite(c.getFd()); err != nil {
-					log.Printf("failed to add write event: %v", err)
+					// 如果事件注册失败，释放缓冲区
+					putBytes(newBuf)
+					slog.Error("failed to add write event", "error", err)
+					return n, err
 				}
+				c.wbufList = append(c.wbufList, newBuf)
 			}
 			return n, nil
 		}
@@ -161,18 +169,27 @@ func (c *Conn) Write(data []byte) (int, error) {
 		if errors.Is(err, core.EAGAIN) || errors.Is(err, core.EINTR) || err == nil {
 			if n < len(*wbuf) {
 				// 部分写入，移动剩余数据到缓冲区开始位置
-				copy(*wbuf, (*wbuf)[n:])
-				*wbuf = (*wbuf)[:len(*wbuf)-n]
+				if n > 0 {
+					newBuf := getBytes(len(*wbuf) - n)
+					copy(*newBuf, (*wbuf)[n:])
+					putBytes(wbuf)
+					c.wbufList[i] = newBuf
+				}
 				// 释放已处理完的缓冲区
 				for j := lastIndex; j < i; j++ {
+					if c.wbufList[j] == nil {
+						continue
+					}
 					putBytes(c.wbufList[j])
+					c.wbufList[j] = nil
 				}
 				// 移动未处理的缓冲区到列表开始位置
 				copy(c.wbufList, c.wbufList[i:])
 				c.wbufList = c.wbufList[:len(c.wbufList)-i]
 				return len(data), nil
 			}
 			putBytes(wbuf)
+			c.wbufList[i] = nil
 			lastIndex = i + 1
 			continue
 		}
@@ -185,14 +202,14 @@ func (c *Conn) Write(data []byte) (int, error) {
 	// 所有数据都已写入
 	c.wbufList = c.wbufList[:0]
 	if err := c.eventLoop.ResetRead(c.getFd()); err != nil {
-		log.Printf("failed to reset read event: %v", err)
+		slog.Error("failed to reset read event", "error", err)
 	}
 	return len(data), nil
 }
 
 func (c *Conn) flush() {
 	if _, err := c.Write(nil); err != nil {
-		log.Printf("failed to flush write buffer: %v", err)
+		slog.Error("failed to flush write buffer", "error", err)
 	}
 }
 
@@ -229,7 +246,7 @@ func handleData(c *Conn, options *Options, rawData []byte) {
 		}
 		return true
 	}); err != nil {
-		log.Printf("failed to add task: %v", err)
+		slog.Error("failed to add task", "error", err)
 		// 释放newBytes since task failed
 		if newBytes != nil {
 			putBytes(newBytes)

multi_event_loops.go

@@ -3,6 +3,7 @@ package pulse
 import (
 	"context"
 	"errors"
+	"io"
 	"log"
 	"log/slog"
 	"net"
@@ -43,6 +44,17 @@ func (m *MultiEventLoop) initDefaultSetting() {
 	if m.options.eventLoopReadBufferSize == 0 {
 		m.options.eventLoopReadBufferSize = defEventLoopReadBufferSize
 	}
+
+	if m.options.maxSocketReadTimes == 0 {
+		if m.options.triggerType == TriggerTypeLevel {
+			// 水平触发模式下使用默认值
+			m.options.maxSocketReadTimes = defMaxSocketReadTimes
+		} else {
+			// 边缘触发模式下不限制读取次数
+			m.options.maxSocketReadTimes = -1
+		}
+		m.options.maxSocketReadTimes = defMaxSocketReadTimes
+	}
 }
 
 func NewMultiEventLoop(ctx context.Context, options ...func(*Options)) (e *MultiEventLoop, err error) {
@@ -133,14 +145,15 @@ func (e *MultiEventLoop) ListenAndServe(addr string) error {
 				if _, err := eventLoop.Poll(0, func(fd int, state core.State, err error) {
 
 					c := safeConns.GetUnsafe(fd)
+					// c := safeConns.Get(fd)
 					// slog.Debug("poll", "fd", fd, "state", state, "err", err)
 					if err != nil {
 						if errors.Is(err, core.EAGAIN) {
 							return
 						}
 						if c != nil {
-							e.options.callback.OnClose(c, err)
 							c.Close()
+							e.options.callback.OnClose(c, err)
 						}
 						return
 					}
@@ -167,7 +180,11 @@ func (e *MultiEventLoop) ListenAndServe(addr string) error {
 }
 
 func (e *MultiEventLoop) doRead(c *Conn, rbuf []byte) {
-	for {
+	for i := 0; ; i++ {
+		if e.options.maxSocketReadTimes > 0 && i >= e.options.maxSocketReadTimes {
+			return
+		}
+
 		// 循环读取数据
 		c.mu.Lock()
 		n, err := core.Read(c.getFd(), rbuf)
@@ -189,6 +206,9 @@ func (e *MultiEventLoop) doRead(c *Conn, rbuf []byte) {
 		}
 
 		if n == 0 {
+			// 如果不是这个错误直接关闭连接
+			c.Close()
+			e.options.callback.OnClose(c, io.EOF)
 			return
 		}
 		if n > 0 {

multi_event_loops_options.go

@@ -11,6 +11,7 @@ var (
 	defTaskMax                 = 30000
 	defTaskInitCount           = 8
 	defEventLoopReadBufferSize = 1024 * 4
+	defMaxSocketReadTimes      = 1
 )
 
 type TaskType int
@@ -38,12 +39,20 @@ type taskConfig struct {
 
 // 边缘触发
 type Options struct {
-	callback                Callback
-	task                    taskConfig
-	level                   slog.Level
-	taskType                TaskType
-	triggerType             core.TriggerType
-	eventLoopReadBufferSize int
+	callback                Callback         // 回调函数
+	task                    taskConfig       // 协程池配置
+	level                   slog.Level       // 日志级别
+	taskType                TaskType         // 任务类型
+	triggerType             core.TriggerType // 触发类型, 水平触发还是边缘触发
+	eventLoopReadBufferSize int              // event loop中读buffer的大小
+	maxSocketReadTimes      int              // socket单次最大读取次数
+}
+
+// 最大读取次数
+func WithMaxSocketReadTimes(maxSocketReadTimes int) func(*Options) {
+	return func(o *Options) {
+		o.maxSocketReadTimes = maxSocketReadTimes
+	}
 }
 
 // 设置回调函数