带有队列的 golang 等待组

在 Golang 中，等待组（WaitGroup）是一个非常流行的同步原语，可以用于等待一组 goroutine 完成执行。然而，某些情况下需要在等待组中添加更多的控制和定制功能，例如队列的方式来处理作业。因此，带有队列的 Golang 等待组（WaitGroup with Queue）应运而生。本文将介绍如何实现一个带有队列的 Golang 等待组。

队列的数据结构

在实现带有队列的 Golang 等待组时，需要定义一个队列的数据结构。这里我们采用 slice 作为队列的内部存储结构，并通过标记头部和尾部的索引来实现循环队列。

type queue struct {
    items []interface{}
    head  int
    tail  int
}

等待组的实现

等待组的实现在 Golang 中非常简单。我们只需要调用 sync.WaitGroup 类型的 Add 函数来增加等待组的计数器，Done() 函数来减少计数器，Wait() 函数来阻塞等待所有 goroutine 执行完成。

wg := &sync.WaitGroup{}
// 增加计数器
wg.Add(1)
// 减少计数器
wg.Done()
// 阻塞等待
wg.Wait()

带有队列的等待组的实现

队列的操作

添加任务时，先将任务放入队列中，并把等待组计数器+1。

func (q *queue) enqueue(item interface{}, wg *sync.WaitGroup) {
    q.items[q.tail] = item
    q.tail = (q.tail + 1) % len(q.items)
    wg.Add(1)
}

执行任务时，从队列中提取任务并执行，执行完后将等待组计数器-1。

func (q *queue) dequeue(fn func(item interface{}), wg *sync.WaitGroup) {
    for q.head != q.tail {
        item := q.items[q.head]
        q.head = (q.head + 1) % len(q.items)
        fn(item)
        wg.Done()
    }
}

带有队列的等待组

在实现等待组时，创建一个 queue 结构体用于存储任务队列。添加任务时，将任务插入队列，并将等待组计数器+1。当所有任务都插入到队列中后，执行 dequeue 函数从队列中获取任务并执行。执行完所有任务后，等待组结束。

type WaitGroupWithQueue struct {
    wg    *sync.WaitGroup
    queue *queue
}

func (wgq *WaitGroupWithQueue) Add(fn func(item interface{}), items ...interface{}) {
    for _, item := range items {
        wgq.queue.enqueue(item, wgq.wg)
    }
}

func (wgq *WaitGroupWithQueue) Wait(fn func(item interface{})) {
    go wgq.queue.dequeue(fn, wgq.wg)
    wgq.wg.Wait()
}

func NewWaitGroupWithQueue(size int) *WaitGroupWithQueue {
    q := queue{items: make([]interface{}, size)}
    return &WaitGroupWithQueue{wg: &sync.WaitGroup{}, queue: &q}
}

使用带有队列的等待组

在使用带有队列的等待组时，首先需要创建 WaitGroupWithQueue 实例，并指定队列的大小。然后，通过 Add 函数添加任务，通过 Wait 函数执行任务。Add 函数支持添加多个任务，Wait 函数支持通过回调函数来处理执行结果。

func main() {
    wgq := NewWaitGroupWithQueue(10)
    for i := 0; i < 10; i++ {
        wgq.Add(func(item interface{}) {
            // 执行任务，例如请求数据接口
        }, i)
    }

    wgq.Wait(func(item interface{}) {
        // 处理任务执行结果，例如打印返回结果
    })
}

这样就完成了一个带有队列的 Golang 等待组的实现。可以通过此结构来处理一系列需要异步处理的任务。