[proxy] concurrent proxy snapshot apply (#6207)

proxy/internal/metrics/metrics.go

@@ -25,6 +25,11 @@ type Metrics struct {
 	backendDuration          metric.Int64Histogram
 	certificateIssueDuration metric.Int64Histogram
 
+	// Management sync metrics.
+	snapshotSyncDuration  metric.Int64Histogram
+	snapshotBatchDuration metric.Int64Histogram
+	addPeerDuration       metric.Int64Histogram
+
 	// L4 service-level metrics.
 	l4Services metric.Int64UpDownCounter
 
@@ -54,6 +59,9 @@ func New(ctx context.Context, meter metric.Meter) (*Metrics, error) {
 	if err := m.initHTTPMetrics(meter); err != nil {
 		return nil, err
 	}
+	if err := m.initSyncMetrics(meter); err != nil {
+		return nil, err
+	}
 	if err := m.initL4Metrics(meter); err != nil {
 		return nil, err
 	}
@@ -126,6 +134,59 @@ func (m *Metrics) initHTTPMetrics(meter metric.Meter) error {
 	return err
 }
 
+func (m *Metrics) initSyncMetrics(meter metric.Meter) error {
+	var err error
+
+	m.snapshotSyncDuration, err = meter.Int64Histogram(
+		"proxy.sync.snapshot.duration.ms",
+		metric.WithUnit("milliseconds"),
+		metric.WithDescription("Duration from management connect until the initial snapshot sync is complete"),
+		metric.WithExplicitBucketBoundaries(100, 250, 500, 1000, 2500, 5000, 10000, 30000, 60000, 120000, 300000),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.snapshotBatchDuration, err = meter.Int64Histogram(
+		"proxy.sync.batch.duration.ms",
+		metric.WithUnit("milliseconds"),
+		metric.WithDescription("Duration to process a single mapping batch during initial snapshot sync"),
+		metric.WithExplicitBucketBoundaries(100, 250, 500, 1000, 2500, 5000, 10000, 30000, 60000, 120000, 300000),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.addPeerDuration, err = meter.Int64Histogram(
+		"proxy.peer.add.duration.ms",
+		metric.WithUnit("milliseconds"),
+		metric.WithDescription("Duration to add a peer for an account (keygen + gRPC CreateProxyPeer + embed.New)"),
+		metric.WithExplicitBucketBoundaries(10, 25, 50, 100, 250, 500, 1000, 2500, 5000, 10000),
+	)
+	return err
+}
+
+// RecordSnapshotSyncDuration records the total time from connect to sync-complete.
+func (m *Metrics) RecordSnapshotSyncDuration(d time.Duration) {
+	m.snapshotSyncDuration.Record(m.ctx, d.Milliseconds())
+}
+
+// RecordSnapshotBatchDuration records the time to process one mapping batch during initial sync.
+func (m *Metrics) RecordSnapshotBatchDuration(d time.Duration) {
+	m.snapshotBatchDuration.Record(m.ctx, d.Milliseconds())
+}
+
+// RecordAddPeerDuration records the time to create a new peer for an account.
+func (m *Metrics) RecordAddPeerDuration(d time.Duration, err error) {
+	result := "success"
+	if err != nil {
+		result = "error"
+	}
+	m.addPeerDuration.Record(m.ctx, d.Milliseconds(), metric.WithAttributes(
+		attribute.String("result", result),
+	))
+}
+
 func (m *Metrics) initL4Metrics(meter metric.Meter) error {
 	var err error
 

proxy/internal/roundtrip/netbird.go

@@ -76,6 +76,11 @@ type clientEntry struct {
 	services          map[ServiceKey]serviceInfo
 	createdAt         time.Time
 	started           bool
+	// ready is closed once the client has been fully initialized.
+	// Callers that find a pending entry wait on this channel before
+	// accessing the client. A nil initErr means success.
+	ready   chan struct{}
+	initErr error
 	// Per-backend in-flight limiting keyed by target host:port.
 	// TODO: clean up stale entries when backend targets change.
 	inflightMu  sync.Mutex
@@ -137,6 +142,11 @@ type NetBird struct {
 	clients        map[types.AccountID]*clientEntry
 	initLogOnce    sync.Once
 	statusNotifier statusNotifier
+
+	// OnAddPeer, when set, is called after AddPeer completes for a new account
+	// (i.e. when a new client was actually created, not when an existing one
+	// was reused). The duration covers keygen + gRPC CreateProxyPeer + embed.New.
+	OnAddPeer func(d time.Duration, err error)
 }
 
 // ClientDebugInfo contains debug information about a client.
@@ -157,6 +167,9 @@ type skipTLSVerifyContextKey struct{}
 // AddPeer registers a service for an account. If the account doesn't have a client yet,
 // one is created by authenticating with the management server using the provided token.
 // Multiple services can share the same client.
+//
+// Client creation (WG keygen, gRPC, embed.New) runs without holding clientsMux
+// so that concurrent AddPeer calls for different accounts execute in parallel.
 func (n *NetBird) AddPeer(ctx context.Context, accountID types.AccountID, key ServiceKey, authToken string, serviceID types.ServiceID) error {
 	si := serviceInfo{serviceID: serviceID}
 
@@ -164,10 +177,23 @@ func (n *NetBird) AddPeer(ctx context.Context, accountID types.AccountID, key Se
 
 	entry, exists := n.clients[accountID]
 	if exists {
+		ready := entry.ready
 		entry.services[key] = si
 		started := entry.started
 		n.clientsMux.Unlock()
 
+		// If the entry is still being initialized by another goroutine, wait.
+		if ready != nil {
+			select {
+			case <-ready:
+			case <-ctx.Done():
+				return ctx.Err()
+			}
+			if entry.initErr != nil {
+				return fmt.Errorf("peer initialization failed: %w", entry.initErr)
+			}
+		}
+
 		n.logger.WithFields(log.Fields{
 			"account_id":  accountID,
 			"service_key": key,
@@ -184,15 +210,43 @@ func (n *NetBird) AddPeer(ctx context.Context, accountID types.AccountID, key Se
 		return nil
 	}
 
-	entry, err := n.createClientEntry(ctx, accountID, key, authToken, si)
+	// Insert a placeholder so other goroutines calling AddPeer for the same
+	// account will wait on the ready channel instead of starting a second
+	// client creation.
+	entry = &clientEntry{
+		services: map[ServiceKey]serviceInfo{key: si},
+		ready:    make(chan struct{}),
+	}
+	n.clients[accountID] = entry
+	n.clientsMux.Unlock()
+
+	createStart := time.Now()
+	created, err := n.createClientEntry(ctx, accountID, key, authToken, si)
+	if n.OnAddPeer != nil {
+		n.OnAddPeer(time.Since(createStart), err)
+	}
 	if err != nil {
+		entry.initErr = err
+		close(entry.ready)
+
+		n.clientsMux.Lock()
+		delete(n.clients, accountID)
 		n.clientsMux.Unlock()
 		return err
 	}
 
-	n.clients[accountID] = entry
+	// Transfer any services that were registered by concurrent AddPeer calls
+	// while we were creating the client.
+	n.clientsMux.Lock()
+	for k, v := range entry.services {
+		created.services[k] = v
+	}
+	created.ready = nil
+	n.clients[accountID] = created
 	n.clientsMux.Unlock()
 
+	close(entry.ready)
+
 	n.logger.WithFields(log.Fields{
 		"account_id":  accountID,
 		"service_key": key,
@@ -200,13 +254,13 @@ func (n *NetBird) AddPeer(ctx context.Context, accountID types.AccountID, key Se
 
 	// Attempt to start the client in the background; if this fails we will
 	// retry on the first request via RoundTrip.
-	go n.runClientStartup(ctx, accountID, entry.client)
+	go n.runClientStartup(ctx, accountID, created.client)
 
 	return nil
 }
 
 // createClientEntry generates a WireGuard keypair, authenticates with management,
-// and creates an embedded NetBird client. Must be called with clientsMux held.
+// and creates an embedded NetBird client.
 func (n *NetBird) createClientEntry(ctx context.Context, accountID types.AccountID, key ServiceKey, authToken string, si serviceInfo) (*clientEntry, error) {
 	serviceID := si.serviceID
 	n.logger.WithFields(log.Fields{