From 8466ab003423cc3a3cc7755eb9561bbaa4a49b87 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Mon, 10 Feb 2020 13:26:18 +0200
Subject: [PATCH] all: allow setting flags via environment vars

Now flags can be set via environment vars with the same names as flags.
Command-line flags override flags set via env vars.

Updates https://github.com/VictoriaMetrics/VictoriaMetrics/issues/311
---
 app/victoria-metrics/main.go      |  3 ++-
 app/victoria-metrics/main_test.go |  3 ++-
 app/vmbackup/main.go              |  3 ++-
 app/vmrestore/main.go             |  3 ++-
 lib/envflag/envflag.go            | 32 +++++++++++++++++++++++++++++++
 5 files changed, 40 insertions(+), 4 deletions(-)
 create mode 100644 lib/envflag/envflag.go

diff --git a/app/victoria-metrics/main.go b/app/victoria-metrics/main.go
index 9eb34fff2..df50376f7 100644
--- a/app/victoria-metrics/main.go
+++ b/app/victoria-metrics/main.go
@@ -9,6 +9,7 @@ import (
 	"github.com/VictoriaMetrics/VictoriaMetrics/app/vmselect"
 	"github.com/VictoriaMetrics/VictoriaMetrics/app/vmstorage"
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/buildinfo"
+	"github.com/VictoriaMetrics/VictoriaMetrics/lib/envflag"
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/fs"
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/httpserver"
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/logger"
@@ -24,7 +25,7 @@ var (
 )
 
 func main() {
-	flag.Parse()
+	envflag.Parse()
 	buildinfo.Init()
 	logger.Init()
 	logger.Infof("starting VictoriaMetrics at %q...", *httpListenAddr)
diff --git a/app/victoria-metrics/main_test.go b/app/victoria-metrics/main_test.go
index 971034b82..7f13ebfc7 100644
--- a/app/victoria-metrics/main_test.go
+++ b/app/victoria-metrics/main_test.go
@@ -21,6 +21,7 @@ import (
 	"github.com/VictoriaMetrics/VictoriaMetrics/app/vminsert"
 	"github.com/VictoriaMetrics/VictoriaMetrics/app/vmselect"
 	"github.com/VictoriaMetrics/VictoriaMetrics/app/vmstorage"
+	"github.com/VictoriaMetrics/VictoriaMetrics/lib/envflag"
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/fs"
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/httpserver"
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/logger"
@@ -146,7 +147,7 @@ func setUp() {
 }
 
 func processFlags() {
-	flag.Parse()
+	envflag.Parse()
 	for _, fv := range []struct {
 		flag  string
 		value string
diff --git a/app/vmbackup/main.go b/app/vmbackup/main.go
index da623721c..eb5f08297 100644
--- a/app/vmbackup/main.go
+++ b/app/vmbackup/main.go
@@ -9,6 +9,7 @@ import (
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/backup/common"
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/backup/fslocal"
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/buildinfo"
+	"github.com/VictoriaMetrics/VictoriaMetrics/lib/envflag"
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/logger"
 )
 
@@ -25,7 +26,7 @@ var (
 
 func main() {
 	flag.Usage = usage
-	flag.Parse()
+	envflag.Parse()
 	buildinfo.Init()
 
 	srcFS, err := newSrcFS()
diff --git a/app/vmrestore/main.go b/app/vmrestore/main.go
index 9bd399e3f..95b16565a 100644
--- a/app/vmrestore/main.go
+++ b/app/vmrestore/main.go
@@ -8,6 +8,7 @@ import (
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/backup/common"
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/backup/fslocal"
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/buildinfo"
+	"github.com/VictoriaMetrics/VictoriaMetrics/lib/envflag"
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/logger"
 )
 
@@ -23,7 +24,7 @@ var (
 
 func main() {
 	flag.Usage = usage
-	flag.Parse()
+	envflag.Parse()
 	buildinfo.Init()
 
 	srcFS, err := newSrcFS()
diff --git a/lib/envflag/envflag.go b/lib/envflag/envflag.go
new file mode 100644
index 000000000..e4a79e9e4
--- /dev/null
+++ b/lib/envflag/envflag.go
@@ -0,0 +1,32 @@
+package envflag
+
+import (
+	"flag"
+	"os"
+)
+
+// Parse parses environment vars and command-line flags.
+//
+// Flags set via command-line override flags set via environment vars.
+//
+// This function must be called instead of flag.Parse() before using any flags in the program.
+func Parse() {
+	flag.Parse()
+
+	// Remember explicitly set command-line flags.
+	flagsSet := make(map[string]bool)
+	flag.Visit(func(f *flag.Flag) {
+		flagsSet[f.Name] = true
+	})
+
+	// Obtain the remaining flag values from environment vars.
+	flag.VisitAll(func(f *flag.Flag) {
+		if flagsSet[f.Name] {
+			// The flag is explicitly set via command-line.
+			return
+		}
+		if v, ok := os.LookupEnv(f.Name); ok {
+			f.Value.Set(v)
+		}
+	})
+}