-
Notifications
You must be signed in to change notification settings - Fork 0
/
root.go
124 lines (110 loc) · 2.83 KB
/
root.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
package cmd
import (
"fmt"
"log"
"os"
"os/exec"
"path"
"path/filepath"
"strconv"
"strings"
"github.com/judwhite/go-svc"
"github.com/samlau0508/imserver/internal/server"
"github.com/samlau0508/imserver/pkg/oklog"
"github.com/spf13/cobra"
"github.com/spf13/viper"
"go.uber.org/zap"
)
var (
cfgFile string
serverOpts = server.NewOptions()
mode string
daemon bool
pidfile string = "imPID"
installDir string
rootCmd = &cobra.Command{
Use: "im",
Short: "IM, a sleek and high-performance instant messaging platform.",
Long: `IM, a sleek and high-performance instant messaging platform. For more details, please refer to the documentation: https://githubim.com`,
CompletionOptions: cobra.CompletionOptions{
DisableDefaultCmd: true,
},
Run: func(cmd *cobra.Command, args []string) {
initServer()
},
}
)
func init() {
homeDir, err := server.GetHomeDir()
if err != nil {
log.Fatal(err)
}
installDir = homeDir
cobra.OnInitialize(initConfig)
rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file")
rootCmd.PersistentFlags().StringVar(&mode, "mode", "debug", "mode")
// 后台运行
rootCmd.PersistentFlags().BoolVarP(&daemon, "daemon", "d", false, "run in daemon mode")
}
func initConfig() {
vp := viper.New()
if cfgFile != "" {
vp.SetConfigFile(cfgFile)
if err := vp.ReadInConfig(); err != nil {
fmt.Println("Using config file:", vp.ConfigFileUsed(), zap.Error(err))
}
}
vp.SetEnvPrefix("im")
vp.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
vp.AutomaticEnv()
// 初始化服务配置
serverOpts.ConfigureWithViper(vp)
}
func initServer() {
logOpts := oklog.NewOptions()
logOpts.Level = serverOpts.Logger.Level
logOpts.LogDir = serverOpts.Logger.Dir
logOpts.LineNum = serverOpts.Logger.LineNum
oklog.Configure(logOpts)
s := server.New(serverOpts)
if daemon {
filePath, _ := filepath.Abs(os.Args[0])
args := os.Args[1:]
newArgs := make([]string, 0)
for _, arg := range args {
if arg == "-d" || arg == "--daemon" {
continue
}
newArgs = append(newArgs, arg)
}
cmd := exec.Command(filePath, newArgs...)
// 将其他命令传入生成出的进程
// cmd.Stdin = os.Stdin // 给新进程设置文件描述符,可以重定向到文件中
// cmd.Stdout = os.Stdout
// cmd.Stderr = os.Stderr
err := cmd.Start() // 开始执行新进程,不等待新进程退出
if err != nil {
log.Fatal(err)
}
err = os.WriteFile(path.Join(installDir, pidfile), []byte(strconv.Itoa(cmd.Process.Pid)), 0644)
if err != nil {
log.Fatal(err)
}
os.Exit(0)
} else {
if err := svc.Run(s); err != nil {
log.Fatal(err)
}
}
}
func addCommand(cmd CMD) {
rootCmd.AddCommand(cmd.CMD())
}
func Execute() {
ctx := &IMContext{}
addCommand(newStopCMD(ctx))
if err := rootCmd.Execute(); err != nil {
fmt.Fprintln(os.Stderr, err)
os.Exit(1)
}
}