-
Notifications
You must be signed in to change notification settings - Fork 1.3k
/
main.go
181 lines (162 loc) · 6.89 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
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
// main.go
//
// This source file is part of the FoundationDB open source project
//
// Copyright 2021 Apple Inc. and the FoundationDB project authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package main
import (
"bufio"
"context"
"fmt"
"io"
"os"
"regexp"
"strings"
"github.com/go-logr/logr"
"github.com/go-logr/zapr"
"github.com/spf13/pflag"
"go.uber.org/zap"
"go.uber.org/zap/zapcore"
"gopkg.in/natefinch/lumberjack.v2"
)
var (
fdbserverPath string
versionFilePath string
sharedBinaryDir string
monitorConfFile string
logPath string
executionModeString string
outputDir string
mainContainerVersion string
additionalEnvFile string
listenAddress string
processCount int
enablePprof bool
)
type executionMode string
const (
executionModeLauncher executionMode = "launcher"
executionModeInit executionMode = "init"
executionModeSidecar executionMode = "sidecar"
)
func initLogger(logPath string) *zap.Logger {
var logWriter io.Writer
if logPath != "" {
lumberjackLogger := &lumberjack.Logger{
Filename: logPath,
MaxSize: 100,
MaxAge: 7,
MaxBackups: 2,
Compress: false,
}
logWriter = io.MultiWriter(os.Stdout, lumberjackLogger)
} else {
logWriter = os.Stdout
}
return zap.New(zapcore.NewCore(zapcore.NewJSONEncoder(zap.NewProductionEncoderConfig()), zapcore.AddSync(logWriter), zapcore.InfoLevel))
}
func main() {
var copyFiles, copyBinaries, copyLibraries, requiredCopyFiles []string
var inputDir, copyPrimaryLibrary, binaryOutputDirectory string
pflag.StringVar(&executionModeString, "mode", "launcher", "Execution mode. Valid options are launcher, sidecar, and init")
pflag.StringVar(&fdbserverPath, "fdbserver-path", "/usr/bin/fdbserver", "Path to the fdbserver binary")
pflag.StringVar(&inputDir, "input-dir", ".", "Directory containing input files")
pflag.StringVar(&monitorConfFile, "input-monitor-conf", "config.json", "Name of the file in the input directory that contains the monitor configuration")
pflag.StringVar(&logPath, "log-path", "", "Name of a file to send logs to. Logs will be sent to stdout in addition the file you pass in this argument. If this is blank, logs will only by sent to stdout")
pflag.StringVar(&outputDir, "output-dir", ".", "Directory to copy files into")
pflag.StringArrayVar(©Files, "copy-file", nil, "A list of files to copy")
pflag.StringArrayVar(©Binaries, "copy-binary", nil, "A list of binaries to copy from /usr/bin")
pflag.StringVar(&versionFilePath, "version-file", "/var/fdb/version", "Path to a file containing the current FDB version")
pflag.StringVar(&sharedBinaryDir, "shared-binary-dir", "/var/fdb/shared-binaries/bin", "A directory containing binaries that are copied from a sidecar process")
pflag.StringVar(&binaryOutputDirectory, "binary-output-dir", "", "A subdirectory within $(output-dir)/bin to store binaries in. This defaults to the value in /var/fdb/version")
pflag.StringArrayVar(©Libraries, "copy-library", nil, "A list of libraries to copy from /usr/lib/fdb/multiversion to $(output-dir)/lib/multiversion")
pflag.StringVar(©PrimaryLibrary, "copy-primary-library", "", "A library to copy from /usr/lib/fdb/multiversion to $(output-dir)/lib. This file will be renamed to libfdb_c.so")
pflag.StringArrayVar(&requiredCopyFiles, "require-not-empty", nil, "When copying this file, exit with an error if the file is empty")
pflag.StringVar(&mainContainerVersion, "main-container-version", "", "For sidecar mode, this specifies the version of the main container. If this is equal to the current container version, no files will be copied")
pflag.StringVar(&additionalEnvFile, "additional-env-file", "", "A file with additional environment variables to use when interpreting the monitor configuration")
pflag.IntVar(&processCount, "process-count", 1, "The number of processes to start")
pflag.BoolVar(&enablePprof, "enable-pprof", false, "Enables /debug/pprof endpoints on the listen address")
pflag.StringVar(&listenAddress, "listen-address", ":8081", "An address and port to listen on")
pflag.Parse()
logger := zapr.NewLogger(initLogger(logPath))
versionBytes, err := os.ReadFile(versionFilePath)
if err != nil {
panic(err)
}
currentContainerVersion := strings.TrimSpace(string(versionBytes))
copyDetails, requiredCopies, err := getCopyDetails(inputDir, copyPrimaryLibrary, binaryOutputDirectory, copyFiles, copyBinaries, copyLibraries, requiredCopyFiles, currentContainerVersion)
if err != nil {
logger.Error(err, "Error getting list of files to copy")
os.Exit(1)
}
mode := executionMode(executionModeString)
switch mode {
case executionModeLauncher:
customEnvironment, err := loadAdditionalEnvironment(logger)
if err != nil {
logger.Error(err, "Error loading additional environment")
os.Exit(1)
}
StartMonitor(context.Background(), logger, fmt.Sprintf("%s/%s", inputDir, monitorConfFile), customEnvironment, processCount, listenAddress, enablePprof, currentContainerVersion)
case executionModeInit:
err = CopyFiles(logger, outputDir, copyDetails, requiredCopies)
if err != nil {
logger.Error(err, "Error copying files")
os.Exit(1)
}
case executionModeSidecar:
if mainContainerVersion != currentContainerVersion {
err = CopyFiles(logger, outputDir, copyDetails, requiredCopies)
if err != nil {
logger.Error(err, "Error copying files")
os.Exit(1)
}
}
logger.Info("Waiting for process to be terminated")
done := make(chan bool)
<-done
default:
logger.Error(nil, "Unknown execution mode", "mode", mode)
os.Exit(1)
}
}
func loadAdditionalEnvironment(logger logr.Logger) (map[string]string, error) {
var customEnvironment = make(map[string]string)
if additionalEnvFile != "" {
environmentPattern := regexp.MustCompile(`export ([A-Za-z0-9_]+)=([^\n]*)`)
file, err := os.Open(additionalEnvFile)
if err != nil {
return nil, err
}
defer func() {
err := file.Close()
if err != nil {
logger.Error(err, "error when closing file")
}
}()
envScanner := bufio.NewScanner(file)
for envScanner.Scan() {
envLine := envScanner.Text()
matches := environmentPattern.FindStringSubmatch(envLine)
if matches == nil || envLine == "" {
logger.Error(nil, "Environment file contains line that we cannot parse", "line", envLine, "environmentPattern", environmentPattern)
continue
}
customEnvironment[matches[1]] = matches[2]
}
}
return customEnvironment, nil
}