diff --git a/combined/cmd/root.go b/combined/cmd/root.go index 917312e57..26d6ceedb 100644 --- a/combined/cmd/root.go +++ b/combined/cmd/root.go @@ -120,7 +120,7 @@ func execute(cmd *cobra.Command, _ []string) error { ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() - err = shutdownServers(ctx, servers.relaySrv, servers.healthcheck, servers.stunServer, servers.mgmtSrv, servers.metricsServer) + err = shutdownServers(ctx, servers.relaySrv, servers.healthcheck, servers.stunServer, servers.mgmtSrv, servers.signalSrv, servers.metricsServer) wg.Wait() return err } @@ -399,7 +399,7 @@ func startServers(wg *sync.WaitGroup, srv *relayServer.Server, httpHealthcheck * } } -func shutdownServers(ctx context.Context, srv *relayServer.Server, httpHealthcheck *healthcheck.Server, stunServer *stun.Server, mgmtSrv mgmtServer.Server, metricsServer *sharedMetrics.Metrics) error { +func shutdownServers(ctx context.Context, srv *relayServer.Server, httpHealthcheck *healthcheck.Server, stunServer *stun.Server, mgmtSrv mgmtServer.Server, signalSrv *signalServer.Server, metricsServer *sharedMetrics.Metrics) error { var errs error if err := httpHealthcheck.Shutdown(ctx); err != nil { @@ -425,6 +425,10 @@ func shutdownServers(ctx context.Context, srv *relayServer.Server, httpHealthche } } + if signalSrv != nil { + signalSrv.Stop() + } + if metricsServer != nil { log.Infof("shutting down metrics server") if err := metricsServer.Shutdown(ctx); err != nil { diff --git a/go.mod b/go.mod index 8e8b7b1d4..543dcc713 100644 --- a/go.mod +++ b/go.mod @@ -40,6 +40,7 @@ require ( github.com/aws/aws-sdk-go-v2/credentials v1.18.10 github.com/aws/aws-sdk-go-v2/service/s3 v1.87.3 github.com/c-robinson/iplib v1.0.3 + github.com/caarlos0/env/v11 v11.4.1 github.com/caddyserver/certmagic v0.21.3 github.com/cilium/ebpf v0.19.0 github.com/coder/websocket v1.8.14 @@ -68,6 +69,7 @@ require ( github.com/google/gopacket v1.1.19 github.com/google/nftables v0.3.0 github.com/gopacket/gopacket v1.4.0 + github.com/grafana/pyroscope-go v1.4.2 github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.0.2-0.20240212192251-757544f21357 github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.3 github.com/hashicorp/go-multierror v1.1.1 @@ -236,6 +238,7 @@ require ( github.com/googleapis/gax-go/v2 v2.21.0 // indirect github.com/goreleaser/chglog v0.7.4 // indirect github.com/gorilla/handlers v1.5.2 // indirect + github.com/grafana/pyroscope-go/godeltaprof v0.1.11 // indirect github.com/hashicorp/errwrap v1.1.0 // indirect github.com/hashicorp/go-cleanhttp v0.5.2 // indirect github.com/hashicorp/go-retryablehttp v0.7.8 // indirect @@ -259,7 +262,7 @@ require ( github.com/josharian/intern v1.0.0 // indirect github.com/kelseyhightower/envconfig v1.4.0 // indirect github.com/kevinburke/ssh_config v1.4.0 // indirect - github.com/klauspost/compress v1.18.3 // indirect + github.com/klauspost/compress v1.18.7 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect github.com/koron/go-ssdp v0.0.4 // indirect github.com/kr/fs v0.1.0 // indirect diff --git a/go.sum b/go.sum index 75a0f1c42..6e5fd0693 100644 --- a/go.sum +++ b/go.sum @@ -106,6 +106,8 @@ github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= github.com/c-robinson/iplib v1.0.3 h1:NG0UF0GoEsrC1/vyfX1Lx2Ss7CySWl3KqqXh3q4DdPU= github.com/c-robinson/iplib v1.0.3/go.mod h1:i3LuuFL1hRT5gFpBRnEydzw8R6yhGkF4szNDIbF8pgo= +github.com/caarlos0/env/v11 v11.4.1 h1:fYwH0sWEsBSMPG7t4e/PEfTFzrWrpjyygXyUnWiSwEw= +github.com/caarlos0/env/v11 v11.4.1/go.mod h1:qupehSf/Y0TUTsxKywqRt/vJjN5nz6vauiYEUUr8P4U= github.com/caddyserver/certmagic v0.21.3 h1:pqRRry3yuB4CWBVq9+cUqu+Y6E2z8TswbhNx1AZeYm0= github.com/caddyserver/certmagic v0.21.3/go.mod h1:Zq6pklO9nVRl3DIFUw9gVUfXKdpc/0qwTUAQMBlfgtI= github.com/caddyserver/zerossl v0.1.3 h1:onS+pxp3M8HnHpN5MMbOMyNjmTheJyWRaZYwn+YTAyA= @@ -327,6 +329,10 @@ github.com/gorilla/handlers v1.5.2 h1:cLTUSsNkgcwhgRqvCNmdbRWG0A3N4F+M2nWKdScwyE github.com/gorilla/handlers v1.5.2/go.mod h1:dX+xVpaxdSw+q0Qek8SSsl3dfMk3jNddUkMzo0GtH0w= github.com/gorilla/mux v1.8.1 h1:TuBL49tXwgrFYWhqrNgrUNEY92u81SPhu7sTdzQEiWY= github.com/gorilla/mux v1.8.1/go.mod h1:AKf9I4AEqPTmMytcMc0KkNouC66V3BtZ4qD5fmWSiMQ= +github.com/grafana/pyroscope-go v1.4.2 h1:0LW5HrUJXgGr9zF5gITP/HaFXN9/LsMiwlgVJAK75l0= +github.com/grafana/pyroscope-go v1.4.2/go.mod h1:Ej13Jr05rRJrjWvrrFhfh6gGYXtfibuukOs3Tl3Y7QQ= +github.com/grafana/pyroscope-go/godeltaprof v0.1.11 h1:el5LYpXissAiCKZ5/6yjlr6mhYVV6Cp5lahTocxraXM= +github.com/grafana/pyroscope-go/godeltaprof v0.1.11/go.mod h1:jl1V8M4cWsXciROCPIDDG7CtjSjT/ECbp6eLVuMxYRI= github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.0.2-0.20240212192251-757544f21357 h1:Fkzd8ktnpOR9h47SXHe2AYPwelXLH2GjGsjlAloiWfo= github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.0.2-0.20240212192251-757544f21357/go.mod h1:w9Y7gY31krpLmrVU5ZPG9H7l9fZuRu5/3R3S3FMtVQ4= github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.3 h1:5ZPtiqj0JL5oKWmcsq4VMaAW5ukBEgSGXEN89zeH1Jo= @@ -413,8 +419,8 @@ github.com/kevinburke/ssh_config v1.4.0 h1:6xxtP5bZ2E4NF5tuQulISpTO2z8XbtH8cg1PW github.com/kevinburke/ssh_config v1.4.0/go.mod h1:q2RIzfka+BXARoNexmF9gkxEX7DmvbW9P4hIVx2Kg4M= github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8= github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= -github.com/klauspost/compress v1.18.3 h1:9PJRvfbmTabkOX8moIpXPbMMbYN60bWImDDU7L+/6zw= -github.com/klauspost/compress v1.18.3/go.mod h1:R0h/fSBs8DE4ENlcrlib3PsXS61voFxhIs2DeRhCvJ4= +github.com/klauspost/compress v1.18.7 h1:aUyZsS4kH3QTKurYhAOwAHxllVPnOthb3vPfnF1Ehjw= +github.com/klauspost/compress v1.18.7/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= github.com/klauspost/cpuid/v2 v2.0.12/go.mod h1:g2LTdtYhdyuGPqyWyv7qRAmj1WBqxuObKfj5c0PQa7c= github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y= github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= diff --git a/management/internals/server/server.go b/management/internals/server/server.go index a1b58fdf1..6d51745a7 100644 --- a/management/internals/server/server.go +++ b/management/internals/server/server.go @@ -23,6 +23,8 @@ import ( "github.com/netbirdio/netbird/management/server/idp" "github.com/netbirdio/netbird/management/server/metrics" "github.com/netbirdio/netbird/management/server/store" + "github.com/netbirdio/netbird/shared/lifecycle" + "github.com/netbirdio/netbird/shared/profiling" "github.com/netbirdio/netbird/util/wsproxy" wsproxyserver "github.com/netbirdio/netbird/util/wsproxy/server" "github.com/netbirdio/netbird/version" @@ -36,6 +38,8 @@ const ( DefaultSelfHostedDomain = "netbird.selfhosted" ContainerKeyBaseServer = "baseServer" + + applicationName = "management" ) type Server interface { @@ -82,6 +86,8 @@ type BaseServer struct { errCh chan error wg sync.WaitGroup cancel context.CancelFunc + + lifecycle.StopHandlers } // Config holds the configuration parameters for creating a new server @@ -117,6 +123,9 @@ func NewServer(cfg *Config) *BaseServer { } s.container[ContainerKeyBaseServer] = s + stopProfiling := profiling.Start(applicationName) + s.OnStop(stopProfiling) + return s } @@ -126,6 +135,14 @@ func (s *BaseServer) AfterInit(fn func(s *BaseServer)) { // Start begins listening for HTTP requests on the configured address func (s *BaseServer) Start(ctx context.Context) error { + if err := s.start(ctx); err != nil { + s.RunStopHandlers() + return err + } + return nil +} + +func (s *BaseServer) start(ctx context.Context) error { srvCtx, cancel := context.WithCancel(ctx) s.cancel = cancel s.errCh = make(chan error, 4) @@ -278,6 +295,7 @@ func (s *BaseServer) setupTLS(ctx context.Context) (bool, error) { func (s *BaseServer) Stop() error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() + defer s.RunStopHandlers() if s.domainCleanupStop != nil { s.domainCleanupStop() } diff --git a/proxy/cmd/proxy/cmd/root.go b/proxy/cmd/proxy/cmd/root.go index 9b180a5c4..765d5c05a 100644 --- a/proxy/cmd/proxy/cmd/root.go +++ b/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 != "" { diff --git a/proxy/cmd/proxy/main.go b/proxy/cmd/proxy/main.go index 16e7e8ac2..6851c6cfc 100644 --- a/proxy/cmd/proxy/main.go +++ b/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() } diff --git a/proxy/internal/metrics/client_metrics_test.go b/proxy/internal/metrics/client_metrics_test.go new file mode 100644 index 000000000..c71e6fb57 --- /dev/null +++ b/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 +} diff --git a/proxy/internal/metrics/metrics.go b/proxy/internal/metrics/metrics.go index 5fd23d934..d7b1797a1 100644 --- a/proxy/internal/metrics/metrics.go +++ b/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 diff --git a/proxy/server.go b/proxy/server.go index 5b652e61c..762ead9b8 100644 --- a/proxy/server.go +++ b/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 diff --git a/proxy/server_test.go b/proxy/server_test.go index 9cef63b95..cf583985f 100644 --- a/proxy/server_test.go +++ b/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(), diff --git a/shared/lifecycle/stop_handlers.go b/shared/lifecycle/stop_handlers.go new file mode 100644 index 000000000..f6ec2688b --- /dev/null +++ b/shared/lifecycle/stop_handlers.go @@ -0,0 +1,57 @@ +package lifecycle + +import ( + "runtime/debug" + "sync" + + log "github.com/sirupsen/logrus" +) + +// StopHandlers collects functions to run once when their owner exits. Embed it +// in a server type to expose OnStop and RunStopHandlers. +type StopHandlers struct { + mu sync.Mutex + stopped bool + handlers []func() +} + +// OnStop registers fn to run once when the owner stops. Handlers run in +// reverse registration order. A handler registered after the owner has +// stopped runs immediately. +func (h *StopHandlers) OnStop(fn func()) { + h.mu.Lock() + stopped := h.stopped + if !stopped { + h.handlers = append(h.handlers, fn) + } + h.mu.Unlock() + + if stopped { + runStopHandler(fn) + } +} + +// RunStopHandlers runs every registered handler once, last registered first. +// Later calls are no-ops, so it can be wired to several exit paths at once. +func (h *StopHandlers) RunStopHandlers() { + h.mu.Lock() + handlers := h.handlers + h.handlers = nil + h.stopped = true + h.mu.Unlock() + + for i := len(handlers) - 1; i >= 0; i-- { + runStopHandler(handlers[i]) + } +} + +// runStopHandler keeps one panicking handler from skipping the ones still +// pending; on the shutdown path there is no second chance to run them. +func runStopHandler(fn func()) { + defer func() { + if r := recover(); r != nil { + log.Errorf("stop handler panicked: %v\n%s", r, debug.Stack()) + } + }() + fn() +} diff --git a/shared/lifecycle/stop_handlers_test.go b/shared/lifecycle/stop_handlers_test.go new file mode 100644 index 000000000..787f39e6d --- /dev/null +++ b/shared/lifecycle/stop_handlers_test.go @@ -0,0 +1,43 @@ +package lifecycle + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestStopHandlers_RunOnceInReverseOrder(t *testing.T) { + var h StopHandlers + var order []string + h.OnStop(func() { order = append(order, "first") }) + h.OnStop(func() { order = append(order, "second") }) + + h.RunStopHandlers() + h.RunStopHandlers() + + assert.Equal(t, []string{"second", "first"}, order, "handlers must run once, last registered first") +} + +func TestStopHandlers_PanicDoesNotSkipRemainingHandlers(t *testing.T) { + var h StopHandlers + var order []string + h.OnStop(func() { order = append(order, "first") }) + h.OnStop(func() { panic("boom") }) + h.OnStop(func() { order = append(order, "third") }) + + h.RunStopHandlers() + + assert.Equal(t, []string{"third", "first"}, order, "handlers around a panicking one must still run") +} + +func TestStopHandlers_LateRegistrationRunsImmediately(t *testing.T) { + var h StopHandlers + h.RunStopHandlers() + + runs := 0 + h.OnStop(func() { runs++ }) + assert.Equal(t, 1, runs, "a handler registered after the stop must run right away") + + h.RunStopHandlers() + assert.Equal(t, 1, runs, "later runs must stay no-ops and must not repeat the handler") +} diff --git a/shared/profiling/profiling.go b/shared/profiling/profiling.go new file mode 100644 index 000000000..1d893048a --- /dev/null +++ b/shared/profiling/profiling.go @@ -0,0 +1,127 @@ +package profiling + +import ( + "errors" + "fmt" + "net/netip" + "net/url" + "os" + "strings" + "sync/atomic" + + "github.com/caarlos0/env/v11" + "github.com/grafana/pyroscope-go" + log "github.com/sirupsen/logrus" +) + +var errNotConfigured = errors.New("pyroscope not configured") + +var started atomic.Bool + +type config struct { + Address string `env:"NB_PYROSCOPE_ADDRESS"` + User string `env:"NB_PYROSCOPE_USER,notEmpty"` + Password string `env:"NB_PYROSCOPE_PASSWORD,notEmpty"` +} + +func Start(applicationName string) func() { + noop := func() {} + + cfg, err := loadConfig() + switch { + case errors.Is(err, errNotConfigured): + log.Info("pyroscope not configured, continuous profiling disabled") + return noop + case err != nil: + log.Errorf("failed to load pyroscope config: %v", err) + return noop + } + + // pprof allows one CPU profile per process, so a second profiler (e.g. the + // signal server inside the combined binary) would only log errors. + if !started.CompareAndSwap(false, true) { + log.Warnf("continuous profiling already running in this process, not starting it for %s", applicationName) + return noop + } + + tags := map[string]string{} + if hostname, err := os.Hostname(); err == nil { + tags["instance"] = hostname + } else { + log.Warnf("failed to resolve hostname for profile tags: %v", err) + } + + profiler, err := pyroscope.Start(pyroscope.Config{ + ApplicationName: applicationName, + ServerAddress: cfg.Address, + BasicAuthUser: cfg.User, + BasicAuthPassword: cfg.Password, + Logger: log.StandardLogger(), + Tags: tags, + ProfileTypes: []pyroscope.ProfileType{ + pyroscope.ProfileCPU, + pyroscope.ProfileAllocObjects, + pyroscope.ProfileAllocSpace, + pyroscope.ProfileInuseObjects, + pyroscope.ProfileInuseSpace, + }, + }) + if err != nil { + started.Store(false) + log.Errorf("failed to start continuous profiling: %v", err) + return noop + } + + return func() { + _ = profiler.Stop() + started.Store(false) + } +} + +func loadConfig() (config, error) { + var cfg config + if err := env.Parse(&cfg); err != nil { + if cfg.Address == "" { + return cfg, errNotConfigured + } + return cfg, fmt.Errorf("failed to parse pyroscope config: %w", err) + } + + if cfg.Address == "" { + return cfg, errNotConfigured + } + if err := validateAddress(cfg.Address); err != nil { + return cfg, err + } + + return cfg, nil +} + +// validateAddress refuses to send the basic-auth credentials in plaintext to +// anything but a loopback or private endpoint. +func validateAddress(address string) error { + u, err := url.Parse(address) + if err != nil { + return fmt.Errorf("invalid pyroscope address %q: %w", address, err) + } + + switch u.Scheme { + case "https": + return nil + case "http": + if isLocalOrPrivate(u.Hostname()) { + return nil + } + return fmt.Errorf("insecure pyroscope address %q: use https for non-local endpoints", address) + default: + return fmt.Errorf("pyroscope address %q must use http or https", address) + } +} + +func isLocalOrPrivate(host string) bool { + if host == "localhost" || strings.HasSuffix(host, ".localhost") { + return true + } + ip, err := netip.ParseAddr(host) + return err == nil && (ip.IsLoopback() || ip.IsPrivate()) +} diff --git a/shared/profiling/profiling_test.go b/shared/profiling/profiling_test.go new file mode 100644 index 000000000..68e56bb4c --- /dev/null +++ b/shared/profiling/profiling_test.go @@ -0,0 +1,202 @@ +package profiling + +import ( + "os" + "testing" + + log "github.com/sirupsen/logrus" + logtest "github.com/sirupsen/logrus/hooks/test" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestStartSkipsSecondProfilerInProcess(t *testing.T) { + clearEnv(t) + t.Setenv("NB_PYROSCOPE_ADDRESS", "http://127.0.0.1:1") + t.Setenv("NB_PYROSCOPE_USER", "user") + t.Setenv("NB_PYROSCOPE_PASSWORD", "token") + + started.Store(true) + t.Cleanup(func() { started.Store(false) }) + hook := logtest.NewGlobal() + t.Cleanup(hook.Reset) + + stop := Start("netbird-second") + stop() + + assert.True(t, started.Load(), "the running profiler must stay marked as started") + entry := hook.LastEntry() + require.NotNil(t, entry, "the skipped start must be logged") + assert.Equal(t, log.WarnLevel, entry.Level) + assert.Contains(t, entry.Message, "already running") +} + +func TestLoadConfig(t *testing.T) { + tests := []struct { + name string + env map[string]string + expected config + errIs error + wantErr bool + }{ + { + name: "address unset disables profiling", + errIs: errNotConfigured, + }, + { + name: "empty address disables profiling", + env: map[string]string{"NB_PYROSCOPE_ADDRESS": ""}, + errIs: errNotConfigured, + }, + { + name: "credentials without address disable profiling", + env: map[string]string{ + "NB_PYROSCOPE_USER": "123456", + "NB_PYROSCOPE_PASSWORD": "token", + }, + errIs: errNotConfigured, + }, + { + name: "address without credentials fails", + env: map[string]string{ + "NB_PYROSCOPE_ADDRESS": "https://profiles-prod-001.grafana.net", + }, + wantErr: true, + }, + { + name: "address with empty credentials fails", + env: map[string]string{ + "NB_PYROSCOPE_ADDRESS": "https://profiles-prod-001.grafana.net", + "NB_PYROSCOPE_USER": "", + "NB_PYROSCOPE_PASSWORD": "", + }, + wantErr: true, + }, + { + name: "address without password fails", + env: map[string]string{ + "NB_PYROSCOPE_ADDRESS": "https://profiles-prod-001.grafana.net", + "NB_PYROSCOPE_USER": "123456", + }, + wantErr: true, + }, + { + name: "full configuration", + env: map[string]string{ + "NB_PYROSCOPE_ADDRESS": "https://profiles-prod-001.grafana.net", + "NB_PYROSCOPE_USER": "123456", + "NB_PYROSCOPE_PASSWORD": "token", + }, + expected: config{ + Address: "https://profiles-prod-001.grafana.net", + User: "123456", + Password: "token", + }, + }, + { + name: "http to loopback is allowed", + env: map[string]string{ + "NB_PYROSCOPE_ADDRESS": "http://127.0.0.1:4040", + "NB_PYROSCOPE_USER": "123456", + "NB_PYROSCOPE_PASSWORD": "token", + }, + expected: config{ + Address: "http://127.0.0.1:4040", + User: "123456", + Password: "token", + }, + }, + { + name: "http to localhost is allowed", + env: map[string]string{ + "NB_PYROSCOPE_ADDRESS": "http://localhost:4040", + "NB_PYROSCOPE_USER": "123456", + "NB_PYROSCOPE_PASSWORD": "token", + }, + expected: config{ + Address: "http://localhost:4040", + User: "123456", + Password: "token", + }, + }, + { + name: "http to private network is allowed", + env: map[string]string{ + "NB_PYROSCOPE_ADDRESS": "http://10.0.0.5:4040", + "NB_PYROSCOPE_USER": "123456", + "NB_PYROSCOPE_PASSWORD": "token", + }, + expected: config{ + Address: "http://10.0.0.5:4040", + User: "123456", + Password: "token", + }, + }, + { + name: "http to public host is rejected", + env: map[string]string{ + "NB_PYROSCOPE_ADDRESS": "http://pyroscope.example.com", + "NB_PYROSCOPE_USER": "123456", + "NB_PYROSCOPE_PASSWORD": "token", + }, + wantErr: true, + }, + { + name: "http to public address is rejected", + env: map[string]string{ + "NB_PYROSCOPE_ADDRESS": "http://203.0.113.10:4040", + "NB_PYROSCOPE_USER": "123456", + "NB_PYROSCOPE_PASSWORD": "token", + }, + wantErr: true, + }, + { + name: "address without scheme is rejected", + env: map[string]string{ + "NB_PYROSCOPE_ADDRESS": "pyroscope.example.com:4040", + "NB_PYROSCOPE_USER": "123456", + "NB_PYROSCOPE_PASSWORD": "token", + }, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + clearEnv(t) + for k, v := range tt.env { + t.Setenv(k, v) + } + + cfg, err := loadConfig() + + switch { + case tt.errIs != nil: + require.ErrorIs(t, err, tt.errIs) + case tt.wantErr: + require.Error(t, err) + require.NotErrorIs(t, err, errNotConfigured) + default: + require.NoError(t, err) + assert.Equal(t, tt.expected, cfg) + } + }) + } +} + +func TestStartWithoutConfigurationIsNoop(t *testing.T) { + clearEnv(t) + + stop := Start("netbird-test") + require.NotNil(t, stop) + stop() +} + +func clearEnv(t *testing.T) { + t.Helper() + + for _, k := range []string{"NB_PYROSCOPE_ADDRESS", "NB_PYROSCOPE_USER", "NB_PYROSCOPE_PASSWORD"} { + t.Setenv(k, "") + require.NoError(t, os.Unsetenv(k)) + } +} diff --git a/signal/cmd/run.go b/signal/cmd/run.go index a36623c6b..42b7d2505 100644 --- a/signal/cmd/run.go +++ b/signal/cmd/run.go @@ -119,6 +119,7 @@ var ( if err != nil { return fmt.Errorf("creating signal server: %v", err) } + defer srv.Stop() proto.RegisterSignalExchangeServer(grpcServer, srv) grpcRootHandler := grpcHandlerFunc(grpcServer, metricsServer.Meter) diff --git a/signal/server/signal.go b/signal/server/signal.go index 7edbb4d34..f991b5d81 100644 --- a/signal/server/signal.go +++ b/signal/server/signal.go @@ -17,6 +17,8 @@ import ( "github.com/netbirdio/signal-dispatcher/dispatcher" + "github.com/netbirdio/netbird/shared/lifecycle" + "github.com/netbirdio/netbird/shared/profiling" "github.com/netbirdio/netbird/shared/signal/proto" "github.com/netbirdio/netbird/signal/metrics" "github.com/netbirdio/netbird/signal/peer" @@ -43,6 +45,8 @@ const ( labelRegistrationNotFound = "not_found" sendTimeout = 10 * time.Second + + applicationName = "signal" ) var ( @@ -51,6 +55,7 @@ var ( // Server an instance of a Signal server type Server struct { + lifecycle.StopHandlers registry *peer.Registry proto.UnimplementedSignalExchangeServer dispatcher *dispatcher.Dispatcher @@ -88,9 +93,17 @@ func NewServer(ctx context.Context, meter metric.Meter, metricsPrefix ...string) sendTimeout: sTimeout, } + stopProfiling := profiling.Start(applicationName) + s.OnStop(stopProfiling) + return s, nil } +// Stop runs the handlers registered with OnStop. +func (s *Server) Stop() { + s.RunStopHandlers() +} + // Send forwards a message to the signal peer func (s *Server) Send(ctx context.Context, msg *proto.EncryptedMessage) (*proto.EncryptedMessage, error) { log.Tracef("received a new message to send from peer [%s] to peer [%s]", msg.Key, msg.RemoteKey)