startimer/timer.go

package startimer

import (
	"context"
	"encoding/json"
	"errors"
	"sort"
	"sync"
	"time"
)

func NewTimer(baseDate time.Time, opts ...TimerOptions) StarTimer {
	timer := StarTimer{
		base:     baseDate,
		nextDate: time.Time{},
		mu:       new(sync.RWMutex),
	}
	for _, opt := range opts {
		var op TimerOption
		opt(&op)
		switch op.idx {
		case 3:
			timer.repeat = append(timer.repeat, &Repeats{
				Repeat: []Repeat{
					{
						Unit:  STAR_YEAR,
						Value: uint32(op.date.Year()),
					},
					{
						Unit:  STAR_MONTH,
						Value: uint32(op.date.Month()),
					},
					{
						Unit:  STAR_DAY,
						Value: uint32(op.date.Day()),
					},
					{
						Unit:  STAR_HOUR,
						Value: uint32(op.date.Hour()),
					},
					{
						Unit:  STAR_MINUTE,
						Value: uint32(op.date.Minute()),
					},
					{
						Unit:  STAR_SECOND,
						Value: uint32(op.date.Second()),
					},
				},
				Every: false,
			})
		case 4:
			timer.repeat = append(timer.repeat, &Repeats{
				Repeat: []Repeat{op.repeat},
				Every:  true,
			})
		}

		if op.repeats != nil {
			timer.repeat = append(timer.repeat, op.repeats)
		}
		if op.tasks != nil {
			timer.tasks = append(timer.tasks, op.tasks)
		}
		if op.runLimit > 0 {
			timer.runLimit = op.runLimit
		}
	}
	return timer
}
func (t *StarTimer) IsRunning() bool {
	t.mu.RLock()
	defer t.mu.RUnlock()
	return t.running
}

func (t *StarTimer) SetRunCountLimit(c int) {
	t.runLimit = c
}

func (t *StarTimer) RunCountLimit() int {
	return t.runLimit
}

func (t *StarTimer) AddTask(task func()) {
	t.tasks = append(t.tasks, task)
}

func (t *StarTimer) SetTasks(tasks []func()) {
	t.tasks = tasks
}

func (t *StarTimer) Stop() error {
	t.mu.Lock()
	defer t.mu.Unlock()
	if !t.running {
		return nil
	}
	t.running = false
	t.stopFn()
	if t.timer != nil {
		t.timer.Stop()
	}
	return nil
}

func (t *StarTimer) NextTimer() time.Time {
	return t.parseNextDate(time.Now(), true)
}

func (t *StarTimer) NextTimerAfterDate(date time.Time) time.Time {
	return t.parseNextDate(date, true)
}

func (t *StarTimer) ExportRepeats() (string, error) {
	var rep []Repeats
	for _, r := range t.repeat {
		rep = append(rep, *r)
	}
	data, err := json.Marshal(rep)
	if err != nil {
		return "", err
	}
	return string(data), nil
}

func (t *StarTimer) ImportRepeats(r string) error {
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.running {
		return errors.New("coonot import repeats to already running timer")
	}
	var rep []Repeats
	err := json.Unmarshal([]byte(r), &rep)
	if err != nil {
		return err
	}
	t.repeat = make([]*Repeats, 0, len(rep))
	for _, v := range rep {
		t.repeat = append(t.repeat, &v)
	}
	return nil
}

func (t *StarTimer) BaseDate() time.Time {
	return t.base
}
func (t *StarTimer) SetBaseDate(date time.Time) {
	t.base = date
}

func (t *StarTimer) ResetWithRepeat(base time.Time, repeat []*Repeats) error {
	t.Stop()
	t.base = base
	t.repeat = repeat
	return t.Run()
}

func (t *StarTimer) Repeats() []*Repeats {
	return t.repeat
}

func (t *StarTimer) Run() error {
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.running {
		return nil
	}
	t.nextDate = t.parseNextDate(time.Now(), false)
	if t.nextDate.Before(time.Now()) {
		return errors.New("Invalid Timer Options,Please Check")
	}
	t.running = true
	t.stopCtx, t.stopFn = context.WithCancel(context.Background())
	go func() {
		for {
			if t.runLimit > 0 && t.runCount >= t.runLimit {
				t.Stop()
			}
			now := time.Now()
			t.mu.Lock()
			t.timer = time.NewTimer(t.nextDate.Sub(now))
			t.mu.Unlock()
			select {
			case <-t.timer.C:
				t.runCount++
				t.nextDate = t.parseNextDate(t.nextDate, false)
				if t.nextDate.Before(now) || t.runLimit > 0 && t.runCount >= t.runLimit {
					t.Stop()
				}
				for _, fn := range t.tasks {
					go fn()
				}
				if !t.running {
					return
				}
			case <-t.stopCtx.Done():
				return
			}
		}
	}()
	return nil
}

func (t *StarTimer) parseNextDate(base time.Time, isMock bool) time.Time {
	if len(t.repeat) == 0 {
		return time.Time{}
	}
	var dates []time.Time
	for _, d := range t.repeat {
		if d == nil {
			continue
		}
		if d.Every {
			dates = append(dates, t.parseEveryNextDate(base, d, isMock)...)
		} else {
			dates = append(dates, t.parseStaticNextDate(base, d, false))
		}
	}
	sort.SliceStable(dates, func(i, j int) bool {
		return dates[i].UnixNano() < dates[j].UnixNano()
	})
	if len(dates) == 0 {
		return time.Time{}
	}
	var bt int64
	if !isMock {
		bt = time.Now().UnixNano()
	} else {
		bt = base.UnixNano()
	}
	for _, v := range dates {
		if v.UnixNano() > bt {
			return v
		}
	}
	return time.Time{}
}

func (t *StarTimer) parseStaticNextDate(base time.Time, r *Repeats, recall bool) time.Time {
	base = time.Date(base.Year(), base.Month(), base.Day(), base.Hour(), base.Minute(), base.Second(), 0, base.Location())
	var targets []time.Time
	var uniqueRepeat [][]Repeat
	selectMap := make(map[Unit][]Repeat)
	{
		for _, d := range r.Repeat {
			selectMap[d.Unit] = append(selectMap[d.Unit], d)
		}
		for key, val := range selectMap {
			if key == STAR_WEEK {
				continue
			}
			defUnikey := make([][]Repeat, 0, 1024)
			for _, vs := range val {
				if len(uniqueRepeat) == 0 {
					defUnikey = append(defUnikey, []Repeat{vs})
				} else {
					for k := range uniqueRepeat {
						tmp := make([]Repeat, len(uniqueRepeat[k])+1)
						copy(tmp, append(uniqueRepeat[k], vs))
						defUnikey = append(defUnikey, tmp)
					}
				}
			}
			if len(defUnikey) > 0 {
				uniqueRepeat = defUnikey
			}
		}
	}
	for _, task := range uniqueRepeat {
		sort.SliceStable(task, func(i, j int) bool {
			return task[i].Unit < task[j].Unit
		})
		target := base
		veryDay, veryMonth := 0, 0 //验证日期
		if !r.Every {              //固定日期
			for _, d := range task {
				switch d.Unit {
				case STAR_SECOND:
					sub := int(d.Value) - target.Second()
					if sub < 0 {
						sub += 60
					}
					target = target.Add(time.Second * time.Duration(sub))
				case STAR_MINUTE:
					sub := int(d.Value) - target.Minute()
					if sub < 0 {
						sub += 60
					}
					target = target.Add(time.Minute * time.Duration(sub))
				case STAR_HOUR:
					sub := int(d.Value) - target.Hour()
					if sub < 0 {
						sub += 24
					}
					target = target.Add(time.Hour * time.Duration(sub))
				case STAR_DAY:
					veryDay = int(d.Value)
					sub := int(d.Value) - target.Day()
					if sub >= 0 {
						target = target.Add(time.Hour * 24 * time.Duration(sub))
						continue
					}
					target = time.Date(target.Year(), target.Month()+1, int(d.Value), target.Hour(), target.Minute(), target.Second(), 0, target.Location())
				case STAR_MONTH:
					veryMonth = int(d.Value)
					sub := int(d.Value) - int(target.Month())
					if sub < 0 {
						sub += 12
					}
					target = target.AddDate(0, sub, 0)
				case STAR_YEAR:
					sub := int(d.Value) - int(target.Year())
					if sub < 0 {
						continue
					}
					target = target.AddDate(sub, 0, 0)
				}
			}
		}
		if target.UnixNano() == base.UnixNano() && veryDay == veryMonth && veryDay == 0 {
			target = target.Add(time.Hour * 24)
		}
		if (veryDay != 0 && target.Day() != veryDay) || (veryMonth != 0 && int(target.Month()) != veryMonth) {
			continue
		}
		if target.UnixNano() == base.UnixNano() && !recall {
			continue
		}
		targets = append(targets, target)

	}
	sort.SliceStable(targets, func(i, j int) bool {
		return targets[i].UnixNano() < targets[j].UnixNano()
	})
	for k, v := range targets {
		if v.UnixNano() > base.UnixNano() || (recall && v.UnixNano() == base.UnixNano()) {
			targets = targets[k:]
			break
		}
	}
	if val, ok := selectMap[STAR_WEEK]; ok {
		if len(targets) > 0 {
			for _, week := range val {
				if int(targets[0].Weekday()) == int(week.Value) {
					return targets[0]
				}
			}
			nextBase := targets[0]
			if nextBase == base {
				nextBase = time.Date(targets[0].Year(), targets[0].Month(), targets[0].Day(), 0, 0, 0, 0, targets[0].Location())
				nextBase = nextBase.Add(time.Hour * 24)
			}
			return t.parseStaticNextDate(nextBase, r, true)
		}
	}
	if len(targets) > 0 {
		return targets[0]
	}
	return time.Time{}
}

func (t *StarTimer) parseEveryNextDate(target time.Time, r *Repeats, isMock bool) []time.Time {
	var res []time.Time
	if r.Every { //定期日期
		for idx, d := range r.Repeat {
			if d.baseDate.Unix() == -62135596800 {
				d.baseDate = t.base
				if !isMock {
					r.Repeat[idx] = d
				}
			}
			if d.baseDate.After(target) {
				res = append(res, d.baseDate)
				continue
			}
			switch d.Unit {
			case STAR_SECOND:
				for {
					d.baseDate = d.baseDate.Add(time.Second * time.Duration(int(d.Value)))
					if d.baseDate.After(target) {
						if !isMock {
							r.Repeat[idx] = d
						}
						break
					}
				}
				res = append(res, d.baseDate)
			case STAR_MINUTE:
				for {
					d.baseDate = d.baseDate.Add(time.Minute * time.Duration(int(d.Value)))
					if d.baseDate.After(target) {
						if !isMock {
							r.Repeat[idx] = d
						}
						break
					}
				}
				res = append(res, d.baseDate)
			case STAR_HOUR:
				for {
					d.baseDate = d.baseDate.Add(time.Hour * time.Duration(int(d.Value)))
					if d.baseDate.After(target) {
						if !isMock {
							r.Repeat[idx] = d
						}
						break
					}
				}
				res = append(res, d.baseDate)
			case STAR_DAY:
				for {
					d.baseDate = d.baseDate.Add(time.Hour * 24 * time.Duration(int(d.Value)))
					if d.baseDate.After(target) {
						if !isMock {
							r.Repeat[idx] = d
						}
						break
					}
				}
				res = append(res, d.baseDate)
			case STAR_MONTH:
				for {
					d.baseDate = d.baseDate.AddDate(0, int(d.Value), 0)
					if d.baseDate.After(target) {
						if !isMock {
							r.Repeat[idx] = d
						}
						break
					}
				}
				res = append(res, d.baseDate)
			case STAR_YEAR:
				for {
					d.baseDate = d.baseDate.AddDate(int(d.Value), 0, 0)
					if d.baseDate.After(target) {
						if !isMock {
							r.Repeat[idx] = d
						}
						break
					}
				}
				res = append(res, d.baseDate)
			}
		}
	}
	return res
}