nezha/cmd/dashboard/main.go

202 lines
5.3 KiB
Go
Raw Normal View History

2019-12-05 09:36:58 -05:00
package main
import (
"context"
"log"
2019-12-08 03:59:58 -05:00
"time"
2019-12-05 09:36:58 -05:00
"github.com/ory/graceful"
2019-12-08 03:59:58 -05:00
"github.com/patrickmn/go-cache"
"github.com/robfig/cron/v3"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
2019-12-05 10:42:20 -05:00
2020-11-10 21:07:45 -05:00
"github.com/naiba/nezha/cmd/dashboard/controller"
"github.com/naiba/nezha/cmd/dashboard/rpc"
"github.com/naiba/nezha/model"
"github.com/naiba/nezha/service/dao"
2019-12-05 09:36:58 -05:00
)
2019-12-08 03:59:58 -05:00
func init() {
shanghai, err := time.LoadLocation("Asia/Shanghai")
if err != nil {
panic(err)
}
// 初始化 dao 包
2019-12-20 10:58:09 -05:00
dao.Conf = &model.Config{}
dao.Cron = cron.New(cron.WithLocation(shanghai))
dao.Crons = make(map[uint64]*model.Cron)
dao.ServerList = make(map[uint64]*model.Server)
dao.SecretToID = make(map[string]uint64)
2019-12-20 10:58:09 -05:00
err = dao.Conf.Read("data/config.yaml")
2019-12-08 03:59:58 -05:00
if err != nil {
panic(err)
}
dao.DB, err = gorm.Open(sqlite.Open("data/sqlite.db"), &gorm.Config{
CreateBatchSize: 200,
})
2019-12-05 09:36:58 -05:00
if err != nil {
panic(err)
}
2019-12-08 10:18:29 -05:00
if dao.Conf.Debug {
dao.DB = dao.DB.Debug()
}
if dao.Conf.GRPCPort == 0 {
dao.Conf.GRPCPort = 5555
}
2019-12-08 03:59:58 -05:00
dao.Cache = cache.New(5*time.Minute, 10*time.Minute)
initSystem()
2019-12-08 10:18:29 -05:00
}
func initSystem() {
dao.DB.AutoMigrate(model.Server{}, model.User{},
model.Notification{}, model.AlertRule{}, model.Monitor{},
model.MonitorHistory{}, model.Cron{}, model.Transfer{})
dao.NewServiceSentinel()
loadServers() //加载服务器列表
loadCrons() //加载计划任务
// 清理 服务请求记录 和 流量记录 的旧数据
2021-07-18 22:37:12 -04:00
_, err := dao.Cron.AddFunc("30 3 * * *", cleanMonitorHistory)
if err != nil {
panic(err)
}
// 流量记录打点
2021-07-18 22:37:12 -04:00
_, err = dao.Cron.AddFunc("0 * * * *", recordTransferHourlyUsage)
if err != nil {
panic(err)
}
}
func recordTransferHourlyUsage() {
dao.ServerLock.Lock()
defer dao.ServerLock.Unlock()
now := time.Now()
nowTrimSeconds := time.Date(now.Year(), now.Month(), now.Day(), now.Hour(), 0, 0, 0, time.Local)
var txs []model.Transfer
for id, server := range dao.ServerList {
tx := model.Transfer{
ServerID: id,
In: server.State.NetInTransfer - uint64(server.PrevHourlyTransferIn),
Out: server.State.NetInTransfer - uint64(server.PrevHourlyTransferIn),
}
server.PrevHourlyTransferIn = int64(server.State.NetInTransfer)
server.PrevHourlyTransferOut = int64(server.State.NetOutTransfer)
tx.CreatedAt = nowTrimSeconds
txs = append(txs, tx)
}
dao.DB.Create(txs)
}
func cleanMonitorHistory() {
dao.DB.Unscoped().Delete(&model.MonitorHistory{}, "created_at < ?", time.Now().AddDate(0, 0, -30))
var allServerKeep time.Time
specialServerKeep := make(map[uint64]time.Time)
var specialServerIDs []uint64
var alerts []model.AlertRule
dao.DB.Find(&alerts)
for i := 0; i < len(alerts); i++ {
for j := 0; j < len(alerts[i].Rules); j++ {
// 是不是流量记录规则
if !alerts[i].Rules[j].IsTransferDurationRule() {
continue
}
dataCouldRemoveBefore := alerts[i].Rules[j].GetTransferDurationStart()
// 判断规则影响的机器范围
if alerts[i].Rules[j].Cover == model.RuleCoverAll {
// 更新全局可以清理的数据点
if allServerKeep.IsZero() || allServerKeep.After(dataCouldRemoveBefore) {
allServerKeep = dataCouldRemoveBefore
}
} else {
// 更新特定机器可以清理数据点
for id := range alerts[i].Rules[j].Ignore {
if specialServerKeep[id].IsZero() || specialServerKeep[id].After(dataCouldRemoveBefore) {
specialServerKeep[id] = dataCouldRemoveBefore
specialServerIDs = append(specialServerIDs, id)
}
}
}
}
}
for id, couldRemove := range specialServerKeep {
dao.DB.Unscoped().Delete(&model.Transfer{}, "id = ? AND created_at < ?", id, couldRemove)
}
2021-07-14 12:07:02 -04:00
if allServerKeep.IsZero() {
dao.DB.Unscoped().Delete(&model.Transfer{}, "id NOT IN (?)", specialServerIDs)
} else {
dao.DB.Unscoped().Delete(&model.Transfer{}, "id NOT IN (?) AND created_at < ?", specialServerIDs, allServerKeep)
}
}
func loadServers() {
2019-12-09 03:02:49 -05:00
var servers []model.Server
dao.DB.Find(&servers)
for _, s := range servers {
2019-12-09 05:14:31 -05:00
innerS := s
2020-12-19 23:18:27 -05:00
innerS.Host = &model.Host{}
innerS.State = &model.HostState{}
2021-01-08 08:04:50 -05:00
dao.ServerList[innerS.ID] = &innerS
dao.SecretToID[innerS.Secret] = innerS.ID
2019-12-09 03:02:49 -05:00
}
2021-01-08 08:04:50 -05:00
dao.ReSortServer()
2019-12-08 03:59:58 -05:00
}
2019-12-05 09:36:58 -05:00
func loadCrons() {
var crons []model.Cron
dao.DB.Find(&crons)
var err error
for i := 0; i < len(crons); i++ {
cr := crons[i]
crIgnoreMap := make(map[uint64]bool)
for j := 0; j < len(cr.Servers); j++ {
crIgnoreMap[cr.Servers[j]] = true
}
cr.CronID, err = dao.Cron.AddFunc(cr.Scheduler, dao.CronTrigger(cr))
if err != nil {
panic(err)
}
dao.Crons[cr.ID] = &cr
}
dao.Cron.Start()
}
2019-12-08 03:59:58 -05:00
func main() {
2021-07-16 06:09:50 -04:00
cleanMonitorHistory()
go rpc.ServeRPC(dao.Conf.GRPCPort)
go rpc.DispatchTask(time.Second * 30)
go dao.AlertSentinelStart()
srv := controller.ServeWeb(dao.Conf.HTTPPort)
graceful.Graceful(func() error {
return srv.ListenAndServe()
}, func(c context.Context) error {
dao.ServerLock.Lock()
defer dao.ServerLock.Unlock()
var txs []model.Transfer
for _, s := range dao.ServerList {
in := s.State.NetInTransfer - uint64(s.PrevHourlyTransferIn)
out := s.State.NetOutTransfer - uint64(s.PrevHourlyTransferOut)
if in > 0 || out > 0 {
tx := model.Transfer{
ServerID: s.ID,
In: in,
Out: out,
}
tx.CreatedAt = time.Now()
txs = append(txs, tx)
}
}
if err := dao.DB.Create(txs).Error; err != nil {
log.Println("流量统计入库", err)
}
srv.Shutdown(c)
return nil
})
2019-12-05 09:36:58 -05:00
}