package queue import ( "context" "math" "sync" "time" ) const minTimerDelay = time.Millisecond const maxTimerDelay = time.Nanosecond * math.MaxInt64 type DelayQueue[T Delayable] struct { enqueuedSignal chan struct{} dequeuedSignal chan struct{} transferChan chan T singleDequeue chan struct{} mutex sync.Mutex priorityQueue *PriorityQueue[T] elmMap map[uint64]T zero T } type Delayable interface { GetDeadline() time.Time GetId() uint64 } func NewDelayQueue[T Delayable](cap int) *DelayQueue[T] { singleDequeue := make(chan struct{}, 1) singleDequeue <- struct{}{} return &DelayQueue[T]{ enqueuedSignal: make(chan struct{}), dequeuedSignal: make(chan struct{}), transferChan: make(chan T), singleDequeue: singleDequeue, elmMap: make(map[uint64]T, 64), priorityQueue: NewPriorityQueue[T](cap, func(src T, dst T) bool { return src.GetDeadline().Before(dst.GetDeadline()) }), } } func (s *DelayQueue[T]) TryDequeue() (T, bool) { s.mutex.Lock() defer s.mutex.Unlock() if elm, ok := s.priorityQueue.Peek(0); ok { delay := elm.GetDeadline().Sub(time.Now()) if delay < minTimerDelay { // 无需延迟,头部元素出队后直接返回 _, _ = s.dequeue() return elm, true } } return s.zero, false } func (s *DelayQueue[T]) TryEnqueue(val T) bool { s.mutex.Lock() defer s.mutex.Unlock() return s.enqueue(val) } func (s *DelayQueue[T]) Dequeue(ctx context.Context) (T, bool) { // 出队锁:避免因重复获取队列头部同一元素降低性能 select { case <-s.singleDequeue: defer func() { s.singleDequeue <- struct{}{} }() case <-ctx.Done(): return s.zero, false } for { // 全局锁:避免入队和出队信号的重置与激活出现并发问题 s.mutex.Lock() if ctx.Err() != nil { s.mutex.Unlock() return s.zero, false } // 接收直接转发的不需要延迟的新元素 select { case elm := <-s.transferChan: s.mutex.Unlock() return elm, true default: } // 延迟时间缺省值为 maxTimerDelay, 表示队列为空 delay := maxTimerDelay if elm, ok := s.priorityQueue.Peek(0); ok { now := time.Now() delay = elm.GetDeadline().Sub(now) if delay < minTimerDelay { // 无需延迟,头部元素出队后直接返回 _, _ = s.dequeue() s.mutex.Unlock() return elm, ok } } // 重置入队信号,避免历史信号干扰 select { case <-s.enqueuedSignal: default: } s.mutex.Unlock() if delay == maxTimerDelay { // 队列为空, 等待新元素 select { case elm := <-s.transferChan: return elm, true case <-s.enqueuedSignal: continue case <-ctx.Done(): return s.zero, false } } else if delay >= minTimerDelay { // 等待时间到期或新元素加入 timer := time.NewTimer(delay) select { case elm := <-s.transferChan: return elm, true case <-s.enqueuedSignal: continue case <-timer.C: continue case <-ctx.Done(): return s.zero, false } } } } func (s *DelayQueue[T]) dequeue() (T, bool) { elm, ok := s.priorityQueue.Dequeue() if !ok { return s.zero, false } delete(s.elmMap, elm.GetId()) select { case s.dequeuedSignal <- struct{}{}: default: } return elm, true } func (s *DelayQueue[T]) enqueue(val T) bool { if ok := s.priorityQueue.Enqueue(val); !ok { return false } s.elmMap[val.GetId()] = val select { case s.enqueuedSignal <- struct{}{}: default: } return true } func (s *DelayQueue[T]) Enqueue(ctx context.Context, val T) bool { for { // 全局锁:避免入队和出队信号的重置与激活出现并发问题 s.mutex.Lock() if _, ok := s.elmMap[val.GetId()]; ok { s.mutex.Unlock() return false } if ctx.Err() != nil { s.mutex.Unlock() return false } // 如果队列未满,入队后直接返回 if !s.priorityQueue.IsFull() { s.enqueue(val) s.mutex.Unlock() return true } // 队列已满,重置出队信号,避免受到历史信号影响 select { case <-s.dequeuedSignal: default: } s.mutex.Unlock() if delay := val.GetDeadline().Sub(time.Now()); delay >= minTimerDelay { // 新元素需要延迟,等待退出信号、出队信号和到期信号 timer := time.NewTimer(delay) select { case <-s.dequeuedSignal: // 收到出队信号,从头开始尝试入队 continue case <-timer.C: // 新元素不再需要延迟 case <-ctx.Done(): return false } } else { // 新元素不需要延迟,等待转发成功信号、出队信号和退出信号 select { case s.transferChan <- val: // 新元素转发成功,直接返回(避免队列满且元素未到期导致新元素长时间无法入队) return true case <-s.dequeuedSignal: // 收到出队信号,从头开始尝试入队 continue case <-ctx.Done(): return false } } } } func (s *DelayQueue[T]) Remove(_ context.Context, elmId uint64) (T, bool) { s.mutex.Lock() defer s.mutex.Unlock() if _, ok := s.elmMap[elmId]; ok { delete(s.elmMap, elmId) return s.priorityQueue.Remove(s.index(elmId)) } return s.zero, false } func (s *DelayQueue[T]) index(elmId uint64) int { for i := 0; i < s.priorityQueue.Len(); i++ { elm, ok := s.priorityQueue.Peek(i) if !ok { continue } if elmId == elm.GetId() { return i } } return -1 }