From 3e2c29a355b58858e53075ebf9879d0c06b17796 Mon Sep 17 00:00:00 2001 From: pascal Date: Mon, 8 Jun 2026 17:08:05 +0200 Subject: [PATCH] simplify --- proxy/internal/roundtrip/netbird.go | 94 ++++++++++-------------- proxy/internal/roundtrip/netbird_test.go | 37 ++++++++-- proxy/mapping_stall_test.go | 5 +- proxy/server.go | 7 +- 4 files changed, 77 insertions(+), 66 deletions(-) diff --git a/proxy/internal/roundtrip/netbird.go b/proxy/internal/roundtrip/netbird.go index dc79ba320..306201d07 100644 --- a/proxy/internal/roundtrip/netbird.go +++ b/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") } } diff --git a/proxy/internal/roundtrip/netbird_test.go b/proxy/internal/roundtrip/netbird_test.go index 7066123be..8aa7a651c 100644 --- a/proxy/internal/roundtrip/netbird_test.go +++ b/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): } diff --git a/proxy/mapping_stall_test.go b/proxy/mapping_stall_test.go index 0ffef1a93..acf313d19 100644 --- a/proxy/mapping_stall_test.go +++ b/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: diff --git a/proxy/server.go b/proxy/server.go index f4586ca16..6d5acfe46 100644 --- a/proxy/server.go +++ b/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) } }