Put pprof endpoint behind auth (#3793)

cmd/armada/main.go

@@ -20,7 +20,6 @@ import (
 	"github.com/armadaproject/armada/internal/common/health"
 	"github.com/armadaproject/armada/internal/common/logging"
 	"github.com/armadaproject/armada/internal/common/profiling"
-	"github.com/armadaproject/armada/internal/common/serve"
 	"github.com/armadaproject/armada/pkg/api"
 )
 
@@ -64,11 +63,9 @@ func main() {
 	})
 
 	// Expose profiling endpoints if enabled.
-	if config.PprofPort != nil {
-		pprofServer := profiling.SetupPprofHttpServer(*config.PprofPort)
-		g.Go(func() error {
-			return serve.ListenAndServe(ctx, pprofServer)
-		})
+	err := profiling.SetupPprof(config.Profiling, ctx, g)
+	if err != nil {
+		log.Fatalf("Pprof setup failed, exiting, %v", err)
 	}
 
 	// TODO This starts a separate HTTP server. Is that intended? Should we have a single mux for everything?

cmd/binoculars/main.go

@@ -19,9 +19,7 @@ import (
 	"github.com/armadaproject/armada/internal/common/armadacontext"
 	gateway "github.com/armadaproject/armada/internal/common/grpc"
 	"github.com/armadaproject/armada/internal/common/health"
-	"github.com/armadaproject/armada/internal/common/logging"
 	"github.com/armadaproject/armada/internal/common/profiling"
-	"github.com/armadaproject/armada/internal/common/serve"
 	api "github.com/armadaproject/armada/pkg/api/binoculars"
 )
 
@@ -47,14 +45,9 @@ func main() {
 	log.Info("Starting...")
 
 	// Expose profiling endpoints if enabled.
-	if config.PprofPort != nil {
-		pprofServer := profiling.SetupPprofHttpServer(*config.PprofPort)
-		go func() {
-			ctx := armadacontext.Background()
-			if err := serve.ListenAndServe(ctx, pprofServer); err != nil {
-				logging.WithStacktrace(ctx, err).Error("pprof server failure")
-			}
-		}()
+	err := profiling.SetupPprof(config.Profiling, armadacontext.Background(), nil)
+	if err != nil {
+		log.Fatalf("Pprof setup failed, exiting, %v", err)
 	}
 
 	stopSignal := make(chan os.Signal, 1)

cmd/executor/main.go

@@ -7,16 +7,14 @@ import (
 	"syscall"
 
 	"github.com/prometheus/client_golang/prometheus"
-	"github.com/sirupsen/logrus"
+	log "github.com/sirupsen/logrus"
 	"github.com/spf13/pflag"
 	"github.com/spf13/viper"
 
 	"github.com/armadaproject/armada/internal/common"
 	"github.com/armadaproject/armada/internal/common/armadacontext"
 	"github.com/armadaproject/armada/internal/common/health"
-	"github.com/armadaproject/armada/internal/common/logging"
 	"github.com/armadaproject/armada/internal/common/profiling"
-	"github.com/armadaproject/armada/internal/common/serve"
 	"github.com/armadaproject/armada/internal/executor"
 	"github.com/armadaproject/armada/internal/executor/configuration"
 )
@@ -41,14 +39,9 @@ func main() {
 	common.LoadConfig(&config, "./config/executor", userSpecifiedConfigs)
 
 	// Expose profiling endpoints if enabled.
-	if config.PprofPort != nil {
-		pprofServer := profiling.SetupPprofHttpServer(*config.PprofPort)
-		go func() {
-			ctx := armadacontext.Background()
-			if err := serve.ListenAndServe(ctx, pprofServer); err != nil {
-				logging.WithStacktrace(ctx, err).Error("pprof server failure")
-			}
-		}()
+	err := profiling.SetupPprof(config.Profiling, armadacontext.Background(), nil)
+	if err != nil {
+		log.Fatalf("Pprof setup failed, exiting, %v", err)
 	}
 
 	mux := http.NewServeMux()
@@ -68,7 +61,7 @@ func main() {
 	)
 	defer shutdownMetricServer()
 
-	shutdown, wg := executor.StartUp(armadacontext.Background(), logrus.NewEntry(logrus.StandardLogger()), config)
+	shutdown, wg := executor.StartUp(armadacontext.Background(), log.NewEntry(log.StandardLogger()), config)
 	go func() {
 		<-shutdownChannel
 		shutdown()

cmd/fakeexecutor/main.go

@@ -5,14 +5,13 @@ import (
 	"os/signal"
 	"syscall"
 
+	log "github.com/sirupsen/logrus"
 	"github.com/spf13/pflag"
 	"github.com/spf13/viper"
 
 	"github.com/armadaproject/armada/internal/common"
 	"github.com/armadaproject/armada/internal/common/armadacontext"
-	"github.com/armadaproject/armada/internal/common/logging"
 	"github.com/armadaproject/armada/internal/common/profiling"
-	"github.com/armadaproject/armada/internal/common/serve"
 	"github.com/armadaproject/armada/internal/executor/configuration"
 	"github.com/armadaproject/armada/internal/executor/fake"
 	"github.com/armadaproject/armada/internal/executor/fake/context"
@@ -38,14 +37,9 @@ func main() {
 	v := common.LoadConfig(&config, "./config/executor", userSpecifiedConfigs)
 
 	// Expose profiling endpoints if enabled.
-	if config.PprofPort != nil {
-		pprofServer := profiling.SetupPprofHttpServer(*config.PprofPort)
-		go func() {
-			ctx := armadacontext.Background()
-			if err := serve.ListenAndServe(ctx, pprofServer); err != nil {
-				logging.WithStacktrace(ctx, err).Error("pprof server failure")
-			}
-		}()
+	err := profiling.SetupPprof(config.Profiling, armadacontext.Background(), nil)
+	if err != nil {
+		log.Fatalf("Pprof setup failed, exiting, %v", err)
 	}
 
 	var nodes []*context.NodeSpec

cmd/lookoutv2/main.go

@@ -14,9 +14,7 @@ import (
 	"github.com/armadaproject/armada/internal/common"
 	"github.com/armadaproject/armada/internal/common/armadacontext"
 	"github.com/armadaproject/armada/internal/common/database"
-	"github.com/armadaproject/armada/internal/common/logging"
 	"github.com/armadaproject/armada/internal/common/profiling"
-	"github.com/armadaproject/armada/internal/common/serve"
 	"github.com/armadaproject/armada/internal/lookoutv2"
 	"github.com/armadaproject/armada/internal/lookoutv2/configuration"
 	"github.com/armadaproject/armada/internal/lookoutv2/gen/restapi"
@@ -127,14 +125,9 @@ func main() {
 	common.LoadConfig(&config, "./config/lookoutv2", userSpecifiedConfigs)
 
 	// Expose profiling endpoints if enabled.
-	if config.PprofPort != nil {
-		pprofServer := profiling.SetupPprofHttpServer(*config.PprofPort)
-		go func() {
-			ctx := armadacontext.Background()
-			if err := serve.ListenAndServe(ctx, pprofServer); err != nil {
-				logging.WithStacktrace(ctx, err).Error("pprof server failure")
-			}
-		}()
+	err := profiling.SetupPprof(config.Profiling, armadacontext.Background(), nil)
+	if err != nil {
+		log.Fatalf("Pprof setup failed, exiting, %v", err)
 	}
 
 	log.SetLevel(log.DebugLevel)

deployment/armada/templates/deployment.yaml

@@ -69,8 +69,8 @@ spec:
             - containerPort: {{ .Values.applicationConfig.httpPort }}
               protocol: TCP
               name: rest
-            {{- if .Values.applicationConfig.pprofPort }}
-            - containerPort: {{ .Values.applicationConfig.pprofPort }}
+            {{- if and .Values.applicationConfig.profiling .Values.applicationConfig.profiling.port }}
+            - containerPort: {{ .Values.applicationConfig.profiling.port }}
               protocol: TCP
               name: pprof
             {{- end }}

deployment/binoculars/templates/deployment.yaml

@@ -54,8 +54,8 @@ spec:
             - containerPort: {{ .Values.applicationConfig.httpPort }}
               protocol: TCP
               name: web
-            {{- if .Values.applicationConfig.pprofPort }}
-            - containerPort: {{ .Values.applicationConfig.pprofPort }}
+            {{- if and .Values.applicationConfig.profiling .Values.applicationConfig.profiling.port }}
+            - containerPort: {{ .Values.applicationConfig.profiling.port }}
               protocol: TCP
               name: pprof
             {{- end }}

deployment/event-ingester/templates/deployment.yaml

@@ -45,8 +45,8 @@ spec:
           resources:
             {{- toYaml .Values.resources | nindent 12 }}
           ports:
-            {{- if .Values.applicationConfig.pprofPort }}
-            - containerPort: {{ .Values.applicationConfig.pprofPort }}
+            {{- if and .Values.applicationConfig.profiling .Values.applicationConfig.profiling.port }}
+            - containerPort: {{ .Values.applicationConfig.profiling.port }}
               protocol: TCP
               name: pprof
             {{- end }}

deployment/executor/templates/deployment.yaml

@@ -52,8 +52,8 @@ spec:
             - containerPort: 9001
               protocol: TCP
               name: metrics
-            {{- if .Values.applicationConfig.pprofPort }}
-            - containerPort: {{ .Values.applicationConfig.pprofPort }}
+            {{- if and .Values.applicationConfig.profiling .Values.applicationConfig.profiling.port }}
+            - containerPort: {{ .Values.applicationConfig.profiling.port }}
               protocol: TCP
               name: pprof
             {{- end }}

deployment/lookout-ingester-v2/templates/deployment.yaml

@@ -45,8 +45,8 @@ spec:
           resources:
             {{- toYaml .Values.resources | nindent 12 }}
           ports:
-            {{- if .Values.applicationConfig.pprofPort }}
-            - containerPort: {{ .Values.applicationConfig.pprofPort }}
+            {{- if and .Values.applicationConfig.profiling .Values.applicationConfig.profiling.port }}
+            - containerPort: {{ .Values.applicationConfig.profiling.port }}
               protocol: TCP
               name: pprof
             {{- end }}

deployment/lookout-v2/templates/deployment.yaml

@@ -48,8 +48,8 @@ spec:
             - containerPort: {{ .Values.applicationConfig.apiPort }}
               protocol: TCP
               name: web
-            {{- if .Values.applicationConfig.pprofPort }}
-            - containerPort: {{ .Values.applicationConfig.pprofPort }}
+            {{- if and .Values.applicationConfig.profiling .Values.applicationConfig.profiling.port }}
+            - containerPort: {{ .Values.applicationConfig.profiling.port }}
               protocol: TCP
               name: pprof
             {{- end }}

deployment/scheduler/templates/scheduler-ingester-deployment.yaml

@@ -45,8 +45,8 @@ spec:
           resources:
             {{- toYaml .Values.ingester.resources | nindent 12 }}
           ports:
-            {{- if .Values.ingester.applicationConfig.pprofPort }}
-            - containerPort: {{ .Values.ingester.applicationConfig.pprofPort }}
+            {{- if and .Values.ingester.applicationConfig.profiling .Values.ingester.applicationConfig.profiling.port }}
+            - containerPort: {{ .Values.ingester.applicationConfig.profiling.port }}
               protocol: TCP
               name: pprof
             {{- end }}

deployment/scheduler/templates/scheduler-statefulset.yaml

@@ -76,8 +76,8 @@ spec:
             - containerPort: {{ .Values.scheduler.applicationConfig.metrics.port }}
               protocol: TCP
               name: metrics
-            {{- if .Values.scheduler.applicationConfig.pprofPort }}
-            - containerPort: {{ .Values.scheduler.applicationConfig.pprofPort }}
+            {{- if and .Values.scheduler.applicationConfig.profiling .Values.scheduler.applicationConfig.profiling.port }}
+            - containerPort: {{ .Values.scheduler.applicationConfig.profiling.port }}
               protocol: TCP
               name: pprof
             {{- end }}

docs/developer/pprof.md

@@ -0,0 +1,14 @@
+# Use of pprof
+
+- Go provides a profiling tool called pprof. It's documented at https://pkg.go.dev/net/http/pprof.
+- If you wish to use this with Armada, enable the profiling socket with the following config (this should be under `applicationConfig` if using the helm charts). This config will listen on the specified port with no auth.
+  ```
+  profiling:
+    port: 6060
+    auth:
+      anonymousAuth: true
+      permissionGroupMapping:
+        pprof: ["everyone"]
+  ```
+- It's possible to put pprof behind auth if you want, see [api.md#authentication](./api.md#authentication) and [oidc.md](./oidc.md).
+- The helm charts do not currently expose the profiling port via a service and ingress. You can use `kubectl port-forward` to access them.

internal/armada/configuration/types.go

@@ -9,6 +9,7 @@ import (
 
 	authconfig "github.com/armadaproject/armada/internal/common/auth/configuration"
 	grpcconfig "github.com/armadaproject/armada/internal/common/grpc/configuration"
+	profilingconfig "github.com/armadaproject/armada/internal/common/profiling/configuration"
 	armadaresource "github.com/armadaproject/armada/internal/common/resource"
 	"github.com/armadaproject/armada/pkg/client"
 )
@@ -19,8 +20,7 @@ type ArmadaConfig struct {
 	GrpcPort    uint16
 	HttpPort    uint16
 	MetricsPort uint16
-	// If non-nil, net/http/pprof endpoints are exposed on localhost on this port.
-	PprofPort *uint16
+	Profiling   *profilingconfig.ProfilingConfig
 
 	CorsAllowedOrigins []string
 	GrpcGatewayPath    string

internal/armada/event/event_test.go

@@ -326,7 +326,7 @@ func TestEventServer_GetJobSetEvents_Permissions(t *testing.T) {
 				err := s.queueRepository.(armadaqueue.QueueRepository).CreateQueue(ctx, q)
 				assert.NoError(t, err)
 
-				principal := auth.NewStaticPrincipal("alice", []string{})
+				principal := auth.NewStaticPrincipal("alice", "test", []string{})
 				ctx := auth.WithPrincipal(armadacontext.Background(), principal)
 				stream := &eventStreamMock{ctx: ctx}
 
@@ -351,7 +351,7 @@ func TestEventServer_GetJobSetEvents_Permissions(t *testing.T) {
 				err := s.queueRepository.(armadaqueue.QueueRepository).CreateQueue(ctx, q)
 				assert.NoError(t, err)
 
-				principal := auth.NewStaticPrincipal("alice", []string{"watch-all-events-group"})
+				principal := auth.NewStaticPrincipal("alice", "test", []string{"watch-all-events-group"})
 				ctx := auth.WithPrincipal(armadacontext.Background(), principal)
 				stream := &eventStreamMock{ctx: ctx}
 
@@ -373,7 +373,7 @@ func TestEventServer_GetJobSetEvents_Permissions(t *testing.T) {
 			err := s.queueRepository.(armadaqueue.QueueRepository).CreateQueue(ctx, q)
 			assert.NoError(t, err)
 
-			principal := auth.NewStaticPrincipal("alice", []string{"watch-events-group", "watch-queue-group"})
+			principal := auth.NewStaticPrincipal("alice", "test", []string{"watch-events-group", "watch-queue-group"})
 			ctx := auth.WithPrincipal(armadacontext.Background(), principal)
 			stream := &eventStreamMock{ctx: ctx}
 

internal/armada/submit/testfixtures/test_fixtures.go

@@ -22,7 +22,7 @@ var (
 	DefaultOwner         = "testUser"
 	DefaultJobset        = "testJobset"
 	DefaultQueue         = queue.Queue{Name: "testQueue"}
-	DefaultPrincipal     = auth.NewStaticPrincipal(DefaultOwner, []string{"groupA"})
+	DefaultPrincipal     = auth.NewStaticPrincipal(DefaultOwner, "test", []string{"groupA"})
 	DefaultContainerPort = v1.ContainerPort{
 		Name:          "testContainerPort",
 		ContainerPort: 8080,

internal/binoculars/configuration/types.go

@@ -3,6 +3,7 @@ package configuration
 import (
 	"github.com/armadaproject/armada/internal/common/auth/configuration"
 	grpcconfig "github.com/armadaproject/armada/internal/common/grpc/configuration"
+	profilingconfig "github.com/armadaproject/armada/internal/common/profiling/configuration"
 )
 
 type BinocularsConfig struct {
@@ -12,8 +13,7 @@ type BinocularsConfig struct {
 	GrpcPort    uint16
 	HttpPort    uint16
 	MetricsPort uint16
-	// If non-nil, net/http/pprof endpoints are exposed on localhost on this port.
-	PprofPort *uint16
+	Profiling   *profilingconfig.ProfilingConfig
 
 	CorsAllowedOrigins []string
 

internal/binoculars/service/cordon_test.go

@@ -38,7 +38,7 @@ var (
 )
 
 func TestCordonNode(t *testing.T) {
-	principal := auth.NewStaticPrincipal("principle", []string{})
+	principal := auth.NewStaticPrincipal("principle", "test", []string{})
 	tests := map[string]struct {
 		additionalLabels map[string]string
 		expectedLabels   map[string]string

internal/common/auth/anonymous.go

@@ -2,12 +2,13 @@ package auth
 
 import "context"
 
+const AnonymousAuthServiceName = "Anonymous"
+
 type AnonymousAuthService struct{}
 
-func (authService *AnonymousAuthService) Name() string {
-	return "Anonymous"
-}
+// Default principal used if no principal can be found in a context.
+var anonymousPrincipal = NewStaticPrincipal("anonymous", AnonymousAuthServiceName, []string{})
 
-func (AnonymousAuthService) Authenticate(ctx context.Context) (Principal, error) {
+func (AnonymousAuthService) Authenticate(ctx context.Context, authHeader string) (Principal, error) {
 	return anonymousPrincipal, nil
 }

internal/common/auth/authorization_test.go

@@ -71,8 +71,8 @@ func TestAuthorizer_AuthorizeQueueAction(t *testing.T) {
 		PriorityFactor: 1,
 	}
 
-	authorizedPrincipal := NewStaticPrincipal("alice", []string{"submit-job-group"})
-	unauthorizedPrincipcal := NewStaticPrincipal("alice", []string{})
+	authorizedPrincipal := NewStaticPrincipal("alice", "test", []string{"submit-job-group"})
+	unauthorizedPrincipcal := NewStaticPrincipal("alice", "test", []string{})
 
 	tests := map[string]struct {
 		ctx                     *armadacontext.Context