feat: add profile for benchmark

grpc/kitex/client/main.go

@@ -17,10 +17,29 @@
 package main
 
 import (
+	"fmt"
+	"os"
+	"runtime/pprof"
+
 	"github.com/cloudwego/kitex-benchmark/runner"
 )
 
 // main is use for routing.
 func main() {
+	if os.Getenv("KITEX_ENABLE_PROFILE") == "1" {
+		fmt.Println("[Kitex profile is enabled]")
+		// start cpu profile
+		cpuProfile, _ := os.Create("output/benchmark-grpc-client-cpu.pprof")
+		defer cpuProfile.Close()
+		_ = pprof.StartCPUProfile(cpuProfile)
+		defer pprof.StopCPUProfile()
+
+		// heap profile after finish
+		heapProfile, _ := os.Create("output/benchmark-grpc-client-mem.pprof")
+		defer func() {
+			_ = pprof.WriteHeapProfile(heapProfile)
+			heapProfile.Close()
+		}()
+	}
 	runner.Main("KITEX", NewKClient)
 }

grpc/kitex/main.go

@@ -21,6 +21,8 @@ import (
 	"fmt"
 	"log"
 	"net"
+	"os"
+	"runtime/pprof"
 
 	"github.com/cloudwego/kitex/server"
 
@@ -51,6 +53,21 @@ func (s *EchoImpl) Echo(ctx context.Context, req *echo.Request) (*echo.Response,
 }
 
 func main() {
+	if os.Getenv("KITEX_ENABLE_PROFILE") == "1" {
+		fmt.Println("[Kitex profile is enabled]")
+		// start cpu profile
+		cpuProfile, _ := os.Create("output/benchmark-grpc-server-cpu.pprof")
+		defer cpuProfile.Close()
+		_ = pprof.StartCPUProfile(cpuProfile)
+		defer pprof.StopCPUProfile()
+
+		// heap profile after finish
+		heapProfile, _ := os.Create("output/benchmark-grpc-server-mem.pprof")
+		defer func() {
+			_ = pprof.WriteHeapProfile(heapProfile)
+			heapProfile.Close()
+		}()
+	}
 	// start pprof server
 	go func() {
 		perf.ServeMonitor(fmt.Sprintf(":%d", port+10000))

scripts/base.sh

@@ -8,6 +8,11 @@ body=(1024)
 concurrent=(100 200 400 600 800 1000)
 qps=(0)
 sleep=0
+# NOTICE: if you want to dump profile, set "enable_profile" to 1
+enable_profile=0
+if [ $enable_profile -eq 1 ]; then
+  export KITEX_ENABLE_PROFILE=1
+fi
 
 CURDIR=$(cd $(dirname $0); pwd)
 
@@ -74,6 +79,15 @@ function kill_pid_listening_on_port() {
         exit 1
     fi
     pids=`lsof -i ":$port" | grep LISTEN | awk '{print $2}' | uniq`
+
+    for p in $pids; do
+        echo "Sending termination signal to $p..."
+        kill -s SIGTERM $p  # 发送 SIGTERM 信号给程序
+    done
+
+    # 给程序一定时间处理信号后再强制终止
+    sleep 2
+
     for p in $pids; do
         echo killing $p...
         kill $p

thrift/kitex/client/main.go

@@ -17,6 +17,9 @@
 package main
 
 import (
+	"fmt"
+	"os"
+	"runtime/pprof"
 	"time"
 
 	"github.com/cloudwego/kitex/client"
@@ -30,6 +33,21 @@ import (
 
 // main is use for routing.
 func main() {
+	if os.Getenv("KITEX_ENABLE_PROFILE") == "1" {
+		fmt.Println("[Kitex profile is enabled]")
+		// start cpu profile
+		cpuProfile, _ := os.Create("output/benchmark-thrift-client-cpu.pprof")
+		defer cpuProfile.Close()
+		_ = pprof.StartCPUProfile(cpuProfile)
+		defer pprof.StopCPUProfile()
+
+		// heap profile after finish
+		heapProfile, _ := os.Create("output/benchmark-thrift-client-mem.pprof")
+		defer func() {
+			_ = pprof.WriteHeapProfile(heapProfile)
+			heapProfile.Close()
+		}()
+	}
 	runner.Main("KITEX", NewThriftKitexClient)
 }
 

thrift/kitex/main.go

@@ -21,6 +21,8 @@ import (
 	"fmt"
 	"log"
 	"net"
+	"os"
+	"runtime/pprof"
 
 	"github.com/cloudwego/kitex/server"
 
@@ -60,6 +62,21 @@ func (s *EchoServerImpl) EchoComplex(ctx context.Context, req *echo.ComplexReque
 }
 
 func main() {
+	if os.Getenv("KITEX_ENABLE_PROFILE") == "1" {
+		fmt.Println("[Kitex profile is enabled]")
+		// start cpu profile
+		cpuProfile, _ := os.Create("output/benchmark-thrift-server-cpu.pprof")
+		defer cpuProfile.Close()
+		_ = pprof.StartCPUProfile(cpuProfile)
+		defer pprof.StopCPUProfile()
+
+		// heap profile after finish
+		heapProfile, _ := os.Create("output/benchmark-thrift-server-mem.pprof")
+		defer func() {
+			_ = pprof.WriteHeapProfile(heapProfile)
+			heapProfile.Close()
+		}()
+	}
 	// start pprof server
 	go func() {
 		perf.ServeMonitor(fmt.Sprintf(":%d", port+10000))