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 }