Fold the rebind counter and traffic flags into one atomic word

connection_manager.go

@@ -105,11 +105,12 @@ func (cm *connectionManager) getInactivityTimeout() time.Duration {
 }
 
 func (cm *connectionManager) In(h *HostInfo) {
-	h.in.Store(true)
+	h.markIn()
 }
 
-func (cm *connectionManager) Out(h *HostInfo) {
-	h.out.Store(true)
+// Out records outbound traffic and reports whether the local network changed since this tunnel last sent.
+func (cm *connectionManager) Out(h *HostInfo) bool {
+	return h.markOut(cm.intf.rebindEpoch.Load())
 }
 
 func (cm *connectionManager) RelayUsed(localIndex uint32) {
@@ -128,8 +129,7 @@ func (cm *connectionManager) RelayUsed(localIndex uint32) {
 // getAndResetTrafficCheck returns if there was any inbound or outbound traffic within the last tick and
 // resets the state for this local index
 func (cm *connectionManager) getAndResetTrafficCheck(h *HostInfo, now time.Time) (bool, bool) {
-	in := h.in.Swap(false)
-	out := h.out.Swap(false)
+	in, out := h.takeTraffic()
 	if in || out {
 		h.lastUsed = now
 	}
@@ -340,7 +340,7 @@ func (cm *connectionManager) makeTrafficDecision(localIndex uint32, now time.Tim
 				"tunnelCheck", m{"state": "alive", "method": "passive"},
 			)
 		}
-		hostinfo.pendingDeletion.Store(false)
+		hostinfo.setPendingDeletion(false)
 
 		if mainHostInfo {
 			decision = tryRehandshake
@@ -363,7 +363,7 @@ func (cm *connectionManager) makeTrafficDecision(localIndex uint32, now time.Tim
 		return decision, hostinfo, primary
 	}
 
-	if hostinfo.pendingDeletion.Load() {
+	if hostinfo.isPendingDeletion() {
 		// We have already sent a test packet and nothing was returned, this hostinfo is dead
 		hostinfo.logger(cm.l).Info("Tunnel status",
 			"tunnelCheck", m{"state": "dead", "method": "active"},
@@ -414,7 +414,7 @@ func (cm *connectionManager) makeTrafficDecision(localIndex uint32, now time.Tim
 		}
 	}
 
-	hostinfo.pendingDeletion.Store(true)
+	hostinfo.setPendingDeletion(true)
 	cm.trafficTimer.Add(hostinfo.localIndexId, cm.pendingDeletionInterval)
 	return decision, hostinfo, nil
 }

connection_manager_test.go

@@ -86,25 +86,25 @@ func Test_NewConnectionManagerTest(t *testing.T) {
 	// We saw traffic out to vpnIp
 	nc.Out(hostinfo)
 	nc.In(hostinfo)
-	assert.False(t, hostinfo.pendingDeletion.Load())
+	assert.False(t, hostinfo.isPendingDeletion())
 	assert.Contains(t, nc.hostMap.Hosts, hostinfo.vpnAddrs[0])
 	assert.Contains(t, nc.hostMap.Indexes, hostinfo.localIndexId)
-	assert.True(t, hostinfo.out.Load())
-	assert.True(t, hostinfo.in.Load())
+	assert.True(t, hostinfo.sentSinceCheck())
+	assert.True(t, (hostinfo.state.Load()&stateIn != 0))
 
 	// Do a traffic check tick, should not be pending deletion but should not have any in/out packets recorded
 	nc.doTrafficCheck(hostinfo.localIndexId, p, nb, out, time.Now())
-	assert.False(t, hostinfo.pendingDeletion.Load())
-	assert.False(t, hostinfo.out.Load())
-	assert.False(t, hostinfo.in.Load())
+	assert.False(t, hostinfo.isPendingDeletion())
+	assert.False(t, hostinfo.sentSinceCheck())
+	assert.False(t, (hostinfo.state.Load()&stateIn != 0))
 
 	// Do another traffic check tick, this host should be pending deletion now
 	nc.Out(hostinfo)
-	assert.True(t, hostinfo.out.Load())
+	assert.True(t, hostinfo.sentSinceCheck())
 	nc.doTrafficCheck(hostinfo.localIndexId, p, nb, out, time.Now())
-	assert.True(t, hostinfo.pendingDeletion.Load())
-	assert.False(t, hostinfo.out.Load())
-	assert.False(t, hostinfo.in.Load())
+	assert.True(t, hostinfo.isPendingDeletion())
+	assert.False(t, hostinfo.sentSinceCheck())
+	assert.False(t, (hostinfo.state.Load()&stateIn != 0))
 	assert.Contains(t, nc.hostMap.Indexes, hostinfo.localIndexId)
 	assert.Contains(t, nc.hostMap.Hosts, hostinfo.vpnAddrs[0])
 
@@ -168,33 +168,33 @@ func Test_NewConnectionManagerTest2(t *testing.T) {
 	// We saw traffic out to vpnIp
 	nc.Out(hostinfo)
 	nc.In(hostinfo)
-	assert.True(t, hostinfo.in.Load())
-	assert.True(t, hostinfo.out.Load())
-	assert.False(t, hostinfo.pendingDeletion.Load())
+	assert.True(t, (hostinfo.state.Load()&stateIn != 0))
+	assert.True(t, hostinfo.sentSinceCheck())
+	assert.False(t, hostinfo.isPendingDeletion())
 	assert.Contains(t, nc.hostMap.Hosts, hostinfo.vpnAddrs[0])
 	assert.Contains(t, nc.hostMap.Indexes, hostinfo.localIndexId)
 
 	// Do a traffic check tick, should not be pending deletion but should not have any in/out packets recorded
 	nc.doTrafficCheck(hostinfo.localIndexId, p, nb, out, time.Now())
-	assert.False(t, hostinfo.pendingDeletion.Load())
-	assert.False(t, hostinfo.out.Load())
-	assert.False(t, hostinfo.in.Load())
+	assert.False(t, hostinfo.isPendingDeletion())
+	assert.False(t, hostinfo.sentSinceCheck())
+	assert.False(t, (hostinfo.state.Load()&stateIn != 0))
 
 	// Do another traffic check tick, this host should be pending deletion now
 	nc.Out(hostinfo)
 	nc.doTrafficCheck(hostinfo.localIndexId, p, nb, out, time.Now())
-	assert.True(t, hostinfo.pendingDeletion.Load())
-	assert.False(t, hostinfo.out.Load())
-	assert.False(t, hostinfo.in.Load())
+	assert.True(t, hostinfo.isPendingDeletion())
+	assert.False(t, hostinfo.sentSinceCheck())
+	assert.False(t, (hostinfo.state.Load()&stateIn != 0))
 	assert.Contains(t, nc.hostMap.Indexes, hostinfo.localIndexId)
 	assert.Contains(t, nc.hostMap.Hosts, hostinfo.vpnAddrs[0])
 
 	// We saw traffic, should no longer be pending deletion
 	nc.In(hostinfo)
 	nc.doTrafficCheck(hostinfo.localIndexId, p, nb, out, time.Now())
-	assert.False(t, hostinfo.pendingDeletion.Load())
-	assert.False(t, hostinfo.out.Load())
-	assert.False(t, hostinfo.in.Load())
+	assert.False(t, hostinfo.isPendingDeletion())
+	assert.False(t, hostinfo.sentSinceCheck())
+	assert.False(t, (hostinfo.state.Load()&stateIn != 0))
 	assert.Contains(t, nc.hostMap.Indexes, hostinfo.localIndexId)
 	assert.Contains(t, nc.hostMap.Hosts, hostinfo.vpnAddrs[0])
 }
@@ -253,31 +253,31 @@ func Test_NewConnectionManager_DisconnectInactive(t *testing.T) {
 	// Do a traffic check tick, in and out should be cleared but should not be pending deletion
 	nc.Out(hostinfo)
 	nc.In(hostinfo)
-	assert.True(t, hostinfo.out.Load())
-	assert.True(t, hostinfo.in.Load())
+	assert.True(t, hostinfo.sentSinceCheck())
+	assert.True(t, (hostinfo.state.Load()&stateIn != 0))
 
 	now := time.Now()
 	decision, _, _ := nc.makeTrafficDecision(hostinfo.localIndexId, now)
 	assert.Equal(t, tryRehandshake, decision)
 	assert.Equal(t, now, hostinfo.lastUsed)
-	assert.False(t, hostinfo.pendingDeletion.Load())
-	assert.False(t, hostinfo.out.Load())
-	assert.False(t, hostinfo.in.Load())
+	assert.False(t, hostinfo.isPendingDeletion())
+	assert.False(t, hostinfo.sentSinceCheck())
+	assert.False(t, (hostinfo.state.Load()&stateIn != 0))
 
 	decision, _, _ = nc.makeTrafficDecision(hostinfo.localIndexId, now.Add(time.Second*5))
 	assert.Equal(t, doNothing, decision)
 	assert.Equal(t, now, hostinfo.lastUsed)
-	assert.False(t, hostinfo.pendingDeletion.Load())
-	assert.False(t, hostinfo.out.Load())
-	assert.False(t, hostinfo.in.Load())
+	assert.False(t, hostinfo.isPendingDeletion())
+	assert.False(t, hostinfo.sentSinceCheck())
+	assert.False(t, (hostinfo.state.Load()&stateIn != 0))
 
 	// Do another traffic check tick, should still not be pending deletion
 	decision, _, _ = nc.makeTrafficDecision(hostinfo.localIndexId, now.Add(time.Second*10))
 	assert.Equal(t, doNothing, decision)
 	assert.Equal(t, now, hostinfo.lastUsed)
-	assert.False(t, hostinfo.pendingDeletion.Load())
-	assert.False(t, hostinfo.out.Load())
-	assert.False(t, hostinfo.in.Load())
+	assert.False(t, hostinfo.isPendingDeletion())
+	assert.False(t, hostinfo.sentSinceCheck())
+	assert.False(t, (hostinfo.state.Load()&stateIn != 0))
 	assert.Contains(t, nc.hostMap.Indexes, hostinfo.localIndexId)
 	assert.Contains(t, nc.hostMap.Hosts, hostinfo.vpnAddrs[0])
 
@@ -285,9 +285,9 @@ func Test_NewConnectionManager_DisconnectInactive(t *testing.T) {
 	decision, _, _ = nc.makeTrafficDecision(hostinfo.localIndexId, now.Add(time.Minute*10))
 	assert.Equal(t, closeTunnel, decision)
 	assert.Equal(t, now, hostinfo.lastUsed)
-	assert.False(t, hostinfo.pendingDeletion.Load())
-	assert.False(t, hostinfo.out.Load())
-	assert.False(t, hostinfo.in.Load())
+	assert.False(t, hostinfo.isPendingDeletion())
+	assert.False(t, hostinfo.sentSinceCheck())
+	assert.False(t, (hostinfo.state.Load()&stateIn != 0))
 	assert.Contains(t, nc.hostMap.Indexes, hostinfo.localIndexId)
 	assert.Contains(t, nc.hostMap.Hosts, hostinfo.vpnAddrs[0])
 }

control.go

@@ -212,7 +212,7 @@ func (c *Control) RebindUDPServer() {
 	c.f.lightHouse.SendUpdate()
 
 	// Let the main interface know that we rebound so that underlying tunnels know to trigger punches from their remotes
-	c.f.rebindCount++
+	c.f.rebindEpoch.Add(1)
 }
 
 // ListHostmapHosts returns details about the actual or pending (handshaking) hostmap by vpn ip

hostmap.go

@@ -262,11 +262,6 @@ type HostInfo struct {
 	// This is used to limit lighthouse re-queries in chatty clients
 	nextLHQuery atomic.Int64
 
-	// lastRebindCount is the other side of Interface.rebindCount, if these values don't match then we need to ask LH
-	// for a punch from the remote end of this tunnel. The goal being to prime their conntrack for our traffic just like
-	// with a handshake
-	lastRebindCount int8
-
 	// lastHandshakeTime records the time the remote side told us about at the stage when the handshake was completed locally
 	// Stage 1 packet will contain it if I am a responder, stage 2 packet if I am an initiator
 	// This is used to avoid an attack where a handshake packet is replayed after some time
@@ -275,8 +270,11 @@ type HostInfo struct {
 	lastRoam       time.Time
 	lastRoamRemote netip.AddrPort
 
-	//TODO: in, out, and others might benefit from being an atomic.Int32. We could collapse connectionManager pendingDeletion, relayUsed, and in/out into this 1 thing
-	in, out, pendingDeletion atomic.Bool
+	// state holds everything the hot paths need to touch per packet, in one word: whether we have seen traffic
+	// each way since the connection manager last looked, whether it has given up on us, and the
+	// Interface.rebindEpoch this tunnel last sent under. Keeping the epoch here means it survives the traffic
+	// bits being cleared, so a tunnel that has not sent since a rebind still notices when it does.
+	state atomic.Uint32
 
 	// lastUsed tracks the last time ConnectionManager checked the tunnel and it was in use.
 	// This value will be behind against actual tunnel utilization in the hot path.
@@ -658,7 +656,7 @@ func (hm *HostMap) unlockedAddHostInfo(hostinfo *HostInfo, f *Interface) {
 	hm.Indexes[hostinfo.localIndexId] = hostinfo
 	hm.RemoteIndexes[hostinfo.remoteIndexId] = hostinfo
 
-	hostinfo.out.Store(true)
+	hostinfo.markOut(f.rebindEpoch.Load())
 	if f.connectionManager != nil { // f.connectionManager is only nil in some unit tests
 		f.connectionManager.trafficTimer.Add(hostinfo.localIndexId, f.connectionManager.checkInterval)
 	}
@@ -759,6 +757,64 @@ func (i *HostInfo) TryPromoteBest(preferredRanges []netip.Prefix, ifce *Interfac
 	}
 }
 
+// Bits within HostInfo.state. Everything above stateEpochShift is the rebind epoch.
+const (
+	stateIn uint32 = 1 << iota
+	stateOut
+	statePendingDeletion
+
+	stateFlags      = stateIn | stateOut | statePendingDeletion
+	stateEpochShift = 3
+)
+
+// markIn records inbound traffic. Reading first keeps the cache line shared on the common path, where the bit
+// is already set.
+func (i *HostInfo) markIn() {
+	if i.state.Load()&stateIn == 0 {
+		i.state.Or(stateIn)
+	}
+}
+
+// markOut records that we sent on this tunnel under the given rebind epoch. It reports whether the epoch moved
+// since our last send, which means the local network changed and we want the far side to punch at us again.
+// The common path is a single load that matches and returns.
+func (i *HostInfo) markOut(epoch uint32) bool {
+	e := epoch << stateEpochShift
+	for {
+		old := i.state.Load()
+		if old&stateOut != 0 && old&^stateFlags == e {
+			return false
+		}
+
+		if i.state.CompareAndSwap(old, old&stateFlags|stateOut|e) {
+			return old&^stateFlags != e
+		}
+	}
+}
+
+// sentSinceCheck reports whether anything has been sent since the connection manager last looked.
+func (i *HostInfo) sentSinceCheck() bool {
+	return i.state.Load()&stateOut != 0
+}
+
+// takeTraffic clears both traffic bits, leaving the epoch alone, and reports what they were.
+func (i *HostInfo) takeTraffic() (in bool, out bool) {
+	old := i.state.And(^(stateIn | stateOut))
+	return old&stateIn != 0, old&stateOut != 0
+}
+
+func (i *HostInfo) setPendingDeletion(v bool) {
+	if v {
+		i.state.Or(statePendingDeletion)
+	} else {
+		i.state.And(^statePendingDeletion)
+	}
+}
+
+func (i *HostInfo) isPendingDeletion() bool {
+	return i.state.Load()&statePendingDeletion != 0
+}
+
 func (i *HostInfo) GetCert() *cert.CachedCertificate {
 	if i.ConnectionState != nil {
 		return i.ConnectionState.peerCert

inside.go

@@ -365,17 +365,12 @@ func (f *Interface) sendNoMetrics(t header.MessageType, st header.MessageSubType
 
 	//l.WithField("trace", string(debug.Stack())).Error("out Header ", &Header{Version, t, st, 0, hostinfo.remoteIndexId, c}, p)
 	out = header.Encode(out, header.Version, t, st, hostinfo.remoteIndexId, c)
-	f.connectionManager.Out(hostinfo)
-
-	// Query our LH if we haven't since the last time we've been rebound, this will cause the remote to punch against
-	// all our addrs and enable a faster roaming.
-	if t != header.CloseTunnel && hostinfo.lastRebindCount != f.rebindCount {
-		//NOTE: there is an update hole if a tunnel isn't used and exactly 256 rebinds occur before the tunnel is
-		// finally used again. This tunnel would eventually be torn down and recreated if this action didn't help.
+	// We rebound since this tunnel last sent, so the local network moved. Ask the lighthouse to have the far side
+	// punch at where we are now, which primes their conntrack the same way a handshake would.
+	if f.connectionManager.Out(hostinfo) && t != header.CloseTunnel {
 		f.lightHouse.QueryServer(hostinfo.vpnAddrs[0])
-		hostinfo.lastRebindCount = f.rebindCount
 		if f.l.Enabled(context.Background(), slog.LevelDebug) {
-			f.l.Debug("Lighthouse update triggered for punch due to rebind counter",
+			f.l.Debug("Lighthouse update triggered for punch due to rebind epoch",
 				"vpnAddrs", hostinfo.vpnAddrs,
 			)
 		}

interface.go

@@ -82,8 +82,10 @@ type Interface struct {
 	sendRecvErrorConfig   recvErrorConfig
 	acceptRecvErrorConfig recvErrorConfig
 
-	// rebindCount is used to decide if an active tunnel should trigger a punch notification through a lighthouse
-	rebindCount int8
+	// rebindEpoch bumps every time the udp listener is rebound, which means the local network moved. Tunnels
+	// compare it against their own copy to decide they need a punch from the far side. Read on every send, only
+	// written on a rebind, so the cache line stays shared across the routines.
+	rebindEpoch atomic.Uint32
 	version     string
 
 	conntrackCacheTimeout time.Duration