package main import ( "context" "fmt" "log" "net" "time" _ "time/tzdata" "github.com/ory/graceful" "github.com/soheilhy/cmux" flag "github.com/spf13/pflag" "golang.org/x/crypto/bcrypt" "github.com/naiba/nezha/cmd/dashboard/controller" "github.com/naiba/nezha/cmd/dashboard/rpc" "github.com/naiba/nezha/model" "github.com/naiba/nezha/proto" "github.com/naiba/nezha/service/singleton" ) type DashboardCliParam struct { Version bool // 当前版本号 ConfigFile string // 配置文件路径 DatebaseLocation string // Sqlite3 数据库文件路径 } var ( dashboardCliParam DashboardCliParam ) func init() { flag.CommandLine.ParseErrorsWhitelist.UnknownFlags = true flag.BoolVarP(&dashboardCliParam.Version, "version", "v", false, "查看当前版本号") flag.StringVarP(&dashboardCliParam.ConfigFile, "config", "c", "data/config.yaml", "配置文件路径") flag.StringVar(&dashboardCliParam.DatebaseLocation, "db", "data/sqlite.db", "Sqlite3数据库文件路径") flag.Parse() // 初始化 dao 包 singleton.InitConfigFromPath(dashboardCliParam.ConfigFile) singleton.InitTimezoneAndCache() singleton.InitDBFromPath(dashboardCliParam.DatebaseLocation) initSystem() } func initSystem() { // 初始化管理员账户 var usersCount int64 if err := singleton.DB.Model(&model.User{}).Count(&usersCount).Error; err != nil { panic(err) } if usersCount == 0 { hash, err := bcrypt.GenerateFromPassword([]byte("admin"), bcrypt.DefaultCost) if err != nil { panic(err) } admin := model.User{ Username: "admin", Password: string(hash), } if err := singleton.DB.Create(&admin).Error; err != nil { panic(err) } } // 启动 singleton 包下的所有服务 singleton.LoadSingleton() // 每天的3:30 对 监控记录 和 流量记录 进行清理 if _, err := singleton.Cron.AddFunc("0 30 3 * * *", singleton.CleanMonitorHistory); err != nil { panic(err) } // 每小时对流量记录进行打点 if _, err := singleton.Cron.AddFunc("0 0 * * * *", singleton.RecordTransferHourlyUsage); err != nil { panic(err) } } func main() { if dashboardCliParam.Version { fmt.Println(singleton.Version) return } l, err := net.Listen("tcp", fmt.Sprintf(":%d", singleton.Conf.ListenPort)) if err != nil { log.Fatal(err) } m := cmux.New(l) grpcL := m.Match(cmux.HTTP2HeaderField("content-type", "application/grpc")) httpL := m.Match(cmux.HTTP1Fast()) // TODO 使用 cmux 在同一端口服务 HTTP 和 gRPC singleton.CleanMonitorHistory() go rpc.ServeRPC(grpcL) serviceSentinelDispatchBus := make(chan model.Monitor) // 用于传递服务监控任务信息的channel go rpc.DispatchTask(serviceSentinelDispatchBus) go rpc.DispatchKeepalive() go singleton.AlertSentinelStart() singleton.NewServiceSentinel(serviceSentinelDispatchBus) srv := controller.ServeWeb() go dispatchReportInfoTask() if err := graceful.Graceful(func() error { return srv.Serve(httpL) }, func(c context.Context) error { log.Println("NEZHA>> Graceful::START") singleton.RecordTransferHourlyUsage() log.Println("NEZHA>> Graceful::END") srv.Shutdown(c) return nil }); err != nil { log.Printf("NEZHA>> ERROR: %v", err) } } func dispatchReportInfoTask() { time.Sleep(time.Second * 15) singleton.ServerLock.RLock() defer singleton.ServerLock.RUnlock() for _, server := range singleton.ServerList { if server == nil || server.TaskStream == nil { continue } server.TaskStream.Send(&proto.Task{ Type: model.TaskTypeReportHostInfo, Data: "", }) } }