pkg/loop: support pyroscope

Author: jmank88

go.mod

@@ -18,6 +18,8 @@ require (
 	github.com/golang-jwt/jwt/v5 v5.2.3
 	github.com/google/go-cmp v0.7.0
 	github.com/google/uuid v1.6.0
+	github.com/grafana/otel-profiling-go v0.5.1
+	github.com/grafana/pyroscope-go v1.2.8
 	github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus v1.0.1
 	github.com/hashicorp/go-hclog v1.6.3
 	github.com/hashicorp/go-plugin v1.7.0
@@ -103,6 +105,7 @@ require (
 	github.com/gofrs/uuid v4.4.0+incompatible // indirect
 	github.com/golang/protobuf v1.5.4 // indirect
 	github.com/google/flatbuffers v25.2.10+incompatible // indirect
+	github.com/grafana/pyroscope-go/godeltaprof v0.1.9 // indirect
 	github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.3.2 // indirect
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.3 // indirect
 	github.com/hako/durafmt v0.0.0-20200710122514-c0fb7b4da026 // indirect

go.sum

@@ -46,6 +46,13 @@ const (
 
 	envPromPort = "CL_PROMETHEUS_PORT"
 
+	envPyroscopeAuthToken                 = "CL_PYROSCOPE_AUTH_TOKEN"
+	envPyroscopeServerAddress             = "CL_PYROSCOPE_SERVER_ADDRESS"
+	envPyroscopeEnvironment               = "CL_PYROSCOPE_ENVIRONMENT"
+	envPyroscopeLinkTracesToProfiles      = "CL_PYROSCOPE_LINK_TRACES_TO_PROFILES"
+	envPyroscopePPROFBlockProfileRate     = "CL_PYROSCOPE_PPROF_BLOCK_PROFILE_RATE"
+	envPyroscopePPROFMutexProfileFraction = "CL_PYROSCOPE_PPROF_MUTEX_PROFILE_FRACTION"
+
 	envTracingEnabled         = "CL_TRACING_ENABLED"
 	envTracingCollectorTarget = "CL_TRACING_COLLECTOR_TARGET"
 	envTracingSamplingRatio   = "CL_TRACING_SAMPLING_RATIO"
@@ -89,6 +96,12 @@ const (
 type EnvConfig struct {
 	AppID string
 
+	ChipIngressEndpoint           string
+	ChipIngressInsecureConnection bool
+
+	CRESettings        string
+	CRESettingsDefault string
+
 	DatabaseURL                          *config.SecretURL
 	DatabaseIdleInTxSessionTimeout       time.Duration
 	DatabaseLockTimeout                  time.Duration
@@ -115,13 +128,14 @@ type EnvConfig struct {
 	MercuryTransmitterReaperMaxAge         time.Duration
 	MercuryVerboseLogging                  bool
 
-	PrometheusPort int //TODO more than just prom
+	PrometheusPort int // also serves pprof routes
 
-	TracingEnabled         bool
-	TracingCollectorTarget string
-	TracingSamplingRatio   float64
-	TracingTLSCertPath     string
-	TracingAttributes      map[string]string
+	PyroscopeAuthToken                 string
+	PyroscopeServerAddress             string
+	PyroscopeEnvironment               string
+	PyroscopeLinkTracesToProfiles      bool
+	PyroscopePPROFBlockProfileRate     int
+	PyroscopePPROFMutexProfileFraction int
 
 	TelemetryEnabled                   bool
 	TelemetryEndpoint                  string
@@ -148,11 +162,11 @@ type EnvConfig struct {
 	TelemetryMetricCompressor          string
 	TelemetryLogCompressor             string
 
-	ChipIngressEndpoint           string
-	ChipIngressInsecureConnection bool
-
-	CRESettings        string
-	CRESettingsDefault string
+	TracingEnabled         bool
+	TracingCollectorTarget string
+	TracingSamplingRatio   float64
+	TracingTLSCertPath     string
+	TracingAttributes      map[string]string
 }
 
 // AsCmdEnv returns a slice of environment variable key/value pairs for an exec.Cmd.
@@ -193,6 +207,13 @@ func (e *EnvConfig) AsCmdEnv() (env []string) {
 
 	add(envPromPort, strconv.Itoa(e.PrometheusPort))
 
+	add(envPyroscopeAuthToken, e.PyroscopeAuthToken)
+	add(envPyroscopeServerAddress, e.PyroscopeServerAddress)
+	add(envPyroscopeEnvironment, e.PyroscopeEnvironment)
+	add(envPyroscopeLinkTracesToProfiles, strconv.FormatBool(e.PyroscopeLinkTracesToProfiles))
+	add(envPyroscopePPROFBlockProfileRate, strconv.Itoa(e.PyroscopePPROFBlockProfileRate))
+	add(envPyroscopePPROFMutexProfileFraction, strconv.Itoa(e.PyroscopePPROFMutexProfileFraction))
+
 	add(envTracingEnabled, strconv.FormatBool(e.TracingEnabled))
 	add(envTracingCollectorTarget, e.TracingCollectorTarget)
 	add(envTracingSamplingRatio, strconv.FormatFloat(e.TracingSamplingRatio, 'f', -1, 64))
@@ -352,6 +373,22 @@ func (e *EnvConfig) parse() error {
 		return fmt.Errorf("failed to parse %s = %q: %w", envPromPort, promPortStr, err)
 	}
 
+	e.PyroscopeAuthToken = os.Getenv(envPyroscopeAuthToken)
+	e.PyroscopeServerAddress = os.Getenv(envPyroscopeServerAddress)
+	e.PyroscopeEnvironment = os.Getenv(envPyroscopeEnvironment)
+	e.PyroscopeLinkTracesToProfiles, err = getBool(envPyroscopeLinkTracesToProfiles)
+	if err != nil {
+		return fmt.Errorf("failed to parse %s: %w", envPyroscopeLinkTracesToProfiles, err)
+	}
+	e.PyroscopePPROFBlockProfileRate, err = getInt(envPyroscopePPROFBlockProfileRate)
+	if err != nil {
+		return fmt.Errorf("failed to parse %s: %w", envPyroscopePPROFBlockProfileRate, err)
+	}
+	e.PyroscopePPROFMutexProfileFraction, err = getInt(envPyroscopePPROFMutexProfileFraction)
+	if err != nil {
+		return fmt.Errorf("failed to parse %s: %w", envPyroscopePPROFMutexProfileFraction, err)
+	}
+
 	e.TracingEnabled, err = getBool(envTracingEnabled)
 	if err != nil {
 		return fmt.Errorf("failed to parse %s: %w", envTracingEnabled, err)

pkg/loop/config.go

@@ -55,6 +55,13 @@ func TestEnvConfig_parse(t *testing.T) {
 
 				envPromPort: "8080",
 
+				envPyroscopeAuthToken:                 "token",
+				envPyroscopeServerAddress:             "http://pyroscope:4040",
+				envPyroscopeEnvironment:               "pyroscope-env",
+				envPyroscopeLinkTracesToProfiles:      "true",
+				envPyroscopePPROFBlockProfileRate:     "42",
+				envPyroscopePPROFMutexProfileFraction: "99",
+
 				envTracingEnabled:           "true",
 				envTracingCollectorTarget:   "some:target",
 				envTracingSamplingRatio:     "1.0",
@@ -160,6 +167,13 @@ var envCfgFull = EnvConfig{
 
 	PrometheusPort: 8080,
 
+	PyroscopeAuthToken:                 "token",
+	PyroscopeServerAddress:             "http://pyroscope:4040",
+	PyroscopeEnvironment:               "pyroscope-env",
+	PyroscopeLinkTracesToProfiles:      true,
+	PyroscopePPROFBlockProfileRate:     42,
+	PyroscopePPROFMutexProfileFraction: 99,
+
 	TracingEnabled:         true,
 	TracingAttributes:      map[string]string{"XYZ": "value"},
 	TracingCollectorTarget: "some:target",
@@ -213,6 +227,13 @@ func TestEnvConfig_AsCmdEnv(t *testing.T) {
 
 	assert.Equal(t, strconv.Itoa(8080), got[envPromPort])
 
+	assert.Equal(t, "token", got[envPyroscopeAuthToken])
+	assert.Equal(t, "http://pyroscope:4040", got[envPyroscopeServerAddress])
+	assert.Equal(t, "pyroscope-env", got[envPyroscopeEnvironment])
+	assert.Equal(t, "true", got[envPyroscopeLinkTracesToProfiles])
+	assert.Equal(t, "42", got[envPyroscopePPROFBlockProfileRate])
+	assert.Equal(t, "99", got[envPyroscopePPROFMutexProfileFraction])
+
 	assert.Equal(t, "true", got[envTracingEnabled])
 	assert.Equal(t, "some:target", got[envTracingCollectorTarget])
 	assert.Equal(t, "1", got[envTracingSamplingRatio])

pkg/loop/config_test.go

@@ -5,9 +5,15 @@ import (
 	"fmt"
 	"os"
 	"os/signal"
+	"path/filepath"
+	"runtime"
+	"runtime/debug"
 	"time"
 
+	otelpyroscope "github.com/grafana/otel-profiling-go"
+	"github.com/grafana/pyroscope-go"
 	"github.com/jmoiron/sqlx"
+	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
 	sdkmetric "go.opentelemetry.io/otel/sdk/metric"
 	semconv "go.opentelemetry.io/otel/semconv/v1.17.0"
@@ -94,6 +100,7 @@ type Server struct {
 	webServer       *webServer
 	checker         *services.HealthChecker
 	LimitsFactory   limits.Factory
+	profiler        *pyroscope.Profiler
 }
 
 func newServer(loggerName string) (*Server, error) {
@@ -221,6 +228,67 @@ func (s *Server) start(opts ...ServerOpt) error {
 		}
 	}
 
+	if addr := s.EnvConfig.PyroscopeServerAddress; addr != "" {
+		runtime.SetBlockProfileRate(s.EnvConfig.PyroscopePPROFBlockProfileRate)
+		runtime.SetMutexProfileFraction(s.EnvConfig.PyroscopePPROFMutexProfileFraction)
+
+		hostname, _ := os.Hostname()
+		var ver, sha, goVer, module string
+		if bi, ok := debug.ReadBuildInfo(); ok {
+			ver = bi.Main.Version
+			sha = bi.Main.Sum
+			if len(sha) > 7 {
+				sha = sha[:7]
+			}
+			goVer = bi.GoVersion
+			module = bi.Main.Path
+		}
+
+		appName, err := os.Executable()
+		if err != nil {
+			s.Logger.Warnf("Failed to get executable name: %v", err)
+			appName = "unknown"
+		} else {
+			appName = filepath.Base(appName)
+		}
+
+		s.profiler, err = pyroscope.Start(pyroscope.Config{
+			ApplicationName: appName,
+			ServerAddress:   s.EnvConfig.PyroscopeServerAddress,
+			AuthToken:       s.EnvConfig.PyroscopeAuthToken,
+
+			Tags: map[string]string{
+				"module":      module,
+				"SHA":         sha,
+				"Version":     ver,
+				"go":          goVer,
+				"Environment": s.EnvConfig.PyroscopeEnvironment,
+				"hostname":    hostname,
+			},
+			ProfileTypes: []pyroscope.ProfileType{
+				// these profile types are enabled by default:
+				pyroscope.ProfileCPU,
+				pyroscope.ProfileAllocObjects,
+				pyroscope.ProfileAllocSpace,
+				pyroscope.ProfileInuseObjects,
+				pyroscope.ProfileInuseSpace,
+
+				// these profile types are optional:
+				pyroscope.ProfileGoroutines,
+				pyroscope.ProfileMutexCount,
+				pyroscope.ProfileMutexDuration,
+				pyroscope.ProfileBlockCount,
+				pyroscope.ProfileBlockDuration,
+			},
+		})
+		if err != nil {
+			return fmt.Errorf("failed to start pyroscope profiler: %w", err)
+		}
+		if tracingConfig.Enabled && s.EnvConfig.PyroscopeLinkTracesToProfiles {
+			otel.SetTracerProvider(otelpyroscope.NewTracerProvider(otel.GetTracerProvider()))
+		}
+	}
+
 	s.webServer = WebServerOpts{}.New(s.Logger, s.EnvConfig.PrometheusPort)
 	if err := s.webServer.Start(ctx); err != nil {
 		return fmt.Errorf("error starting prometheus server: %w", err)
@@ -291,6 +359,9 @@ func (s *Server) Stop() {
 	}
 	s.Logger.ErrorIfFn(s.checker.Close, "Failed to close health checker")
 	s.Logger.ErrorIfFn(s.webServer.Close, "Failed to close web server")
+	if s.profiler != nil {
+		s.Logger.ErrorIfFn(s.profiler.Stop, "Failed to stop pyroscope profiler")
+	}
 	if err := s.Logger.Sync(); err != nil {
 		fmt.Println("Failed to sync logger:", err)
 	}