diff --git a/client/internal/routemanager/manager.go b/client/internal/routemanager/manager.go index 40bcce84f..ef69b81a4 100644 --- a/client/internal/routemanager/manager.go +++ b/client/internal/routemanager/manager.go @@ -246,17 +246,15 @@ func (m *DefaultManager) ReconcilePeerAllowedIPs(peerKey string) error { return nil } - prefixes := m.allowedIPsRefCounter.KeysMatching(func(out string) bool { - return out == peerKey - }) - - var merr *multierror.Error - for _, prefix := range prefixes { - if err := m.wgInterface.AddAllowedIP(peerKey, prefix); err != nil { - merr = multierror.Append(merr, fmt.Errorf("add allowed IP %s for peer %s: %w", prefix, peerKey, err)) - } - } - return nberrors.FormatErrorOrNil(merr) + return m.allowedIPsRefCounter.ReapplyMatching( + func(out string) bool { return out == peerKey }, + func(prefix netip.Prefix) error { + if err := m.wgInterface.AddAllowedIP(peerKey, prefix); err != nil { + return fmt.Errorf("add allowed IP %s for peer %s: %w", prefix, peerKey, err) + } + return nil + }, + ) } // Init sets up the routing diff --git a/client/internal/routemanager/refcounter/refcounter.go b/client/internal/routemanager/refcounter/refcounter.go index 02fc835cb..917120275 100644 --- a/client/internal/routemanager/refcounter/refcounter.go +++ b/client/internal/routemanager/refcounter/refcounter.go @@ -94,21 +94,24 @@ func (rm *Counter[Key, I, O]) Get(key Key) (Ref[O], bool) { return ref, ok } -// KeysMatching returns every key whose stored Out satisfies the predicate. The predicate is -// evaluated under the counter's lock, so keep it cheap and side-effect free. It is used to -// enumerate the keys (e.g. allowed-IP prefixes) associated with a given Out value (e.g. a peer -// key) without exposing the internal map. -func (rm *Counter[Key, I, O]) KeysMatching(pred func(out O) bool) []Key { +// ReapplyMatching calls apply for every key whose stored Out satisfies pred, holding the +// counter lock for the whole pass. Running apply under the lock keeps it atomic with respect +// to Increment/Decrement: a prefix dropped to zero is removed from the map (and had its +// RemoveFunc called) before this pass observes it, so a stale key can never be re-applied. +// pred and apply are invoked under the lock, so they must not call back into the counter. +func (rm *Counter[Key, I, O]) ReapplyMatching(pred func(out O) bool, apply func(key Key) error) error { rm.mu.Lock() defer rm.mu.Unlock() - var keys []Key + var merr *multierror.Error for key, ref := range rm.refCountMap { if pred(ref.Out) { - keys = append(keys, key) + if err := apply(key); err != nil { + merr = multierror.Append(merr, err) + } } } - return keys + return nberrors.FormatErrorOrNil(merr) } // Increment increments the reference count for the given key. diff --git a/client/internal/routemanager/refcounter/refcounter_test.go b/client/internal/routemanager/refcounter/refcounter_test.go index 686782fd5..79a99c388 100644 --- a/client/internal/routemanager/refcounter/refcounter_test.go +++ b/client/internal/routemanager/refcounter/refcounter_test.go @@ -8,10 +8,10 @@ import ( "github.com/stretchr/testify/require" ) -// TestKeysMatching verifies KeysMatching returns exactly the keys whose stored Out satisfies the -// predicate — the primitive ReconcilePeerAllowedIPs relies on to enumerate a single peer's routed -// prefixes without touching the others. -func TestKeysMatching(t *testing.T) { +// TestReapplyMatching verifies ReapplyMatching invokes apply for exactly the keys whose stored +// Out satisfies the predicate (no duplicates for multiply-referenced keys) — the primitive +// ReconcilePeerAllowedIPs relies on to re-apply a single peer's routed prefixes. +func TestReapplyMatching(t *testing.T) { rc := New[netip.Prefix, string, string]( func(_ netip.Prefix, peerKey string) (string, error) { return peerKey, nil }, func(netip.Prefix, string) error { return nil }, @@ -21,23 +21,27 @@ func TestKeysMatching(t *testing.T) { peerA2 := netip.MustParsePrefix("10.1.0.0/24") peerB1 := netip.MustParsePrefix("10.2.0.0/24") + for prefix, peer := range map[netip.Prefix]string{peerA1: "peerA", peerA2: "peerA", peerB1: "peerB"} { + _, err := rc.Increment(prefix, peer) + require.NoError(t, err) + } + // a second reference must not make the key applied twice _, err := rc.Increment(peerA1, "peerA") require.NoError(t, err) - _, err = rc.Increment(peerA2, "peerA") + + var applied []netip.Prefix + err = rc.ReapplyMatching( + func(out string) bool { return out == "peerA" }, + func(key netip.Prefix) error { applied = append(applied, key); return nil }, + ) require.NoError(t, err) - _, err = rc.Increment(peerB1, "peerB") + assert.ElementsMatch(t, []netip.Prefix{peerA1, peerA2}, applied) + + var none []netip.Prefix + err = rc.ReapplyMatching( + func(out string) bool { return out == "missing" }, + func(key netip.Prefix) error { none = append(none, key); return nil }, + ) require.NoError(t, err) - - // a second reference to peerA1 must not duplicate it in the result - _, err = rc.Increment(peerA1, "peerA") - require.NoError(t, err) - - keysA := rc.KeysMatching(func(out string) bool { return out == "peerA" }) - assert.ElementsMatch(t, []netip.Prefix{peerA1, peerA2}, keysA) - - keysB := rc.KeysMatching(func(out string) bool { return out == "peerB" }) - assert.ElementsMatch(t, []netip.Prefix{peerB1}, keysB) - - keysNone := rc.KeysMatching(func(out string) bool { return out == "missing" }) - assert.Empty(t, keysNone) + assert.Empty(t, none) }