timer.go

// This code is based on work from the goim project
//
// Source: https://github.com/Terry-Mao/goim/blob/master/pkg/time/timer.go
// Original author: Terry.Mao
// Licensed under the MIT License: https://opensource.org/licenses/MIT
package xcron

import (
	"sync"
	"time"
)

const (
	timerFormat      = "2006-01-02 15:04:05"
	infiniteDuration = time.Duration(1<<63 - 1)
)

// TimerData timer data.
type TimerData struct {
	expire time.Time
	index  int
	next   *TimerData
	job    Job
}

// Delay delay duration.
func (td *TimerData) Delay() time.Duration {
	return time.Until(td.expire)
}

// ExpireString expire string.
func (td *TimerData) ExpireString() string {
	return td.expire.Format(timerFormat)
}

// Timer timer.
type Timer struct {
	lock   sync.Mutex
	free   *TimerData
	timers []*TimerData
	signal *time.Timer
	num    int
	ch     chan Job
}

// NewTimer new a timer.
// A heap must be initialized before any of the heap operations
// can be used. Init is idempotent with respect to the heap invariants
// and may be called whenever the heap invariants may have been invalidated.
// Its complexity is O(n) where n = h.Len().
func NewTimer(num int, ch chan Job) (t *Timer) {
	t = new(Timer)
	t.ch = ch
	t.init(num)
	return t
}

// Init init the timer.
func (t *Timer) Init(num int, ch chan Job) {
	t.ch = ch
	t.init(num)
}

func (t *Timer) init(num int) {
	t.signal = time.NewTimer(infiniteDuration)
	t.timers = make([]*TimerData, 0, num)
	t.num = num
	t.grow()
	go t.start()
}

func (t *Timer) grow() {
	var (
		i   int
		td  *TimerData
		tds = make([]TimerData, t.num)
	)
	t.free = &(tds[0])
	td = t.free
	for i = 1; i < t.num; i++ {
		td.next = &(tds[i])
		td = td.next
	}
	td.next = nil
}

// get get a free timer data.
func (t *Timer) get() (td *TimerData) {
	if td = t.free; td == nil {
		t.grow()
		td = t.free
	}
	t.free = td.next
	return
}

// put put back a timer data.
func (t *Timer) put(td *TimerData) {
	td.job = nil
	td.next = t.free
	t.free = td
}

// Add add the element x onto the heap. The complexity is
// O(log(n)) where n = h.Len().
func (t *Timer) Add(expire time.Time, job Job) (td *TimerData) {
	t.lock.Lock()
	td = t.get()
	td.expire = expire
	td.job = job
	t.add(td)
	t.lock.Unlock()
	return
}

// Del removes the element at index i from the heap.
// The complexity is O(log(n)) where n = h.Len().
func (t *Timer) Del(td *TimerData) {
	t.lock.Lock()
	t.del(td)
	t.put(td)
	t.lock.Unlock()
}

// Push pushes the element x onto the heap. The complexity is
// O(log(n)) where n = h.Len().
func (t *Timer) add(td *TimerData) {
	var d time.Duration
	td.index = len(t.timers)
	// add to the minheap last node
	t.timers = append(t.timers, td)
	t.up(td.index)
	if td.index == 0 {
		// if first node, signal start goroutine
		d = td.Delay()
		t.signal.Reset(d)
	}
}

func (t *Timer) del(td *TimerData) {
	var (
		i    = td.index
		last = len(t.timers) - 1
	)
	if i < 0 || i > last || t.timers[i] != td {
		// already remove, usually by expire
		return
	}
	if i != last {
		t.swap(i, last)
		t.down(i, last)
		t.up(i)
	}
	// remove item is the last node
	t.timers[last].index = -1 // for safety
	t.timers = t.timers[:last]
}

// Set update timer data.
func (t *Timer) Set(td *TimerData, expire time.Time) {
	t.lock.Lock()
	t.del(td)
	td.expire = expire
	t.add(td)
	t.lock.Unlock()
}

// start start the timer.
func (t *Timer) start() {
	for {
		t.expire()
		<-t.signal.C
	}
}

// expire removes the minimum element (according to Less) from the heap.
// The complexity is O(log(n)) where n = max.
// It is equivalent to Del(0).
func (t *Timer) expire() {
	var (
		job Job
		td  *TimerData
		d   time.Duration
	)
	t.lock.Lock()
	for {
		if len(t.timers) == 0 {
			d = infiniteDuration
			break
		}
		td = t.timers[0]
		if d = td.Delay(); d > 0 {
			break
		}
		job = td.job
		// let caller put back
		t.del(td)
		t.lock.Unlock()
		t.ch <- job
		t.lock.Lock()
	}
	t.signal.Reset(d)
	t.lock.Unlock()
}

func (t *Timer) up(j int) {
	for {
		i := (j - 1) / 2 // parent
		if i >= j || !t.less(j, i) {
			break
		}
		t.swap(i, j)
		j = i
	}
}

func (t *Timer) down(i, n int) {
	for {
		j1 := 2*i + 1
		if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
			break
		}
		j := j1 // left child
		if j2 := j1 + 1; j2 < n && !t.less(j1, j2) {
			j = j2 // = 2*i + 2  // right child
		}
		if !t.less(j, i) {
			break
		}
		t.swap(i, j)
		i = j
	}
}

func (t *Timer) less(i, j int) bool {
	return t.timers[i].expire.Before(t.timers[j].expire)
}

func (t *Timer) swap(i, j int) {
	t.timers[i], t.timers[j] = t.timers[j], t.timers[i]
	t.timers[i].index = i
	t.timers[j].index = j
}