support pyroscope

Signed-off-by: 童剑 <[email protected]>
bufferflies · Dec 30, 2024 · 55c3422 · 55c3422
1 parent 5999d64
commit 55c3422
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 0 deletions.
diff --git a/cmd/pd-server/main.go b/cmd/pd-server/main.go
@@ -232,6 +232,7 @@ func start(cmd *cobra.Command, args []string, services ...string) {
 	grpcprometheus.EnableHandlingTimeHistogram()
 
 	metricutil.Push(&cfg.Metric)
+	metricutil.EnablePyroscope()
 
 	err = join.PrepareJoinCluster(cfg)
 	if err != nil {

diff --git a/go.mod b/go.mod
@@ -63,6 +63,7 @@ require (
 	golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d
 	google.golang.org/grpc v1.62.1
 	gotest.tools/gotestsum v1.7.0
+	github.com/grafana/pyroscope-go v1.2.0
 )
 
 require (

diff --git a/go.sum b/go.sum
@@ -245,6 +245,7 @@ github.com/gorilla/mux v1.7.4/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB7
 github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/gorilla/websocket v1.5.1 h1:gmztn0JnHVt9JZquRuzLw3g4wouNVzKL15iLr/zn/QY=
 github.com/gorilla/websocket v1.5.1/go.mod h1:x3kM2JMyaluk02fnUJpQuwD2dCS5NDG2ZHL0uE0tcaY=
+github.com/grafana/pyroscope-go v1.2.0/go.mod h1:2GHr28Nr05bg2pElS+dDsc98f3JTUh2f6Fz1hWXrqwk=
 github.com/grpc-ecosystem/go-grpc-middleware v1.4.0 h1:UH//fgunKIs4JdUbpDl1VZCDaL56wXCB/5+wF6uHfaI=
 github.com/grpc-ecosystem/go-grpc-middleware v1.4.0/go.mod h1:g5qyo/la0ALbONm6Vbp88Yd8NsDy6rZz+RcrMPxvld8=
 github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0 h1:Ovs26xHkKqVztRpIrF/92BcuyuQ/YW4NSIpoGtfXNho=

diff --git a/pkg/utils/metricutil/metricutil.go b/pkg/utils/metricutil/metricutil.go
@@ -15,7 +15,10 @@
 package metricutil
 
 import (
+	"github.com/grafana/pyroscope-go"
+	"go.uber.org/zap"
 	"os"
+	"runtime"
 	"time"
 	"unicode"
 
@@ -102,3 +105,28 @@ func instanceName() string {
 	}
 	return hostname
 }
+
+// EnablePyroscope enables pyroscope if pyroscope is enabled.
+func EnablePyroscope() {
+	if os.Getenv("PYROSCOPE_SERVER_ADDRESS") != "" {
+		runtime.SetMutexProfileFraction(5)
+		runtime.SetBlockProfileRate(5)
+		_, err := pyroscope.Start(pyroscope.Config{
+			ApplicationName:   "pd",
+			ServerAddress:     os.Getenv("PYROSCOPE_SERVER_ADDRESS"),
+			Logger:            pyroscope.StandardLogger,
+			AuthToken:         os.Getenv("PYROSCOPE_AUTH_TOKEN"),
+			TenantID:          os.Getenv("PYROSCOPE_TENANT_ID"),
+			BasicAuthUser:     os.Getenv("PYROSCOPE_BASIC_AUTH_USER"),
+			BasicAuthPassword: os.Getenv("PYROSCOPE_BASIC_AUTH_PASSWORD"),
+			ProfileTypes: []pyroscope.ProfileType{
+				pyroscope.ProfileCPU,
+				pyroscope.ProfileAllocSpace,
+			},
+			UploadRate: 30 * time.Second,
+		})
+		if err != nil {
+			log.Fatal("fail to start pyroscope", zap.Error(err))
+		}
+	}
+}