手搓流的有序输入输出队列

80 阅读2分钟

流式场景中输入和输出大多情况都要求流是有序的,真实的例子有视频片段、语音片段的输入,这里需要解决的问题和TCP的接收窗口非常类似,我们需要一个队列作为缓冲,这个队列需要有两个能力

  • 提前到达的包需要在队列中等待,直到它前一个包被处理
  • 每一个包的等待都有超时时间,时间到后放弃这个包

实现原理是维护一个指针始终指向下一个可以弹出的位置,并维护一个时钟计时,每当指针所指的位置有包到达或超时,就移动指针

golang的channel让阻塞等待和超时的实现变得非常简单

const (
   maxWaitTime = 2 * time.Second
)

type Index interface {
   Index() int32
}

type RecvQueue[T Index] struct {
   mu        sync.Mutex
   cond      *sync.Cond
   signal    chan struct{}
   messages  map[int32]T
   nextIndex int32

   c    chan T
   stop atomic.Bool
   done atomic.Bool
}

func NewRecvQueue[T Index](startIndex int32) *RecvQueue[T] {
   queue := &RecvQueue[T]{
      messages:  make(map[int32]T),
      nextIndex: startIndex,
      signal:    make(chan struct{}, 100),
      c:         make(chan T, 100),
   }
   queue.cond = sync.NewCond(&queue.mu)
   queue.run()
   return queue
}

func (q *RecvQueue[T]) run() {
   go safely.RunVoid(func() {
      defer close(q.c)
      for !q.stop.Load() {
         msg, ok := q.pop()
         if !ok {
            continue
         }
         q.c <- msg
      }
   })
}
func (q *RecvQueue[T]) Add(msg T) {
   q.mu.Lock()
   defer q.mu.Unlock()
   if q.stop.Load() || q.done.Load() {
      return
   }
   // 过期的消息不处理
   if q.nextIndex > msg.Index() {
      return
   }
   q.messages[msg.Index()] = msg
   // 唤醒在空队列中等待的协程
   if len(q.messages) == 1 {
      q.cond.Signal()
      return
   }
   if next := q.messages[q.nextIndex]; !reflect.ValueOf(&next).Elem().IsZero() {
      q.signal <- struct{}{}
   }
}

func (q *RecvQueue[T]) Pop() <-chan T {
   return q.c
}

// Done 标记队列不再接收消息,等待消息队列处理完毕后关闭
func (q *RecvQueue[T]) Done() {
   defer q.Stop()
   q.done.Store(true)
   tick := time.NewTicker(100 * time.Millisecond)
   defer tick.Stop()
   for {
      select {
      case <-tick.C:
         q.mu.Lock()
         if len(q.messages) == 0 {
            q.mu.Unlock()
            return
         }
         q.mu.Unlock()
      }
   }
}

// Stop 立即停止接收和处理消息
func (q *RecvQueue[T]) Stop() {
   q.stop.Store(true)
   q.cond.Signal()
}

func (q *RecvQueue[T]) pop() (T, bool) {
   // 队列为空时阻塞等待
   q.waitOnEmptyQueue()
   // 可能消息已经到达队列,防止signal通知丢失,因此先尝试获取下一个消息
   if msg, ok := q.getNext(); ok {
      return msg, true
   }
   // 当队列不为空时,等待一段时间后继续尝试获取下一个消息
   timeout := time.After(maxWaitTime)
   for {
      select {
      case <-timeout:
         q.mu.Lock()
         q.nextIndex++
         q.mu.Unlock()
         var msg T
         return msg, false
      case <-q.signal:
         if msg, ok := q.getNext(); ok {
            return msg, ok
         }
      }
   }
}

func (q *RecvQueue[T]) waitOnEmptyQueue() {
   q.mu.Lock()
   defer q.mu.Unlock()
   if len(q.messages) == 0 {
      q.cond.Wait()
   }
}

func (q *RecvQueue[T]) getNext() (T, bool) {
   q.mu.Lock()
   defer q.mu.Unlock()
   msg, ok := q.messages[q.nextIndex]
   if ok {
      delete(q.messages, msg.Index())
      q.nextIndex++
   }
   return msg, ok
}