From 758ec183833a7a6525a4b21820d47495e6c60e91 Mon Sep 17 00:00:00 2001
From: Zain Budhwani <99770260+zbud-msft@users.noreply.github.com>
Date: Tue, 12 Mar 2024 14:28:27 -0700
Subject: [PATCH] Call flag.Parse() to parse global flags like -logtostderr
 (#198)

---
 telemetry/telemetry.go      | 54 ++++++++++++++++++++++++++++++++++++-
 telemetry/telemetry_test.go | 21 ++++++++-------
 2 files changed, 65 insertions(+), 10 deletions(-)

diff --git a/telemetry/telemetry.go b/telemetry/telemetry.go
index b32677e5..b779a44e 100644
--- a/telemetry/telemetry.go
+++ b/telemetry/telemetry.go
@@ -11,6 +11,7 @@ import (
 	"time"
 	"os"
 	"os/signal"
+	"strings"
 	"syscall"
 	"sync"
 	log "github.com/golang/glog"
@@ -50,8 +51,20 @@ func main() {
 }
 
 func runTelemetry(args []string) error {
+	/* Glog flags like -logtostderr have to be part of the global flagset.
+	   Because we use a custom flagset to avoid the use of global var and improve
+	   testability, we have to parse cmd line args two different times such that
+	   in the first parse, cmd line args will contain global flags and flag.Parse() will be called.
+	   The second parse, cmd line args will contain flags only relevant to telemetry, and our custom flagset will
+	   call Parse().
+	*/
+	glogFlags, telemetryFlags := parseOSArgs()
+	os.Args = glogFlags
+	flag.Parse() // glog flags will be populated after global flag parse
+
+	os.Args = telemetryFlags
 	fs := flag.NewFlagSet(args[0], flag.ExitOnError)
-	telemetryCfg, cfg, err := setupFlags(fs)
+	telemetryCfg, cfg, err := setupFlags(fs) // telemetry flags will be populated after second parse
 	if err != nil {
 		return err
 	}
@@ -72,6 +85,45 @@ func runTelemetry(args []string) error {
 	return nil
 }
 
+func getGlogFlagsMap() map[string] bool {
+	// glog flags: https://pkg.go.dev/github.com/golang/glog
+	return map[string]bool {
+		"-alsologtostderr":  true,
+		"-log_backtrace_at": true,
+		"-log_dir":          true,
+		"-log_link":         true,
+		"-logbuflevel":      true,
+		"-logtostderr":      true,
+		"-stderrthreshold":  true,
+		"-v":                true,
+		"-vmodule":          true,
+	}
+}
+
+func parseOSArgs() ([]string, []string) {
+	glogFlags := []string{os.Args[0]}
+	telemetryFlags := []string{os.Args[0]}
+	glogFlagsMap := getGlogFlagsMap()
+
+	for _, arg := range os.Args[1:] {
+		if strings.HasPrefix(arg, "-v") { // only flag in both glog and telemetry
+			glogFlags = append(glogFlags, arg)
+			telemetryFlags = append(telemetryFlags, arg)
+			continue
+		}
+		isGlogFlag := false
+		if glogFlagsMap[arg] {
+			glogFlags = append(glogFlags, arg)
+			isGlogFlag = true
+			continue
+		}
+		if !isGlogFlag {
+			telemetryFlags = append(telemetryFlags, arg)
+		}
+	}
+	return glogFlags, telemetryFlags
+}
+
 func setupFlags(fs *flag.FlagSet) (*TelemetryConfig, *gnmi.Config, error) {
 	telemetryCfg := &TelemetryConfig {
 		UserAuth:              gnmi.AuthTypes{"password": false, "cert": false, "jwt": false},
diff --git a/telemetry/telemetry_test.go b/telemetry/telemetry_test.go
index e416f6c6..ecb25c54 100644
--- a/telemetry/telemetry_test.go
+++ b/telemetry/telemetry_test.go
@@ -23,13 +23,7 @@ import (
 )
 
 func TestRunTelemetry(t *testing.T) {
-	telemetryCfg := &TelemetryConfig{}
-	cfg := &gnmi.Config{}
-
-	patches := gomonkey.ApplyFunc(setupFlags, func(*flag.FlagSet) (*TelemetryConfig, *gnmi.Config, error) {
-		return telemetryCfg, cfg, nil
-	})
-	patches.ApplyFunc(startGNMIServer, func(_ *TelemetryConfig, _ *gnmi.Config, serverControlSignal <-chan int, wg *sync.WaitGroup) {
+	patches := gomonkey.ApplyFunc(startGNMIServer, func(_ *TelemetryConfig, _ *gnmi.Config, serverControlSignal <-chan int, wg *sync.WaitGroup) {
 		defer wg.Done()
 	})
 	patches.ApplyFunc(signalHandler, func(serverControlSignal chan<- int, wg *sync.WaitGroup, sigchannel <-chan os.Signal) {
@@ -37,11 +31,20 @@ func TestRunTelemetry(t *testing.T) {
 	})
 	defer patches.Reset()
 
-	args := []string{"test"}
-	err := runTelemetry(args)
+	args := []string{"telemetry", "-logtostderr",  "-port", "50051", "-v=2"}
+	os.Args = args
+	err := runTelemetry(os.Args)
 	if err != nil {
 		t.Errorf("Expected err to be nil, but got %v", err)
 	}
+	vflag := flag.Lookup("v")
+	if vflag.Value.String() != "2" {
+		t.Errorf("Expected v to be 2")
+	}
+	logtostderrflag := flag.Lookup("logtostderr")
+	if logtostderrflag.Value.String() != "true" {
+		t.Errorf("Expected logtostderr to be true")
+	}
 }
 
 func TestFlags(t *testing.T) {