TUN-2921: Rework address selection logic to avoid corner cases

edgediscovery/allregions/region.go

@@ -2,6 +2,8 @@ package allregions
 
 import (
 	"net"
+
+	"github.com/sirupsen/logrus"
 )
 
 // Region contains cloudflared edge addresses. The edge is partitioned into several regions for
@@ -56,6 +58,10 @@ func (r Region) GetUnusedIP(excluding *net.TCPAddr) *net.TCPAddr {
 
 // Use the address, assigning it to a proxy connection.
 func (r Region) Use(addr *net.TCPAddr, connID int) {
+	if addr == nil {
+		logrus.Errorf("Attempted to use nil address for connection %d", connID)
+		return
+	}
 	r.connFor[addr] = InUse(connID)
 }
 

edgediscovery/allregions/regions.go

@@ -86,21 +86,28 @@ func (rs *Regions) AddrUsedBy(connID int) *net.TCPAddr {
 // GetUnusedAddr gets an unused addr from the edge, excluding the given addr. Prefer to use addresses
 // evenly across both regions.
 func (rs *Regions) GetUnusedAddr(excluding *net.TCPAddr, connID int) *net.TCPAddr {
-	var addr *net.TCPAddr
 	if rs.region1.AvailableAddrs() > rs.region2.AvailableAddrs() {
-		addr = rs.region1.GetUnusedIP(excluding)
+		return getAddrs(excluding, connID, &rs.region1, &rs.region2)
-		rs.region1.Use(addr, connID)
-	} else {
-		addr = rs.region2.GetUnusedIP(excluding)
-		rs.region2.Use(addr, connID)
 	}
 
-	if addr == nil {
+	return getAddrs(excluding, connID, &rs.region2, &rs.region1)
-		return nil
+}
+
+// getAddrs tries to grab address form `first` region, then `second` region
+// this is an unrolled loop over 2 element array
+func getAddrs(excluding *net.TCPAddr, connID int, first *Region, second *Region) *net.TCPAddr {
+	addr := first.GetUnusedIP(excluding)
+	if addr != nil {
+		first.Use(addr, connID)
+		return addr
+	}
+	addr = second.GetUnusedIP(excluding)
+	if addr != nil {
+		second.Use(addr, connID)
+		return addr
 	}
 
-	// Mark the address as used and return it
+	return nil
-	return addr
 }
 
 // AvailableAddrs returns how many edge addresses aren't used.

edgediscovery/allregions/regions_test.go

@@ -58,6 +58,7 @@ func TestRegions_Giveback_Region1(t *testing.T) {
 	rs.GiveBack(&addr0)
 	assert.Equal(t, &addr0, rs.GetUnusedAddr(nil, 3))
 }
+
 func TestRegions_Giveback_Region2(t *testing.T) {
 	rs := makeRegions()
 	rs.region1.Use(&addr0, 0)

edgediscovery/edgediscovery.go

@@ -100,7 +100,7 @@ func (ed *Edge) GetAddr(connID int) (*net.TCPAddr, error) {
 		logger.Debug("No addresses left to give proxy connection")
 		return nil, errNoAddressesLeft
 	}
-	logger.Debug("Giving connection its new address")
+	logger.Debugf("Giving connection its new address %s", addr)
 	return addr, nil
 }
 
@@ -120,9 +120,10 @@ func (ed *Edge) GetDifferentAddr(connID int) (*net.TCPAddr, error) {
 	addr := ed.regions.GetUnusedAddr(oldAddr, connID)
 	if addr == nil {
 		logger.Debug("No addresses left to give proxy connection")
+		// note: if oldAddr were not nil, it will become available on the next iteration
 		return nil, errNoAddressesLeft
 	}
-	logger.Debug("Giving connection its new address")
+	logger.Debugf("Giving connection its new address %s", addr)
 	return addr, nil
 }
 

edgediscovery/edgediscovery_test.go

@@ -83,6 +83,30 @@ func TestGetAddrForRPC(t *testing.T) {
 	assert.Equal(t, 4, edge.AvailableAddrs())
 }
 
+func TestOnePerRegion(t *testing.T) {
+	l := logrus.New()
+
+	// Make an edge with only one address
+	edge := MockEdge(l, []*net.TCPAddr{&addr0, &addr1})
+
+	// Use the only address
+	const connID = 0
+	a1, err := edge.GetAddr(connID)
+	assert.NoError(t, err)
+	assert.NotNil(t, a1)
+
+	// if the first address is bad, get the second one
+	a2, err := edge.GetDifferentAddr(connID)
+	assert.NoError(t, err)
+	assert.NotNil(t, a2)
+	assert.NotEqual(t, a1, a2)
+
+	// now that second one is bad, get the first one again
+	a3, err := edge.GetDifferentAddr(connID)
+	assert.NoError(t, err)
+	assert.Equal(t, a1, a3)
+}
+
 func TestOnlyOneAddrLeft(t *testing.T) {
 	l := logrus.New()
 
@@ -98,6 +122,11 @@ func TestOnlyOneAddrLeft(t *testing.T) {
 	// If that edge address is "bad", there's no alternative address.
 	_, err = edge.GetDifferentAddr(connID)
 	assert.Error(t, err)
+
+	// previously bad address should become available again on next iteration.
+	addr, err = edge.GetDifferentAddr(connID)
+	assert.NoError(t, err)
+	assert.NotNil(t, addr)
 }
 
 func TestNoAddrsLeft(t *testing.T) {

origin/supervisor.go

@@ -274,7 +274,7 @@ func (s *Supervisor) startTunnel(ctx context.Context, index int, connectedSignal
 		s.tunnelErrors <- tunnelError{index: index, addr: addr, err: err}
 	}()
 
-	addr, err = s.edgeIPs.GetAddr(index)
+	addr, err = s.edgeIPs.GetDifferentAddr(index)
 	if err != nil {
 		return
 	}