流式场景中输入和输出大多情况都要求流是有序的,真实的例子有视频片段、语音片段的输入,这里需要解决的问题和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
}