Go 实战单队列到优先级队列实现图文示例

type JobQueue struct {
    mu sync.Mutex //队列的操作需要并发安全
    jobList *list.List //List是golang库的双向队列实现，每个元素都是一个job
    noticeChan chan struct{} //入队一个job就往该channel中放入一个消息，以供消费者消费
}

/**
 * 队列的Push操作
 */
func (queue *JobQueue) PushJob(job Job) {
    queue.jobList.PushBack(job) //将job加到队尾
    queue.noticeChan <- struct{}{}
}

queue.noticeChan <- struct{}{}

/**
 * 弹出队列的第一个元素
 */
func (queue *JobQueue) PopJob() Job {
    queue.mu.Lock()
    defer queue.mu.Unlock()
    /**
     * 说明在队列中没有元素了
     */
    if queue.jobList.Len() == 0 {
        return nil
    }
    elements := queue.jobList.Front() //获取队里的第一个元素
    return queue.jobList.Remove(elements).(Job) //将元素从队列中移除并返回
}

func (queue *JobQueue) WaitJob() <-chan struct{} {
    return queue.noticeChan
}

type BaseJob struct {
    Err error
    DoneChan chan struct{} //当作业完成时，或者作业被取消时，通知调用者
    Ctx context.Context
    cancelFunc context.CancelFunc
}

/**
 * 作业执行完毕，关闭DoneChan，所有监听DoneChan的接收者都能收到关闭的信号
 */
func (job *BaseJob) Done() {
    close(job.DoneChan)
}

/**
 * 等待job执行完成
 */
func (job *BaseJob) WaitDone()  {
    select {
    case <-job.DoneChan:
        return
    }
}

type SquareJob struct {
    *BaseJob
    x int
}

func (s *SquareJob) Execute() error {
    result := s.x * s.x
    fmt.Println("the result is ", result)
    return nil
}

type WorkerManager struct {
    queue *JobQueue
    closeChan chan struct{}
}

 func (m *WorkerManager) StartWork() error {
    fmt.Println("Start to Work")
    for {
        select {
            case <-m.closeChan:
                return nil
            case <-m.queue.noticeChan:
                job := m.queue.PopJob()
                m.ConsumeJob(job)
        }
    }
    return nil
}
func (m *WorkerManager) ConsumeJob(job Job) {
    defer func() {
        job.Done()
    }()
    job.Execute()
}

func main() {
    //初始化一个队列
    queue := &JobQueue{
        jobList: list.New(),
        noticeChan: make(chan struct{}, 10),
    }
    //初始化一个消费worker
    workerManger := NewWorkerManager(queue)
    // worker开始监听队列
    go workerManger.StartWork()
    // 构造SquareJob
    job := &SquareJob{
        BaseJob: &BaseJob{
            DoneChan: make(chan struct{}, 1),
        },
        x: 5,
    }
    //压入队列尾部
    queue.PushJob(job)
    //等待job执行完成
    job.WaitDone()
    print("The End")
}

type BaseJob struct {
    Err error
    DoneChan chan struct{} //当作业完成时，或者作业被取消时，通知调用者
    Ctx context.Context
    cancelFunc context.CancelFunc
    priority int //工作单元的优先级
}

var queues = make([]*JobQueue, 10, 100)

var priorityIdx map[int][int]//该map的key是优先级，value代表的是queues切片的索引

type PriorityQueue struct {
    mu sync.Mutex
    noticeChan chan struct{}
    queues []*JobQueue
    priorityIdx map[int]int
}
//原来的JobQueue会变成如下这样：
type JobQueue struct {
    priority int //代表该队列是哪种优先级的队列
    jobList *list.List //List是golang库的双向队列实现，每个元素都是一个job
}

func (priorityQueue *PriorityQueue) Push(job Job) {
    priorityQueue.mu.Lock()
    defer priorityQueue.mu.Unlock()
    //先根据job的优先级找要入队的队列
    var idx int
    var ok bool
    //从优先级-切片索引的map中查找该优先级的队列是否存在
    if idx, ok = priorityQueue.priorityIdx[job.Priority()]; !ok {
        //如果不存在该优先级的队列，则需要初始化一个队列，并返回该队列在切片中的索引位置
        idx = priorityQueue.addPriorityQueue(job.Priority)
    }
    //根据获取到的切片索引idx，找到具体的队列
    queue := priority.queues[idx]
    //将job推送到队列的队尾
    queue.JobList.PushBack(job)
    //队列job个数+1
    priorityQueue.Size++
    //如果队列job个数超过队列的最大容量，则从优先级最低的队列中移除工作单元
    if priorityQueue.size > priorityQueue.capacity {
        priorityQueue.RemoveLeastPriorityJob()
    }else {
        //通知新进来一个job
        priorityQueue.noticeChan <- struct{}{}
    }
}

func (priorityQueue *PriorityQueue) addPriorityQueue(priority int) int {
    n := len(priorityQueue.queues)
    //通过二分查找找到priority应插入的切片索引
    pos := sort.Search(n, func(i int) bool {
        return priority < priorityQueue.priority
    })
    //更新映射表中优先级和切片索引的对应关系
    for i := pos; i < n; i++ {
        priorityQueue.priorityIdx[priorityQueue.queues[i].priority] = i + 1
    }
    tail := make([]*jobQueue, n-pos)
    copy(tail, priorityQueue.queues[pos:])
    //初始化一个新的优先级队列，并将该元素放到切片的pos位置中
    priorityQueue.queues = append(priorityQueue.queues[0:pos], newJobQueue(priority))
    //将高于priority优先级的元素也拼接到切片后面
    priorityQueue.queues = append(priorityQueue.queues, tail...) 
    return pos
}

func main() {
    //初始化一个队列
    queue := &PriorityQueue{
        noticeChan: make(chan struct{}, cap),
        capacity: cap,
        priorityIdx: make(map[int]int),
        size: 0,
    }
    //初始化一个消费worker
    workerManger := NewWorkerManager(queue)
    // worker开始监听队列
    go workerManger.StartWork()
    // 构造SquareJob
    job := &SquareJob{
        BaseJob: &BaseJob{
            DoneChan: make(chan struct{}, 1),
        },
        x: 5,
        priority: 10,
    }
    //压入队列尾部
    queue.PushJob(job)
    //等待job执行完成
    job.WaitDone()
    print("The End")
}

type WorkerManager struct {
    queue *PriorityQueue
    closeChans []chan struct{}
}

 func (m *WorkerManager) StartWork(n int) error {
    fmt.Println("Start to Work")
    for i := 0; i < n; i++ {
        m.createWorker();
    }
    return nil
}
func (m *WorkerManager) createWorker() {
    closeChan := make(chan struct{})
    //每个协程，就是一个worker
    go func(closeChan chan struct{}) {
        var job Job
        for {
                select {
                    case <-m.closeChan:
                        return nil
                    case <-m.queue.noticeChan:
                        job := m.queue.PopJob()
                        m.ConsumeJob(job)
                }   
        }
    }(closeChan)
    m.closeChanMu.Lock()
    defer m.closeChanMu.Unlock()
    m.closeChans = append(m.closeChans, closeChan)
    return nil
}
func (m *WorkerManager) ConsumeJob(job Job) {
    defer func() {
        job.Done()
    }()
    job.Execute()
}