new-api/service/channel_monitor_runner.go

package service

import (
	"log"
	"sync"
	"time"

	"github.com/QuantumNous/new-api/model"
)

var (
	channelMonitorRunnerOnce sync.Once
	channelMonitorStopMu     sync.Mutex
	channelMonitorStopCh     chan struct{}
	channelMonitorStopped    bool

	channelMonitorRunningMu sync.Mutex
	channelMonitorRunning   = map[int]struct{}{}
	channelMonitorWg        sync.WaitGroup
)

func StartChannelMonitorRunner() {
	channelMonitorRunnerOnce.Do(func() {
		channelMonitorStopMu.Lock()
		defer channelMonitorStopMu.Unlock()

		channelMonitorStopCh = make(chan struct{})
		channelMonitorStopped = false
		go runChannelMonitorLoop(channelMonitorStopCh)
	})
}

func StopChannelMonitorRunner() {
	channelMonitorStopMu.Lock()
	if channelMonitorStopCh != nil && !channelMonitorStopped {
		close(channelMonitorStopCh)
		channelMonitorStopped = true
	}
	channelMonitorStopMu.Unlock()

	channelMonitorWg.Wait()
}

func runChannelMonitorLoop(stopCh <-chan struct{}) {
	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			runDueChannelMonitors()
		case <-stopCh:
			return
		}
	}
}

func runDueChannelMonitors() {
	monitors, err := model.GetEnabledChannelMonitors()
	if err != nil {
		log.Printf("channel monitor runner: get monitors error: %v", err)
		return
	}

	now := time.Now()
	for _, monitor := range monitors {
		if !channelMonitorDue(&monitor, now) {
			continue
		}
		if !markChannelMonitorRunning(monitor.Id) {
			continue
		}

		monitorCopy := monitor
		channelMonitorStopMu.Lock()
		if channelMonitorStopped {
			channelMonitorStopMu.Unlock()
			unmarkChannelMonitorRunning(monitorCopy.Id)
			continue
		}
		channelMonitorWg.Add(1)
		channelMonitorStopMu.Unlock()
		go func() {
			defer channelMonitorWg.Done()
			defer unmarkChannelMonitorRunning(monitorCopy.Id)
			defer func() {
				if r := recover(); r != nil {
					log.Printf("channel monitor runner: monitor %d panicked: %v", monitorCopy.Id, r)
				}
			}()
			if err := RunMonitorCheck(&monitorCopy); err != nil {
				log.Printf("channel monitor runner: monitor %d check error: %v", monitorCopy.Id, err)
			}
		}()
	}
}

func channelMonitorDue(monitor *model.ChannelMonitor, now time.Time) bool {
	if monitor.LastCheckedAt == nil {
		return true
	}
	interval := monitor.CheckInterval
	if interval <= 0 {
		interval = 300
	}
	return !monitor.LastCheckedAt.Add(time.Duration(interval) * time.Second).After(now)
}

func markChannelMonitorRunning(monitorID int) bool {
	channelMonitorRunningMu.Lock()
	defer channelMonitorRunningMu.Unlock()

	if _, ok := channelMonitorRunning[monitorID]; ok {
		return false
	}
	channelMonitorRunning[monitorID] = struct{}{}
	return true
}

func unmarkChannelMonitorRunning(monitorID int) {
	channelMonitorRunningMu.Lock()
	defer channelMonitorRunningMu.Unlock()

	delete(channelMonitorRunning, monitorID)
}