diff --git a/internal/cmd/branch/vtctld/throttler.go b/internal/cmd/branch/vtctld/throttler.go index ed7228d9..1770f3ca 100644 --- a/internal/cmd/branch/vtctld/throttler.go +++ b/internal/cmd/branch/vtctld/throttler.go @@ -2,6 +2,7 @@ package vtctld import ( "fmt" + "time" "github.com/planetscale/cli/internal/cmdutil" ps "github.com/planetscale/cli/internal/planetscale" @@ -141,18 +142,33 @@ func ThrottlerCheckCmd(ch *cmdutil.Helper) *cobra.Command { // keyspace. func ThrottlerUpdateConfigCmd(ch *cmdutil.Helper) *cobra.Command { var flags struct { - keyspace string - enabled bool - threshold float64 + keyspace string + enabled bool + threshold float64 + throttleApp string + throttleAppRatio float64 + throttleAppDuration time.Duration + unthrottleApp string + appName string + appMetrics []string } cmd := &cobra.Command{ Use: "update-config ", Short: "Update the throttler configuration for a keyspace", - Long: "Update the tablet throttler configuration for a keyspace. The throttler is " + - "enabled or disabled with --enabled; this flag is required because there is no " + - "separate \"leave unchanged\" state.", + Long: "Update the tablet throttler configuration for a keyspace. " + + "Omit --enabled to leave the keyspace enable state unchanged. " + + "Flag behavior mirrors vtctldclient UpdateThrottlerConfig: " + + "--throttle-app and --unthrottle-app are mutually exclusive; " + + "--app-name and --app-metrics are required together.", Args: cmdutil.RequiredArgs("database", "branch"), + PreRunE: func(cmd *cobra.Command, args []string) error { + // Mirrors vtctldclient validateUpdateThrottlerConfig. + if cmd.Flags().Changed("app-name") && flags.appName == "" { + return fmt.Errorf("--app-name must not be empty") + } + return nil + }, RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() database, branch := args[0], args[1] @@ -172,11 +188,31 @@ func ThrottlerUpdateConfigCmd(ch *cmdutil.Helper) *cobra.Command { Database: database, Branch: branch, Keyspace: flags.keyspace, - Enabled: flags.enabled, + } + if cmd.Flags().Changed("enabled") { + enabled := flags.enabled + req.Enabled = &enabled } if cmd.Flags().Changed("threshold") { req.Threshold = &flags.threshold } + if flags.throttleApp != "" { + ratio := flags.throttleAppRatio + expireAt := time.Now().Add(flags.throttleAppDuration).UTC().Format(time.RFC3339) + req.Apps = []ps.VtctldThrottledAppConfig{{ + Name: flags.throttleApp, + Ratio: &ratio, + ExpireAt: expireAt, + }} + } + if flags.unthrottleApp != "" { + req.UnthrottleApps = []string{flags.unthrottleApp} + } + if flags.appName != "" { + req.AppCheckedMetrics = map[string][]string{ + flags.appName: flags.appMetrics, + } + } data, err := client.Vtctld.UpdateThrottlerConfig(ctx, req) if err != nil { @@ -189,10 +225,17 @@ func ThrottlerUpdateConfigCmd(ch *cmdutil.Helper) *cobra.Command { } cmd.Flags().StringVar(&flags.keyspace, "keyspace", "", "Keyspace whose throttler config to update") - cmd.Flags().BoolVar(&flags.enabled, "enabled", false, "Enable (true) or disable (false) the throttler for the keyspace") - cmd.Flags().Float64Var(&flags.threshold, "threshold", 0, "Replication lag threshold in seconds for the default check (defaults to 5.0 server-side when omitted)") + cmd.Flags().BoolVar(&flags.enabled, "enabled", false, "Enable (true) or disable (false) the throttler for the keyspace. Omit to leave unchanged.") + cmd.Flags().Float64Var(&flags.threshold, "threshold", 0, "Replication lag threshold in seconds for the default check") + cmd.Flags().StringVar(&flags.throttleApp, "throttle-app", "", "App name to throttle (e.g. \"rowstreamer\")") + cmd.Flags().Float64Var(&flags.throttleAppRatio, "throttle-app-ratio", 1.0, "Ratio to throttle the app specified by --throttle-app (0.00-1.00)") + cmd.Flags().DurationVar(&flags.throttleAppDuration, "throttle-app-duration", time.Hour, "Duration after which the --throttle-app rule expires") + cmd.Flags().StringVar(&flags.unthrottleApp, "unthrottle-app", "", "App name whose throttled-app rule should be removed") + cmd.Flags().StringVar(&flags.appName, "app-name", "", "App name for which to assign checked metrics (requires --app-metrics)") + cmd.Flags().StringSliceVar(&flags.appMetrics, "app-metrics", nil, "Metrics to check for --app-name (e.g. lag,loadavg)") cmd.MarkFlagRequired("keyspace") // nolint:errcheck - cmd.MarkFlagRequired("enabled") // nolint:errcheck + cmd.MarkFlagsMutuallyExclusive("unthrottle-app", "throttle-app") + cmd.MarkFlagsRequiredTogether("app-name", "app-metrics") return cmd } diff --git a/internal/cmd/branch/vtctld/throttler_test.go b/internal/cmd/branch/vtctld/throttler_test.go index e2d734e6..61422547 100644 --- a/internal/cmd/branch/vtctld/throttler_test.go +++ b/internal/cmd/branch/vtctld/throttler_test.go @@ -129,7 +129,8 @@ func TestThrottlerUpdateConfig(t *testing.T) { UpdateThrottlerConfigFn: func(ctx context.Context, req *ps.VtctldUpdateThrottlerConfigRequest) (json.RawMessage, error) { c.Assert(req.Organization, qt.Equals, org) c.Assert(req.Keyspace, qt.Equals, "commerce") - c.Assert(req.Enabled, qt.IsTrue) + c.Assert(req.Enabled, qt.IsNotNil) + c.Assert(*req.Enabled, qt.IsTrue) c.Assert(req.Threshold, qt.IsNotNil) c.Assert(*req.Threshold, qt.Equals, 2.5) return json.RawMessage(`{}`), nil @@ -154,8 +155,8 @@ func TestThrottlerUpdateConfig_DisableOmitsThreshold(t *testing.T) { svc := &mock.VtctldService{ UpdateThrottlerConfigFn: func(ctx context.Context, req *ps.VtctldUpdateThrottlerConfigRequest) (json.RawMessage, error) { - c.Assert(req.Enabled, qt.IsFalse) - // Threshold stays nil when not provided so the server keeps its default. + c.Assert(req.Enabled, qt.IsNotNil) + c.Assert(*req.Enabled, qt.IsFalse) c.Assert(req.Threshold, qt.IsNil) return json.RawMessage(`{}`), nil }, @@ -170,15 +171,82 @@ func TestThrottlerUpdateConfig_DisableOmitsThreshold(t *testing.T) { c.Assert(svc.UpdateThrottlerConfigFnInvoked, qt.IsTrue) } -func TestThrottlerUpdateConfig_RequiresEnabled(t *testing.T) { +func TestThrottlerUpdateConfig_KeyspaceOnlyAllowed(t *testing.T) { c := qt.New(t) - svc := &mock.VtctldService{} + svc := &mock.VtctldService{ + UpdateThrottlerConfigFn: func(ctx context.Context, req *ps.VtctldUpdateThrottlerConfigRequest) (json.RawMessage, error) { + c.Assert(req.Keyspace, qt.Equals, "commerce") + c.Assert(req.Enabled, qt.IsNil) + c.Assert(req.Threshold, qt.IsNil) + c.Assert(req.Apps, qt.HasLen, 0) + c.Assert(req.UnthrottleApps, qt.HasLen, 0) + c.Assert(req.AppCheckedMetrics, qt.IsNil) + return json.RawMessage(`{}`), nil + }, + } ch, _ := newThrottlerTestHelper("my-org", svc) cmd := ThrottlerUpdateConfigCmd(ch) cmd.SetArgs([]string{"my-db", "my-branch", "--keyspace", "commerce"}) err := cmd.Execute() - c.Assert(err, qt.IsNotNil) - c.Assert(svc.UpdateThrottlerConfigFnInvoked, qt.IsFalse) + c.Assert(err, qt.IsNil) + c.Assert(svc.UpdateThrottlerConfigFnInvoked, qt.IsTrue) +} + +func TestThrottlerUpdateConfig_ThrottleAppWithMetrics(t *testing.T) { + c := qt.New(t) + + svc := &mock.VtctldService{ + UpdateThrottlerConfigFn: func(ctx context.Context, req *ps.VtctldUpdateThrottlerConfigRequest) (json.RawMessage, error) { + c.Assert(req.Enabled, qt.IsNil) + c.Assert(len(req.Apps), qt.Equals, 1) + c.Assert(req.Apps[0].Name, qt.Equals, "rowstreamer") + c.Assert(req.Apps[0].Ratio, qt.IsNotNil) + c.Assert(*req.Apps[0].Ratio, qt.Equals, 0.5) + c.Assert(req.Apps[0].ExpireAt, qt.Not(qt.Equals), "") + c.Assert(req.AppCheckedMetrics, qt.DeepEquals, map[string][]string{ + "rowstreamer": {"lag", "loadavg"}, + }) + return json.RawMessage(`{}`), nil + }, + } + + ch, _ := newThrottlerTestHelper("my-org", svc) + + cmd := ThrottlerUpdateConfigCmd(ch) + cmd.SetArgs([]string{"my-db", "my-branch", + "--keyspace", "commerce", + "--throttle-app", "rowstreamer", + "--throttle-app-ratio", "0.5", + "--throttle-app-duration", "96h", + "--app-name", "rowstreamer", + "--app-metrics", "lag,loadavg", + }) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.UpdateThrottlerConfigFnInvoked, qt.IsTrue) +} + +func TestThrottlerUpdateConfig_UnthrottleApp(t *testing.T) { + c := qt.New(t) + + svc := &mock.VtctldService{ + UpdateThrottlerConfigFn: func(ctx context.Context, req *ps.VtctldUpdateThrottlerConfigRequest) (json.RawMessage, error) { + c.Assert(req.Enabled, qt.IsNil) + c.Assert(req.UnthrottleApps, qt.DeepEquals, []string{"rowstreamer"}) + return json.RawMessage(`{}`), nil + }, + } + + ch, _ := newThrottlerTestHelper("my-org", svc) + + cmd := ThrottlerUpdateConfigCmd(ch) + cmd.SetArgs([]string{"my-db", "my-branch", + "--keyspace", "commerce", + "--unthrottle-app", "rowstreamer", + }) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.UpdateThrottlerConfigFnInvoked, qt.IsTrue) } diff --git a/internal/planetscale/vtctld_general.go b/internal/planetscale/vtctld_general.go index 5ea3ecfe..28026dc1 100644 --- a/internal/planetscale/vtctld_general.go +++ b/internal/planetscale/vtctld_general.go @@ -164,14 +164,19 @@ type VtctldUpdateThrottlerConfigRequest struct { Keyspace string `json:"keyspace"` // Enabled controls whether the tablet throttler is enabled for the keyspace. - // It is required: the API has no tri-state, so omitting it would disable the - // throttler. - Enabled bool `json:"enabled"` - // Threshold is the threshold for the default throttler check (replication lag - // in seconds). It must be >= 0; the server defaults to 5.0 when omitted. + // Omit to leave the enable state unchanged. + Enabled *bool `json:"enabled,omitempty"` + // Threshold is the threshold for the default throttler check (replication + // lag in seconds). It must be >= 0. Omit to leave the existing threshold + // unchanged. Threshold *float64 `json:"threshold,omitempty"` // Apps configures zero or more per-app throttling rules. Apps []VtctldThrottledAppConfig `json:"apps,omitempty"` + // UnthrottleApps names apps whose throttled-app rules should be removed. + UnthrottleApps []string `json:"unthrottle_apps,omitempty"` + // AppCheckedMetrics assigns metrics to check per app name + // (e.g. {"rowstreamer": ["lag", "loadavg"]}). + AppCheckedMetrics map[string][]string `json:"app_checked_metrics,omitempty"` } type vtctldService struct { diff --git a/internal/planetscale/vtctld_general_test.go b/internal/planetscale/vtctld_general_test.go index bfc2285e..738b23fa 100644 --- a/internal/planetscale/vtctld_general_test.go +++ b/internal/planetscale/vtctld_general_test.go @@ -520,7 +520,7 @@ func TestVtctld_UpdateThrottlerConfig(t *testing.T) { Database: "my-db", Branch: "my-branch", Keyspace: "commerce", - Enabled: true, + Enabled: boolPtr(true), Threshold: float64Ptr(2.5), Apps: []VtctldThrottledAppConfig{ {Name: "online-ddl", Ratio: float64Ptr(0.5)}, @@ -562,7 +562,7 @@ func TestVtctld_UpdateThrottlerConfig_DisableSendsEnabledFalse(t *testing.T) { Database: "my-db", Branch: "my-branch", Keyspace: "commerce", - Enabled: false, + Enabled: boolPtr(false), }) c.Assert(err, qt.IsNil) }