Add kubernetes enrichment option with pyroscope (#51521)

* feat: add kubernetes enrichment option for pyroscope Signed-off-by: Evan Freed <evan.freed@goteleport.com> * fix license and refactor Signed-off-by: Evan Freed <evan.freed@goteleport.com> * Update lib/service/pyroscope.go Co-authored-by: rosstimothy <39066650+rosstimothy@users.noreply.github.com> * Update lib/service/pyroscope.go Co-authored-by: rosstimothy <39066650+rosstimothy@users.noreply.github.com> * refactor and tests Signed-off-by: Evan Freed <evan.freed@goteleport.com> * review comments Signed-off-by: Evan Freed <evan.freed@goteleport.com> * Update lib/service/pyroscope.go Co-authored-by: rosstimothy <39066650+rosstimothy@users.noreply.github.com> * fix Signed-off-by: Evan Freed <evan.freed@goteleport.com> * Update lib/service/pyroscope.go Co-authored-by: rosstimothy <39066650+rosstimothy@users.noreply.github.com> * review comments for tests Signed-off-by: Evan Freed <evan.freed@goteleport.com> * Update lib/service/pyroscope_test.go Co-authored-by: rosstimothy <39066650+rosstimothy@users.noreply.github.com> * logger Signed-off-by: Evan Freed <evan.freed@goteleport.com> * replace logger Signed-off-by: Evan Freed <evan.freed@goteleport.com> --------- Signed-off-by: Evan Freed <evan.freed@goteleport.com> Co-authored-by: rosstimothy <39066650+rosstimothy@users.noreply.github.com>
gravitational · Feb 3, 2025 · c642532 · c642532
1 parent 1ef1f84
commit c642532
Show file tree

Hide file tree

Showing 2 changed files with 124 additions and 10 deletions.
diff --git a/lib/service/pyroscope.go b/lib/service/pyroscope.go
@@ -24,6 +24,7 @@ import (
 	"time"
 
 	"github.com/grafana/pyroscope-go"
+	"github.com/gravitational/trace"
 
 	"github.com/gravitational/teleport"
 )
@@ -59,22 +60,21 @@ func (l pyroscopeLogger) Errorf(format string, args ...interface{}) {
 	l.l.Error(fmt.Sprintf(format, args...))
 }
 
-// initPyroscope instruments Teleport to run with continuous profiling for Pyroscope
-func (process *TeleportProcess) initPyroscope(address string) {
+// createPyroscopeConfig generates the Pyroscope configuration for the Teleport process.
+func createPyroscopeConfig(ctx context.Context, logger *slog.Logger, address string) (pyroscope.Config, error) {
 	if address == "" {
-		return
+		return pyroscope.Config{}, trace.BadParameter("pyroscope address is empty")
 	}
 
 	hostname, err := os.Hostname()
 	if err != nil {
 		hostname = "unknown"
 	}
 
-	// Build pyroscope config
 	config := pyroscope.Config{
 		ApplicationName: teleport.ComponentTeleport,
 		ServerAddress:   address,
-		Logger:          pyroscope.Logger(pyroscopeLogger{l: slog.Default()}),
+		Logger:          pyroscope.Logger(pyroscopeLogger{l: logger}),
 		Tags: map[string]string{
 			"host":    hostname,
 			"version": teleport.Version,
@@ -84,38 +84,59 @@ func (process *TeleportProcess) initPyroscope(address string) {
 
 	// Evaluate if profile configuration is customized
 	if p := getPyroscopeProfileTypesFromEnv(); len(p) == 0 {
-		slog.InfoContext(process.ExitContext(), "No profile types enabled, using default")
+		logger.InfoContext(ctx, "No profile types enabled, using default")
 	} else {
 		config.ProfileTypes = p
+		logger.InfoContext(ctx, "Pyroscope will configure profiles from env")
 	}
 
 	var uploadRate *time.Duration
 	if rate := os.Getenv("TELEPORT_PYROSCOPE_UPLOAD_RATE"); rate != "" {
 		parsedRate, err := time.ParseDuration(rate)
 		if err != nil {
-			slog.InfoContext(process.ExitContext(), "invalid TELEPORT_PYROSCOPE_UPLOAD_RATE, ignoring value", "provided_value", rate, "error", err)
+			logger.InfoContext(ctx, "invalid TELEPORT_PYROSCOPE_UPLOAD_RATE, ignoring value", "provided_value", rate, "error", err)
 		} else {
 			uploadRate = &parsedRate
 		}
 	} else {
-		slog.InfoContext(process.ExitContext(), "TELEPORT_PYROSCOPE_UPLOAD_RATE not specified, using default")
+		logger.InfoContext(ctx, "TELEPORT_PYROSCOPE_UPLOAD_RATE not specified, using default")
 	}
 
 	// Set UploadRate or fall back to defaults
 	if uploadRate != nil {
 		config.UploadRate = *uploadRate
 	}
 
+	if value, isSet := os.LookupEnv("TELEPORT_PYROSCOPE_KUBE_COMPONENT"); isSet {
+		config.Tags["component"] = value
+	}
+
+	if value, isSet := os.LookupEnv("TELEPORT_PYROSCOPE_KUBE_NAMESPACE"); isSet {
+		config.Tags["namespace"] = value
+	}
+
+	return config, nil
+}
+
+// initPyroscope instruments Teleport to run with continuous profiling for Pyroscope
+func (process *TeleportProcess) initPyroscope(address string) {
+	logger := process.logger.With(teleport.ComponentKey, "pyroscope")
+	config, err := createPyroscopeConfig(process.ExitContext(), logger, address)
+	if err != nil {
+		logger.ErrorContext(process.ExitContext(), "failed to create Pyroscope config", "address", address, "error", err)
+		return
+	}
+
 	profiler, err := pyroscope.Start(config)
 	if err != nil {
-		slog.ErrorContext(process.ExitContext(), "error starting pyroscope profiler", "error", err)
+		logger.ErrorContext(process.ExitContext(), "error starting pyroscope profiler", "address", address, "error", err)
 	} else {
 		process.OnExit("pyroscope.profiler", func(payload any) {
 			profiler.Flush(payload == nil)
 			_ = profiler.Stop()
 		})
 	}
-	slog.InfoContext(process.ExitContext(), "Pyroscope has successfully started")
+	logger.InfoContext(process.ExitContext(), "Pyroscope has successfully started")
 }
 
 // getPyroscopeProfileTypesFromEnv sets the profile types based on environment variables.

diff --git a/lib/service/pyroscope_test.go b/lib/service/pyroscope_test.go
@@ -0,0 +1,93 @@
+// Teleport
+// Copyright (C) 2025 Gravitational, Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+package service
+
+import (
+	"context"
+	"log/slog"
+	"testing"
+
+	"github.com/grafana/pyroscope-go"
+	"github.com/stretchr/testify/require"
+)
+
+func TestPyroscopeConfig(t *testing.T) {
+	tests := []struct {
+		name           string
+		address        string
+		envVars        map[string]string
+		want           pyroscope.Config
+		errorAssertion require.ErrorAssertionFunc
+	}{
+		{
+			name:    "No address configured",
+			envVars: map[string]string{},
+			want: pyroscope.Config{
+				Tags: map[string]string{},
+			},
+			errorAssertion: require.Error,
+		},
+		{
+			name:    "No environment vars configured",
+			address: "127.0.0.1",
+			envVars: map[string]string{},
+			want: pyroscope.Config{
+				Tags: map[string]string{},
+			},
+			errorAssertion: require.NoError,
+		},
+		{
+			name:    "Environment vars configured",
+			address: "127.0.0.1",
+			envVars: map[string]string{
+				"TELEPORT_PYROSCOPE_PROFILE_CPU_ENABLED":        "true",
+				"TELEPORT_PYROSCOPE_PROFILE_GOROUTINES_ENABLED": "true",
+				"TELEPORT_PYROSCOPE_PROFILE_MEMORY_ENABLED":     "true",
+				"TELEPORT_PYROSCOPE_KUBE_COMPONENT":             "auth",
+				"TELEPORT_PYROSCOPE_KUBE_NAMESPACE":             "test-namespace",
+			},
+			want: pyroscope.Config{
+				ProfileTypes: []pyroscope.ProfileType{
+					pyroscope.ProfileAllocObjects,
+					pyroscope.ProfileAllocSpace,
+					pyroscope.ProfileInuseObjects,
+					pyroscope.ProfileInuseSpace,
+					pyroscope.ProfileCPU,
+					pyroscope.ProfileGoroutines,
+				},
+				Tags: map[string]string{
+					"component": "auth",
+					"namespace": "test-namespace",
+				},
+			},
+			errorAssertion: require.NoError,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			for k, v := range tt.envVars {
+				t.Setenv(k, v)
+			}
+			got, err := createPyroscopeConfig(context.Background(), slog.Default(), tt.address)
+			tt.errorAssertion(t, err)
+
+			require.Equal(t, tt.want.ProfileTypes, got.ProfileTypes)
+			require.Subset(t, got.Tags, tt.want.Tags)
+		})
+	}
+}