From 5abeda0df701feb21e099e279f8fa4081c8321bf Mon Sep 17 00:00:00 2001 From: mlsmaycon Date: Mon, 31 Aug 2026 17:07:23 +0200 Subject: [PATCH] [client] Add embed knobs for running many clients in one process Running several embedded clients in one process is already how the proxy works, and both of the costs below scale with the number of clients. Every engine retains its latest management sync response so that GetLatestSyncResponse can read it back. Retaining it pins a decoded copy of the whole network map for the lifetime of the client, roughly 680 KB per client against a 2000 peer network map. A process holding many clients that never read the response back pays that for nothing, so DisableSyncResponsePersistence lets a caller opt out. Persistence stays on by default. Status runs health probes against every STUN and TURN server and takes the engine lock, which is too expensive to poll at high frequency or across many clients. StatusSnapshot returns the same recorder state without the probes. Reading that state from another module also needs the peer status constants and the per-peer state type, which were only partly exported. --- client/embed/embed.go | 37 ++++++++++-- client/embed/embed_test.go | 120 ++++++++++++++++++++++++++++++++++++- 2 files changed, 150 insertions(+), 7 deletions(-) diff --git a/client/embed/embed.go b/client/embed/embed.go index 5a3d11f24..6b1d69935 100644 --- a/client/embed/embed.go +++ b/client/embed/embed.go @@ -36,6 +36,10 @@ var ( ) const ( + // PeerStatusIdle indicates the peer is in disconnected state. + PeerStatusIdle = peer.StatusIdle + // PeerStatusConnecting indicates the peer is in connecting state. + PeerStatusConnecting = peer.StatusConnecting // PeerStatusConnected indicates the peer is in connected state. PeerStatusConnected = peer.StatusConnected ) @@ -43,6 +47,10 @@ const ( // PeerConnStatus is a peer's connection status. type PeerConnStatus = peer.ConnStatus +// PeerState is the status recorder's view of one remote peer, as carried in +// the Peers field of the value returned by Status and StatusSnapshot. +type PeerState = peer.State + // Client manages a netbird embedded client instance. type Client struct { deviceName string @@ -53,6 +61,8 @@ type Client struct { jwtToken string connect *internal.ConnectClient recorder *peer.Status + + disableSyncPersistence bool } // Options configures a new Client. @@ -115,6 +125,12 @@ type Options struct { DNSLabels []string // Performance configures the tunnel's buffer pool cap and batch size. Performance Performance + // DisableSyncResponsePersistence stops the client from retaining the latest + // management sync response. That response is only ever read back through + // GetLatestSyncResponse, and retaining it pins a decoded copy of the whole + // network map for the lifetime of the client. Set this when many clients + // share one process and none of them read the sync response back. + DisableSyncResponsePersistence bool } // Performance configures the embedded client's tunnel memory/throughput knobs. @@ -251,11 +267,12 @@ func New(opts Options) (*Client, error) { } return &Client{ - deviceName: opts.DeviceName, - setupKey: opts.SetupKey, - jwtToken: opts.JWTToken, - config: config, - recorder: peer.NewRecorder(config.ManagementURL.String()), + deviceName: opts.DeviceName, + setupKey: opts.SetupKey, + jwtToken: opts.JWTToken, + config: config, + recorder: peer.NewRecorder(config.ManagementURL.String()), + disableSyncPersistence: opts.DisableSyncResponsePersistence, }, nil } @@ -288,7 +305,7 @@ func (c *Client) Start(startCtx context.Context) error { return fmt.Errorf("login: %w", err) } client := internal.NewConnectClient(ctx, c.config, c.recorder) - client.SetSyncResponsePersistence(true) + client.SetSyncResponsePersistence(!c.disableSyncPersistence) // either startup error (permanent backoff err) or nil err (successful engine up) // TODO: make after-startup backoff err available @@ -500,6 +517,14 @@ func (c *Client) Status() (peer.FullStatus, error) { return c.recorder.GetFullStatus(), nil } +// StatusSnapshot returns the client's current status without running health +// probes. Status probes every STUN and TURN server and takes the engine lock, +// which is too expensive to call at high frequency or across many clients +// sharing one process. +func (c *Client) StatusSnapshot() peer.FullStatus { + return c.recorder.GetFullStatus() +} + // GetLatestSyncResponse returns the latest sync response from the management server. func (c *Client) GetLatestSyncResponse() (*mgmProto.SyncResponse, error) { engine, err := c.getEngine() diff --git a/client/embed/embed_test.go b/client/embed/embed_test.go index 4ff5c9978..0f135efe8 100644 --- a/client/embed/embed_test.go +++ b/client/embed/embed_test.go @@ -6,8 +6,9 @@ import ( "testing" "time" - "go.uber.org/mock/gomock" "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel" + "go.uber.org/mock/gomock" "google.golang.org/grpc" "github.com/netbirdio/netbird/management/internals/controllers/network_map/controller" @@ -29,6 +30,8 @@ import ( "github.com/netbirdio/netbird/management/server/telemetry" "github.com/netbirdio/netbird/management/server/types" mgmtProto "github.com/netbirdio/netbird/shared/management/proto" + signalProto "github.com/netbirdio/netbird/shared/signal/proto" + signalServer "github.com/netbirdio/netbird/signal/server" "github.com/netbirdio/netbird/util" ) @@ -166,3 +169,118 @@ func startManagement(t *testing.T, signalAddr string) string { return lis.Addr().String() } + +// startSignal starts a signal server that serves the SignalExchange service, so +// an embedded client can get past WaitStreamConnected and finish Engine.Start. +func startSignal(t *testing.T) string { + t.Helper() + + lis, err := net.Listen("tcp", "localhost:0") + require.NoError(t, err) + + s := grpc.NewServer() + srv, err := signalServer.NewServer(context.Background(), otel.Meter("")) + require.NoError(t, err) + signalProto.RegisterSignalExchangeServer(s, srv) + + go func() { + if err := s.Serve(lis); err != nil { + t.Error(err) + } + }() + t.Cleanup(s.Stop) + + return lis.Addr().String() +} + +// TestClientSyncResponsePersistence checks that DisableSyncResponsePersistence +// controls whether the engine retains the latest management sync response, which +// is observable through GetLatestSyncResponse. +func TestClientSyncResponsePersistence(t *testing.T) { + tests := []struct { + name string + disable bool + persisted bool + }{ + {name: "retained by default", disable: false, persisted: true}, + {name: "dropped when disabled", disable: true, persisted: false}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + signalAddr := startSignal(t) + mgmAddr := startManagement(t, signalAddr) + + wgPort := 0 + client, err := New(Options{ + DeviceName: "embed-persistence-test", + SetupKey: testSetupKey, + ManagementURL: "http://" + mgmAddr, + WireguardPort: &wgPort, + DisableSyncResponsePersistence: tc.disable, + }) + require.NoError(t, err, "embed client creation must succeed") + + startCtx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + require.NoError(t, client.Start(startCtx), "client must start") + + t.Cleanup(func() { + stopCtx, stopCancel := context.WithTimeout(context.Background(), 15*time.Second) + defer stopCancel() + if err := client.Stop(stopCtx); err != nil { + t.Logf("stop client: %v", err) + } + }) + + if !tc.persisted { + _, err := client.GetLatestSyncResponse() + require.Error(t, err, "no sync response may be retained when persistence is disabled") + return + } + + require.Eventually(t, func() bool { + resp, err := client.GetLatestSyncResponse() + return err == nil && resp.GetNetworkMap() != nil + }, 30*time.Second, 200*time.Millisecond, "the sync response and its network map should be retained by default") + }) + } +} + +// TestClientStatusSnapshot checks that StatusSnapshot reports a started client's +// state without going through the health probes Status runs. +func TestClientStatusSnapshot(t *testing.T) { + signalAddr := startSignal(t) + mgmAddr := startManagement(t, signalAddr) + + mgmtURL := "http://" + mgmAddr + wgPort := 0 + client, err := New(Options{ + DeviceName: "embed-status-snapshot-test", + SetupKey: testSetupKey, + ManagementURL: mgmtURL, + WireguardPort: &wgPort, + }) + require.NoError(t, err, "embed client creation must succeed") + + // Safe before Start: the recorder exists from New, and no engine is needed. + require.Empty(t, client.StatusSnapshot().LocalPeerState.IP, "an unstarted client has no overlay address") + + startCtx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + require.NoError(t, client.Start(startCtx), "client must start") + + t.Cleanup(func() { + stopCtx, stopCancel := context.WithTimeout(context.Background(), 15*time.Second) + defer stopCancel() + if err := client.Stop(stopCtx); err != nil { + t.Logf("stop client: %v", err) + } + }) + + require.Eventually(t, func() bool { + return client.StatusSnapshot().LocalPeerState.IP != "" + }, 30*time.Second, 200*time.Millisecond, "a started client should report its overlay address") + + require.Equal(t, mgmtURL, client.StatusSnapshot().ManagementState.URL, "snapshot should carry the management URL") +}