From f4ae66a2fd338a500235a5fdd026c1fd065f43f6 Mon Sep 17 00:00:00 2001 From: James <83447078+xiaofan-luan@users.noreply.github.com> Date: Fri, 9 Oct 2026 23:04:42 +0800 Subject: [PATCH] enhance: add startup flag for existing native profiling (#53997) Milvus already provides jemalloc native heap profiling and `/debug/jemalloc/pprof/heap`. This PR only adds an opt-in startup flag, using the existing subprocess path so jemalloc receives its configuration before allocator initialization. ```sh milvus run querynode --native-profile curl -f http://localhost:9091/debug/jemalloc/pprof/heap -o native.heap ``` The flag enables profiling and allocation sampling while preserving other `MALLOC_CONF` options, including the sampling interval. Existing builds, HTTP routing, dump behavior, and dependencies are reused. The legacy HTTP endpoint requires a Linux `with_jemalloc` build with profiling support and remains subject to existing pprof/management routing; enabling management authentication hides legacy DefaultServeMux routes. Fixes #53985 Validation: focused CLI tests and Go vet passed; the startup-only revision also passed a Linux `go build -mod=readonly -tags dynamic,with_jemalloc ./cmd` and gosec on `cmd/milvus`. Formatting and diff checks passed. No service or profiling workload was started. Signed-off-by: xiaofanluan Co-authored-by: xiaofanluan --- cmd/milvus/help.go | 2 + cmd/milvus/main_entry.go | 7 +++ cmd/milvus/native_profile.go | 100 ++++++++++++++++++++++++++++++ cmd/milvus/native_profile_test.go | 87 ++++++++++++++++++++++++++ 4 files changed, 196 insertions(+) create mode 100644 cmd/milvus/native_profile.go create mode 100644 cmd/milvus/native_profile_test.go diff --git a/cmd/milvus/help.go b/cmd/milvus/help.go index d784e8cf71..10b5d74eaf 100644 --- a/cmd/milvus/help.go +++ b/cmd/milvus/help.go @@ -33,6 +33,8 @@ milvus run [server type] [flags] Start the datacoord server. -alias '' Set alias + --native-profile + Enable the existing jemalloc heap profiler (requires a Linux build with jemalloc profiling support). ` stopLine = ` milvus stop [server type] [flags] diff --git a/cmd/milvus/main_entry.go b/cmd/milvus/main_entry.go index 1817bd86fe..4f71bbcc8a 100644 --- a/cmd/milvus/main_entry.go +++ b/cmd/milvus/main_entry.go @@ -17,6 +17,7 @@ package milvus import ( + "fmt" "log" "os" "os/exec" @@ -53,6 +54,11 @@ func Main(args []string) { streamingutil.SetStreamingServiceEnabled() defer asan.LsanDoLeakCheck() + args, profileEnv, err := prepareNativeProfile(args, os.Environ()) + if err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } // execute command as a subprocess if the command contains "--run-with-subprocess" if subArgs, runSubprocess := subprocessArgs(args); runSubprocess { @@ -60,6 +66,7 @@ func Main(args []string) { /* #nosec G204 */ cmd := exec.Command(subArgs[0], subArgs[1:]...) //nolint:gosec // args are from os.Args, not user input + cmd.Env = profileEnv cmd.Stdout = os.Stdout cmd.Stderr = os.Stderr diff --git a/cmd/milvus/native_profile.go b/cmd/milvus/native_profile.go new file mode 100644 index 0000000000..408e54d42f --- /dev/null +++ b/cmd/milvus/native_profile.go @@ -0,0 +1,100 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package milvus + +import ( + "flag" + "io" + "slices" + "strconv" + "strings" +) + +// prepareNativeProfile consumes the opt-in startup flag and configures the +// existing subprocess. A nil environment keeps normal startup unchanged. +func prepareNativeProfile(args, environ []string) ([]string, []string, error) { + command, role := "", -1 + for i := 1; i < len(args); i++ { + if args[i] == "--run-with-subprocess" { + continue + } + if command == "" { + command = args[i] + } else { + role = i + break + } + } + if command != "run" || role < 0 { + return args, nil, nil + } + flags := flag.NewFlagSet("milvus", flag.ContinueOnError) + flags.SetOutput(io.Discard) + flags.String("alias", "", "") + for _, name := range []string{"rootcoord", "querycoord", "datacoord", "indexcoord", "querynode", "datanode", "proxy", "streamingnode"} { + flags.Bool(name, false, "") + } + var flagArgs []string + var positions []int + for i := role + 1; i < len(args); i++ { + if args[i] != "--run-with-subprocess" { + flagArgs = append(flagArgs, args[i]) + positions = append(positions, i) + } + } + enabled := false + var nativeErr error + consumed := make(map[int]bool) + flags.BoolFunc("native-profile", "", func(value string) error { + // BoolFunc runs after the flag token is consumed. Remove only tokens + // parsed as this flag, preserving alias values and unparsed arguments. + consumed[positions[len(flagArgs)-len(flags.Args())-1]] = true + enabled, nativeErr = strconv.ParseBool(value) + return nativeErr + }) + parseErr := flags.Parse(flagArgs) + if nativeErr != nil { + return nil, nil, parseErr + } + if len(consumed) == 0 { + return args, nil, nil + } + clean := make([]string, 0, len(args)) + for i, arg := range args { + if !consumed[i] { + clean = append(clean, arg) + } + } + // Let the existing role parser handle help and invalid flags. + if parseErr != nil || !enabled { + return clean, nil, nil + } + conf := "" + for _, entry := range environ { + if value, ok := strings.CutPrefix(entry, "MALLOC_CONF="); ok { + conf = value + } + } + if conf != "" { + conf += "," + } + conf += "prof:true,prof_active:true,prof_thread_active_init:true" + if !slices.Contains(clean, "--run-with-subprocess") { + clean = append(clean, "--run-with-subprocess") + } + return clean, append(append([]string(nil), environ...), "MALLOC_CONF="+conf), nil +} diff --git a/cmd/milvus/native_profile_test.go b/cmd/milvus/native_profile_test.go new file mode 100644 index 0000000000..cf39cc6b11 --- /dev/null +++ b/cmd/milvus/native_profile_test.go @@ -0,0 +1,87 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package milvus + +import ( + "reflect" + "testing" +) + +func TestNativeProfileDefaultAndHelp(t *testing.T) { + option := "--native-profile" + environ := []string{"MALLOC_CONF=prof:true,lg_prof_sample:17", "KEEP=value"} + for _, tc := range []struct { + args, want []string + }{ + {args: []string{"milvus"}}, + {args: []string{"milvus", "run", "querynode", "--run-with-subprocess"}}, + {args: []string{"milvus", "run", "querynode", "--alias", option}}, + {args: []string{"milvus", "run", "querynode", "-alias", "-native-profile"}}, + {args: []string{"milvus", "run", "querynode", "--", option}}, + {args: []string{"milvus", "run", "querynode", "positional", option}}, + {args: []string{"milvus", "run", "querynode", "--help", option}}, + {args: []string{"milvus", "run", "querynode", "--unknown", option}}, + {args: []string{"milvus", "stop", "querynode", option}}, + {args: []string{"milvus", "run", "querynode", option, "--help", option}, want: []string{"milvus", "run", "querynode", "--help", option}}, + {args: []string{"milvus", "run", "querynode", "--native-profile=false"}, want: []string{"milvus", "run", "querynode"}}, + {args: []string{"milvus", "run", "querynode", option, "--native-profile=false"}, want: []string{"milvus", "run", "querynode"}}, + } { + before := append([]string(nil), tc.args...) + args, env, err := prepareNativeProfile(tc.args, environ) + want := tc.want + if want == nil { + want = before + } + if err != nil || env != nil || !reflect.DeepEqual(args, want) || !reflect.DeepEqual(tc.args, before) { + t.Fatalf("default/help %v = %v, %v, %v", tc.args, args, env, err) + } + } +} + +func TestNativeProfilePreparesSubprocess(t *testing.T) { + environ := []string{"KEEP=value", "MALLOC_CONF=background_thread:true,prof:false,prof_active:false,prof_thread_active_init:false,lg_prof_sample:17,prof_final:true"} + beforeEnv := append([]string(nil), environ...) + wantEnv := append(append([]string(nil), environ...), environ[1]+",prof:true,prof_active:true,prof_thread_active_init:true") + for _, tc := range []struct{ args, want []string }{ + {[]string{"milvus", "run", "querynode", "--native-profile", "--alias", "node"}, []string{"milvus", "run", "querynode", "--alias", "node", "--run-with-subprocess"}}, + {[]string{"milvus", "run", "querynode", "--run-with-subprocess", "--native-profile=true"}, []string{"milvus", "run", "querynode", "--run-with-subprocess"}}, + {[]string{"milvus", "--run-with-subprocess", "run", "querynode", "-native-profile"}, []string{"milvus", "--run-with-subprocess", "run", "querynode"}}, + {[]string{"milvus", "run", "standalone", "--native-profile"}, []string{"milvus", "run", "standalone", "--run-with-subprocess"}}, + } { + before := append([]string(nil), tc.args...) + clean, env, err := prepareNativeProfile(tc.args, environ) + if err != nil || !reflect.DeepEqual(clean, tc.want) || !reflect.DeepEqual(env, wantEnv) { + t.Fatalf("subprocess %v = %v, %v, %v", tc.args, clean, env, err) + } + if !reflect.DeepEqual(tc.args, before) || !reflect.DeepEqual(environ, beforeEnv) { + t.Fatal("profiling modified caller args/environment") + } + } +} + +func TestNativeProfileInvalidValue(t *testing.T) { + for _, option := range []string{"--native-profile=", "--native-profile=garbage", "--native-profile=/profiles"} { + _, env, err := prepareNativeProfile([]string{"milvus", "run", "querynode", option}, nil) + if err == nil || env != nil { + t.Fatalf("invalid value %q accepted: %v, %v", option, env, err) + } + } + _, _, err := prepareNativeProfile([]string{"milvus", "run", "querynode", "--native-profile", "--native-profile=bad"}, nil) + if err == nil { + t.Fatal("invalid value after an earlier profiling option must fail") + } +}