-
Notifications
You must be signed in to change notification settings - Fork 0
/
mainhelper.go
155 lines (132 loc) · 3.57 KB
/
mainhelper.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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
package util
import (
"encoding/json"
"log"
"log/slog"
"os"
"os/exec"
"os/signal"
"strconv"
"strings"
"syscall"
"time"
)
// Main helper provides boilerplate utilities for running a main()
// and loading initial config.
var startupTime = time.Now()
// FindConfig is a simple loader for a config file.
func FindConfig(base string, s string) []byte {
basecfg := os.Getenv(base)
if basecfg != "" {
return []byte(basecfg)
}
fb, err := os.ReadFile("./" + base + s)
if err == nil {
return fb
}
fb, err = os.ReadFile("/" + base + "/" + base + s)
if err == nil {
return fb
}
// Also look in the .ssh directory - this is mainly for secrets.
fb, err = os.ReadFile(os.Getenv("HOME") + "/.ssh/" + base + s)
if err == nil {
return fb
}
return nil
}
// MainStart is an opinionated startup - configures build in components.
// 'base' is the name of the config - for example 'mds'
// If it is set as an environment variable - it is expected to be a json config.
// Otherwise, a file /$base/$base.json or ./$base.json will be loaded.
// Other env variables of type string may be merged into the config.
//
// - Will init slog with a json handler
//
// Larger binaries should use viper - which provides support for:
// - ini, json, yaml, java properties
// - remote providers (with encryption) - built in etcd3, consul, firestore
func MainStart(base string, out interface{}) error {
basecfg := FindConfig(base, ".json")
if basecfg != nil {
err := json.Unmarshal([]byte(basecfg), out)
if err != nil {
return err
}
}
// Quick hack to load environment variables into the config struct.
envl := os.Environ()
envm := map[string]string{}
for _, k := range envl {
kv := strings.SplitN(k, "=", 2)
if len(kv) == 2 {
envm[kv[0]] = kv[1]
}
}
envb, err := json.Marshal(envm)
if err != nil {
log.Println("Failed to overlay env", envl, err, envb)
}
return json.Unmarshal(envb, out)
}
// Main config helper - base implementation for minimal deps CLI.
//
// Larger binaries should use viper - which provides support for:
// - ini, json, yaml, java properties
// - remote providers (with encryption) - built in etcd3, consul, firestore
func GetString(key string) string {
return os.Getenv(key)
}
// MainEnd should be the last call in main(). The app is expected to get all the config
// from file or env variables - if the command line arguments are not empty: exec the remaining
// and wait to complete - else wait for a signal.
func MainEnd() {
sigCh := make(chan os.Signal, 1)
signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
var cmd string
var argv []string
posArgs := os.Args
if len(posArgs) == 1 {
for {
sig := <-sigCh
slog.Info("Exit", "sig", sig, "running", time.Since(startupTime))
d := GetString("DRAIN_TIMEOUT")
if d == "" {
d = "1"
}
di, _ := strconv.Atoi(d)
time.AfterFunc(time.Second*time.Duration(di), func() {
os.Exit(0)
})
// Testing force exit timing
// return
}
}
// If it has extra args, exec the command
if len(posArgs) > 2 {
cmd, argv = posArgs[1], posArgs[2:]
} else {
cmd = posArgs[1]
}
c := exec.Command(cmd, argv...)
c.Stdout = os.Stdout
c.Stderr = os.Stderr
c.Stdin = os.Stdin
c.Env = os.Environ()
if err := c.Start(); err != nil {
slog.Error("failed to start subprocess", "cmd", cmd, "args", argv, "err", err)
os.Exit(c.ProcessState.ExitCode())
}
go func() {
sig := <-sigCh
if err := c.Process.Signal(sig); err != nil {
log.Printf("failed to signal process: %v", err)
}
}()
if err := c.Wait(); err != nil {
if v, ok := err.(*exec.ExitError); ok {
ec := v.ExitCode()
os.Exit(ec)
}
}
}