cmd: apply pprof flags in Before so all commands can be profiled (#3266)

Author: serprex

flow/cmd/worker.go

@@ -3,14 +3,8 @@ package cmd
 import (
 	"context"
 	"fmt"
-	"log"
 	"log/slog"
-	"net/http"
-	//nolint:gosec
-	_ "net/http/pprof"
 	"os"
-	"runtime"
-	"time"
 
 	"go.temporal.io/sdk/client"
 	temporalotel "go.temporal.io/sdk/contrib/opentelemetry"
@@ -31,10 +25,8 @@ type WorkerSetupOptions struct {
 	TemporalNamespace                  string
 	TemporalMaxConcurrentActivities    int
 	TemporalMaxConcurrentWorkflowTasks int
-	EnableProfiling                    bool
 	EnableOtelMetrics                  bool
 	UseMaintenanceTaskQueue            bool
-	PprofPort                          int // Port for pprof HTTP server
 }
 
 type WorkerSetupResponse struct {
@@ -51,35 +43,7 @@ func (w *WorkerSetupResponse) Close(ctx context.Context) {
 	}
 }
 
-func setupPprof(opts *WorkerSetupOptions) {
-	// Set default pprof port if not specified
-	pprofPort := opts.PprofPort
-
-	// Enable mutex and block profiling
-	runtime.SetMutexProfileFraction(5)
-	runtime.SetBlockProfileRate(5)
-
-	// Start HTTP server with pprof endpoints
-	go func() {
-		pprofAddr := fmt.Sprintf(":%d", pprofPort)
-		slog.Info("Starting pprof HTTP server on " + pprofAddr)
-		server := &http.Server{
-			Addr:         pprofAddr,
-			ReadTimeout:  1 * time.Minute,
-			WriteTimeout: 11 * time.Minute,
-		}
-
-		if err := server.ListenAndServe(); err != nil {
-			log.Fatalf("Failed to start pprof HTTP server: %v", err)
-		}
-	}()
-}
-
 func WorkerSetup(ctx context.Context, opts *WorkerSetupOptions) (*WorkerSetupResponse, error) {
-	if opts.EnableProfiling {
-		setupPprof(opts)
-	}
-
 	conn, err := internal.GetCatalogConnectionPoolFromEnv(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("unable to create catalog connection pool: %w", err)
@@ -115,11 +79,10 @@ func WorkerSetup(ctx context.Context, opts *WorkerSetupOptions) (*WorkerSetupRes
 	}
 	taskQueue := internal.PeerFlowTaskQueueName(queueId)
 	slog.Info(
-		fmt.Sprintf("Creating temporal worker for queue %v: %v workflow workers %v activity workers",
-			taskQueue,
-			opts.TemporalMaxConcurrentWorkflowTasks,
-			opts.TemporalMaxConcurrentActivities,
-		),
+		"Creating temporal worker",
+		slog.String("taskQueue", taskQueue),
+		slog.Int("workflowConcurrency", opts.TemporalMaxConcurrentWorkflowTasks),
+		slog.Int("activityConcurrency", opts.TemporalMaxConcurrentActivities),
 	)
 	w := worker.New(c, taskQueue, worker.Options{
 		EnableSessionWorker:                    true,

flow/main.go

@@ -2,12 +2,17 @@ package main
 
 import (
 	"context"
+	"fmt"
 	"log"
 	"log/slog"
+	"net/http"
+	//nolint:gosec
+	_ "net/http/pprof"
 	"os"
 	"os/signal"
 	"runtime"
 	"syscall"
+	"time"
 
 	"github.com/urfave/cli/v3"
 	"go.temporal.io/sdk/worker"
@@ -42,7 +47,7 @@ func main() {
 		Sources: cli.EnvVars("ENABLE_OTEL_METRICS"),
 	}
 
-	pprofPortFlag := &cli.IntFlag{
+	pprofPortFlag := &cli.Uint16Flag{
 		Name:    "pprof-port",
 		Value:   6060,
 		Usage:   "Port for pprof HTTP server",
@@ -133,20 +138,45 @@ func main() {
 
 	app := &cli.Command{
 		Name: "PeerDB Flows CLI",
+		Before: func(ctx context.Context, clicmd *cli.Command) (context.Context, error) {
+			if clicmd.Bool(profilingFlag.Name) {
+				// Enable mutex and block profiling
+				runtime.SetMutexProfileFraction(5)
+				runtime.SetBlockProfileRate(5)
+				pprofPort := clicmd.Uint16(pprofPortFlag.Name)
+				pprofAddr := fmt.Sprintf(":%d", pprofPort)
+
+				// Start HTTP server with pprof endpoints
+				go func() {
+					slog.Info("Starting pprof HTTP server", slog.String("address", pprofAddr))
+
+					server := &http.Server{
+						Addr:         pprofAddr,
+						ReadTimeout:  1 * time.Minute,
+						WriteTimeout: 11 * time.Minute,
+					}
+					if err := server.ListenAndServe(); err != nil {
+						log.Fatalf("Failed to start pprof HTTP server: %v", err)
+					}
+				}()
+			}
+			return nil, nil
+		},
+		Flags: []cli.Flag{
+			profilingFlag,
+			pprofPortFlag,
+		},
 		Commands: []*cli.Command{
 			{
 				Name: "worker",
 				Action: func(ctx context.Context, clicmd *cli.Command) error {
-					temporalHostPort := clicmd.String("temporal-host-port")
 					res, err := cmd.WorkerSetup(ctx, &cmd.WorkerSetupOptions{
-						TemporalHostPort:                   temporalHostPort,
-						EnableProfiling:                    clicmd.Bool("enable-profiling"),
-						EnableOtelMetrics:                  clicmd.Bool("enable-otel-metrics"),
-						TemporalNamespace:                  clicmd.String("temporal-namespace"),
-						TemporalMaxConcurrentActivities:    clicmd.Int("temporal-max-concurrent-activities"),
-						TemporalMaxConcurrentWorkflowTasks: clicmd.Int("temporal-max-concurrent-workflow-tasks"),
+						TemporalHostPort:                   clicmd.String(temporalHostPortFlag.Name),
+						TemporalNamespace:                  clicmd.String(temporalNamespaceFlag.Name),
+						TemporalMaxConcurrentActivities:    clicmd.Int(temporalMaxConcurrentActivitiesFlag.Name),
+						TemporalMaxConcurrentWorkflowTasks: clicmd.Int(temporalMaxConcurrentWorkflowTasksFlag.Name),
 						UseMaintenanceTaskQueue:            clicmd.Bool(useMaintenanceTaskQueueFlag.Name),
-						PprofPort:                          clicmd.Int(pprofPortFlag.Name),
+						EnableOtelMetrics:                  clicmd.Bool(otelMetricsFlag.Name),
 					})
 					if err != nil {
 						return err
@@ -156,23 +186,20 @@ func main() {
 				},
 				Flags: []cli.Flag{
 					temporalHostPortFlag,
-					profilingFlag,
-					otelMetricsFlag,
-					pprofPortFlag,
 					temporalNamespaceFlag,
 					temporalMaxConcurrentActivitiesFlag,
 					temporalMaxConcurrentWorkflowTasksFlag,
 					useMaintenanceTaskQueueFlag,
+					otelMetricsFlag,
 				},
 			},
 			{
 				Name: "snapshot-worker",
 				Action: func(ctx context.Context, clicmd *cli.Command) error {
-					temporalHostPort := clicmd.String("temporal-host-port")
 					res, err := cmd.SnapshotWorkerMain(ctx, &cmd.SnapshotWorkerOptions{
-						EnableOtelMetrics: clicmd.Bool("enable-otel-metrics"),
-						TemporalHostPort:  temporalHostPort,
-						TemporalNamespace: clicmd.String("temporal-namespace"),
+						TemporalHostPort:  clicmd.String(temporalHostPortFlag.Name),
+						TemporalNamespace: clicmd.String(temporalNamespaceFlag.Name),
+						EnableOtelMetrics: clicmd.Bool(otelMetricsFlag.Name),
 					})
 					if err != nil {
 						return err
@@ -181,21 +208,21 @@ func main() {
 					return res.Worker.Run(worker.InterruptCh())
 				},
 				Flags: []cli.Flag{
-					otelMetricsFlag,
 					temporalHostPortFlag,
 					temporalNamespaceFlag,
+					otelMetricsFlag,
 				},
 			},
 			{
 				Name: "api",
 				Flags: []cli.Flag{
-					&cli.UintFlag{
+					&cli.Uint16Flag{
 						Name:    "port",
 						Aliases: []string{"p"},
 						Value:   8110,
 					},
 					// gateway port is the port that the grpc-gateway listens on
-					&cli.UintFlag{
+					&cli.Uint16Flag{
 						Name:  "gateway-port",
 						Value: 8111,
 					},
@@ -204,13 +231,11 @@ func main() {
 					otelMetricsFlag,
 				},
 				Action: func(ctx context.Context, clicmd *cli.Command) error {
-					temporalHostPort := clicmd.String("temporal-host-port")
-
 					return cmd.APIMain(ctx, &cmd.APIServerParams{
-						Port:              uint16(clicmd.Uint("port")),
-						TemporalHostPort:  temporalHostPort,
-						GatewayPort:       uint16(clicmd.Uint("gateway-port")),
-						TemporalNamespace: clicmd.String("temporal-namespace"),
+						Port:              clicmd.Uint16("port"),
+						GatewayPort:       clicmd.Uint16("gateway-port"),
+						TemporalHostPort:  clicmd.String(temporalHostPortFlag.Name),
+						TemporalNamespace: clicmd.String(temporalNamespaceFlag.Name),
 						EnableOtelMetrics: clicmd.Bool(otelMetricsFlag.Name),
 					})
 				},