[management,signal,proxy] add pyroscope profiling (#7536)

proxy/cmd/proxy/cmd/root.go

@@ -14,6 +14,7 @@ import (
 	"golang.org/x/crypto/acme"
 
 	"github.com/netbirdio/netbird/shared/management/domain"
+	"github.com/netbirdio/netbird/shared/profiling"
 
 	"github.com/netbirdio/netbird/client/embed"
 	"github.com/netbirdio/netbird/proxy"
@@ -30,6 +31,8 @@ const (
 	// how many buffers each receive/TUN worker eagerly allocates. Zero
 	// (unset) keeps the platform default.
 	envMaxBatchSize = "NB_PROXY_MAX_BATCH_SIZE"
+
+	applicationName = "proxy"
 )
 
 const DefaultManagementURL = "https://api.netbird.io:443"
@@ -160,6 +163,9 @@ func runServer(cmd *cobra.Command, args []string) error {
 
 	logger.Infof("configured log level: %s", level)
 
+	stopProfiling := profiling.Start(applicationName)
+	defer stopProfiling()
+
 	var wgPool, wgBatch uint64
 	var perf embed.Performance
 	if raw := os.Getenv(envPreallocatedBuffers); raw != "" {

proxy/cmd/proxy/main.go

@@ -4,6 +4,7 @@ import (
 	"net/http"
 	// nolint:gosec
 	_ "net/http/pprof"
+	"os"
 	"runtime"
 
 	log "github.com/sirupsen/logrus"
@@ -26,9 +27,13 @@ var (
 )
 
 func main() {
-	go func() {
-		log.Println(http.ListenAndServe("localhost:6060", nil))
-	}()
+	if pprofAddr := os.Getenv("NB_PPROF_ADDR"); pprofAddr != "" {
+		log.Infof("pprof enabled, listening on: %s", pprofAddr)
+		go func() {
+			log.Println(http.ListenAndServe(pprofAddr, nil))
+		}()
+	}
+
 	cmd.SetVersionInfo(Version, Commit, BuildDate, GoVersion)
 	cmd.Execute()
 }

proxy/internal/metrics/client_metrics_test.go

@@ -0,0 +1,49 @@
+package metrics_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	sdkmetric "go.opentelemetry.io/otel/sdk/metric"
+	"go.opentelemetry.io/otel/sdk/metric/metricdata"
+
+	"github.com/netbirdio/netbird/proxy/internal/metrics"
+)
+
+func TestRegisterClientObserver(t *testing.T) {
+	reader := sdkmetric.NewManualReader()
+	provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader))
+	m, err := metrics.New(context.Background(), provider.Meter("test"))
+	require.NoError(t, err)
+
+	clients := 2
+	require.NoError(t, m.RegisterClientObserver(func() int { return clients }))
+
+	var rm metricdata.ResourceMetrics
+	require.NoError(t, reader.Collect(context.Background(), &rm))
+	assert.Equal(t, int64(2), gaugeValue(t, rm, "proxy.clients.count"), "gauge must report the current client count")
+
+	clients = 1
+	require.NoError(t, reader.Collect(context.Background(), &rm))
+	assert.Equal(t, int64(1), gaugeValue(t, rm, "proxy.clients.count"), "gauge must follow the client count on the next collection")
+}
+
+func gaugeValue(t *testing.T, rm metricdata.ResourceMetrics, name string) int64 {
+	t.Helper()
+
+	for _, sm := range rm.ScopeMetrics {
+		for _, mtr := range sm.Metrics {
+			if mtr.Name != name {
+				continue
+			}
+			gauge, ok := mtr.Data.(metricdata.Gauge[int64])
+			require.True(t, ok, "%s must be an int64 gauge", name)
+			require.Len(t, gauge.DataPoints, 1, "%s must have a single data point", name)
+			return gauge.DataPoints[0].Value
+		}
+	}
+	t.Fatalf("gauge %s not found", name)
+	return 0
+}

proxy/internal/metrics/metrics.go

@@ -196,6 +196,21 @@ func (m *Metrics) RecordAddPeerDuration(d time.Duration, err error) {
 	))
 }
 
+// RegisterClientObserver reports the number of embedded clients as a gauge.
+// clientCount runs on every collection cycle, so it must stay cheap.
+func (m *Metrics) RegisterClientObserver(clientCount func() int) error {
+	_, err := m.meter.Int64ObservableGauge(
+		"proxy.clients.count",
+		metric.WithUnit("1"),
+		metric.WithDescription("Current number of embedded NetBird clients running on the netbird proxy"),
+		metric.WithInt64Callback(func(_ context.Context, o metric.Int64Observer) error {
+			o.Observe(int64(clientCount()))
+			return nil
+		}),
+	)
+	return err
+}
+
 func (m *Metrics) initL4Metrics(meter metric.Meter) error {
 	var err error
 

proxy/server.go

@@ -362,6 +362,13 @@ func (s *Server) Start(ctx context.Context) error {
 		return err
 	}
 
+	startupOK := false
+	defer func() {
+		if !startupOK {
+			s.cleanupFailedStart()
+		}
+	}()
+
 	// Management client must be initialised BEFORE the middleware manager —
 	// initMiddlewareManager passes s.mgmtClient into the builtin FactoryContext
 	// that the limit-check / limit-record middlewares pull from. Reversed
@@ -374,7 +381,9 @@ func (s *Server) Start(ctx context.Context) error {
 	runCtx, runCancel := context.WithCancel(ctx)
 	s.runCancel = runCancel
 
-	s.initNetBirdClient()
+	if err := s.initNetBirdClient(); err != nil {
+		return err
+	}
 	// Create health checker before the mapping worker so it can track
 	// management connectivity from the first stream connection.
 	s.healthChecker = health.NewChecker(s.Logger, s.netbird)
@@ -395,18 +404,6 @@ func (s *Server) Start(ctx context.Context) error {
 		return err
 	}
 
-	startupOK := false
-	defer func() {
-		if startupOK {
-			return
-		}
-		if s.geoRaw != nil {
-			if closeErr := s.geoRaw.Close(); closeErr != nil {
-				s.Logger.Debugf("close geolocation on startup failure: %v", closeErr)
-			}
-		}
-	}()
-
 	s.auth = auth.NewMiddleware(s.Logger, s.mgmtClient, s.geo)
 	s.accessLog = accesslog.NewLogger(s.mgmtClient, s.Logger, s.TrustedProxies)
 
@@ -475,14 +472,7 @@ func (s *Server) Stop(ctx context.Context) error {
 		go func() {
 			defer close(done)
 			s.gracefulShutdown()
-			if s.runCancel != nil {
-				s.runCancel()
-			}
-			if s.mgmtConn != nil {
-				if err := s.mgmtConn.Close(); err != nil {
-					s.Logger.Debugf("management connection close: %v", err)
-				}
-			}
+			s.releaseRunResources()
 		}()
 
 		select {
@@ -497,6 +487,27 @@ func (s *Server) Stop(ctx context.Context) error {
 	return s.runErr
 }
 
+// cleanupFailedStart releases what a failed Start already brought up. It
+// skips the drain and pre-stop delay because nothing has served yet, and
+// consumes stopOnce so a later Stop stays a no-op.
+func (s *Server) cleanupFailedStart() {
+	s.stopOnce.Do(func() {
+		s.shutdownServices()
+		s.releaseRunResources()
+	})
+}
+
+func (s *Server) releaseRunResources() {
+	if s.runCancel != nil {
+		s.runCancel()
+	}
+	if s.mgmtConn != nil {
+		if err := s.mgmtConn.Close(); err != nil {
+			s.Logger.Debugf("management connection close: %v", err)
+		}
+	}
+}
+
 // waitAndStop blocks until ctx is cancelled or a background goroutine
 // reports a fatal error, then drains and stops. Used by ListenAndServe.
 func (s *Server) waitAndStop(ctx context.Context) error {
@@ -568,7 +579,7 @@ func (s *Server) initManagementClient() error {
 // initNetBirdClient builds the multi-tenant embedded NetBird client used
 // for outbound RoundTripping and (when --private is on) per-account
 // inbound listeners.
-func (s *Server) initNetBirdClient() {
+func (s *Server) initNetBirdClient() error {
 	s.netbird = roundtrip.NewNetBird(s.ctx, s.ID, s.ProxyURL, roundtrip.ClientConfig{
 		MgmtAddr:     s.ManagementAddress,
 		WGPort:       s.WireguardPort,
@@ -581,6 +592,10 @@ func (s *Server) initNetBirdClient() {
 		BlockInbound: !s.Private,
 	}, s.Logger, s, s.mgmtClient)
 	s.netbird.OnAddPeer = s.meter.RecordAddPeerDuration
+	if err := s.meter.RegisterClientObserver(s.netbird.ClientCount); err != nil {
+		return fmt.Errorf("register client metrics: %w", err)
+	}
+	return nil
 }
 
 // initReverseProxy builds the meter-instrumented reverse proxy. MultiTransport

proxy/server_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/otel/metric/noop"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/connectivity"
 
 	"github.com/netbirdio/netbird/proxy/internal/auth"
 	proxymetrics "github.com/netbirdio/netbird/proxy/internal/metrics"
@@ -106,6 +107,25 @@ func TestStartFailsWithoutManagement(t *testing.T) {
 	assert.Contains(t, err.Error(), "already started", "error must explain why the call was rejected")
 }
 
+func TestStartFailureReleasesManagementConnection(t *testing.T) {
+	srv := New(t.Context(), Config{
+		Logger:               quietLifecycleLogger(),
+		ListenAddr:           "127.0.0.1:0",
+		ManagementAddress:    "https://127.0.0.1:1",
+		CertificateDirectory: t.TempDir(),
+		CertificateFile:      "missing.crt",
+		CertificateKeyFile:   "missing.key",
+	})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	err := srv.Start(ctx)
+	require.Error(t, err, "Start must fail on the missing certificate")
+	require.NotNil(t, srv.mgmtConn, "the management connection is created before the certificate step")
+	assert.Equal(t, connectivity.Shutdown, srv.mgmtConn.GetState(), "a failed Start must close the management connection it opened")
+}
+
 func TestStopIsIdempotent(t *testing.T) {
 	srv := &Server{
 		Logger:    quietLifecycleLogger(),