github.com/fumiama/NanoBot@v0.0.0-20231122134259-c22d8183efca/single.go (about)

     1  package nano
     2  
     3  import (
     4  	"github.com/RomiChan/syncx"
     5  )
     6  
     7  // Option 配置项
     8  type Option[K comparable] func(*Single[K])
     9  
    10  // Single 反并发
    11  type Single[K comparable] struct {
    12  	group syncx.Map[K, struct{}]
    13  	key   func(ctx *Ctx) K
    14  	post  func(ctx *Ctx)
    15  }
    16  
    17  // WithKeyFn 指定反并发的 Key
    18  func WithKeyFn[K comparable](fn func(ctx *Ctx) K) Option[K] {
    19  	return func(s *Single[K]) {
    20  		s.key = fn
    21  	}
    22  }
    23  
    24  // WithPostFn 指定反并发拦截后的操作
    25  func WithPostFn[K comparable](fn func(ctx *Ctx)) Option[K] {
    26  	return func(s *Single[K]) {
    27  		s.post = fn
    28  	}
    29  }
    30  
    31  // NewSingle 创建反并发中间件
    32  func NewSingle[K comparable](op ...Option[K]) *Single[K] {
    33  	s := Single[K]{}
    34  	for _, option := range op {
    35  		option(&s)
    36  	}
    37  	return &s
    38  }
    39  
    40  // Apply 为指定 Engine 添加反并发功能
    41  func (s *Single[K]) Apply(engine *Engine) {
    42  	engine.UseMidHandler(func(ctx *Ctx) bool {
    43  		if s.key == nil {
    44  			return true
    45  		}
    46  		key := s.key(ctx)
    47  		if _, ok := s.group.Load(key); ok {
    48  			if s.post != nil {
    49  				defer s.post(ctx)
    50  			}
    51  			return false
    52  		}
    53  		s.group.Store(key, struct{}{})
    54  		ctx.State["__single-key__"] = key
    55  		return true
    56  	})
    57  
    58  	engine.UsePostHandler(func(ctx *Ctx) {
    59  		s.group.Delete(ctx.State["__single-key__"].(K))
    60  	})
    61  }