From 7a243ae5b1968727e3f06ffd15c5c40ad8ab0026 Mon Sep 17 00:00:00 2001 From: Brad Fitzpatrick Date: Mon, 18 Oct 2021 14:38:03 -0700 Subject: [PATCH] wgengine/magicsock: finish TODO to speed up peerMap.forEachEndpointWithDiscoKey Now that peerMap tracks the set of nodes for a DiscoKey. Updates #3088 Change-Id: I927bf2bdfd2b8126475f6b6acc44bc799fcb489f Signed-off-by: Brad Fitzpatrick --- wgengine/magicsock/magicsock.go | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go index 440322685..4160eaadf 100644 --- a/wgengine/magicsock/magicsock.go +++ b/wgengine/magicsock/magicsock.go @@ -153,12 +153,17 @@ func (m *peerMap) forEachEndpoint(f func(ep *endpoint)) { // forEachEndpointWithDiscoKey invokes f on every endpoint in m // that has the provided DiscoKey. func (m *peerMap) forEachEndpointWithDiscoKey(dk tailcfg.DiscoKey, f func(ep *endpoint)) { - // TODO(bradfitz): once byDiscoKey is a set of endpoints, then range - // over that instead. - for _, pi := range m.byNodeKey { - if pi.ep.discoKey == dk { - f(pi.ep) + for nk := range m.nodesOfDisco[dk] { + pi, ok := m.byNodeKey[nk] + if !ok { + // Unexpected. Data structures would have to + // be out of sync. But we don't have a logger + // here to log [unexpected], so just skip. + // Maybe log later once peerMap is merged back + // into Conn. + continue } + f(pi.ep) } }