/
main.go
149 lines (125 loc) · 3.83 KB
/
main.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
package main
import (
"os"
"os/signal"
"syscall"
log "github.com/sirupsen/logrus"
"github.com/spf13/pflag"
"github.com/spf13/viper"
"k8s.io/apimachinery/pkg/util/clock"
"github.com/armadaproject/armada/internal/common"
"github.com/armadaproject/armada/internal/common/armadacontext"
"github.com/armadaproject/armada/internal/common/database"
"github.com/armadaproject/armada/internal/common/logging"
"github.com/armadaproject/armada/internal/common/profiling"
"github.com/armadaproject/armada/internal/common/serve"
"github.com/armadaproject/armada/internal/lookoutv2"
"github.com/armadaproject/armada/internal/lookoutv2/configuration"
"github.com/armadaproject/armada/internal/lookoutv2/gen/restapi"
"github.com/armadaproject/armada/internal/lookoutv2/pruner"
"github.com/armadaproject/armada/internal/lookoutv2/schema"
)
const (
CustomConfigLocation string = "config"
MigrateDatabase = "migrateDatabase"
PruneDatabase = "pruneDatabase"
)
func init() {
pflag.StringSlice(
CustomConfigLocation,
[]string{},
"Fully qualified path to application configuration file (for multiple config files repeat this arg or separate paths with commas)",
)
pflag.Bool(MigrateDatabase, false, "Migrate database instead of running server")
pflag.Bool(PruneDatabase, false, "Prune database of old jobs instead of running server")
pflag.Parse()
}
func makeContext() (*armadacontext.Context, func()) {
ctx := armadacontext.Background()
ctx, cancel := armadacontext.WithCancel(ctx)
c := make(chan os.Signal, 1)
signal.Notify(c, os.Interrupt, syscall.SIGTERM)
go func() {
select {
case <-c:
cancel()
case <-ctx.Done():
}
}()
return ctx, func() {
signal.Stop(c)
cancel()
}
}
func migrate(ctx *armadacontext.Context, config configuration.LookoutV2Config) {
db, err := database.OpenPgxPool(config.Postgres)
if err != nil {
panic(err)
}
migrations, err := schema.LookoutMigrations()
if err != nil {
panic(err)
}
err = database.UpdateDatabase(ctx, db, migrations)
if err != nil {
panic(err)
}
}
func prune(ctx *armadacontext.Context, config configuration.LookoutV2Config) {
db, err := database.OpenPgxConn(config.Postgres)
if err != nil {
panic(err)
}
if config.PrunerConfig.Timeout <= 0 {
panic("timeout must be greater than 0")
}
if config.PrunerConfig.ExpireAfter <= 0 {
panic("expireAfter must be greater than 0")
}
if config.PrunerConfig.BatchSize <= 0 {
panic("batchSize must be greater than 0")
}
log.Infof("expireAfter: %v, batchSize: %v, timeout: %v",
config.PrunerConfig.ExpireAfter, config.PrunerConfig.BatchSize, config.PrunerConfig.Timeout)
ctxTimeout, cancel := armadacontext.WithTimeout(ctx, config.PrunerConfig.Timeout)
defer cancel()
err = pruner.PruneDb(ctxTimeout, db, config.PrunerConfig.ExpireAfter, config.PrunerConfig.BatchSize, clock.RealClock{})
if err != nil {
panic(err)
}
}
func main() {
common.ConfigureLogging()
common.BindCommandlineArguments()
var config configuration.LookoutV2Config
userSpecifiedConfigs := viper.GetStringSlice(CustomConfigLocation)
common.LoadConfig(&config, "./config/lookoutv2", userSpecifiedConfigs)
// Expose profiling endpoints if enabled.
if config.PprofPort != nil {
pprofServer := profiling.SetupPprofHttpServer(*config.PprofPort)
go func() {
ctx := armadacontext.Background()
if err := serve.ListenAndServe(ctx, pprofServer); err != nil {
logging.WithStacktrace(ctx, err).Error("pprof server failure")
}
}()
}
log.SetLevel(log.DebugLevel)
ctx, cleanup := makeContext()
defer cleanup()
if viper.GetBool(MigrateDatabase) {
log.Info("Migrating database")
migrate(ctx, config)
return
}
if viper.GetBool(PruneDatabase) {
log.Info("Pruning database")
prune(ctx, config)
return
}
restapi.UIConfig = config.UIConfig
if err := lookoutv2.Serve(config); err != nil {
log.Error(err)
os.Exit(1)
}
}