定时器和时间轮系列(二): 优先级队列和延时队列

// Element for priority queue element
type Element struct {
    // Value for element
    Value interface{}




    // priority of the element, make it private to avoid change from element
    priority int64




    // index of the element in the slice
    index int




    // pq is the refer of priority queue
    pq *priorityQueue
}




// ...




// PriorityQueue for priority queue trait
type PriorityQueue interface {
    // Add element to the PriorityQueue, it will return the element witch been added
    Add(v interface{}, priority int64) *Element




    // Peek return the lowest priority element
    Peek() *Element




    // Pop return the lowest priority element and remove it
    Pop() *Element




    // Remove will remove the element from the priority queue
    Remove(v *Element) error




    // Update the element in the priority queue with the new priority
    Update(v *Element, priority int64) error




    // Size return the element size of queue
    Size() int
}




// elems is an slice of Elements, it implement the heap.Interface
type elems = []*Element




// heapi implement the heap.Interface
type heapi struct {
    pq *priorityQueue
}




// ...




// Push the value at the end of slice, implement heap.Push
func (h *heapi) Push(x interface{}) {
    h.pq.e = append(h.pq.e, x.(*Element))
}




// Pop the value at the last position of slice, implement heap.Pop
func (h *heapi) Pop() interface{} {
    old := h.pq.e
    n := len(old)
    if n == 0 {
        return nil
    }




    // set element to nil for GC
    x := old[n-1]
    old[n-1] = nil
    h.pq.e = old[0 : n-1]
    return x
}




// priorityQueue is a implement by min heap, the 0th element is the lowest value
type priorityQueue struct {
    e elems
    h *heapi
}




// NewPriorityQueue construct a PriorityQueue
func NewPriorityQueue(size int) PriorityQueue {
    pq := &priorityQueue{
        e: make(elems, 0, size),
    }
    pq.h = &heapi{
        pq: pq,
    }
    return pq
}




// Add element to the PriorityQueue, it will return the element witch been added
func (pq *priorityQueue) Add(x interface{}, priority int64) *Element {
    e := &Element{
        Value:    x,
        priority: priority,
        index:    len(pq.e),
        pq:       pq,
    }
    heap.Push(pq.h, e)
    return e
}




// Pop return the lowest priority element and remove it
func (pq *priorityQueue) Pop() *Element {
    if len(pq.e) == 0 {
        return nil
    }




    x := heap.Pop(pq.h)
    e := x.(*Element)
    e.index = -1
    e.pq = nil
    return e
}




// Remove will remove the element from the priority queue
func (pq *priorityQueue) Remove(e *Element) error {
    if e.pq != pq {
        return fmt.Errorf("PriorityQueue.Remove: QueueMatchFailed: Element[%v], Queue[%v]", e.pq, pq)
    }




    if e.index < 0 || e.index >= len(pq.e) {
        return fmt.Errorf("PriorityQueue.Remove: OutOfIndex: Index[%v], Len[%v]", e.index, len(pq.e))
    }
    if e.priority != pq.e[e.index].priority {
        return fmt.Errorf("PriorityQueue.Remove: PriorityMatchFailed: Element[%v], Queue[%v]", e.priority, pq.e[e.index].priority)
    }




    heap.Remove(pq.h, e.index)
    e.index = -1
    e.pq = nil
    return nil
}

// DelayQueue is an blocking queue of *Delay* elements, the element
// can only been taken when its delay has expired. The head of the queue
// is the element whose delay expired most recent in the queue.
type DelayQueue interface {
    // Offer insert the element into the current DelayQueue,
    // if the expiration is blow the current min expiration, the item will
    // been fired first.
    Offer(elem interface{}, expireation int64)




    // Poll starts an infinite loop, it will continually waits for an element to
    // been fired, and send the element to the output Chan.
    Poll(ctx context.Context)




    // Chan return the output chan, when the element is fired the element
    // will send to the channel.
    Chan() <-chan interface{}




    // Size return the element count in the queue
    Size() int
}

// Offer implement the DelayQueue.Offer
func (q *delayQueue) Offer(element interface{}, expireation int64) {
    _push := func() (*Element, int) {
        q.mu.Lock()
        defer q.mu.Unlock()




        e := q.pq.Add(element, expireation)
        return e, e.Index()
    }
    _, index := _push()




    // there is no concurrent protection, EX:
    // 1. goroutine1 add element with expireation 100
    // 2. goroutine2 add element with expireation 50
    // 3. the both goroutine get the element index 0
    // 4. goroutine2 cas the sleeping state to 0, and send the wakeup signal
    // 5. pool wakeup and update the fired point, cas the sleeping state to 1
    // 6. goroutine1 cas the sleeping state to 0, and send the wakeup signal
    // 7. pool wakeup and update the fired point, cas the sleeping state to 1
    // because the pool always update the fired point to the min expireation, so there is no problem(always update to 50)
    if index == 0 {
        // the element is the first element(with the earliest expireation), we
        // need week up the Pool loop to update the fired point
        if atomic.CompareAndSwapInt32(&q.sleeping, 1, 0) {
            // if we change the sleeping state from sleep to weekup success, send the signal to wakepupC
            q.wakeupC <- struct{}{}
        }
    }
}

// Poll implement the DelayQueue.Pool
func (q *delayQueue) Poll(ctx context.Context) {
    defer func() {
        // reset the state to wakeup
        atomic.StoreInt32(&q.sleeping, 0)
    }()




    // an infinite loop
    // 1. wakeup at the min expiration
    // 2. send to the C
    for poll(ctx, q) {
    }
}




var (
    poll func(ctx context.Context, q *delayQueue) bool
)




// the inner implement of poll, split from Poll for test
// return true if been wakeup or fired, false to shutdown the loop
func pollImpl(ctx context.Context, q *delayQueue) bool {
    n := q.T.Now()




    q.mu.Lock()
    item := q.pq.Peek()
    if item == nil || item.Priority() > n {
        // No item left, change the sleeping state to 1
        atomic.StoreInt32(&q.sleeping, 1)
    }
    q.mu.Unlock()




    // we have got the min expiration item, it maybe nil for empty pq
    if item == nil {
        // wait for wakeup (new item Offer into the queue)
        select {
        case <-ctx.Done():
            return false
        case <-q.wakeupC:
            return true
        }
    }




    // have item, wait for the fired point
    delta := item.Priority() - n
    if delta <= 0 {
        // the item need fired, send the value to the output channel
        select {
        case q.C <- item.Value:
            // the element is fired
            q.mu.Lock()
            _ = q.pq.Remove(item)
            q.mu.Unlock()
            return true
        case <-ctx.Done():
            return false
        }
    }




    // the item is pending, wait for fired or new min element add
    select {
    case <-q.wakeupC:
        return true
    case <-time.After(time.Duration(delta) * time.Millisecond):
        // we doesn't fired the item at there, go to next loop and the item will been fired because delta <= 0
        if atomic.SwapInt32(&q.sleeping, 0) == 0 {
            // if the old state is wakeup, the maybe an signal in wakeupC,
            // so we drain it the unblock the caller
            select {
            case <-q.wakeupC:
            default:
            }
        }
        return true
    case <-ctx.Done():
        return false
    }
}