Merge pull request #20 from antlabs/dev

Dev

Author: guonaihong

common_options.go

@@ -283,28 +283,46 @@ func WithClientOnCloseFunc(onClose func(c *Conn, err error)) ClientOption {
 // 18.1 配置服务端Callback相关方法在io event loop中执行
 func WithServerCallbackInEventLoop() ServerOption {
 	return func(o *ConnOption) {
-		o.callbackInEventLoop = true
+		o.runInGoTask = "io"
 	}
 }
 
 // 18.2 配置服务端Callback相关方法在io event loop中执行
 func WithClientCallbackInEventLoop() ClientOption {
 	return func(o *DialOption) {
-		o.callbackInEventLoop = true
+		o.runInGoTask = "io"
 	}
 }
 
 // 19.1 配置服务端使用stream模式处理请求，该模式一个连接会独占一个go程，如果你的请求对时序有要求，可以使用这个模式
 func WithServerStreamMode() ServerOption {
 	return func(o *ConnOption) {
-		o.runInGoStrategy = taskStrategyStream
+		o.runInGoTask = "stream"
 	}
 }
 
 // 19.2 配置客户端使用stream模式处理请求，该模式一个连接会独占一个go程，如果你的请求对时序有要求，可以使用这个模式
 func WithClientStreamMode() ClientOption {
 	return func(o *DialOption) {
-		o.runInGoStrategy = taskStrategyStream
+		o.runInGoTask = "stream"
+	}
+}
+
+// 20.1 配置自定义task, 需要确保传入的值是有效的，不然会panic
+func WithServerCustomTaskMode(taskName string) ServerOption {
+	return func(o *ConnOption) {
+		if len(taskName) > 0 {
+			o.runInGoTask = taskName
+		}
+	}
+}
+
+// 20.2 配置自定义task, 需要确保传入的值是有效的，不然会panic
+func WithClientCustomTaskMode(taskName string) ClientOption {
+	return func(o *DialOption) {
+		if len(taskName) > 0 {
+			o.runInGoTask = taskName
+		}
 	}
 }
 

config.go

@@ -36,8 +36,7 @@ type Config struct {
 	maxDelayWriteDuration           time.Duration     // 最大延迟时间, 默认值是10ms
 	subProtocols                    []string          // 设置支持的子协议
 	multiEventLoop                  *MultiEventLoop   // 事件循环
-	callbackInEventLoop             bool              // 在event loop中运行websocket OnOpen, OnMessage, OnClose 回调函数
-	runInGoStrategy                 taskStrategy      //运行业务OnMessage的策略
+	runInGoTask                     string            // 运行业务OnMessage的策略, 现在greatws集成三种OnMessage运行模式，分别是io, task
 }
 
 // func (c *Config) useIoUring() bool {
@@ -55,8 +54,9 @@ func (c *Config) defaultSetting() {
 	c.windowsMultipleTimesPayloadSize = 1.0
 	c.delayWriteInitBufferSize = 8 * 1024
 	c.maxDelayWriteDuration = 10 * time.Millisecond
-	c.runInGoStrategy = taskStrategyBind
+	// c.runInGoStrategy = taskStrategyBind
 	c.tcpNoDelay = true
 	// 对于text消息，默认不检查text是utf8字符
 	c.utf8Check = func(b []byte) bool { return true }
+	c.runInGoTask = "unstream"
 }

conn_core.go

@@ -23,7 +23,6 @@ import (
 	"math/rand"
 	"sync/atomic"
 	"time"
-	"unsafe"
 
 	"github.com/antlabs/wsutil/bytespool"
 	"github.com/antlabs/wsutil/enum"
@@ -63,14 +62,14 @@ const (
 type conn struct {
 	fd                   int64              // 文件描述符fd
 	rbuf                 *[]byte            // 读缓冲区
-	rr                   int                // rbuf读索引
-	rw                   int                // rbuf写索引
+	rr                   int                // rbuf读索引，rfc标准里面有超过4个字节的大包，所以索引只能用int类型
+	rw                   int                // rbuf写索引，rfc标准里面有超过4个字节的大包，所以索引只能用int类型
+	wbuf                 *[]byte            // 写缓冲区, 当直接Write失败时，会将数据写入缓冲区
 	lenAndMaskSize       int                // payload长度和掩码的长度
-	lastPayloadLen       int64              // 上一次读取的payload长度
 	rh                   frame.FrameHeader  // frame头部
 	fragmentFramePayload *[]byte            // 存放分片帧的缓冲区
 	fragmentFrameHeader  *frame.FrameHeader // 存放分段帧的头部
-	closed               int32              // 是否关闭
+	lastPayloadLen       int32              // 上一次读取的payload长度, TODO启用
 	curState             frameState         // 保存当前状态机的状态
 	client               bool               // 客户端为true，服务端为false
 }
@@ -79,56 +78,14 @@ func (c *Conn) getLogger() *slog.Logger {
 	return c.multiEventLoop.Logger
 }
 
-// 获取当前绑定的go程
-func (c *Conn) getCurrBindGo() *businessGo {
-	return (*businessGo)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&c.currBindGo))))
-}
-
-// 设置当前绑定的go程
-func (c *Conn) setCurrBindGo(b *businessGo) {
-	atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&c.currBindGo)), unsafe.Pointer(b))
-}
-
-// addTask 进行任务调度， 任务调试分为三种模式， 1. stream模式， 2. go模式， 3. io模式
-func (c *Conn) addTask(ts taskStrategy, f func() bool) {
+func (c *Conn) addTask(f func() bool) {
 	if c.isClosed() {
 		return
 	}
 
-	if c.callbackInEventLoop {
-		c.multiEventLoop.runInIo.addTask(ts, f)
-		return
-	}
-	if ts == taskStrategyStream {
-		c.streamGo.addTask(ts, f)
-		return
-	}
-
-	var err error
-	retry := 2
-	if c.parent.localTask.isFull() {
-		retry = 1
-	}
-
-	for i := 0; i < retry; i++ {
-		err = c.parent.localTask.addTask(c, ts, f)
-		if err == nil || retry == 1 {
-			break
-		}
-		if err == ErrTaskQueueFull {
-			if c.parent.localTask.addGoWithSteal(c.getCurrBindGo()) {
-				c.parent.localTask.rebindGo(c)
-			}
-			continue
-		}
-
-	}
-
-	if err == ErrTaskQueueFull {
-		// 负载高走自己专用chan
-		c.getCurrBindGo().taskChan <- f
-		// 负载低走公共chan, TODO， 需要找一个判断高/低雷负载的条件
-		// c.parent.localTask.public <- f
+	err := c.task.AddTask(f)
+	if err != nil {
+		c.getLogger().Error("addTask", "err", err.Error())
 	}
 
 }
@@ -298,7 +255,7 @@ func (c *Conn) readPayload() (f frame.Frame2, success bool, err error) {
 	if needRead > 0 {
 		return
 	}
-	c.lastPayloadLen = c.rh.PayloadLen
+	c.lastPayloadLen = int32(c.rh.PayloadLen)
 	// 普通frame
 	newBuf := GetPayloadBytes(int(c.rh.PayloadLen))
 	copy(*newBuf, (*c.rbuf)[c.rr:c.rr+int(c.rh.PayloadLen)])
@@ -359,7 +316,7 @@ func (c *Conn) processCallback(f frame.Frame2) (err error) {
 				c.fragmentFramePayload = nil
 
 				// 进入业务协程执行
-				c.addTask(c.runInGoStrategy, func() (exit bool) {
+				c.addTask(func() (exit bool) {
 					if fragmentFrameHeader.GetRsv1() && decompression {
 						tempBuf, err := decode(*fragmentFramePayload)
 						if err != nil {
@@ -421,7 +378,7 @@ func (c *Conn) processCallback(f frame.Frame2) (err error) {
 		// payloadPtr.Store(f.Payload)
 
 		// text或者binary进入业务协程执行
-		c.addTask(c.runInGoStrategy, func() bool {
+		c.addTask(func() bool {
 			return c.processCallbackData(f, payload, rsv1, decompression, needMask, maskKey)
 		})
 
@@ -487,8 +444,7 @@ func (c *Conn) processCallback(f frame.Frame2) (err error) {
 				// 进入业务协程执行
 				payload := f.Payload
 				// here
-				c.addTask(c.runInGoStrategy, func() bool {
-
+				c.addTask(func() bool {
 					return c.processPing(f, payload)
 				})
 				return
@@ -500,7 +456,7 @@ func (c *Conn) processCallback(f frame.Frame2) (err error) {
 		}
 
 		// 进入业务协程执行
-		c.addTask(c.runInGoStrategy, func() bool {
+		c.addTask(func() bool {
 			c.Callback.OnMessage(c, f.Opcode, nil)
 			return false
 		})

conn_unix.go

@@ -27,6 +27,7 @@ import (
 	"time"
 	"unsafe"
 
+	"github.com/antlabs/greatws/task/driver"
 	"github.com/antlabs/wsutil/bytespool"
 	"github.com/antlabs/wsutil/enum"
 	"golang.org/x/sys/unix"
@@ -62,17 +63,15 @@ var (
 type Conn struct {
 	conn
 
-	wbuf          *[]byte     // 写缓冲区, 当直接Write失败时，会将数据写入缓冲区
-	mu            sync.Mutex  // 锁
-	*Config                   // 配置
-	parent        *EventLoop  // event loop
-	currBindGo    *businessGo // 绑定模式下，当前绑定的go程
-	streamGo      *taskStream // stream模式下，当前绑定的go程
-	rtime         *time.Timer // 控制读超时
-	wtime         *time.Timer // 控制写超时
-	closeConnOnce sync.Once   // 关闭一次
-	onCloseOnce   Once        // 保证只调用一次OnClose函数
-
+	mu            sync.Mutex          // 锁
+	*Config                           // 配置
+	parent        *EventLoop          // event loop
+	task          driver.TaskExecutor // 任务，该任务会进协程池里面执行
+	rtime         *time.Timer         // 控制读超时
+	wtime         *time.Timer         // 控制写超时
+	closeConnOnce sync.Once           // 关闭一次
+	onCloseOnce   Once                // 保证只调用一次OnClose函数
+	closed        int32               // 是否关闭
 }
 
 func newConn(fd int64, client bool, conf *Config) *Conn {
@@ -83,14 +82,10 @@ func newConn(fd int64, client bool, conf *Config) *Conn {
 		},
 		// 初始化不分配内存，只有在需要的时候才分配
 		Config: conf,
-
 		parent: conf.multiEventLoop.getEventLoop(int(fd)),
 	}
 
-	if conf.runInGoStrategy == taskStrategyStream {
-		c.streamGo = newTaskStream()
-	}
-
+	c.task = c.parent.localTask.newTask(conf.runInGoTask)
 	if conf.readTimeout > 0 {
 		c.setReadDeadline(time.Now().Add(conf.readTimeout))
 	}
@@ -113,15 +108,15 @@ func (c *Conn) closeInnerWithOnClose(err error, onClose bool) {
 		if err != nil {
 			err = io.EOF
 		}
-		switch c.Config.runInGoStrategy {
-		case taskStrategyBind:
-			if c.getCurrBindGo() != nil {
-				c.currBindGo.subBinConnCount()
-			}
-		case taskStrategyStream:
-			c.streamGo.close()
-
-		}
+		// switch c.Config.runInGoStrategy {
+		// case taskStrategyBind:
+		// 	if c.getCurrBindGo() != nil {
+		// 		c.currBindGo.subBinConnCount()
+		// 	}
+		// case taskStrategyStream:
+		// 	c.streamGo.close()
+
+		// }
 		fd := c.getFd()
 		c.getLogger().Debug("close conn", slog.Int64("fd", int64(fd)))
 		c.parent.del(c)

event_loop.go

@@ -34,7 +34,7 @@ type EventLoop struct {
 	*apiState          // 每个平台对应的异步io接口/epoll/kqueue/iouring(暂时不加，除非io-uring性能超过epoll才加回来)
 	shutdown  bool
 	parent    *MultiEventLoop
-	localTask task
+	localTask selectTasks
 }
 
 // 初始化函数
@@ -44,9 +44,14 @@ func CreateEventLoop(setSize int, flag evFlag, parent *MultiEventLoop) (e *Event
 		maxFd:   -1,
 		parent:  parent,
 	}
-	e.localTask.taskConfig = e.parent.configTask.taskConfig
-	e.localTask.taskMode = e.parent.configTask.taskMode
-	e.localTask.init()
+
+	// 初始化任务池
+	e.localTask = newSelectTask(parent.configTask.initCount, parent.configTask.min, parent.configTask.max)
+
+	// TODO+
+	// e.localTask.taskConfig = e.parent.configTask.taskConfig
+	// e.localTask.taskMode = e.parent.configTask.taskMode
+	// e.localTask.init()
 	err = e.apiCreate(flag)
 	return e, err
 }

multi_event_loops.go

@@ -19,18 +19,32 @@ import (
 	"runtime"
 	"sync/atomic"
 	"time"
+
+	_ "github.com/antlabs/greatws/task/io"
+	_ "github.com/antlabs/greatws/task/stream"
+	_ "github.com/antlabs/greatws/task/unstream"
 )
 
+type taskConfig struct {
+	initCount int // 初始化的协程数
+	min       int // 最小协程数
+	max       int // 最大协程数
+}
+
 type multiEventLoopOption struct {
 	numLoops int //起多少个event loop
-	// 只是配置的作用，不是真正的任务池， fd是绑定到某个事件循环上的，
+
+	// 为何不设计全局池, 现在的做法是
+	// fd是绑定到某个事件循环上的，
 	// 任务池是绑定到某个事件循环上的，所以这里的任务池也绑定到对应的localTask上
 	// 如果设计全局任务池，那么概念就会很乱，容易出错，也会临界区竞争
-	configTask       task
-	taskMode         taskMode
+	configTask taskConfig
+	// taskMode         taskMode
 	level            slog.Level //控制日志等级
 	maxEventNum      int        //每次epoll/kqueue返回时，一次最多处理多少事件
 	parseInParseLoop *bool      //在解析循环中运行websocket OnOpen, OnMessage, OnClose 回调函数
+
+	selectTask selectTasks // 任务池
 }
 
 type MultiEventLoop struct {
@@ -39,8 +53,7 @@ type MultiEventLoop struct {
 	loops     []*EventLoop
 	parseLoop *taskParse
 
-	runInIo taskIo
-	flag    evFlag // 是否使用io_uring，目前没有使用
+	flag evFlag // 是否使用io_uring，目前没有使用
 
 	stat // 统计信息
 	*slog.Logger
@@ -52,7 +65,7 @@ var (
 	defMaxEventNum   = 256
 	defTaskMin       = 50
 	defTaskMax       = 30000
-	defTaskInitCount = 1000
+	defTaskInitCount = 8
 	defNumLoops      = runtime.NumCPU()
 )
 
@@ -120,11 +133,7 @@ func NewMultiEventLoop(opts ...EvOption) (e *MultiEventLoop, err error) {
 		o(m)
 	}
 	m.initDefaultSetting()
-	// m.startOk = make(chan struct{}, 1)
-	// 设置任务池模式(tps, 或者流量模式)
-	m.configTask.taskMode = m.taskMode
 
-	m.configTask.init()
 	if *m.parseInParseLoop {
 		m.parseLoop = newTaskParse()
 	}

multi_event_loops_option.go

@@ -51,11 +51,11 @@ func WithBusinessGoNum(initCount, min, max int) EvOption {
 }
 
 // 设置business go程池 对流量压测友好的模式
-func WithBusinessGoTrafficMode() EvOption {
-	return func(e *MultiEventLoop) {
-		e.taskMode = trafficMode
-	}
-}
+// func WithBusinessGoTrafficMode() EvOption {
+// 	return func(e *MultiEventLoop) {
+// 		e.taskMode = trafficMode
+// 	}
+// }
 
 // 设置日志级别
 func WithLogLevel(level slog.Level) EvOption {