[cli] [servenv] migrate --service_map and pprof flags to pflag (#…

…11179)

* unexport CheckServiceMap

Signed-off-by: Andrew Mason <andrew@planetscale.com>

* [cli] [servenv] Migrate `service_map` to pflag

Relates to #11144.

Also converts the flag from our custom `StringListValue` to pflag-native
`StringSlice`.

Signed-off-by: Andrew Mason <andrew@planetscale.com>

* [cli] [servenv] Migrate `pprof` flag to `pflag`

Relates to #11144.

Signed-off-by: Andrew Mason <andrew@planetscale.com>

* fix tests

Signed-off-by: Andrew Mason <andrew@planetscale.com>

go/cmd/mysqlctld/mysqlctld.go

@@ -20,12 +20,11 @@ limitations under the License.
 package main
 
 import (
+	"context"
 	"flag"
 	"os"
 	"time"
 
-	"context"
-
 	"vitess.io/vitess/go/exit"
 	"vitess.io/vitess/go/vt/dbconfigs"
 	"vitess.io/vitess/go/vt/log"
@@ -50,10 +49,11 @@ var (
 
 func init() {
 	servenv.RegisterDefaultFlags()
+	servenv.RegisterDefaultSocketFileFlags()
 	servenv.RegisterFlags()
 	servenv.RegisterGRPCServerFlags()
-	servenv.RegisterDefaultSocketFileFlags()
 	servenv.RegisterGRPCServerAuthFlags()
+	servenv.RegisterServiceMapFlag()
 }
 
 func main() {

go/cmd/vtcombo/main.go

@@ -77,6 +77,7 @@ func init() {
 	servenv.RegisterFlags()
 	servenv.RegisterGRPCServerFlags()
 	servenv.RegisterGRPCServerAuthFlags()
+	servenv.RegisterServiceMapFlag()
 }
 
 func startMysqld(uid uint32) (*mysqlctl.Mysqld, *mysqlctl.Mycnf) {

go/cmd/vtctld/main.go

@@ -28,6 +28,7 @@ func init() {
 	servenv.RegisterFlags()
 	servenv.RegisterGRPCServerFlags()
 	servenv.RegisterGRPCServerAuthFlags()
+	servenv.RegisterServiceMapFlag()
 }
 
 // used at runtime by plug-ins

go/cmd/vtgate/vtgate.go

@@ -54,6 +54,7 @@ func init() {
 	servenv.RegisterFlags()
 	servenv.RegisterGRPCServerFlags()
 	servenv.RegisterGRPCServerAuthFlags()
+	servenv.RegisterServiceMapFlag()
 }
 
 // CheckCellFlags will check validation of cell and cells_to_watch flag

go/cmd/vtgateclienttest/main.go

@@ -31,6 +31,7 @@ func init() {
 	servenv.RegisterFlags()
 	servenv.RegisterGRPCServerFlags()
 	servenv.RegisterGRPCServerAuthFlags()
+	servenv.RegisterServiceMapFlag()
 }
 
 func main() {

go/cmd/vttablet/vttablet.go

@@ -59,6 +59,7 @@ func init() {
 	servenv.RegisterFlags()
 	servenv.RegisterGRPCServerFlags()
 	servenv.RegisterGRPCServerAuthFlags()
+	servenv.RegisterServiceMapFlag()
 }
 
 func main() {

go/cmd/vttestserver/main.go

@@ -205,6 +205,7 @@ func (t *topoFlags) buildTopology() (*vttestpb.VTTestTopology, error) {
 // Annoying, but in unit tests, parseFlags gets called multiple times per process
 // (anytime startCluster is called), so we need to guard against the second test
 // to run failing with, for example:
+//
 //	flag redefined: log_rotate_max_size
 var flagsOnce sync.Once
 
@@ -217,6 +218,7 @@ func parseFlags() (env vttest.Environment, err error) {
 		servenv.RegisterFlags()
 		servenv.RegisterGRPCServerFlags()
 		servenv.RegisterGRPCServerAuthFlags()
+		servenv.RegisterServiceMapFlag()
 		servenv.ParseFlags("vttestserver")
 
 		// Move all pflag flags back to the goflag CommandLine.

go/flags/endtoend/vtctld.txt

@@ -125,7 +125,7 @@ Usage of vtctld:
       --pid_file string                                                  If set, the process will write its pid to the named file, and delete it on graceful shutdown.
       --pool_hostname_resolve_interval duration                          if set force an update to all hostnames and reconnect if changed, defaults to 0 (disabled) (default 0s)
       --port int                                                         port for the server
-      --pprof string                                                     enable profiling
+      --pprof strings                                                    enable profiling
       --proxy_tablets                                                    Setting this true will make vtctld proxy the tablet status instead of redirecting to them
       --purge_logs_interval duration                                     how often try to remove old logs (default 1h0m0s)
       --query-log-stream-handler string                                  URL handler for streaming queries log (default "/debug/querylog")
@@ -181,7 +181,7 @@ Usage of vtctld:
       --schema_change_replicas_timeout duration                          how long to wait for replicas to receive the schema change (default 10s)
       --schema_change_user string                                        The user who submits this schema change.
       --security_policy string                                           the name of a registered security policy to use for controlling access to URLs - empty means allow all for anyone (built-in policies: deny-all, read-only)
-      --service_map StringList                                           comma separated list of services to enable (or disable if prefixed with '-') Example: grpc-queryservice
+      --service_map strings                                              comma separated list of services to enable (or disable if prefixed with '-') Example: grpc-queryservice
       --serving_state_grace_period duration                              how long to pause after broadcasting health to vtgate, before enforcing a new serving state (default 0s)
       --shutdown_grace_period float                                      how long to wait (in seconds) for queries and transactions to complete during graceful shutdown.
       --sql-max-length-errors int                                        truncate queries in error logs to the given length (default unlimited)

go/flags/endtoend/vtexplain.txt

@@ -121,7 +121,7 @@ Usage of vtexplain:
       --planner-version string                                       Sets the query planner version to use when generating the explain output. Valid values are V3 and Gen4
       --planner_version string                                       Deprecated flag. Use planner-version instead
       --pool_hostname_resolve_interval duration                      if set force an update to all hostnames and reconnect if changed, defaults to 0 (disabled) (default 0s)
-      --pprof string                                                 enable profiling
+      --pprof strings                                                enable profiling
       --proxy_protocol                                               Enable HAProxy PROXY protocol on MySQL listener socket
       --purge_logs_interval duration                                 how often try to remove old logs (default 1h0m0s)
       --query-log-stream-handler string                              URL handler for streaming queries log (default "/debug/querylog")
@@ -169,7 +169,6 @@ Usage of vtexplain:
       --schema_change_signal                                         Enable the schema tracker; requires queryserver-config-schema-change-signal to be enabled on the underlying vttablets for this to work (default true)
       --schema_change_signal_user string                             User to be used to send down query to vttablet to retrieve schema changes
       --security_policy string                                       the name of a registered security policy to use for controlling access to URLs - empty means allow all for anyone (built-in policies: deny-all, read-only)
-      --service_map StringList                                       comma separated list of services to enable (or disable if prefixed with '-') Example: grpc-queryservice
       --serving_state_grace_period duration                          how long to pause after broadcasting health to vtgate, before enforcing a new serving state (default 0s)
       --shards int                                                   Number of shards per keyspace. Passing --ks-shard-map/--ks-shard-map-file causes this flag to be ignored. (default 2)
       --shutdown_grace_period float                                  how long to wait (in seconds) for queries and transactions to complete during graceful shutdown.

go/flags/endtoend/vtgate.txt

@@ -123,7 +123,7 @@ Usage of vtgate:
       --planner-version string                                           Sets the default planner to use when the session has not changed it. Valid values are: V3, Gen4, Gen4Greedy and Gen4Fallback. Gen4Fallback tries the gen4 planner and falls back to the V3 planner if the gen4 fails.
       --planner_version string                                           Deprecated flag. Use planner-version instead
       --port int                                                         port for the server
-      --pprof string                                                     enable profiling
+      --pprof strings                                                    enable profiling
       --proxy_protocol                                                   Enable HAProxy PROXY protocol on MySQL listener socket
       --purge_logs_interval duration                                     how often try to remove old logs (default 1h0m0s)
       --querylog-buffer-size int                                         Maximum number of buffered query logs before throttling log output (default 10)
@@ -136,7 +136,7 @@ Usage of vtgate:
       --schema_change_signal                                             Enable the schema tracker; requires queryserver-config-schema-change-signal to be enabled on the underlying vttablets for this to work (default true)
       --schema_change_signal_user string                                 User to be used to send down query to vttablet to retrieve schema changes
       --security_policy string                                           the name of a registered security policy to use for controlling access to URLs - empty means allow all for anyone (built-in policies: deny-all, read-only)
-      --service_map StringList                                           comma separated list of services to enable (or disable if prefixed with '-') Example: grpc-queryservice
+      --service_map strings                                              comma separated list of services to enable (or disable if prefixed with '-') Example: grpc-queryservice
       --sql-max-length-errors int                                        truncate queries in error logs to the given length (default unlimited)
       --sql-max-length-ui int                                            truncate queries in debug UIs to the given length (default 512) (default 512)
       --srv_topo_cache_refresh duration                                  how frequently to refresh the topology for cached entries (default 1s)

go/flags/endtoend/vtgr.txt

@@ -31,14 +31,13 @@ Usage of vtgr:
       --mysql_server_version string                MySQL server version to advertise.
       --pid_file string                            If set, the process will write its pid to the named file, and delete it on graceful shutdown.
       --ping_tablet_timeout duration               time to wait when we ping a tablet (default 2s)
-      --pprof string                               enable profiling
+      --pprof strings                              enable profiling
       --purge_logs_interval duration               how often try to remove old logs (default 1h0m0s)
       --refresh_interval duration                  Refresh interval to load tablets. (default 10s)
       --remote_operation_timeout duration          time to wait for a remote operation (default 30s)
       --scan_interval duration                     Scan interval to diagnose and repair. (default 3s)
       --scan_repair_timeout duration               Time to wait for a Diagnose and repair operation. (default 3s)
       --security_policy string                     the name of a registered security policy to use for controlling access to URLs - empty means allow all for anyone (built-in policies: deny-all, read-only)
-      --service_map StringList                     comma separated list of services to enable (or disable if prefixed with '-') Example: grpc-queryservice
       --stats_backend string                       The name of the registered push-based monitoring/stats backend to use
       --stats_combine_dimensions string            List of dimensions to be combined into a single "all" value in exported stats vars
       --stats_common_tags string                   Comma-separated list of common tags for the stats backend. It provides both label and values. Example: label1:value1,label2:value2

go/flags/endtoend/vttablet.txt

@@ -300,7 +300,7 @@ Usage of vttablet:
       --pitr_gtid_lookup_timeout duration                                PITR restore parameter: timeout for fetching gtid from timestamp. (default 1m0s)
       --pool_hostname_resolve_interval duration                          if set force an update to all hostnames and reconnect if changed, defaults to 0 (disabled) (default 0s)
       --port int                                                         port for the server
-      --pprof string                                                     enable profiling
+      --pprof strings                                                    enable profiling
       --pt-osc-path string                                               override default pt-online-schema-change binary full path
       --publish_retry_interval duration                                  how long vttablet waits to retry publishing the tablet record (default 30s)
       --purge_logs_interval duration                                     how often try to remove old logs (default 1h0m0s)
@@ -362,7 +362,7 @@ Usage of vttablet:
       --s3_backup_tls_skip_verify_cert                                   skip the 'certificate is valid' check for SSL connections.
       --sanitize_log_messages                                            Remove potentially sensitive information in tablet INFO, WARNING, and ERROR log messages such as query parameters.
       --security_policy string                                           the name of a registered security policy to use for controlling access to URLs - empty means allow all for anyone (built-in policies: deny-all, read-only)
-      --service_map StringList                                           comma separated list of services to enable (or disable if prefixed with '-') Example: grpc-queryservice
+      --service_map strings                                              comma separated list of services to enable (or disable if prefixed with '-') Example: grpc-queryservice
       --serving_state_grace_period duration                              how long to pause after broadcasting health to vtgate, before enforcing a new serving state (default 0s)
       --shard_sync_retry_delay duration                                  delay between retries of updates to keep the tablet and its shard record in sync (default 30s)
       --shutdown_grace_period float                                      how long to wait (in seconds) for queries and transactions to complete during graceful shutdown.

go/vt/servenv/grpc_server.go

@@ -321,7 +321,7 @@ func GRPCCheckServiceMap(name string) bool {
 	}
 
 	// then check ServiceMap
-	return CheckServiceMap("grpc", name)
+	return checkServiceMap("grpc", name)
 }
 
 func authenticatingStreamInterceptor(srv any, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {

go/vt/servenv/pprof.go

@@ -17,7 +17,6 @@ limitations under the License.
 package servenv
 
 import (
-	"flag"
 	"fmt"
 	"io"
 	"os"
@@ -31,11 +30,13 @@ import (
 	"sync/atomic"
 	"syscall"
 
+	"github.com/spf13/pflag"
+
 	"vitess.io/vitess/go/vt/log"
 )
 
 var (
-	pprofFlag = flag.String("pprof", "", "enable profiling")
+	pprofFlag []string
 )
 
 type profmode string
@@ -63,15 +64,14 @@ type profile struct {
 	waitSig bool
 }
 
-func parseProfileFlag(pf string) (*profile, error) {
-	if pf == "" {
+func parseProfileFlag(pf []string) (*profile, error) {
+	if len(pf) == 0 {
 		return nil, nil
 	}
 
 	var p profile
 
-	items := strings.Split(pf, ",")
-	switch items[0] {
+	switch pf[0] {
 	case "cpu":
 		p.mode = profileCPU
 	case "mem", "mem=heap":
@@ -93,10 +93,10 @@ func parseProfileFlag(pf string) (*profile, error) {
 	case "goroutine":
 		p.mode = profileGoroutine
 	default:
-		return nil, fmt.Errorf("unknown profile mode: %q", items[0])
+		return nil, fmt.Errorf("unknown profile mode: %q", pf[0])
 	}
 
-	for _, kv := range items[1:] {
+	for _, kv := range pf[1:] {
 		var err error
 		fields := strings.SplitN(kv, "=", 2)
 
@@ -299,7 +299,7 @@ func (prof *profile) init() (start func(), stop func()) {
 }
 
 func pprofInit() {
-	prof, err := parseProfileFlag(*pprofFlag)
+	prof, err := parseProfileFlag(pprofFlag)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -338,5 +338,8 @@ func pprofInit() {
 }
 
 func init() {
+	OnParse(func(fs *pflag.FlagSet) {
+		fs.StringSliceVar(&pprofFlag, "pprof", pprofFlag, "enable profiling")
+	})
 	OnInit(pprofInit)
 }

go/vt/servenv/pprof_test.go

@@ -6,9 +6,9 @@
 package servenv
 
 import (
-	"flag"
 	"os/signal"
 	"reflect"
+	"strings"
 	"syscall"
 	"testing"
 	"time"
@@ -45,7 +45,11 @@ func TestParseProfileFlag(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.arg, func(t *testing.T) {
-			got, err := parseProfileFlag(tt.arg)
+			var profileFlag []string
+			if tt.arg != "" {
+				profileFlag = strings.Split(tt.arg, ",")
+			}
+			got, err := parseProfileFlag(profileFlag)
 			if (err != nil) != tt.wantErr {
 				t.Errorf("parseProfileFlag() error = %v, wantErr %v", err, tt.wantErr)
 				return
@@ -60,7 +64,7 @@ func TestParseProfileFlag(t *testing.T) {
 // with waitSig, we should start with profiling off and toggle on-off-on-off
 func TestPProfInitWithWaitSig(t *testing.T) {
 	signal.Reset(syscall.SIGUSR1)
-	flag.Set("pprof", "cpu,waitSig")
+	pprofFlag = strings.Split("cpu,waitSig", ",")
 
 	pprofInit()
 	time.Sleep(1 * time.Second)
@@ -86,7 +90,7 @@ func TestPProfInitWithWaitSig(t *testing.T) {
 // without waitSig, we should start with profiling on and toggle off-on-off
 func TestPProfInitWithoutWaitSig(t *testing.T) {
 	signal.Reset(syscall.SIGUSR1)
-	flag.Set("pprof", "cpu")
+	pprofFlag = strings.Split("cpu", ",")
 
 	pprofInit()
 	time.Sleep(1 * time.Second)

go/vt/servenv/service_map.go

@@ -17,14 +17,13 @@ limitations under the License.
 package servenv
 
 import (
-	"flag"
+	"github.com/spf13/pflag"
 
-	"vitess.io/vitess/go/flagutil"
 	"vitess.io/vitess/go/vt/log"
 )
 
 var (
-	serviceMapFlag flagutil.StringListValue
+	serviceMapFlag []string
 
 	// serviceMap is the used version of the service map.
 	// init() functions can add default values to it (using InitServiceMap).
@@ -33,11 +32,14 @@ var (
 	serviceMap = make(map[string]bool)
 )
 
-func init() {
-	flag.Var(&serviceMapFlag, "service_map", "comma separated list of services to enable (or disable if prefixed with '-') Example: grpc-queryservice")
-	OnInit(func() {
-		updateServiceMap()
+// RegisterServiceMapFlag registers an OnParse hook to install the
+// `--service_map` flag for a given cmd. It must be called before ParseFlags or
+// ParseFlagsWithArgs.
+func RegisterServiceMapFlag() {
+	OnParse(func(fs *pflag.FlagSet) {
+		fs.StringSliceVar(&serviceMapFlag, "service_map", serviceMapFlag, "comma separated list of services to enable (or disable if prefixed with '-') Example: grpc-queryservice")
 	})
+	OnInit(updateServiceMap)
 }
 
 // InitServiceMap will set the default value for a protocol/name to be served.
@@ -57,9 +59,9 @@ func updateServiceMap() {
 	}
 }
 
-// CheckServiceMap returns if we should register a RPC service
+// checkServiceMap returns if we should register a RPC service
 // (and also logs how to enable / disable it)
-func CheckServiceMap(protocol, name string) bool {
+func checkServiceMap(protocol, name string) bool {
 	if serviceMap[protocol+"-"+name] {
 		log.Infof("Registering %v for %v, disable it with -%v-%v service_map parameter", name, protocol, protocol, name)
 		return true