From 535574400e2384a37d598acfa9cf7240c5a738c9 Mon Sep 17 00:00:00 2001
From: Rohan Sreerama <rsreerama@influxdata.com>
Date: Thu, 12 Sep 2024 09:13:17 -0700
Subject: [PATCH] fix(profiling): fixes cpu profiling flagging

---
 cmd/wirey/main.go | 17 -----------------
 cmd/wirey/root.go | 17 +++++++++++++++++
 2 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/cmd/wirey/main.go b/cmd/wirey/main.go
index 7244caa..cfa688a 100644
--- a/cmd/wirey/main.go
+++ b/cmd/wirey/main.go
@@ -1,27 +1,10 @@
 package main
 
 import (
-	"flag"
-	"log"
-	"net/http"
 	_ "net/http/pprof"
 )
 
 func main() {
-	enableProfiling := flag.Bool("enable-profiling", false, "Enable pprof profiling")
-	profileAddr := flag.String("profile-addr", "localhost:6060", "Address for pprof profiling")
-
-	flag.Parse()
-
-	if *enableProfiling {
-		go func() {
-			log.Printf("Starting pprof server on %s\n", *profileAddr)
-			if err := http.ListenAndServe(*profileAddr, nil); err != nil {
-				log.Fatalf("Error starting pprof server: %v\n", err)
-			}
-		}()
-	}
-
 	// Execute main application logic
 	Execute()
 }
diff --git a/cmd/wirey/root.go b/cmd/wirey/root.go
index 241af82..85280b9 100644
--- a/cmd/wirey/root.go
+++ b/cmd/wirey/root.go
@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"math/rand"
 	"net"
+	"net/http"
+	_ "net/http/pprof"
 	"os"
 	"path/filepath"
 	"strings"
@@ -116,6 +118,17 @@ func backendFactory() (backend.Backend, error) {
 	httpBackend := viper.GetString("http")
 	//httpPortBackend := viper.GetInt("http-port")
 	discoverConf := viper.GetString("discover")
+	enableProfiling := viper.GetBool("enable-profiling")
+
+	if enableProfiling {
+		go func() {
+			profileAddr := viper.GetString("profile-addr")
+			log.Printf("Starting pprof server on %s\n", profileAddr)
+			if err := http.ListenAndServe(profileAddr, nil); err != nil {
+				log.Fatalf("Error starting pprof server: %v\n", err)
+			}
+		}()
+	}
 
 	// discover
 	discoverHosts := []string{}
@@ -237,6 +250,10 @@ func init() {
 	pflags.StringSlice("allowedips", nil, "array of allowed ips")
 	pflags.String("log-level", "info", "logging level to be used panic, fatal, error, trace, debug, warn, info")
 
+	// profiling flags
+	pflags.Bool("enable-profiling", false, "Enable pprof profiling")
+	pflags.String("profile-addr", "localhost:6060", "Address for pprof profiling")
+
 	rootCmd.MarkFlagRequired("endpoint")
 	rootCmd.MarkFlagRequired("ipaddr")