simplify

proxy/internal/roundtrip/netbird.go

@@ -209,33 +209,30 @@ func (n *NetBird) AddPeer(ctx context.Context, accountID types.AccountID, key Se
 		return nil
 	}
 
-	entry, unlock, err := n.createClientLocked(ctx, accountID, key, authToken, si)
+	lifecycle := n.accountLifecycle(accountID)
+	lifecycle.Lock()
+
+	entry, err := n.createClientLocked(ctx, accountID, key, authToken, si)
 	if entry == nil || err != nil {
-		unlock()
+		lifecycle.Unlock()
 		return err
 	}
 
-	// runClientStartup inherits the lifecycle lock so a new client.Start for
-	// this account waits behind any in-flight teardown.
 	go func() {
-		defer unlock()
+		defer lifecycle.Unlock()
 		n.runClientStartup(accountID, entry.client)
 	}()
 
 	return nil
 }
 
-// createClientLocked acquires the account's lifecycle lock, re-checks for a
-// concurrently-created client, and creates a new one. It returns with the lock
-// HELD; the caller must call the returned unlock once startup is done (or
-// immediately on a nil entry / error). A nil entry means another caller won
-// the race and the service was registered against the existing client.
-func (n *NetBird) createClientLocked(ctx context.Context, accountID types.AccountID, key ServiceKey, authToken string, si serviceInfo) (*clientEntry, func(), error) {
-	lifecycle := n.accountLifecycle(accountID)
-	lifecycle.Lock()
-
+// createClientLocked re-checks for a concurrently-created client and creates a
+// new one. The caller must hold the account's lifecycle lock. A nil entry means
+// another caller won the race and the service was registered against the
+// existing client.
+func (n *NetBird) createClientLocked(ctx context.Context, accountID types.AccountID, key ServiceKey, authToken string, si serviceInfo) (*clientEntry, error) {
 	if n.registerExistingClient(accountID, key, si) {
-		return nil, lifecycle.Unlock, nil
+		return nil, nil
 	}
 
 	createStart := time.Now()
@@ -244,7 +241,7 @@ func (n *NetBird) createClientLocked(ctx context.Context, accountID types.Accoun
 		n.OnAddPeer(time.Since(createStart), err)
 	}
 	if err != nil {
-		return nil, lifecycle.Unlock, err
+		return nil, err
 	}
 
 	n.clientsMux.Lock()
@@ -256,7 +253,7 @@ func (n *NetBird) createClientLocked(ctx context.Context, accountID types.Accoun
 		"service_key": key,
 	}).Info("created new client for account")
 
-	return entry, lifecycle.Unlock, nil
+	return entry, nil
 }
 
 // registerExistingClient registers the service against an already-present
@@ -489,6 +486,15 @@ func (n *NetBird) notifyClientReady(accountID types.AccountID, client *embed.Cli
 // RemovePeer unregisters a service from an account. The client is only stopped
 // when no services are using it anymore.
 func (n *NetBird) RemovePeer(ctx context.Context, accountID types.AccountID, key ServiceKey) error {
+	lifecycle := n.accountLifecycle(accountID)
+	lifecycle.Lock()
+	transferred := false
+	defer func() {
+		if !transferred {
+			lifecycle.Unlock()
+		}
+	}()
+
 	n.clientsMux.Lock()
 
 	entry, exists := n.clients[accountID]
@@ -511,17 +517,8 @@ func (n *NetBird) RemovePeer(ctx context.Context, accountID types.AccountID, key
 	delete(entry.services, key)
 
 	stopClient := len(entry.services) == 0
-	var client *embed.Client
-	var transport, insecureTransport *http.Transport
-	var inbound any
-	var stopHandler func(types.AccountID, any)
 	if stopClient {
 		n.logger.WithField("account_id", accountID).Info("stopping client, no more services")
-		client = entry.client
-		transport = entry.transport
-		insecureTransport = entry.insecureTransport
-		inbound = entry.inbound
-		stopHandler = n.stopHandler
 		delete(n.clients, accountID)
 	} else {
 		n.logger.WithFields(log.Fields{
@@ -535,53 +532,36 @@ func (n *NetBird) RemovePeer(ctx context.Context, accountID types.AccountID, key
 	n.notifyDisconnect(ctx, accountID, key, si.serviceID)
 
 	if stopClient {
-		go n.stopClientAsync(accountID, clientTeardown{
-			client:            client,
-			transport:         transport,
-			insecureTransport: insecureTransport,
-			inbound:           inbound,
-			stopHandler:       stopHandler,
-		})
+		transferred = true
+		go n.stopClientLocked(accountID, lifecycle, entry)
 	}
 
 	return nil
 }
 
-// clientTeardown bundles the resources released when an account's last service
-// is removed.
-type clientTeardown struct {
-	client            *embed.Client
-	transport         *http.Transport
-	insecureTransport *http.Transport
-	inbound           any
-	stopHandler       func(types.AccountID, any)
-}
-
-// stopClientAsync releases a client's resources off the caller's goroutine so a
+// stopClientLocked releases a client's resources off the caller's goroutine so a
 // slow client.Stop cannot wedge the mapping receive loop (which calls RemovePeer
-// synchronously). The account's lifecycle mutex is held so a new client.Start
-// for the same account waits for this teardown.
-func (n *NetBird) stopClientAsync(accountID types.AccountID, td clientTeardown) {
-	lifecycle := n.accountLifecycle(accountID)
-	lifecycle.Lock()
+// synchronously). It unlocks lifecycle when done so a new client.Start for the
+// same account waits for this teardown.
+func (n *NetBird) stopClientLocked(accountID types.AccountID, lifecycle *sync.Mutex, entry *clientEntry) {
 	defer lifecycle.Unlock()
 
-	if td.inbound != nil && td.stopHandler != nil {
-		td.stopHandler(accountID, td.inbound)
+	if entry.inbound != nil && n.stopHandler != nil {
+		n.stopHandler(accountID, entry.inbound)
 	}
-	if td.transport != nil {
-		td.transport.CloseIdleConnections()
+	if entry.transport != nil {
+		entry.transport.CloseIdleConnections()
 	}
-	if td.insecureTransport != nil {
-		td.insecureTransport.CloseIdleConnections()
+	if entry.insecureTransport != nil {
+		entry.insecureTransport.CloseIdleConnections()
 	}
-	if td.client == nil {
+	if entry.client == nil {
 		return
 	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), clientStopTimeout)
 	defer cancel()
-	if err := td.client.Stop(ctx); err != nil {
+	if err := entry.client.Stop(ctx); err != nil {
 		n.logger.WithField("account_id", accountID).WithError(err).Warn("failed to stop netbird client")
 	}
 }

proxy/internal/roundtrip/netbird_test.go

@@ -23,6 +23,18 @@ func (m *mockMgmtClient) CreateProxyPeer(_ context.Context, _ *proto.CreateProxy
 	return &proto.CreateProxyPeerResponse{Success: true}, nil
 }
 
+// signalMgmtClient closes entered the first time CreateProxyPeer is called, so
+// tests can detect AddPeer reaching client creation.
+type signalMgmtClient struct {
+	entered chan struct{}
+	once    sync.Once
+}
+
+func (m *signalMgmtClient) CreateProxyPeer(_ context.Context, _ *proto.CreateProxyPeerRequest, _ ...grpc.CallOption) (*proto.CreateProxyPeerResponse, error) {
+	m.once.Do(func() { close(m.entered) })
+	return &proto.CreateProxyPeerResponse{Success: true}, nil
+}
+
 type mockStatusNotifier struct {
 	mu       sync.Mutex
 	statuses []statusCall
@@ -422,6 +434,13 @@ func TestNetBird_RemovePeer_TeardownIsAsync(t *testing.T) {
 // TestNetBird_AddPeer_WaitsForTeardown proves the lifecycle lock serialises a
 // new client bringup behind an in-flight teardown for the same account, so a
 // slow client.Stop can never race a new client.Start for that account.
+//
+// It targets the handoff race specifically: AddPeer is launched immediately
+// after RemovePeer returns, WITHOUT waiting for the teardown goroutine to start.
+// This only passes if RemovePeer acquires the lifecycle lock synchronously
+// (before returning) and hands it to the teardown goroutine — if the goroutine
+// acquired the lock itself, AddPeer could win the lock in this window and start
+// a replacement client while the old teardown is still pending.
 func TestNetBird_AddPeer_WaitsForTeardown(t *testing.T) {
 	nb := NewNetBird(context.Background(), "test-proxy", "invalid.test", ClientConfig{
 		MgmtAddr: "http://invalid.test:9999",
@@ -430,10 +449,12 @@ func TestNetBird_AddPeer_WaitsForTeardown(t *testing.T) {
 	accountID := types.AccountID("acct-serialize")
 	key := DomainServiceKey("svc.example")
 
-	teardownEntered := make(chan struct{})
+	addEntered := make(chan struct{})
 	releaseTeardown := make(chan struct{})
 	nb.SetClientLifecycle(nil, func(types.AccountID, any) {
-		close(teardownEntered)
+		// Block teardown until released. If AddPeer ever reaches createClientEntry
+		// (signalled via the mgmt client below) while we hold the lock, the lock
+		// failed to serialise and the test fails before we release.
 		<-releaseTeardown
 	})
 
@@ -445,9 +466,13 @@ func TestNetBird_AddPeer_WaitsForTeardown(t *testing.T) {
 	}
 	nb.clientsMux.Unlock()
 
-	require.NoError(t, nb.RemovePeer(context.Background(), accountID, key))
-	<-teardownEntered
+	// createClientEntry calls CreateProxyPeer; closing addEntered there tells us
+	// AddPeer got past the lifecycle lock and into client creation.
+	nb.mgmtClient = &signalMgmtClient{entered: addEntered}
 
+	require.NoError(t, nb.RemovePeer(context.Background(), accountID, key))
+
+	// Launch AddPeer with NO synchronisation against the teardown goroutine.
 	addReturned := make(chan struct{})
 	go func() {
 		_ = nb.AddPeer(context.Background(), accountID, DomainServiceKey("svc2.example"), "key-2", types.ServiceID("svc-2"))
@@ -455,8 +480,10 @@ func TestNetBird_AddPeer_WaitsForTeardown(t *testing.T) {
 	}()
 
 	select {
+	case <-addEntered:
+		t.Fatal("AddPeer entered client creation while teardown held the lifecycle lock — handoff race not closed")
 	case <-addReturned:
-		t.Fatal("AddPeer returned while teardown still held the lifecycle lock — not serialised")
+		t.Fatal("AddPeer completed while teardown held the lifecycle lock — not serialised")
 	case <-time.After(300 * time.Millisecond):
 	}
 

proxy/mapping_stall_test.go

@@ -174,10 +174,11 @@ func TestMappingStream_StallsWhenApplyBlocks(t *testing.T) {
 	// THE DEADLOCK: while the first batch is parked in CreateProxyPeer, the
 	// single-threaded loop cannot advance. The second batch is never pulled,
 	// even though it is already available on the stream. Give it ample time.
+	// deliveredCount is atomic; syncDone is intentionally not read here because
+	// the loop goroutine owns it (reading it from the test would race).
 	time.Sleep(500 * time.Millisecond)
 	assert.Equal(t, int32(1), stream.deliveredCount(),
 		"loop must NOT consume the second batch while the first is blocked in apply — proxy is stuck")
-	assert.False(t, syncDone, "initial sync cannot complete while the loop is wedged")
 
 	select {
 	case <-loopDone:
@@ -260,10 +261,10 @@ func TestMappingStream_StallsWhenRemoveBlocks(t *testing.T) {
 	}
 
 	// THE DEADLOCK: the loop is parked in the blocked remove and cannot advance.
+	// syncDone is owned by the loop goroutine, so it is not read here.
 	time.Sleep(500 * time.Millisecond)
 	assert.Equal(t, int32(1), stream.deliveredCount(),
 		"loop must NOT consume the second batch while the first remove is blocked — proxy is stuck")
-	assert.False(t, syncDone, "initial sync cannot complete while the loop is wedged on a remove")
 
 	select {
 	case <-loopDone:

proxy/server.go

@@ -1489,10 +1489,13 @@ func (s *Server) mappingBatchWatchdog() time.Duration {
 // if processing exceeds the watchdog so the caller reconnects and resyncs
 // instead of wedging silently.
 func (s *Server) processMappingsGuarded(ctx context.Context, mappings []*proto.ProxyMapping) error {
+	batchCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	done := make(chan struct{})
 	go func() {
 		defer close(done)
-		s.processMappings(ctx, mappings)
+		s.processMappings(batchCtx, mappings)
 	}()
 
 	watchdog := s.mappingBatchWatchdog()
@@ -1505,7 +1508,7 @@ func (s *Server) processMappingsGuarded(ctx context.Context, mappings []*proto.P
 	case <-ctx.Done():
 		return ctx.Err()
 	case <-timer.C:
-		s.Logger.Errorf("processing mapping batch exceeded %s, reconnecting to resync", watchdog)
+		s.Logger.Errorf("processing mapping batch exceeded %s, cancelling and reconnecting to resync", watchdog)
 		return fmt.Errorf("mapping batch processing stalled after %s", watchdog)
 	}
 }