72 lines
1.6 KiB
Go
72 lines
1.6 KiB
Go
|
package gtq
|
||
|
|
|
||
|
|
import (
|
||
|
|
"gopkg.in/eapache/queue.v1"
|
||
|
|
"sync"
|
||
|
|
)
|
||
|
|
|
||
|
|
// NewQueue initializes and returns an new Queue.
|
||
|
|
func NewQueue() (q *Queue) {
|
||
|
|
return &Queue{
|
||
|
|
queue: queue.New(),
|
||
|
|
Done: NewStat(0),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Queue wraps eapache/queue with locking for thread safety. Also bundles in
|
||
|
|
// Stats, used by schedulers to help pick tasks to execute.
|
||
|
|
type Queue struct {
|
||
|
|
queue *queue.Queue
|
||
|
|
mutex sync.Mutex
|
||
|
|
// Stat which tracks task removal rate from the queue.
|
||
|
|
Done *Stat
|
||
|
|
}
|
||
|
|
|
||
|
|
// Add puts an element on the end of the queue.
|
||
|
|
func (q *Queue) Add(task func()) {
|
||
|
|
q.mutex.Lock()
|
||
|
|
defer q.mutex.Unlock()
|
||
|
|
q.queue.Add(task)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Get returns the element at index i in the queue.
|
||
|
|
// This method accepts both positive and negative index values.
|
||
|
|
// Index 0 refers to the first element, and index -1 refers to the last.
|
||
|
|
func (q *Queue) Get(i int) (task func()) {
|
||
|
|
q.mutex.Lock()
|
||
|
|
defer q.mutex.Unlock()
|
||
|
|
if (q.queue.Length() < 1 || i > q.queue.Length()-1) || (i < -1) {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
return q.queue.Get(i).(func())
|
||
|
|
}
|
||
|
|
|
||
|
|
// Length returns the number of elements currently stored in the queue.
|
||
|
|
func (q *Queue) Length() int {
|
||
|
|
q.mutex.Lock()
|
||
|
|
defer q.mutex.Unlock()
|
||
|
|
return q.queue.Length()
|
||
|
|
}
|
||
|
|
|
||
|
|
// Peek returns the element at the head of the queue.
|
||
|
|
func (q *Queue) Peek() (task func()) {
|
||
|
|
q.mutex.Lock()
|
||
|
|
defer q.mutex.Unlock()
|
||
|
|
if q.Length() < 1 {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
return q.queue.Peek().(func())
|
||
|
|
}
|
||
|
|
|
||
|
|
// Remove removes and returns the element from the front of the queue.
|
||
|
|
func (q *Queue) Remove() (task func()) {
|
||
|
|
q.mutex.Lock()
|
||
|
|
defer q.mutex.Unlock()
|
||
|
|
if q.queue.Length() < 1 {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
task = q.queue.Remove().(func())
|
||
|
|
q.Done.Add(1)
|
||
|
|
return task
|
||
|
|
}
|