From e33275c99af03e3c4694cea97137974313780a90 Mon Sep 17 00:00:00 2001 From: Andrew Mason Date: Wed, 7 Sep 2022 18:52:30 -0400 Subject: [PATCH] [cli] [servenv] migrate `--service_map` and `pprof` flags to `pflag` (#11179) * unexport CheckServiceMap Signed-off-by: Andrew Mason * [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 * [cli] [servenv] Migrate `pprof` flag to `pflag` Relates to #11144. Signed-off-by: Andrew Mason * fix tests Signed-off-by: Andrew Mason --- go/cmd/mysqlctld/mysqlctld.go | 6 +++--- go/cmd/vtcombo/main.go | 1 + go/cmd/vtctld/main.go | 1 + go/cmd/vtgate/vtgate.go | 1 + go/cmd/vtgateclienttest/main.go | 1 + go/cmd/vttablet/vttablet.go | 1 + go/cmd/vttestserver/main.go | 2 ++ go/flags/endtoend/vtctld.txt | 4 ++-- go/flags/endtoend/vtexplain.txt | 3 +-- go/flags/endtoend/vtgate.txt | 4 ++-- go/flags/endtoend/vtgr.txt | 3 +-- go/flags/endtoend/vttablet.txt | 4 ++-- go/vt/servenv/grpc_server.go | 2 +- go/vt/servenv/pprof.go | 21 ++++++++++++--------- go/vt/servenv/pprof_test.go | 12 ++++++++---- go/vt/servenv/service_map.go | 20 +++++++++++--------- 16 files changed, 50 insertions(+), 36 deletions(-) diff --git a/go/cmd/mysqlctld/mysqlctld.go b/go/cmd/mysqlctld/mysqlctld.go index e5d43028a08..b9f5059bb89 100644 --- a/go/cmd/mysqlctld/mysqlctld.go +++ b/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() { diff --git a/go/cmd/vtcombo/main.go b/go/cmd/vtcombo/main.go index a5ea155d46c..212a0b17048 100644 --- a/go/cmd/vtcombo/main.go +++ b/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) { diff --git a/go/cmd/vtctld/main.go b/go/cmd/vtctld/main.go index 610e90d86c7..ba40cbffa41 100644 --- a/go/cmd/vtctld/main.go +++ b/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 diff --git a/go/cmd/vtgate/vtgate.go b/go/cmd/vtgate/vtgate.go index 2d4ad8fb730..f7d88d12115 100644 --- a/go/cmd/vtgate/vtgate.go +++ b/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 diff --git a/go/cmd/vtgateclienttest/main.go b/go/cmd/vtgateclienttest/main.go index a06c26d72b4..c101db77c95 100644 --- a/go/cmd/vtgateclienttest/main.go +++ b/go/cmd/vtgateclienttest/main.go @@ -31,6 +31,7 @@ func init() { servenv.RegisterFlags() servenv.RegisterGRPCServerFlags() servenv.RegisterGRPCServerAuthFlags() + servenv.RegisterServiceMapFlag() } func main() { diff --git a/go/cmd/vttablet/vttablet.go b/go/cmd/vttablet/vttablet.go index c97b341c3cf..22bb0c47dd3 100644 --- a/go/cmd/vttablet/vttablet.go +++ b/go/cmd/vttablet/vttablet.go @@ -59,6 +59,7 @@ func init() { servenv.RegisterFlags() servenv.RegisterGRPCServerFlags() servenv.RegisterGRPCServerAuthFlags() + servenv.RegisterServiceMapFlag() } func main() { diff --git a/go/cmd/vttestserver/main.go b/go/cmd/vttestserver/main.go index 17e3afeafb8..4ec96d7ed99 100644 --- a/go/cmd/vttestserver/main.go +++ b/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. diff --git a/go/flags/endtoend/vtctld.txt b/go/flags/endtoend/vtctld.txt index 263542a0c11..55fa4f97920 100644 --- a/go/flags/endtoend/vtctld.txt +++ b/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) diff --git a/go/flags/endtoend/vtexplain.txt b/go/flags/endtoend/vtexplain.txt index 3843c653448..556e81914da 100644 --- a/go/flags/endtoend/vtexplain.txt +++ b/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. diff --git a/go/flags/endtoend/vtgate.txt b/go/flags/endtoend/vtgate.txt index 0a66876244c..cfa3e3c7a1e 100644 --- a/go/flags/endtoend/vtgate.txt +++ b/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) diff --git a/go/flags/endtoend/vtgr.txt b/go/flags/endtoend/vtgr.txt index d632b7c089b..62029d377c3 100644 --- a/go/flags/endtoend/vtgr.txt +++ b/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 diff --git a/go/flags/endtoend/vttablet.txt b/go/flags/endtoend/vttablet.txt index f8accd737b4..2481f8bc38e 100644 --- a/go/flags/endtoend/vttablet.txt +++ b/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. diff --git a/go/vt/servenv/grpc_server.go b/go/vt/servenv/grpc_server.go index f60852a1e6d..177f0f26b4e 100644 --- a/go/vt/servenv/grpc_server.go +++ b/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 { diff --git a/go/vt/servenv/pprof.go b/go/vt/servenv/pprof.go index 8119cf50a47..d431dd9239b 100644 --- a/go/vt/servenv/pprof.go +++ b/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) } diff --git a/go/vt/servenv/pprof_test.go b/go/vt/servenv/pprof_test.go index fdae777aba2..8f602d0f529 100644 --- a/go/vt/servenv/pprof_test.go +++ b/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) diff --git a/go/vt/servenv/service_map.go b/go/vt/servenv/service_map.go index 4616f058b99..69d5cd0bcbc 100644 --- a/go/vt/servenv/service_map.go +++ b/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