From 2e53a4b2a3d8d077ef0e332a6a30bada21ed88e3 Mon Sep 17 00:00:00 2001
From: HeyJavaBean <liyun.339@bytedance.com>
Date: Thu, 26 Sep 2024 17:35:12 +0800
Subject: [PATCH] feat/add_profile

---
 grpc/kitex/client/main.go   | 15 +++++++++++++++
 grpc/kitex/main.go          | 15 +++++++++++++++
 scripts/base.sh             | 14 ++++++++++++++
 thrift/kitex/client/main.go | 15 +++++++++++++++
 thrift/kitex/main.go        | 15 +++++++++++++++
 5 files changed, 74 insertions(+)

diff --git a/grpc/kitex/client/main.go b/grpc/kitex/client/main.go
index b3b3e0a..04cfd35 100644
--- a/grpc/kitex/client/main.go
+++ b/grpc/kitex/client/main.go
@@ -22,5 +22,20 @@ 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-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)
 }
diff --git a/grpc/kitex/main.go b/grpc/kitex/main.go
index cd80a16..f033ddc 100644
--- a/grpc/kitex/main.go
+++ b/grpc/kitex/main.go
@@ -51,6 +51,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))
diff --git a/scripts/base.sh b/scripts/base.sh
index ab83c61..08a8763 100755
--- a/scripts/base.sh
+++ b/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
diff --git a/thrift/kitex/client/main.go b/thrift/kitex/client/main.go
index 4079c02..7cc4a2c 100644
--- a/thrift/kitex/client/main.go
+++ b/thrift/kitex/client/main.go
@@ -30,6 +30,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)
 }
 
diff --git a/thrift/kitex/main.go b/thrift/kitex/main.go
index 233f483..b85b308 100644
--- a/thrift/kitex/main.go
+++ b/thrift/kitex/main.go
@@ -60,6 +60,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))