feat: add channel monitor runner
Implement monitor probe execution, history recording, redirect-safe target checks, and a background runner for due monitors.
This commit is contained in:
zizi
parent
a63300fc1f
commit
8f91ab1eb6
@ -32,6 +32,11 @@ func (m *ChannelMonitor) Update() error {
|
|||||||
return DB.Save(m).Error
|
return DB.Save(m).Error
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (m *ChannelMonitor) UpdateLastCheckedAt(checkedAt time.Time) error {
|
||||||
|
m.LastCheckedAt = &checkedAt
|
||||||
|
return DB.Model(m).Update("last_checked_at", checkedAt).Error
|
||||||
|
}
|
||||||
|
|
||||||
func GetChannelMonitorByID(id int) (*ChannelMonitor, error) {
|
func GetChannelMonitorByID(id int) (*ChannelMonitor, error) {
|
||||||
var monitor ChannelMonitor
|
var monitor ChannelMonitor
|
||||||
err := DB.First(&monitor, "id = ?", id).Error
|
err := DB.First(&monitor, "id = ?", id).Error
|
||||||
|
|||||||
251
service/channel_monitor_checker.go
Normal file
251
service/channel_monitor_checker.go
Normal file
@ -0,0 +1,251 @@
|
|||||||
|
package service
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/QuantumNous/new-api/common"
|
||||||
|
"github.com/QuantumNous/new-api/model"
|
||||||
|
)
|
||||||
|
|
||||||
|
const channelMonitorResponseSnippetLimit = 500
|
||||||
|
|
||||||
|
type channelMonitorCheckOptions struct {
|
||||||
|
client *http.Client
|
||||||
|
validateURL func(string) error
|
||||||
|
now func() time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
func RunMonitorCheck(monitor *model.ChannelMonitor) error {
|
||||||
|
return runMonitorCheckWithOptions(monitor, channelMonitorCheckOptions{})
|
||||||
|
}
|
||||||
|
|
||||||
|
func runMonitorCheckWithOptions(monitor *model.ChannelMonitor, options channelMonitorCheckOptions) error {
|
||||||
|
if monitor == nil {
|
||||||
|
return errors.New("channel monitor is nil")
|
||||||
|
}
|
||||||
|
|
||||||
|
checkNow := options.now
|
||||||
|
if checkNow == nil {
|
||||||
|
checkNow = time.Now
|
||||||
|
}
|
||||||
|
|
||||||
|
template := SelectChannelMonitorRequestTemplate(monitor.RequestTemplates, nil)
|
||||||
|
models, modelParseErr := channelMonitorModels(monitor)
|
||||||
|
|
||||||
|
var checkErrs []error
|
||||||
|
if modelParseErr != nil {
|
||||||
|
checkErrs = append(checkErrs, modelParseErr)
|
||||||
|
}
|
||||||
|
for _, modelName := range models {
|
||||||
|
if err := checkChannelMonitorModel(monitor, modelName, template, options, checkNow); err != nil {
|
||||||
|
checkErrs = append(checkErrs, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
now := checkNow()
|
||||||
|
if err := monitor.UpdateLastCheckedAt(now); err != nil {
|
||||||
|
checkErrs = append(checkErrs, fmt.Errorf("update monitor last_checked_at: %w", err))
|
||||||
|
}
|
||||||
|
return errors.Join(checkErrs...)
|
||||||
|
}
|
||||||
|
|
||||||
|
func channelMonitorModels(monitor *model.ChannelMonitor) ([]string, error) {
|
||||||
|
models := make([]string, 0, 1)
|
||||||
|
primaryModel := strings.TrimSpace(monitor.PrimaryModel)
|
||||||
|
if primaryModel != "" {
|
||||||
|
models = append(models, primaryModel)
|
||||||
|
}
|
||||||
|
|
||||||
|
rawExtraModels := strings.TrimSpace(monitor.ExtraModels)
|
||||||
|
if rawExtraModels == "" {
|
||||||
|
return models, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
var extraModels []string
|
||||||
|
if err := common.Unmarshal([]byte(rawExtraModels), &extraModels); err != nil {
|
||||||
|
return models, fmt.Errorf("invalid extra_models json array ignored: %w", err)
|
||||||
|
}
|
||||||
|
for _, modelName := range extraModels {
|
||||||
|
modelName = strings.TrimSpace(modelName)
|
||||||
|
if modelName != "" {
|
||||||
|
models = append(models, modelName)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return models, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func checkChannelMonitorModel(monitor *model.ChannelMonitor, modelName string, template string, options channelMonitorCheckOptions, now func() time.Time) error {
|
||||||
|
start := now()
|
||||||
|
requestBody, err := buildChannelMonitorRequestBody(monitor, modelName, template)
|
||||||
|
record := &model.ChannelMonitorHistory{
|
||||||
|
MonitorId: monitor.Id,
|
||||||
|
ChannelId: monitor.ChannelId,
|
||||||
|
ModelName: modelName,
|
||||||
|
RequestBody: requestBody,
|
||||||
|
CheckedAt: start,
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return insertChannelMonitorErrorHistory(record, fmt.Errorf("build request body: %w", err))
|
||||||
|
}
|
||||||
|
|
||||||
|
validateURL := options.validateURL
|
||||||
|
if validateURL == nil {
|
||||||
|
validateURL = ValidateChannelMonitorTargetURL
|
||||||
|
}
|
||||||
|
if err := validateURL(monitor.TargetUrl); err != nil {
|
||||||
|
return insertChannelMonitorErrorHistory(record, fmt.Errorf("SSRF validation failed: %w", err))
|
||||||
|
}
|
||||||
|
|
||||||
|
headers, err := parseChannelMonitorCustomHeaders(monitor.CustomHeaders)
|
||||||
|
if err != nil {
|
||||||
|
return insertChannelMonitorErrorHistory(record, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
timeout := channelMonitorTimeout(monitor.Timeout)
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, monitor.TargetUrl, bytes.NewReader([]byte(requestBody)))
|
||||||
|
if err != nil {
|
||||||
|
return insertChannelMonitorErrorHistory(record, fmt.Errorf("create request: %w", err))
|
||||||
|
}
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
for key, value := range headers {
|
||||||
|
req.Header.Set(key, value)
|
||||||
|
}
|
||||||
|
|
||||||
|
client := options.client
|
||||||
|
if client == nil {
|
||||||
|
client = newChannelMonitorHTTPClient(timeout, validateURL)
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := client.Do(req)
|
||||||
|
record.LatencyMs = int(time.Since(start).Milliseconds())
|
||||||
|
if err != nil {
|
||||||
|
return insertChannelMonitorErrorHistory(record, fmt.Errorf("request failed: %w", err))
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
record.StatusCode = resp.StatusCode
|
||||||
|
snippet, readErr := readChannelMonitorResponseSnippet(resp.Body)
|
||||||
|
record.ResponseSnippet = snippet
|
||||||
|
if readErr != nil {
|
||||||
|
return insertChannelMonitorErrorHistory(record, fmt.Errorf("read response: %w", readErr))
|
||||||
|
}
|
||||||
|
|
||||||
|
record.Status = channelMonitorStatus(resp.StatusCode, snippet, record.LatencyMs, timeout)
|
||||||
|
if record.Status == "failed" && resp.StatusCode >= 200 && resp.StatusCode < 300 && strings.TrimSpace(snippet) == "" {
|
||||||
|
record.ErrorMessage = "empty response body"
|
||||||
|
}
|
||||||
|
if err := record.Insert(); err != nil {
|
||||||
|
return fmt.Errorf("insert monitor history: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildChannelMonitorRequestBody(monitor *model.ChannelMonitor, modelName string, template string) (string, error) {
|
||||||
|
if strings.TrimSpace(monitor.CustomBody) != "" {
|
||||||
|
return replaceChannelMonitorTemplateVariables(monitor.CustomBody, modelName, template), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
body := map[string]any{
|
||||||
|
"model": modelName,
|
||||||
|
"messages": []map[string]string{
|
||||||
|
{"role": "user", "content": template},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
encoded, err := common.Marshal(body)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return string(encoded), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func replaceChannelMonitorTemplateVariables(raw string, modelName string, template string) string {
|
||||||
|
replacer := strings.NewReplacer(
|
||||||
|
"{{model}}", modelName,
|
||||||
|
"{{model_name}}", modelName,
|
||||||
|
"${model}", modelName,
|
||||||
|
"${model_name}", modelName,
|
||||||
|
"{{message}}", template,
|
||||||
|
"{{template}}", template,
|
||||||
|
"${message}", template,
|
||||||
|
"${template}", template,
|
||||||
|
)
|
||||||
|
return replacer.Replace(raw)
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseChannelMonitorCustomHeaders(raw string) (map[string]string, error) {
|
||||||
|
raw = strings.TrimSpace(raw)
|
||||||
|
if raw == "" {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
var headers map[string]string
|
||||||
|
if err := common.Unmarshal([]byte(raw), &headers); err != nil {
|
||||||
|
return nil, fmt.Errorf("invalid custom_headers json object: %w", err)
|
||||||
|
}
|
||||||
|
return headers, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func newChannelMonitorHTTPClient(timeout time.Duration, validateURL func(string) error) *http.Client {
|
||||||
|
return &http.Client{
|
||||||
|
Timeout: timeout,
|
||||||
|
CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
||||||
|
if len(via) >= 10 {
|
||||||
|
return errors.New("stopped after 10 redirects")
|
||||||
|
}
|
||||||
|
if validateURL == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if err := validateURL(req.URL.String()); err != nil {
|
||||||
|
return fmt.Errorf("redirect target rejected: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func channelMonitorTimeout(timeoutSeconds int) time.Duration {
|
||||||
|
if timeoutSeconds <= 0 {
|
||||||
|
timeoutSeconds = 30
|
||||||
|
}
|
||||||
|
return time.Duration(timeoutSeconds) * time.Second
|
||||||
|
}
|
||||||
|
|
||||||
|
func channelMonitorStatus(statusCode int, snippet string, latencyMs int, timeout time.Duration) string {
|
||||||
|
if statusCode < 200 || statusCode >= 300 {
|
||||||
|
return "failed"
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(snippet) == "" {
|
||||||
|
return "failed"
|
||||||
|
}
|
||||||
|
if latencyMs >= int(timeout.Milliseconds()) {
|
||||||
|
return "degraded"
|
||||||
|
}
|
||||||
|
return "pass"
|
||||||
|
}
|
||||||
|
|
||||||
|
func readChannelMonitorResponseSnippet(reader io.Reader) (string, error) {
|
||||||
|
body, err := io.ReadAll(io.LimitReader(reader, channelMonitorResponseSnippetLimit))
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return string(body), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func insertChannelMonitorErrorHistory(record *model.ChannelMonitorHistory, err error) error {
|
||||||
|
record.Status = "error"
|
||||||
|
record.ErrorMessage = err.Error()
|
||||||
|
if insertErr := record.Insert(); insertErr != nil {
|
||||||
|
return fmt.Errorf("%w; insert monitor history: %v", err, insertErr)
|
||||||
|
}
|
||||||
|
return err
|
||||||
|
}
|
||||||
182
service/channel_monitor_checker_test.go
Normal file
182
service/channel_monitor_checker_test.go
Normal file
@ -0,0 +1,182 @@
|
|||||||
|
package service
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/QuantumNous/new-api/common"
|
||||||
|
"github.com/QuantumNous/new-api/model"
|
||||||
|
"github.com/glebarez/sqlite"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
"gorm.io/gorm"
|
||||||
|
)
|
||||||
|
|
||||||
|
func setupChannelMonitorTestDB(t *testing.T) *gorm.DB {
|
||||||
|
t.Helper()
|
||||||
|
|
||||||
|
common.UsingSQLite = true
|
||||||
|
common.UsingMySQL = false
|
||||||
|
common.UsingPostgreSQL = false
|
||||||
|
|
||||||
|
previousDB := model.DB
|
||||||
|
previousLogDB := model.LOG_DB
|
||||||
|
dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
|
||||||
|
db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NoError(t, db.AutoMigrate(&model.ChannelMonitor{}, &model.ChannelMonitorHistory{}))
|
||||||
|
|
||||||
|
model.DB = db
|
||||||
|
model.LOG_DB = db
|
||||||
|
|
||||||
|
t.Cleanup(func() {
|
||||||
|
model.DB = previousDB
|
||||||
|
model.LOG_DB = previousLogDB
|
||||||
|
sqlDB, err := db.DB()
|
||||||
|
if err == nil {
|
||||||
|
_ = sqlDB.Close()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
return db
|
||||||
|
}
|
||||||
|
|
||||||
|
func allowTestChannelMonitorURL(string) error {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRunMonitorCheckPassWithCustomTemplateBodyAndHeaders(t *testing.T) {
|
||||||
|
db := setupChannelMonitorTestDB(t)
|
||||||
|
|
||||||
|
var receivedBody string
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
require.Equal(t, "ok", r.Header.Get("X-Test"))
|
||||||
|
body, err := io.ReadAll(r.Body)
|
||||||
|
require.NoError(t, err)
|
||||||
|
receivedBody = string(body)
|
||||||
|
_, _ = w.Write([]byte(`{"id":"chatcmpl_test"}`))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
monitor := &model.ChannelMonitor{
|
||||||
|
Name: "custom",
|
||||||
|
ChannelId: 7,
|
||||||
|
TargetUrl: server.URL,
|
||||||
|
PrimaryModel: "m-custom",
|
||||||
|
RequestTemplates: `["custom prompt"]`,
|
||||||
|
CustomBody: `{"model":"{{model}}","input":"{{template}}"}`,
|
||||||
|
CustomHeaders: `{"X-Test":"ok"}`,
|
||||||
|
Timeout: 5,
|
||||||
|
Status: 1,
|
||||||
|
}
|
||||||
|
require.NoError(t, db.Create(monitor).Error)
|
||||||
|
|
||||||
|
err := runMonitorCheckWithOptions(monitor, channelMonitorCheckOptions{
|
||||||
|
client: server.Client(),
|
||||||
|
validateURL: allowTestChannelMonitorURL,
|
||||||
|
})
|
||||||
|
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Contains(t, receivedBody, `"model":"m-custom"`)
|
||||||
|
require.Contains(t, receivedBody, `"input":"custom prompt"`)
|
||||||
|
require.NotNil(t, monitor.LastCheckedAt)
|
||||||
|
|
||||||
|
var history model.ChannelMonitorHistory
|
||||||
|
require.NoError(t, db.First(&history, "monitor_id = ?", monitor.Id).Error)
|
||||||
|
require.Equal(t, monitor.Id, history.MonitorId)
|
||||||
|
require.Equal(t, 7, history.ChannelId)
|
||||||
|
require.Equal(t, "m-custom", history.ModelName)
|
||||||
|
require.Equal(t, "pass", history.Status)
|
||||||
|
require.Equal(t, 200, history.StatusCode)
|
||||||
|
require.GreaterOrEqual(t, history.LatencyMs, 0)
|
||||||
|
require.Empty(t, history.ErrorMessage)
|
||||||
|
require.Contains(t, history.RequestBody, `"model":"m-custom"`)
|
||||||
|
require.Contains(t, history.ResponseSnippet, "chatcmpl_test")
|
||||||
|
require.False(t, history.CheckedAt.IsZero())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRunMonitorCheckRecordsFailedStatusCode(t *testing.T) {
|
||||||
|
db := setupChannelMonitorTestDB(t)
|
||||||
|
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusInternalServerError)
|
||||||
|
_, _ = w.Write([]byte(`{"error":"bad"}`))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
monitor := &model.ChannelMonitor{
|
||||||
|
Name: "failed",
|
||||||
|
TargetUrl: server.URL,
|
||||||
|
PrimaryModel: "m-failed",
|
||||||
|
Timeout: 5,
|
||||||
|
Status: 1,
|
||||||
|
}
|
||||||
|
require.NoError(t, db.Create(monitor).Error)
|
||||||
|
|
||||||
|
err := runMonitorCheckWithOptions(monitor, channelMonitorCheckOptions{
|
||||||
|
client: server.Client(),
|
||||||
|
validateURL: allowTestChannelMonitorURL,
|
||||||
|
})
|
||||||
|
|
||||||
|
require.NoError(t, err)
|
||||||
|
var history model.ChannelMonitorHistory
|
||||||
|
require.NoError(t, db.First(&history, "monitor_id = ?", monitor.Id).Error)
|
||||||
|
require.Equal(t, "failed", history.Status)
|
||||||
|
require.Equal(t, http.StatusInternalServerError, history.StatusCode)
|
||||||
|
require.Contains(t, history.ResponseSnippet, "bad")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRunMonitorCheckRecordsSSRFFailure(t *testing.T) {
|
||||||
|
db := setupChannelMonitorTestDB(t)
|
||||||
|
|
||||||
|
monitor := &model.ChannelMonitor{
|
||||||
|
Name: "ssrf",
|
||||||
|
TargetUrl: "http://127.0.0.1:1/v1/chat/completions",
|
||||||
|
PrimaryModel: "m-ssrf",
|
||||||
|
Timeout: 5,
|
||||||
|
Status: 1,
|
||||||
|
}
|
||||||
|
require.NoError(t, db.Create(monitor).Error)
|
||||||
|
|
||||||
|
err := RunMonitorCheck(monitor)
|
||||||
|
|
||||||
|
require.Error(t, err)
|
||||||
|
var history model.ChannelMonitorHistory
|
||||||
|
require.NoError(t, db.First(&history, "monitor_id = ?", monitor.Id).Error)
|
||||||
|
require.Equal(t, "error", history.Status)
|
||||||
|
require.Contains(t, history.ErrorMessage, "SSRF validation failed")
|
||||||
|
require.Contains(t, history.RequestBody, "m-ssrf")
|
||||||
|
require.Equal(t, 0, history.StatusCode)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRunMonitorCheckRecordsInvalidCustomHeaders(t *testing.T) {
|
||||||
|
db := setupChannelMonitorTestDB(t)
|
||||||
|
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
t.Fatal("request should not be sent when custom headers are invalid")
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
monitor := &model.ChannelMonitor{
|
||||||
|
Name: "headers",
|
||||||
|
TargetUrl: server.URL,
|
||||||
|
PrimaryModel: "m-headers",
|
||||||
|
CustomHeaders: `{"X-Test":123}`,
|
||||||
|
Timeout: 5,
|
||||||
|
Status: 1,
|
||||||
|
}
|
||||||
|
require.NoError(t, db.Create(monitor).Error)
|
||||||
|
|
||||||
|
err := runMonitorCheckWithOptions(monitor, channelMonitorCheckOptions{
|
||||||
|
client: server.Client(),
|
||||||
|
validateURL: allowTestChannelMonitorURL,
|
||||||
|
})
|
||||||
|
|
||||||
|
require.Error(t, err)
|
||||||
|
var history model.ChannelMonitorHistory
|
||||||
|
require.NoError(t, db.First(&history, "monitor_id = ?", monitor.Id).Error)
|
||||||
|
require.Equal(t, "error", history.Status)
|
||||||
|
require.Contains(t, history.ErrorMessage, "invalid custom_headers json object")
|
||||||
|
}
|
||||||
109
service/channel_monitor_runner.go
Normal file
109
service/channel_monitor_runner.go
Normal file
@ -0,0 +1,109 @@
|
|||||||
|
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{}{}
|
||||||
|
)
|
||||||
|
|
||||||
|
func StartChannelMonitorRunner() {
|
||||||
|
channelMonitorRunnerOnce.Do(func() {
|
||||||
|
channelMonitorStopMu.Lock()
|
||||||
|
defer channelMonitorStopMu.Unlock()
|
||||||
|
|
||||||
|
channelMonitorStopCh = make(chan struct{})
|
||||||
|
channelMonitorStopped = false
|
||||||
|
go runChannelMonitorLoop(channelMonitorStopCh)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func StopChannelMonitorRunner() {
|
||||||
|
channelMonitorStopMu.Lock()
|
||||||
|
defer channelMonitorStopMu.Unlock()
|
||||||
|
|
||||||
|
if channelMonitorStopCh != nil && !channelMonitorStopped {
|
||||||
|
close(channelMonitorStopCh)
|
||||||
|
channelMonitorStopped = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
go func() {
|
||||||
|
defer unmarkChannelMonitorRunning(monitorCopy.Id)
|
||||||
|
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)
|
||||||
|
}
|
||||||
Loading…
x
Reference in New Issue
Block a user