github.com/go-board/x-go@v0.1.2-0.20220610024734-db1323f6cb15/xcontainer/queue/delay_queue.go (about)

     1  package queue
     2  
     3  import (
     4  	"math"
     5  	"time"
     6  
     7  	"github.com/go-board/x-go/types"
     8  	"github.com/go-board/x-go/xcontainer/priority_queue"
     9  )
    10  
    11  type DelayTask struct {
    12  	data interface{}
    13  	at   time.Time
    14  }
    15  
    16  func (t DelayTask) Compare(o types.Comparable) types.Ordering {
    17  	oo := o.(*DelayTask)
    18  	if t.at.Before(oo.at) {
    19  		return types.OrderingLess
    20  	}
    21  	if t.at.Equal(oo.at) {
    22  		return types.OrderingEqual
    23  	}
    24  	return types.OrderingGreater
    25  }
    26  
    27  func NewTimedTask(v interface{}, at time.Time) *DelayTask {
    28  	return &DelayTask{data: v, at: at}
    29  }
    30  
    31  func NewDelayTask(v interface{}, d time.Duration) *DelayTask {
    32  	return NewTimedTask(v, time.Now().Add(d))
    33  }
    34  
    35  type DelayQueue struct {
    36  	q      *priority_queue.PriorityQueue
    37  	notify chan struct{}
    38  }
    39  
    40  func NewDelayQueue() *DelayQueue {
    41  	return &DelayQueue{
    42  		q:      priority_queue.NewComparablePriorityQueue(false),
    43  		notify: make(chan struct{}, 1),
    44  	}
    45  }
    46  
    47  // Push new data into delay queue.
    48  func (q *DelayQueue) Push(task *DelayTask) {
    49  	q.q.Push(task)
    50  	// this should not block
    51  	select {
    52  	case q.notify <- struct{}{}:
    53  	default:
    54  	}
    55  }
    56  
    57  // Pop try to pop nearest expired data.
    58  func (q *DelayQueue) Pop() (interface{}, bool) {
    59  	val, _, ok := q.popNearest()
    60  	return val, ok
    61  }
    62  
    63  func (q *DelayQueue) popNearest() (interface{}, time.Duration, bool) {
    64  	task := q.q.Pop()
    65  	if task == nil {
    66  		return nil, time.Duration(math.MaxInt64), false
    67  	}
    68  	t := task.(*DelayTask)
    69  	now := time.Now()
    70  	duration := t.at.Sub(now)
    71  	if !t.at.After(now) {
    72  		return t.data, 0, true
    73  	}
    74  	q.q.Push(task)
    75  	return nil, duration, false
    76  }
    77  
    78  // BlockPop must get a data otherwise wait for data ready.
    79  func (q *DelayQueue) BlockPop() interface{} {
    80  	for {
    81  		v, duration, ok := q.popNearest()
    82  		if ok {
    83  			// drain notification if possible, otherwise we may get old notification
    84  			select {
    85  			case <-q.notify:
    86  			default:
    87  			}
    88  			return v
    89  		}
    90  		// block until timeout or new element pushed
    91  		select {
    92  		case <-time.NewTimer(duration).C:
    93  		case <-q.notify:
    94  		}
    95  	}
    96  }