367 lines
8.6 KiB
Go
367 lines
8.6 KiB
Go
package main
|
||
|
||
import (
|
||
"fmt"
|
||
"os"
|
||
"os/exec"
|
||
"os/signal"
|
||
"path/filepath"
|
||
"strconv"
|
||
"strings"
|
||
"syscall"
|
||
"time"
|
||
"version/pkg/handler"
|
||
|
||
client "backend-service/pkg/api"
|
||
"backend-service/pkg/machine"
|
||
"backend-service/pkg/service"
|
||
"backend-service/pkg/version"
|
||
"bash_go_service/shared/pkg/constants"
|
||
"bash_go_service/shared/pkg/logger"
|
||
|
||
"github.com/spf13/cobra"
|
||
"github.com/spf13/viper"
|
||
"golang.org/x/sys/unix"
|
||
)
|
||
|
||
var rootCmd = &cobra.Command{
|
||
Use: "bash_service",
|
||
Short: "北溟 BASH 后台服务",
|
||
Long: `欢迎使用北溟云计算服务，这是BASH增强服务，如出现问题请联系管理员`,
|
||
Run: func(cmd *cobra.Command, args []string) {
|
||
if err := run(); err != nil {
|
||
logger.Error("Application failed: %v", err)
|
||
os.Exit(1)
|
||
}
|
||
},
|
||
}
|
||
|
||
var daemonCmd = &cobra.Command{
|
||
Use: "daemon",
|
||
Short: "Run the application as a daemon",
|
||
Run: func(cmd *cobra.Command, args []string) {
|
||
if err := checkExistingProcess(); err != nil {
|
||
logger.Error("Error starting daemon.")
|
||
return
|
||
}
|
||
pid := os.Getpid()
|
||
pidStr := []byte(fmt.Sprintf("%d", pid))
|
||
if err := os.WriteFile(constants.PidFilePath, pidStr, constants.LogFileMode); err != nil {
|
||
logger.Error("Failed to write PID file: %v", err)
|
||
os.Exit(1)
|
||
}
|
||
runDaemon()
|
||
},
|
||
}
|
||
|
||
var forceCmd = &cobra.Command{
|
||
Use: "force",
|
||
Short: "Force stop the daemon and restart",
|
||
Run: func(cmd *cobra.Command, args []string) {
|
||
logger.Debug("Force start...")
|
||
if err := forceStopDaemon(); err != nil {
|
||
logger.Error("Error stopping daemon.")
|
||
return
|
||
}
|
||
if err := run(); err != nil {
|
||
logger.Error("Application failed: %v", err)
|
||
os.Exit(1)
|
||
}
|
||
},
|
||
}
|
||
|
||
var quitCmd = &cobra.Command{
|
||
Use: "quit",
|
||
Short: "Force stop the daemon and quit",
|
||
Run: func(cmd *cobra.Command, args []string) {
|
||
logger.Debug("Force quit...")
|
||
if err := forceStopDaemon(); err != nil {
|
||
logger.Error("Error stopping daemon.")
|
||
return
|
||
}
|
||
},
|
||
}
|
||
|
||
var updateCmd = &cobra.Command{
|
||
Use: "update",
|
||
Short: "Update the application",
|
||
Run: func(cmd *cobra.Command, args []string) {
|
||
logger.Debug("Update command executed.")
|
||
need, err := version.CheckUpdate()
|
||
if err != nil {
|
||
logger.Error("Error checking update: %v", err)
|
||
return
|
||
}
|
||
if need {
|
||
logger.Info("Update needed. Restarting the application...")
|
||
|
||
executable, err := os.Executable()
|
||
if err != nil {
|
||
logger.Error("Error getting executable path: %v", err)
|
||
return
|
||
}
|
||
|
||
env := os.Environ()
|
||
|
||
cmd := exec.Command(executable)
|
||
cmd.Env = env
|
||
|
||
err = cmd.Start()
|
||
if err != nil {
|
||
logger.Error("Error restarting the application: %v", err)
|
||
return
|
||
}
|
||
|
||
logger.Info("Application restarted successfully.")
|
||
os.Exit(0)
|
||
}
|
||
},
|
||
}
|
||
|
||
func init() {
|
||
cobra.OnInitialize(initConfig, checkAndUpdate)
|
||
rootCmd.AddCommand(daemonCmd)
|
||
rootCmd.AddCommand(forceCmd)
|
||
rootCmd.AddCommand(quitCmd)
|
||
rootCmd.AddCommand(updateCmd)
|
||
}
|
||
|
||
func checkAndUpdate() {
|
||
exePath, err := os.Executable()
|
||
if err != nil {
|
||
logger.Error("Error getting executable path: %v", err)
|
||
return
|
||
}
|
||
info, err := handler.LoadUpdateInfo(exePath)
|
||
if err != nil {
|
||
logger.Debug("Error loading update info: %v", err)
|
||
return
|
||
}
|
||
|
||
// 删除info文件，下次再试
|
||
err = handler.RemoveUpdateInfo(exePath)
|
||
if err != nil {
|
||
logger.Error("Error removing update info: %v", err)
|
||
return
|
||
}
|
||
|
||
exeDir := filepath.Dir(exePath)
|
||
// 先更新so文件
|
||
logger.Debug("Updating so file...")
|
||
err = handler.ExecuteUpdate(exeDir, info.SoFileNewPath, info.SoFileOldPath, info.SoTargetPath, false)
|
||
if err != nil {
|
||
logger.Error("Error updating so file: %v", err)
|
||
return
|
||
}
|
||
// 更新exe程序
|
||
logger.Debug("Updating exe file...")
|
||
forceStopDaemon()
|
||
err = handler.ExecuteUpdate(exeDir, info.ExeFileNewPath, info.ExeFileOldPath, info.ExeTargetPath, true)
|
||
if err != nil {
|
||
logger.Error("Error updating exe file: %v", err)
|
||
return
|
||
}
|
||
}
|
||
|
||
func main() {
|
||
if err := rootCmd.Execute(); err != nil {
|
||
logger.Error("Error executing command: %v", err)
|
||
os.Exit(1)
|
||
}
|
||
}
|
||
|
||
func initConfig() {
|
||
var configDir string
|
||
if isDebug() {
|
||
configDir = constants.ConfigPath
|
||
logger.Debug("Debug mode, using config path: %s", configDir)
|
||
} else {
|
||
exePath, err := os.Executable()
|
||
if err != nil {
|
||
logger.Error("Error getting executable path: %v", err)
|
||
return
|
||
}
|
||
exeDir := filepath.Dir(exePath)
|
||
configDir = filepath.Join(exeDir, constants.ConfigPath)
|
||
}
|
||
viper.AddConfigPath(configDir)
|
||
viper.SetConfigName("config")
|
||
viper.SetConfigType("yaml")
|
||
if err := viper.ReadInConfig(); err != nil {
|
||
logger.Error("Error initializing config: %v", err)
|
||
}
|
||
logger.UpdateLogLevel()
|
||
logger.Debug("Initializing...")
|
||
}
|
||
|
||
func run() error {
|
||
info, err := collectAndSendMachineInfo()
|
||
if err != nil {
|
||
return fmt.Errorf("collect and send machine info: %w", err)
|
||
}
|
||
showWelcomeMessage(info)
|
||
if err := handleDaemonProcess(); err != nil {
|
||
logger.Debug("Handle daemon process failed: %s", err)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func forceStopDaemon() error {
|
||
pidBytes, err := os.ReadFile(constants.PidFilePath)
|
||
if err != nil {
|
||
if os.IsNotExist(err) {
|
||
return nil
|
||
}
|
||
return fmt.Errorf("read pid file: %w", err)
|
||
}
|
||
pid, err := strconv.Atoi(strings.TrimSpace(string(pidBytes)))
|
||
if err != nil {
|
||
os.Remove(constants.PidFilePath)
|
||
return nil
|
||
}
|
||
process, err := os.FindProcess(pid)
|
||
if err != nil {
|
||
os.Remove(constants.PidFilePath)
|
||
return nil
|
||
}
|
||
if err := process.Signal(syscall.SIGTERM); err != nil {
|
||
if err == unix.ESRCH {
|
||
os.Remove(constants.PidFilePath)
|
||
return nil
|
||
}
|
||
if err := process.Kill(); err != nil {
|
||
return fmt.Errorf("kill process: %w", err)
|
||
}
|
||
}
|
||
for i := 0; i < 10; i++ {
|
||
if err := unix.Kill(pid, 0); err != nil {
|
||
if err == unix.ESRCH {
|
||
os.Remove(constants.PidFilePath)
|
||
return nil
|
||
}
|
||
}
|
||
time.Sleep(100 * time.Millisecond)
|
||
}
|
||
return fmt.Errorf("failed to stop daemon process")
|
||
}
|
||
|
||
func collectAndSendMachineInfo() (*client.Result, error) {
|
||
info := machine.CollectMachineInfo()
|
||
result, err := client.SendMachineInfo(info)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("send machine info: %w", err)
|
||
}
|
||
return &result, nil
|
||
}
|
||
|
||
func handleDaemonProcess() error {
|
||
exePath, err := os.Executable()
|
||
if err != nil {
|
||
return fmt.Errorf("get executable path: %w", err)
|
||
}
|
||
logger.Debug("Exe path: %s", filepath.Dir(exePath))
|
||
if err := checkExistingProcess(); err != nil {
|
||
return err
|
||
}
|
||
return startDaemonProcess()
|
||
}
|
||
|
||
func checkExistingProcess() error {
|
||
pidFile := constants.PidFilePath
|
||
pidBytes, err := os.ReadFile(pidFile)
|
||
if err != nil {
|
||
logger.Debug("Pid file not exist.")
|
||
if os.IsNotExist(err) {
|
||
return nil
|
||
}
|
||
return fmt.Errorf("read pid file: %w", err)
|
||
}
|
||
pid, err := strconv.Atoi(strings.TrimSpace(string(pidBytes)))
|
||
if err != nil {
|
||
logger.Debug("Pid file is broken.")
|
||
os.Remove(pidFile)
|
||
return nil
|
||
}
|
||
process, err := os.FindProcess(pid)
|
||
if err != nil {
|
||
logger.Debug("Process not found.")
|
||
os.Remove(pidFile)
|
||
return nil
|
||
} else {
|
||
logger.Debug("Process found, pid is: %d", process.Pid)
|
||
}
|
||
if err := unix.Kill(pid, 0); err != nil {
|
||
logger.Debug("Kill process error: %v", err)
|
||
if err == unix.ESRCH {
|
||
logger.Debug("Process not exist.")
|
||
os.Remove(pidFile)
|
||
return nil
|
||
}
|
||
}
|
||
logger.Info("Daemon already running with PID: %d", pid)
|
||
return fmt.Errorf("daemon already running")
|
||
}
|
||
|
||
func formatGoTime(format string) string {
|
||
replacer := strings.NewReplacer(
|
||
"%Y", "2006",
|
||
"%m", "01",
|
||
"%d", "02",
|
||
"%H", "15",
|
||
"%M", "04",
|
||
"%S", "05",
|
||
)
|
||
return replacer.Replace(format)
|
||
}
|
||
|
||
func startDaemonProcess() error {
|
||
if err := os.MkdirAll(constants.LogFilePath, constants.LogFileMode); err != nil {
|
||
return fmt.Errorf("create logs directory: %w", err)
|
||
}
|
||
goTimeFormat := formatGoTime(constants.LogNameFormate)
|
||
logFileName := time.Now().Format(goTimeFormat)
|
||
logFilePath := filepath.Join(constants.LogFilePath, logFileName)
|
||
logFile, err := os.OpenFile(logFilePath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, constants.LogFileMode)
|
||
if err != nil {
|
||
return fmt.Errorf("open log file: %w", err)
|
||
}
|
||
cmd := exec.Command(os.Args[0], constants.DaemonFlag)
|
||
cmd.Stdout = logFile
|
||
cmd.Stderr = logFile
|
||
cmd.SysProcAttr = &syscall.SysProcAttr{
|
||
Setsid: true,
|
||
}
|
||
if err := cmd.Start(); err != nil {
|
||
return fmt.Errorf("start daemon: %w", err)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func runDaemon() {
|
||
// daemon的默认日志等级
|
||
logger.SetLevel(logger.DEBUG)
|
||
logger.Info("Daemon started with PID: %d", os.Getpid())
|
||
sigChan := make(chan os.Signal, 1)
|
||
signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT, syscall.SIGQUIT)
|
||
go func() {
|
||
sig := <-sigChan
|
||
logger.Info("Received signal: %v", sig)
|
||
os.Remove(constants.PidFilePath)
|
||
os.Exit(0)
|
||
}()
|
||
service.Run()
|
||
}
|
||
|
||
func showWelcomeMessage(info *client.Result) {
|
||
fmt.Println(info.Msg)
|
||
}
|
||
|
||
func isDebug() bool {
|
||
for _, arg := range os.Args[1:] {
|
||
if arg == constants.DebugFlag {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|