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 <[email protected]>
Co-authored-by: xiaofanluan <[email protected]>

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]

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

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
+}

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")
+	}
+}