github.com/ngicks/gokugen@v0.0.5/scheduler/queue.go (about)

     1  package scheduler
     2  
     3  import (
     4  	"errors"
     5  	"sync"
     6  
     7  	"github.com/ngicks/gokugen/heap"
     8  )
     9  
    10  var (
    11  	ErrMax = errors.New("heap size exceeds max limit")
    12  )
    13  
    14  // TaskQueue is priority queue.
    15  // Task with least scheduled time has most priority.
    16  type TaskQueue interface {
    17  	// Len returns number of elements stored in underlying heap.
    18  	// The complexity of typical implementation is O(1)
    19  	Len() int
    20  	// Push pushes *Task into heap.
    21  	// Push may return ErrMax if max is set, namely max != 0.
    22  	// The complexity is O(log n) where n = t.Len()
    23  	Push(task *Task) error
    24  	// Peek peeks a min element in heap without any change.
    25  	// The complexity of typical implementation is O(1)
    26  	Peek() *Task
    27  	// Pop pops a min element from heap if given `checker` function is nil or `checker` returns true.
    28  	// `checker` will be called with a min element.
    29  	// Argument to `checker` would be nil if heap is empty.
    30  	// Pop returns nil when heap is empty, given `checker` function is nil or `checker` returns false.
    31  	Pop(checker func(*Task) bool) *Task
    32  	// Exclude excludes elements from underlying heap if filter returns true.
    33  	// Heap is scanned in given range, namely [start,end).
    34  	// Wider range is, longer it will hold lock. Range size must be chosen wisely.
    35  	// If filter func is nil, it is no-op.
    36  	// The complexity of execlusion is O(n) where n = end-start,
    37  	// and restoration of heap invariants is O(m) where m = new len.
    38  	// Number of exclued elements has significant performance effect much more than restoration of invariants.
    39  	Exclude(filter func(ent *Task) bool, start, end int) (removed []*Task)
    40  }
    41  
    42  type internalHeap interface {
    43  	Len() int
    44  	Peek() *Task
    45  	Pop() *Task
    46  	Push(ele *Task)
    47  	Remove(i int) *Task
    48  	Fix(i int)
    49  	Exclude(filter func(ent *Task) bool, start, end int) (removed []*Task)
    50  }
    51  
    52  func isUnlimited(u uint) bool {
    53  	return u == 0
    54  }
    55  
    56  func less(i, j *Task) bool {
    57  	return i.scheduledTime.Before(j.scheduledTime)
    58  }
    59  
    60  type SyncTaskQueue struct {
    61  	mu   sync.Mutex
    62  	max  uint
    63  	heap internalHeap
    64  }
    65  
    66  func NewSyncQueue(max uint) *SyncTaskQueue {
    67  	return &SyncTaskQueue{heap: heap.NewHeap(less), max: max}
    68  }
    69  
    70  func (q *SyncTaskQueue) Len() int {
    71  	q.mu.Lock()
    72  	defer q.mu.Unlock()
    73  	return q.heap.Len()
    74  }
    75  
    76  func (q *SyncTaskQueue) Push(task *Task) error {
    77  	q.mu.Lock()
    78  	defer q.mu.Unlock()
    79  
    80  	if !isUnlimited(q.max) && uint(q.heap.Len()) >= q.max {
    81  		return ErrMax
    82  	}
    83  
    84  	q.heap.Push(task)
    85  	return nil
    86  }
    87  
    88  func (q *SyncTaskQueue) Peek() *Task {
    89  	q.mu.Lock()
    90  	defer q.mu.Unlock()
    91  	return q.heap.Peek()
    92  }
    93  
    94  func (q *SyncTaskQueue) Pop(checker func(*Task) bool) *Task {
    95  	q.mu.Lock()
    96  	defer q.mu.Unlock()
    97  	if checker != nil && checker(q.heap.Peek()) {
    98  		return q.heap.Pop()
    99  	} else {
   100  		return nil
   101  	}
   102  }
   103  
   104  func (q *SyncTaskQueue) Exclude(filter func(ent *Task) bool, start, end int) (removed []*Task) {
   105  	q.mu.Lock()
   106  	defer q.mu.Unlock()
   107  	return q.heap.Exclude(filter, start, end)
   108  }
   109  
   110  type UnsafeTaskQueue struct {
   111  	max  uint
   112  	heap internalHeap
   113  }
   114  
   115  func NewUnsafeQueue(max uint) *UnsafeTaskQueue {
   116  	return &UnsafeTaskQueue{heap: heap.NewHeap(less), max: max}
   117  }
   118  
   119  func (q *UnsafeTaskQueue) Push(task *Task) error {
   120  	if !isUnlimited(q.max) && uint(q.heap.Len()) >= q.max {
   121  		return ErrMax
   122  	}
   123  	q.heap.Push(task)
   124  	return nil
   125  }
   126  
   127  func (q *UnsafeTaskQueue) Len() int {
   128  	return q.heap.Len()
   129  }
   130  
   131  func (q *UnsafeTaskQueue) Peek() *Task {
   132  	return q.heap.Peek()
   133  }
   134  
   135  func (q *UnsafeTaskQueue) Pop(checker func(*Task) bool) *Task {
   136  	if checker == nil || checker(q.heap.Peek()) {
   137  		return q.heap.Pop()
   138  	} else {
   139  		return nil
   140  	}
   141  }
   142  
   143  func (q *UnsafeTaskQueue) Exclude(filter func(ent *Task) bool, start, end int) (removed []*Task) {
   144  	return q.heap.Exclude(filter, start, end)
   145  }