synctv/internal/sysnotify/sysnotify.go

package sysnotify

import (
	"errors"
	"os"
	"sync"

	log "github.com/sirupsen/logrus"

	"github.com/zijiren233/gencontainer/pqueue"
	"github.com/zijiren233/gencontainer/rwmap"
)

var sysNotify SysNotify

func Init() {
	sysNotify.Init()
}

func RegisterSysNotifyTask(priority int, task *Task) error {
	return sysNotify.RegisterSysNotifyTask(priority, task)
}

func WaitCbk() {
	sysNotify.WaitCbk()
}

type SysNotify struct {
	c         chan os.Signal
	taskGroup rwmap.RWMap[NotifyType, *taskQueue]
	once      sync.Once
}

type NotifyType int

const (
	NotifyTypeEXIT NotifyType = iota + 1
	NotifyTypeRELOAD
)

type taskQueue struct {
	notifyTaskQueue *pqueue.PQueue[*Task]
	notifyTaskLock  sync.Mutex
}

type Task struct {
	Task       func() error
	Name       string
	NotifyType NotifyType
}

func NewSysNotifyTask(name string, notifyType NotifyType, task func() error) *Task {
	return &Task{
		Name:       name,
		NotifyType: notifyType,
		Task:       task,
	}
}

func runTask(tq *taskQueue) {
	tq.notifyTaskLock.Lock()
	defer tq.notifyTaskLock.Unlock()
	for tq.notifyTaskQueue.Len() > 0 {
		_, task := tq.notifyTaskQueue.Pop()
		func() {
			defer func() {
				if err := recover(); err != nil {
					log.Errorf("task: %s panic has returned: %v", task.Name, err)
				}
			}()
			log.Infof("task: %s running", task.Name)
			if err := task.Task(); err != nil {
				log.Errorf("task: %s an error occurred: %v", task.Name, err)
			}
			log.Infof("task: %s done", task.Name)
		}()
	}
}

func (sn *SysNotify) RegisterSysNotifyTask(priority int, task *Task) error {
	if task == nil || task.Task == nil {
		return errors.New("task is nil")
	}
	if task.NotifyType == 0 {
		panic("task notify type is 0")
	}
	tasks, _ := sn.taskGroup.LoadOrStore(task.NotifyType, &taskQueue{
		notifyTaskQueue: pqueue.NewMinPriorityQueue[*Task](),
	})
	tasks.notifyTaskLock.Lock()
	defer tasks.notifyTaskLock.Unlock()
	tasks.notifyTaskQueue.Push(priority, task)
	return nil
}

func (sn *SysNotify) waitCbk() {
	log.Info("wait sys notify")
	for s := range sn.c {
		log.Infof("receive sys notify: %v", s)
		switch parseSysNotifyType(s) {
		case NotifyTypeEXIT:
			tq, ok := sn.taskGroup.Load(NotifyTypeEXIT)
			if ok {
				log.Info("task: NotifyTypeEXIT running...")
				runTask(tq)
			}
			return
		case NotifyTypeRELOAD:
			tq, ok := sn.taskGroup.Load(NotifyTypeRELOAD)
			if ok {
				log.Info("task: NotifyTypeRELOAD running...")
				runTask(tq)
			}
		}
	}
	log.Info("task: all done")
}

func (sn *SysNotify) WaitCbk() {
	sn.once.Do(sn.waitCbk)
}