new-api-custom/controller/uptime_kuma.go

package controller

import (
	"context"
	"errors"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
	"github.com/QuantumNous/new-api/service"

	"github.com/gin-gonic/gin"
	"golang.org/x/sync/errgroup"
)

const (
	requestTimeout   = 30 * time.Second
	httpTimeout      = 10 * time.Second
	uptimeKeySuffix  = "_24"
	apiStatusPath    = "/api/status-page/"
	apiHeartbeatPath = "/api/status-page/heartbeat/"
)

type Monitor struct {
	Metrics         *perfmetrics.GroupResult `json:"metrics,omitempty"`
	Name            string                   `json:"name"`
	Uptime          float64                  `json:"uptime"`
	Status          int                      `json:"status"`
	Group           string                   `json:"group,omitempty"`
	Ping            int                      `json:"ping,omitempty"`
	LastChecked     string                   `json:"last_checked,omitempty"`
	History         []MonitorSample          `json:"history,omitempty"`
	CacheStats      *model.GroupCacheStats   `json:"cache_stats,omitempty"`
	ProbeStatus     int                      `json:"probe_status"`
	ChannelDisabled bool                     `json:"channel_disabled"`
}

type MonitorSample struct {
	Status int    `json:"status"`
	Ping   int    `json:"ping,omitempty"`
	Time   string `json:"time"`
}

type UptimeGroupResult struct {
	CategoryName string    `json:"categoryName"`
	Monitors     []Monitor `json:"monitors"`
}

func getAndDecode(ctx context.Context, client *http.Client, url string, dest any) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return err
	}

	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return errors.New("non-200 status")
	}

	return common.DecodeJson(resp.Body, dest)
}

func fetchGroupData(ctx context.Context, client *http.Client, groupConfig map[string]any) UptimeGroupResult {
	url, _ := groupConfig["url"].(string)
	slug, _ := groupConfig["slug"].(string)
	categoryName, _ := groupConfig["categoryName"].(string)

	result := UptimeGroupResult{
		CategoryName: categoryName,
		Monitors:     []Monitor{},
	}

	if url == "" || slug == "" {
		return result
	}

	baseURL := strings.TrimSuffix(url, "/")

	var statusData struct {
		PublicGroupList []struct {
			ID          int    `json:"id"`
			Name        string `json:"name"`
			MonitorList []struct {
				ID   int    `json:"id"`
				Name string `json:"name"`
			} `json:"monitorList"`
		} `json:"publicGroupList"`
	}

	var heartbeatData struct {
		HeartbeatList map[string][]struct {
			Status int    `json:"status"`
			Ping   int    `json:"ping"`
			Time   string `json:"time"`
		} `json:"heartbeatList"`
		UptimeList map[string]float64 `json:"uptimeList"`
	}

	g, gCtx := errgroup.WithContext(ctx)
	g.Go(func() error {
		return getAndDecode(gCtx, client, baseURL+apiStatusPath+slug, &statusData)
	})
	g.Go(func() error {
		return getAndDecode(gCtx, client, baseURL+apiHeartbeatPath+slug, &heartbeatData)
	})

	if g.Wait() != nil {
		return result
	}

	for _, pg := range statusData.PublicGroupList {
		if len(pg.MonitorList) == 0 {
			continue
		}

		for _, m := range pg.MonitorList {
			monitor := Monitor{
				Name:   m.Name,
				Group:  pg.Name,
				Status: -1,
			}

			monitorID := strconv.Itoa(m.ID)

			if uptime, exists := heartbeatData.UptimeList[monitorID+uptimeKeySuffix]; exists {
				monitor.Uptime = uptime
			}

			if heartbeats, exists := heartbeatData.HeartbeatList[monitorID]; exists && len(heartbeats) > 0 {
				latest := heartbeats[len(heartbeats)-1]
				monitor.Status = latest.Status
				monitor.Ping = latest.Ping
				monitor.LastChecked = latest.Time
				monitor.History = make([]MonitorSample, 0, len(heartbeats))
				for _, heartbeat := range heartbeats {
					monitor.History = append(monitor.History, MonitorSample{
						Status: heartbeat.Status,
						Ping:   heartbeat.Ping,
						Time:   heartbeat.Time,
					})
				}
			}

			result.Monitors = append(result.Monitors, monitor)
		}
	}

	return result
}

var groupCacheSnapshot struct {
	sync.Mutex
	updated  time.Time
	channels string
	stats    map[string]model.GroupCacheStats
	metrics  map[string]perfmetrics.GroupResult
}

const requestStatusFreshness = 2 * time.Hour

func currentGroupMonitors(ctx context.Context, results []UptimeGroupResult) ([]UptimeGroupResult, error) {
	channels, err := model.GetStatusChannels(ctx)
	if err != nil {
		return nil, err
	}
	usable := service.GetUserUsableGroups("")
	active := make(map[string]bool)
	order := make([]string, 0)
	for _, channel := range channels {
		for _, group := range strings.Split(channel.Group, ",") {
			group = strings.TrimSpace(group)
			if _, visible := usable[group]; !visible {
				continue
			}
			if _, exists := active[group]; !exists {
				order = append(order, group)
			}
			active[group] = active[group] || channel.Status == common.ChannelStatusEnabled
		}
	}
	encoded, err := common.Marshal(channels)
	if err != nil {
		return nil, err
	}
	groupCacheSnapshot.Lock()
	defer groupCacheSnapshot.Unlock()
	if time.Since(groupCacheSnapshot.updated) >= time.Minute || groupCacheSnapshot.channels != string(encoded) {
		now := time.Now()
		stats, statsErr := model.GetGroupCacheStats(ctx, channels, now.Add(-24*time.Hour), now)
		if statsErr != nil {
			return nil, statsErr
		}
		groupCacheSnapshot.stats = stats
		names := []string{}
		for _, item := range model.GetPricing() {
			names = append(names, item.ModelName)
		}
		performance, metricsErr := perfmetrics.QueryGroups(names, 24)
		if metricsErr != nil {
			return nil, metricsErr
		}
		groupCacheSnapshot.metrics = make(map[string]perfmetrics.GroupResult)
		for _, group := range performance.Groups {
			groupCacheSnapshot.metrics[group.Group] = group
		}
		groupCacheSnapshot.updated = now
		groupCacheSnapshot.channels = string(encoded)
	}
	byName := make(map[string]Monitor)
	for _, result := range results {
		for _, monitor := range result.Monitors {
			if previous, exists := byName[monitor.Name]; !exists || (monitor.Status == 1 && previous.Status != 1) {
				byName[monitor.Name] = monitor
			}
		}
	}
	monitors := make([]Monitor, 0, len(order))
	for _, group := range order {
		monitor, exists := byName[group]
		if !exists {
			monitor = Monitor{Name: group, Status: -1}
		}
		monitor.Group = group
		monitor.ProbeStatus = monitor.Status
		monitor.ChannelDisabled = !active[group]
		if !active[group] {
			monitor.Status = 0
		}
		stats := groupCacheSnapshot.stats[group]
		monitor.CacheStats = &stats
		applyRequestMetrics(&monitor, groupCacheSnapshot.metrics[group])
		monitors = append(monitors, monitor)
	}
	return []UptimeGroupResult{{CategoryName: "Groups", Monitors: monitors}}, nil
}

func GetUptimeKumaStatus(c *gin.Context) {
	ctx, cancel := context.WithTimeout(c.Request.Context(), requestTimeout)
	defer cancel()

	results, err := currentGroupMonitors(ctx, nil)
	if err != nil {
		common.SysError("group status statistics: " + err.Error())
		c.JSON(http.StatusServiceUnavailable, gin.H{"success": false, "message": "Status data is temporarily unavailable"})
		return
	}
	c.JSON(http.StatusOK, gin.H{"success": true, "message": "", "data": results})
}

func requestStatus(rate float64) int {
	if rate >= 90 {
		return 1
	}
	if rate >= 70 {
		return 2
	}
	return 0
}

func applyRequestMetrics(monitor *Monitor, metrics perfmetrics.GroupResult) {
	monitor.Status = -1
	monitor.Uptime = 0
	monitor.Ping = 0
	monitor.LastChecked = ""
	monitor.History = nil
	if len(metrics.Series) > 0 {
		monitor.Metrics = &metrics
		monitor.Uptime = metrics.SuccessRate / 100
		latest := metrics.Series[len(metrics.Series)-1]
		if time.Since(time.Unix(latest.Ts, 0)) <= requestStatusFreshness {
			monitor.Status = requestStatus(latest.SuccessRate)
		}
		for _, point := range metrics.Series {
			monitor.History = append(monitor.History, MonitorSample{
				Status: requestStatus(point.SuccessRate), Ping: int(point.AvgTtftMs),
				Time: time.Unix(point.Ts, 0).UTC().Format("2006-01-02 15:04:05"),
			})
		}
		monitor.LastChecked = monitor.History[len(monitor.History)-1].Time
	}
	if monitor.ChannelDisabled {
		monitor.Status = 0
	}
}