github.com/binbinly/pkg@v0.0.11-0.20240321014439-f4fbf666eb0f/transport/ws/router.go (about)

     1  package ws
     2  
     3  import (
     4  	"math"
     5  	"sync"
     6  	"time"
     7  )
     8  
     9  const abortIndex int8 = math.MaxInt8 / 2 // 路由中间件嵌套最大值
    10  
    11  // HandlerFunc 路由处理方法
    12  type HandlerFunc func(c *Context)
    13  
    14  // HandlerChain 执行链
    15  type HandlerChain []HandlerFunc
    16  
    17  // RouterGroup 路由组
    18  type RouterGroup struct {
    19  	Handlers HandlerChain
    20  	engine   *Engine
    21  }
    22  
    23  // Engine 路由引擎
    24  type Engine struct {
    25  	RouterGroup
    26  
    27  	tree map[string]HandlerChain
    28  	pool sync.Pool
    29  }
    30  
    31  // Context 上下文对象
    32  type Context struct {
    33  	Req      *Request
    34  	handlers HandlerChain
    35  	index    int8
    36  }
    37  
    38  // NewEngine get a new engine(tcp.Handler)
    39  func NewEngine() *Engine {
    40  	engine := &Engine{
    41  		RouterGroup: RouterGroup{
    42  			Handlers: nil,
    43  		},
    44  		tree: make(map[string]HandlerChain),
    45  	}
    46  	engine.RouterGroup.engine = engine
    47  	engine.pool.New = func() any {
    48  		return engine.allocateContext()
    49  	}
    50  	return engine
    51  }
    52  
    53  // Next 执行下一个
    54  func (c *Context) Next() {
    55  	c.index++
    56  	// 调用 c.Abort() 的时候不会往后执行
    57  	for c.index < int8(len(c.handlers)) {
    58  		c.handlers[c.index](c)
    59  		c.index++
    60  	}
    61  }
    62  
    63  // Abort 退出路由
    64  func (c *Context) Abort() {
    65  	c.index = abortIndex
    66  }
    67  
    68  // Reset 重置路由
    69  func (c *Context) Reset() {
    70  	c.handlers = nil
    71  	c.index = -1
    72  }
    73  
    74  // Done always returns nil (chan which will wait forever),
    75  // if you want to abort your work when the connection was closed
    76  // you should use Request.Context().Done() instead.
    77  func (c *Context) Done() <-chan struct{} {
    78  	return nil
    79  }
    80  
    81  // Err always returns nil, maybe you want to use Request.Context().Err() instead.
    82  func (c *Context) Err() error {
    83  	return nil
    84  }
    85  
    86  // Deadline always returns that there is no deadline (ok==false),
    87  // maybe you want to use Request.Context().Deadline() instead.
    88  func (c *Context) Deadline() (deadline time.Time, ok bool) {
    89  	return
    90  }
    91  
    92  // Value returns the value associated with this context for key, or nil
    93  // if no value is associated with key. Successive calls to Value with
    94  // the same key returns the same result.
    95  func (c *Context) Value(any) any {
    96  	return c.Req
    97  }
    98  
    99  // allocateContext
   100  func (e *Engine) allocateContext() *Context {
   101  	return &Context{}
   102  }
   103  
   104  // Start 执行入口
   105  func (e *Engine) Start(req *Request) {
   106  	c := e.pool.Get().(*Context)
   107  	c.Req = req
   108  	c.Reset()
   109  	handlers := e.getHandlers(c.Req)
   110  	if handlers != nil {
   111  		c.handlers = handlers
   112  		c.Next()
   113  	}
   114  }
   115  
   116  func (e *Engine) addRoute(event string, handlers HandlerChain) {
   117  	e.tree[event] = handlers
   118  }
   119  
   120  func (e *Engine) getHandlers(req *Request) HandlerChain {
   121  	handlers, ok := e.tree[req.Event()]
   122  	if !ok {
   123  		return nil
   124  	}
   125  	return handlers
   126  }
   127  
   128  // Use set common middleware
   129  func (e *Engine) Use(middleware ...HandlerFunc) {
   130  	e.RouterGroup.Use(middleware...)
   131  }
   132  
   133  // Use 加载中间件
   134  func (g *RouterGroup) Use(middleware ...HandlerFunc) {
   135  	g.Handlers = append(g.Handlers, middleware...)
   136  }
   137  
   138  // AddRoute specific middleware
   139  func (g *RouterGroup) AddRoute(event string, handlers ...HandlerFunc) {
   140  	handlers = g.mergeHandlers(handlers)
   141  	g.engine.addRoute(event, handlers)
   142  }
   143  
   144  // merge specific and common middleware
   145  func (g *RouterGroup) mergeHandlers(handlers HandlerChain) HandlerChain {
   146  	finalSize := len(g.Handlers) + len(handlers)
   147  	mergedHandlers := make(HandlerChain, finalSize)
   148  	copy(mergedHandlers, g.Handlers)
   149  	copy(mergedHandlers[len(g.Handlers):], handlers)
   150  	return mergedHandlers
   151  }