diff --git a/cmd/build_memstats_test.go b/cmd/build_memstats_test.go new file mode 100644 index 0000000000..46ea7b97a1 --- /dev/null +++ b/cmd/build_memstats_test.go @@ -0,0 +1,67 @@ +// Copyright (c) 2026 Lark Technologies Pte. Ltd. +// SPDX-License-Identifier: MIT + +package cmd + +import ( + "context" + "runtime" + "testing" + + "github.com/larksuite/cli/internal/cmdutil" +) + +// TestBuild_DefaultNoCompletionLeak verifies that, without any call to +// SetFlagCompletionsEnabled, repeated cmd.Build invocations do not leak +// *cobra.Command instances into cobra's package-global flag-completion map. +// +// This guards the new default (completions disabled) — if someone flips the +// zero-value back to "enabled", the per-Build memory growth observed under +// `scripts/bench_build` would resurface in production hot paths that build +// the root command without serving a completion request. +func TestBuild_DefaultNoCompletionLeak(t *testing.T) { + if cmdutil.FlagCompletionsEnabled() { + t.Fatalf("precondition: FlagCompletionsEnabled() = true, want false (state polluted by another test)") + } + + snap := func() (heapMB float64, objs uint64) { + runtime.GC() + runtime.GC() + runtime.GC() + var m runtime.MemStats + runtime.ReadMemStats(&m) + return float64(m.HeapAlloc) / 1024 / 1024, m.HeapObjects + } + + // Warm one-time caches (registry JSON decode, embed reads) so the first + // Build's lazy allocations don't skew the per-iteration delta. + _ = Build(context.Background(), cmdutil.InvocationContext{}) + baseMB, baseObj := snap() + + const N = 20 + for range N { + _ = Build(context.Background(), cmdutil.InvocationContext{}) + } + mb, obj := snap() + + deltaMB := mb - baseMB + deltaObj := int64(obj) - int64(baseObj) + perBuildKB := deltaMB * 1024 / float64(N) + perBuildObj := deltaObj / int64(N) + + t.Logf("%d builds: +%.2f MB, +%d objects (%.1f KB/build, %d objs/build)", + N, deltaMB, deltaObj, perBuildKB, perBuildObj) + + // With completions disabled (the default), per-Build retained growth + // should be minimal. Threshold is conservative: the previously observed + // leak with completions enabled was ~hundreds of KB and thousands of + // objects per Build, well above this bound. + const maxKBPerBuild = 50.0 + const maxObjsPerBuild = 500 + if perBuildKB > maxKBPerBuild { + t.Errorf("per-build heap growth = %.1f KB, want <= %.1f KB (completion registration may be leaking)", perBuildKB, maxKBPerBuild) + } + if perBuildObj > maxObjsPerBuild { + t.Errorf("per-build object growth = %d, want <= %d", perBuildObj, maxObjsPerBuild) + } +} diff --git a/cmd/root.go b/cmd/root.go index 1f630bac27..ea69e30f09 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -158,7 +158,7 @@ func isCompletionCommand(args []string) bool { // configureFlagCompletions enables cmdutil.RegisterFlagCompletion only when // the invocation will actually serve a __complete request. func configureFlagCompletions(args []string) { - cmdutil.SetFlagCompletionsDisabled(!isCompletionCommand(args)) + cmdutil.SetFlagCompletionsEnabled(isCompletionCommand(args)) } // handleRootError dispatches a command error to the appropriate handler diff --git a/cmd/root_test.go b/cmd/root_test.go index fe718841fa..cbf3ccf386 100644 --- a/cmd/root_test.go +++ b/cmd/root_test.go @@ -198,7 +198,7 @@ func TestRootLong_AgentSkillsLinkTargetsReadmeSection(t *testing.T) { } func TestConfigureFlagCompletions(t *testing.T) { - t.Cleanup(func() { cmdutil.SetFlagCompletionsDisabled(false) }) + t.Cleanup(func() { cmdutil.SetFlagCompletionsEnabled(false) }) tests := []struct { name string @@ -213,10 +213,10 @@ func TestConfigureFlagCompletions(t *testing.T) { } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { - cmdutil.SetFlagCompletionsDisabled(!tc.wantDisabled) + cmdutil.SetFlagCompletionsEnabled(tc.wantDisabled) configureFlagCompletions(tc.args) - if got := cmdutil.FlagCompletionsDisabled(); got != tc.wantDisabled { - t.Fatalf("FlagCompletionsDisabled() = %v, want %v", got, tc.wantDisabled) + if got := !cmdutil.FlagCompletionsEnabled(); got != tc.wantDisabled { + t.Fatalf("FlagCompletionsEnabled() = %v, want disabled=%v", !got, tc.wantDisabled) } }) } diff --git a/internal/cmdutil/completion.go b/internal/cmdutil/completion.go index c16160955a..60ecd732a7 100644 --- a/internal/cmdutil/completion.go +++ b/internal/cmdutil/completion.go @@ -11,26 +11,27 @@ import ( // Cobra keeps completion callbacks in a package-global map keyed by // *pflag.Flag with no removal path, so registrations made for a *cobra.Command -// outlive the command itself. Skip registration when the current invocation -// will not serve a completion request. -var flagCompletionsDisabled atomic.Bool +// outlive the command itself. Default to disabled (zero value = false) and let +// callers that actually serve a completion request opt in via +// SetFlagCompletionsEnabled(true). +var flagCompletionsEnabled atomic.Bool -// SetFlagCompletionsDisabled switches RegisterFlagCompletion between -// registering and no-op. Typically set once at process start. -func SetFlagCompletionsDisabled(disabled bool) { - flagCompletionsDisabled.Store(disabled) +// SetFlagCompletionsEnabled toggles whether RegisterFlagCompletion actually +// registers callbacks with cobra. Typically set once at process start. +func SetFlagCompletionsEnabled(enabled bool) { + flagCompletionsEnabled.Store(enabled) } -// FlagCompletionsDisabled reports the current switch state. -func FlagCompletionsDisabled() bool { - return flagCompletionsDisabled.Load() +// FlagCompletionsEnabled reports the current switch state. +func FlagCompletionsEnabled() bool { + return flagCompletionsEnabled.Load() } // RegisterFlagCompletion wraps (*cobra.Command).RegisterFlagCompletionFunc // and honors the package switch. The underlying error is swallowed to match // the `_ = cmd.RegisterFlagCompletionFunc(...)` style already used here. func RegisterFlagCompletion(cmd *cobra.Command, flagName string, fn cobra.CompletionFunc) { - if flagCompletionsDisabled.Load() { + if !flagCompletionsEnabled.Load() { return } _ = cmd.RegisterFlagCompletionFunc(flagName, fn) diff --git a/internal/cmdutil/completion_test.go b/internal/cmdutil/completion_test.go index a20d290ba7..b802944c58 100644 --- a/internal/cmdutil/completion_test.go +++ b/internal/cmdutil/completion_test.go @@ -12,18 +12,18 @@ import ( "github.com/spf13/cobra" ) -func TestSetFlagCompletionsDisabled_RoundTrip(t *testing.T) { - t.Cleanup(func() { SetFlagCompletionsDisabled(false) }) +func TestSetFlagCompletionsEnabled_RoundTrip(t *testing.T) { + t.Cleanup(func() { SetFlagCompletionsEnabled(false) }) - if FlagCompletionsDisabled() { - t.Fatal("expected default false") + if FlagCompletionsEnabled() { + t.Fatal("expected default false (completions disabled by default)") } - SetFlagCompletionsDisabled(true) - if !FlagCompletionsDisabled() { + SetFlagCompletionsEnabled(true) + if !FlagCompletionsEnabled() { t.Fatal("expected true after Set(true)") } - SetFlagCompletionsDisabled(false) - if FlagCompletionsDisabled() { + SetFlagCompletionsEnabled(false) + if FlagCompletionsEnabled() { t.Fatal("expected false after Set(false)") } } @@ -31,8 +31,8 @@ func TestSetFlagCompletionsDisabled_RoundTrip(t *testing.T) { // When disabled, a *cobra.Command must be collectable after the caller drops // its reference — i.e. the wrapper did not touch cobra's global map. func TestRegisterFlagCompletion_Disabled_DoesNotRetainCommand(t *testing.T) { - SetFlagCompletionsDisabled(true) - t.Cleanup(func() { SetFlagCompletionsDisabled(false) }) + SetFlagCompletionsEnabled(false) + t.Cleanup(func() { SetFlagCompletionsEnabled(false) }) const N = 5 var collected atomic.Int32 @@ -58,7 +58,8 @@ func TestRegisterFlagCompletion_Disabled_DoesNotRetainCommand(t *testing.T) { // When enabled, the registered completion must be reachable via cobra. func TestRegisterFlagCompletion_Enabled_DoesRegister(t *testing.T) { - SetFlagCompletionsDisabled(false) + SetFlagCompletionsEnabled(true) + t.Cleanup(func() { SetFlagCompletionsEnabled(false) }) cmd := &cobra.Command{Use: "x"} cmd.Flags().String("foo", "", "") diff --git a/shortcuts/common/runner_flag_completion_test.go b/shortcuts/common/runner_flag_completion_test.go index 35ce532fe9..3c96dbd164 100644 --- a/shortcuts/common/runner_flag_completion_test.go +++ b/shortcuts/common/runner_flag_completion_test.go @@ -16,8 +16,8 @@ import ( // the per-flag enum completion (runner.go:879) and the auto-injected --format // completion (runner.go:895). func TestShortcutMount_FlagCompletionsRegistered(t *testing.T) { - t.Cleanup(func() { cmdutil.SetFlagCompletionsDisabled(false) }) - cmdutil.SetFlagCompletionsDisabled(false) + t.Cleanup(func() { cmdutil.SetFlagCompletionsEnabled(false) }) + cmdutil.SetFlagCompletionsEnabled(true) f, _, _, _ := cmdutil.TestFactory(t, nil) parent := &cobra.Command{Use: "root"} @@ -68,8 +68,8 @@ func TestShortcutMount_FlagCompletionsRegistered(t *testing.T) { // TestShortcutMount_FlagCompletionsDisabled verifies the switch actually // prevents the two registrations from landing in cobra's global map. func TestShortcutMount_FlagCompletionsDisabled(t *testing.T) { - t.Cleanup(func() { cmdutil.SetFlagCompletionsDisabled(false) }) - cmdutil.SetFlagCompletionsDisabled(true) + t.Cleanup(func() { cmdutil.SetFlagCompletionsEnabled(false) }) + cmdutil.SetFlagCompletionsEnabled(false) f, _, _, _ := cmdutil.TestFactory(t, nil) parent := &cobra.Command{Use: "root"}