DB1000N | Add app. auto-restart upon successful update (arriven#317)

* DB1000N | Add app. auto-restart upon successful update * DB1000N | Refactor flags values * DB1000N | Fix build for Windows, fix linting issues * DB1000N | Fix linting issues * DB1000N | Fix build for Windows, Fix linters issues * DB1000N | Fix MD linters issues * DB1000N | Fix backward compatibility with Go 1.16 * DB1000N | Allow restarts only on Linux & Darwin systems * DB1000N | For god sake fix for the test compilation on Windows
DmitryZ1 · Mar 13, 2022 · 1aab748 · 1aab748
1 parent 24a092a
commit 1aab748
Show file tree

Hide file tree

Showing 9 changed files with 367 additions and 85 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -5,8 +5,7 @@ WORKDIR /build
 COPY go.mod .
 RUN go mod download && go mod verify
 COPY . .
-# use -s -w to strip extra debug data
-RUN make LDFLAGS="-s -w" build_encrypted
+RUN make build_encrypted
 
 FROM alpine:3.11.3 as advanced
 

diff --git a/main.go b/main.go
@@ -43,10 +43,14 @@ import (
 	"github.com/Arriven/db1000n/src/utils/templates"
 )
 
+const (
+	DefaultUpdateCheckFrequency = 24 * time.Hour
+)
+
 func main() {
 	log.SetOutput(os.Stdout)
 	log.SetFlags(log.Ldate | log.Lmicroseconds | log.Lshortfile | log.LUTC)
-	log.Printf("DB1000n [Version: %s]", ota.Version)
+	log.Printf("DB1000n [Version: %s][PID=%d]\n", ota.Version, os.Getpid())
 
 	configPaths := flag.String("c", utils.GetEnvStringDefault("CONFIG", "https://raw.githubusercontent.com/db1000n-coordinators/LoadTestConfig/main/config.v0.7.json"), "path to config files, separated by a comma, each path can be a web endpoint")
 	backupConfig := flag.String("b", config.DefaultConfig, "raw backup config in case the primary one is unavailable")
@@ -60,7 +64,10 @@ func main() {
 	configFormat := flag.String("format", utils.GetEnvStringDefault("CONFIG_FORMAT", "json"), "config format")
 	prometheusOn := flag.Bool("prometheus_on", utils.GetEnvBoolDefault("PROMETHEUS_ON", false), "Start metrics exporting via HTTP and pushing to gateways (specified via <prometheus_gateways>)")
 	prometheusPushGateways := flag.String("prometheus_gateways", utils.GetEnvStringDefault("PROMETHEUS_GATEWAYS", ""), "Comma separated list of prometheus push gateways")
-	doSelfUpdate := flag.Bool("enable-self-update", utils.GetEnvBoolDefault("ENABLE_SELF_UPDATE", false), "Enable the application automatic updates on the startup")
+	doAutoUpdate := flag.Bool("enable-self-update", utils.GetEnvBoolDefault("ENABLE_SELF_UPDATE", false), "Enable the application automatic updates on the startup")
+	doRestartOnUpdate := flag.Bool("restart-on-update", utils.GetEnvBoolDefault("RESTART_ON_UPDATE", true), "Allows application to restart upon successful update (ignored if auto-update is disabled)")
+	skipUpdateCheckOnStart := flag.Bool("skip-update-check-on-start", utils.GetEnvBoolDefault("SKIP_UPDATE_CHECK_ON_START", false), "Allows to skip the update check at the startup (usually set automatically by the previous version)")
+	autoUpdateCheckFrequency := flag.Duration("self-update-check-frequency", utils.GetEnvDurationDefault("SELF_UPDATE_CHECK_FREQUENCY", DefaultUpdateCheckFrequency), "How often to run auto-update checks")
 
 	flag.Parse()
 
@@ -70,8 +77,8 @@ func main() {
 		return
 	}
 
-	if *doSelfUpdate {
-		ota.DoSelfUpdate()
+	if *doAutoUpdate {
+		go watchUpdates(*doRestartOnUpdate, *skipUpdateCheckOnStart, *autoUpdateCheckFrequency)
 	}
 
 	setUpPprof(*pprof, *debug)
@@ -106,9 +113,13 @@ func main() {
 	}
 
 	go func() {
-		// Wait for sigterm
 		sigs := make(chan os.Signal, 1)
-		signal.Notify(sigs, syscall.SIGTERM)
+		signal.Notify(sigs,
+			syscall.SIGTERM,
+			syscall.SIGABRT,
+			syscall.SIGHUP,
+			syscall.SIGINT,
+		)
 		<-sigs
 		log.Println("Terminating")
 		cancel()
@@ -137,3 +148,53 @@ func setUpPprof(pprof string, debug bool) {
 		log.Println(http.ListenAndServe(pprof, mux))
 	}()
 }
+
+func watchUpdates(doRestartOnUpdate, skipUpdateCheckOnStart bool, autoUpdateCheckFrequency time.Duration) {
+	if !skipUpdateCheckOnStart {
+		runUpdate(doRestartOnUpdate)
+	} else {
+		log.Printf("Version update on startup is skipped, next update check is scheduled in %s",
+			autoUpdateCheckFrequency)
+	}
+
+	periodicalUpdateChecker := time.NewTicker(autoUpdateCheckFrequency)
+	defer periodicalUpdateChecker.Stop()
+
+	for range periodicalUpdateChecker.C {
+		runUpdate(doRestartOnUpdate)
+	}
+}
+
+//nolint:nestif // The nested if linter is disabled as it would add unnecessary function splitting. This function is quite obvious
+func runUpdate(doRestartOnUpdate bool) {
+	log.Println("Running a check for a newer version...")
+
+	isUpdateFound, newVersion, changeLog, err := ota.DoAutoUpdate()
+	if err != nil {
+		log.Printf("Auto-Update is failed: %s", err)
+
+		return
+	}
+
+	if isUpdateFound {
+		log.Printf("Newer version of the application is found [version=%s]\n", newVersion)
+		log.Printf("What's new:\n%s", changeLog)
+
+		if doRestartOnUpdate {
+			log.Println("Auto restart is enabled, restarting the application to run a new version")
+
+			additionalArgs := []string{
+				"-skip-update-check-on-start",
+			}
+
+			if err = ota.Restart(additionalArgs...); err != nil {
+				log.Printf("Failed to restart the application after the update to the new version: %s", err)
+				log.Printf("Restart the application manually to apply changes!\n")
+			}
+		} else {
+			log.Println("Auto restart is disabled, restart the application manually to apply changes!")
+		}
+	} else {
+		log.Println("We are running the latest version, OK!")
+	}
+}
diff --git a/src/utils/ota/README.md b/src/utils/ota/README.md
@@ -4,19 +4,12 @@ Lots of maintainers run their needles on a bare metal machines.
 As long as this project is so frequently updated, it might be
 a good idea to let them update it without the hassle.
 
-## TODO
-
-_TO BE CONFIRMED WITH THE CODE OWNER!_
-
-- [?] Enable automatic time-based version check
-- [?] Enable push-based version check
-- [?] Enable application self-restart after it downloaded the update
-- [?] Disable OTA updates for needles running inside the Docker container
+- [V] Enabled automatic time-based version check
+- [V] Enabled application self-restart after it downloaded the update
 
 ## Description
 
-Support for the application self-update by downloading
-the latest release from the official repository.
+Support for the application self-update by downloading the latest release from the official repository.
 
 ```text
 Stay strong, be the first in line!
@@ -27,26 +20,76 @@ target in the Makefile.
 
 ## Usage
 
+### Available flags
+
+```bash
+  -enable-self-update
+        Enable the application automatic updates on the startup
+  -restart-on-update
+        Allows application to restart upon the successful update (ignored if auto-update is disabled) (default true)
+  -self-update-check-frequency duration
+        How often to run auto-update checks (default 24h0m0s)
+  -skip-update-check-on-start
+        Allows to skip the update check at the startup (usually set automatically by the previous version) (default false)
+```
+
+The default behavior if the self-update enabled:
+
+```bash
+* Check for the update
+* If update is available - download it
+   * If auto-restart is enabled
+      * Notify the user that a newer version is available
+      * Fork-Exec a new process (will have a different PID), add a flag to skip the version check upon startup
+      * Stop the current process
+   * If auto-restart is disabled - notify user that manual restart is required
+* If update is NOT available - schedule the next check
+```
+
+### Examples
+
 To update your needle, start it with a flag `-enable-self-update`
 
 ```sh
 ./db1000n -enable-self-update
 ```
 
-### Update example
+#### Advanced options
+
+Start the needle with the **self-update & self-restart**
 
 ```bash
-$ make build
-CGO_ENABLED=0 go build -ldflags="-s -w -X 'github.com/Arriven/db1000n/ota.Version=v0.6.4'" -o db1000n -a ./main.go
 $ ./db1000n -enable-self-update
-1970/01/01 00:00:00.00000 main.go:76: DB1000n [Version: v0.6.4]
-1970/01/01 00:00:00.00000 ota.go:30: Successfully updated to version 0.6.5
-1970/01/01 00:00:00.00000 ota.go:31: Release note:
- ## What's Changed
-* User friendly logs by @Arriven in https://github.com/Arriven/db1000n/pull/271
-...
-
-**Full Changelog**: https://github.com/Arriven/db1000n/compare/v0.6.4...v0.6.5
-1970/01/01 00:00:00.00000 config.go:36: Loading config from "https://raw.githubusercontent.com/db1000n-coordinators/LoadTestConfig/main/config.json"
-1970/01/01 00:00:00.00000 config.go:97: New config received, applying
+0000/00/00 00:00:00 main.go:82: DB1000n [Version: v0.6.4][PID=75259]
+0000/00/00 00:00:00 main.go:166: Running a check for a newer version...
+0000/00/00 00:00:00 main.go:176: Newer version of the application is found [0.7.0]
+0000/00/00 00:00:00 main.go:177: What's new:
+* Added some great improvements
+* Added some spectacular bugs
+0000/00/00 00:00:00 main.go:180: Auto restart is enabled, restarting the application to run a new version
+0000/00/00 00:00:00 restart.go:45: new process has been started successfully [old_pid=75259,new_pid=75262]
+
+# NOTE: Process 75259 exited, Process 75262 has started with a flag to skip version check on the startup
+
+0000/00/00 00:00:00 main.go:82: DB1000n [Version: v0.7.0][PID=75262]
+0000/00/00 00:00:00 main.go:155: Version update on startup is skipped, next update check is scheduled in 24h0m0s
+```
+
+Start the needle with the self-update but do not restart the process upon update (`systemd` friendly)
+
+```bash
+$ ./db1000n -enable-self-update -self-update-check-frequency=5m -restart-on-update=false
+0000/00/00 00:00:00 main.go:82: DB1000n [Version: v0.6.4][PID=75320]
+0000/00/00 00:00:00 main.go:166: Running a check for a newer version...
+0000/00/00 00:00:00 main.go:176: Newer version of the application is found [0.7.0]
+0000/00/00 00:00:00 main.go:177: What's new:
+* Added some great improvements
+* Added some spectacular bugs
+0000/00/00 00:00:00 main.go:191: Auto restart is disabled, restart the application manually to apply changes!
 ```
+
+## References
+
+1. Graceful restart with zero downtime for TCP connection -  [https://github.com/Scalingo/go-graceful-restart-example](https://github.com/Scalingo/go-graceful-restart-example)
+1. Graceful restart with zero downtime for TCP connection (two variants) [https://github.com/rcrowley/goagain](https://github.com/rcrowley/goagain)
+1. Graceful restart with zero downtime for TCP connection (alternative) [https://grisha.org/blog/2014/06/03/graceful-restart-in-golang](https://grisha.org/blog/2014/06/03/graceful-restart-in-golang)
diff --git a/src/utils/ota/ota.go b/src/utils/ota/ota.go
@@ -2,80 +2,44 @@
 package ota
 
 import (
-	"bufio"
 	"fmt"
-	"log"
-	"os"
 
 	"github.com/blang/semver"
 	"github.com/rhysd/go-github-selfupdate/selfupdate"
 )
 
-// DoSelfUpdate updates the app to the latest version
-func DoSelfUpdate() {
+// DoAutoUpdate updates the app to the latest version.
+func DoAutoUpdate() (updateFound bool, newVersion, changeLog string, err error) {
 	v, err := semver.ParseTolerant(Version)
 	if err != nil {
-		log.Println("Binary version validation failed:", err)
+		err = fmt.Errorf("binary version validation failed: %w", err)
 
 		return
 	}
 
 	latest, err := selfupdate.UpdateSelf(v, Repository)
 	if err != nil {
-		log.Println("Binary update failed:", err)
+		err = fmt.Errorf("binary update failed: %w", err)
 
 		return
 	}
 
-	if latest.Version.Equals(v) {
-		// latest version is the same as current version. It means current binary is up-to-date.
-		log.Println("Current binary is the latest version", Version)
-	} else {
-		log.Println("Successfully updated to version", latest.Version)
-		log.Println("Release note:\n", latest.ReleaseNotes)
+	if !latest.Version.Equals(v) {
+		updateFound = true
+		newVersion = latest.Version.String()
+		changeLog = latest.ReleaseNotes
 	}
-}
-
-// ConfirmAndSelfUpdate ask for user confirmation to do a self-update
-func ConfirmAndSelfUpdate() {
-	latest, found, err := selfupdate.DetectLatest(Repository)
-	if err != nil {
-		log.Println("Error occurred while detecting version:", err)
-
-		return
-	}
-
-	if v := semver.MustParse(Version); !found || latest.Version.LTE(v) {
-		log.Println("Current version is the latest")
 
-		return
-	}
-
-	fmt.Print("Do you want to update to", latest.Version, "? (y/n): ") //nolint:forbidigo // Here we actually write to console and expect user input
-
-	input, err := bufio.NewReader(os.Stdin).ReadString('\n')
-	if err != nil || (input != "y\n" && input != "n\n") {
-		log.Println("Invalid input")
-
-		return
-	}
-
-	if input == "n\n" {
-		return
-	}
-
-	exe, err := os.Executable()
-	if err != nil {
-		log.Println("Could not locate executable path")
-
-		return
-	}
-
-	if err := selfupdate.UpdateTo(latest.AssetURL, exe); err != nil {
-		log.Println("Error occurred while updating binary:", err)
+	return
+}
 
-		return
+func MockAutoUpdate(shouldUpdateBeFound bool) (updateFound bool, newVersion, changeLog string, err error) {
+	updateFound = shouldUpdateBeFound
+	if updateFound {
+		newVersion = "brand-new-version"
+		changeLog = "something-was-changed"
+		err = nil
 	}
 
-	log.Println("Successfully updated to version", latest.Version)
+	return
 }
diff --git a/src/utils/ota/restart.go b/src/utils/ota/restart.go
@@ -0,0 +1,10 @@
+//go:build windows || netbsd || solaris || aix || dragonfly || freebsd || (illumos && !linux && !darwin)
+// +build windows netbsd solaris aix dragonfly freebsd illumos,!linux,!darwin
+
+package ota
+
+import "errors"
+
+func Restart(extraArgs ...string) error {
+	return errors.New("restart on the Windows system is not available")
+}