YCMediaKit/internal/pkg/unit/keepalive/keepalive.go

package keepalive

import (
	"log"
	"time"
	"ycmediakit/internal/pkg/util"
)

type Supervisor struct {
	Cfg        *Config
	HandlePath string
	guarder    Guarder
	fastMap    *util.RwMap[string, time.Time]
}

type Guarder interface {
	CreateKey(string) (string, bool)
	Expire(string) bool
}

func NewSupervisor(cfg *Config) *Supervisor {
	s := &Supervisor{
		Cfg:     cfg,
		fastMap: util.NewRwMap[string, time.Time](),
	}
	if s.Cfg.Enable {
		go s.start()
	}
	return s
}

func (s *Supervisor) start() {
	i := time.Duration(s.Cfg.PollInterval) * time.Minute
	for range time.Tick(i) {
		go s.inspect()
	}
}

func (s *Supervisor) inspect() {
	g := s.guarder
	if g == nil {
		return
	}
	nowTime := time.Now()
	pollTime := time.Duration(s.Cfg.PollInterval) * time.Minute
	s.fastMap.ModifyRange(func(key string, value time.Time) {
		if nowTime.Sub(value) >= pollTime {
			g.Expire(key)
			delete(s.fastMap.Map, key)
		}
	})
}

func (s *Supervisor) Register(handlePath string, i interface{}) {
	g, ok := i.(Guarder)
	if !ok {
		log.Panicf("%v is not implements Guarder\n", i)
	}
	if s.guarder == nil {
		s.HandlePath = handlePath
		s.guarder = g
	}
}

func (s *Supervisor) CheckIn(oriUrl string) {
	g := s.guarder
	if g == nil {
		return
	}
	key, ok := g.CreateKey(oriUrl)
	if !ok {
		return
	}
	oldTime, ok := s.fastMap.Get(key)
	newTime := time.Now()
	if ok && newTime.Sub(oldTime) < time.Duration(s.Cfg.CheckInInterval)*time.Minute {
		return
	}
	s.Activate(key)
}

func (s *Supervisor) Activate(key string) {
	s.fastMap.Set(key, time.Now())
}

func (s *Supervisor) DeActivate(key string) {
	if !s.fastMap.Has(key) {
		return
	}
	s.fastMap.Delete(key)
}