TUN-6007: Implement new edge discovery algorithm

edgediscovery/allregions/address.go

@@ -0,0 +1,64 @@
+package allregions
+
+// Region contains cloudflared edge addresses. The edge is partitioned into several regions for
+// redundancy purposes.
+type AddrSet map[*EdgeAddr]UsedBy
+
+// AddrUsedBy finds the address used by the given connection in this region.
+// Returns nil if the connection isn't using any IP.
+func (a AddrSet) AddrUsedBy(connID int) *EdgeAddr {
+	for addr, used := range a {
+		if used.Used && used.ConnID == connID {
+			return addr
+		}
+	}
+	return nil
+}
+
+// AvailableAddrs counts how many unused addresses this region contains.
+func (a AddrSet) AvailableAddrs() int {
+	n := 0
+	for _, usedby := range a {
+		if !usedby.Used {
+			n++
+		}
+	}
+	return n
+}
+
+// GetUnusedIP returns a random unused address in this region.
+// Returns nil if all addresses are in use.
+func (a AddrSet) GetUnusedIP(excluding *EdgeAddr) *EdgeAddr {
+	for addr, usedby := range a {
+		if !usedby.Used && addr != excluding {
+			return addr
+		}
+	}
+	return nil
+}
+
+// Use the address, assigning it to a proxy connection.
+func (a AddrSet) Use(addr *EdgeAddr, connID int) {
+	if addr == nil {
+		return
+	}
+	a[addr] = InUse(connID)
+}
+
+// GetAnyAddress returns an arbitrary address from the region.
+func (a AddrSet) GetAnyAddress() *EdgeAddr {
+	for addr := range a {
+		return addr
+	}
+	return nil
+}
+
+// GiveBack the address, ensuring it is no longer assigned to an IP.
+// Returns true if the address is in this region.
+func (a AddrSet) GiveBack(addr *EdgeAddr) (ok bool) {
+	if _, ok := a[addr]; !ok {
+		return false
+	}
+	a[addr] = Unused()
+	return true
+}

edgediscovery/allregions/address_test.go

@@ -0,0 +1,247 @@
+package allregions
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestAddrSet_AddrUsedBy(t *testing.T) {
+	type args struct {
+		connID int
+	}
+	tests := []struct {
+		name    string
+		addrSet AddrSet
+		args    args
+		want    *EdgeAddr
+	}{
+		{
+			name: "happy trivial test",
+			addrSet: AddrSet{
+				&addr0: InUse(0),
+			},
+			args: args{connID: 0},
+			want: &addr0,
+		},
+		{
+			name: "sad trivial test",
+			addrSet: AddrSet{
+				&addr0: InUse(0),
+			},
+			args: args{connID: 1},
+			want: nil,
+		},
+		{
+			name: "sad test",
+			addrSet: AddrSet{
+				&addr0: InUse(0),
+				&addr1: InUse(1),
+				&addr2: InUse(2),
+			},
+			args: args{connID: 3},
+			want: nil,
+		},
+		{
+			name: "happy test",
+			addrSet: AddrSet{
+				&addr0: InUse(0),
+				&addr1: InUse(1),
+				&addr2: InUse(2),
+			},
+			args: args{connID: 1},
+			want: &addr1,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := tt.addrSet.AddrUsedBy(tt.args.connID); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("Region.AddrUsedBy() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestAddrSet_AvailableAddrs(t *testing.T) {
+	tests := []struct {
+		name    string
+		addrSet AddrSet
+		want    int
+	}{
+		{
+			name: "contains addresses",
+			addrSet: AddrSet{
+				&addr0: InUse(0),
+				&addr1: Unused(),
+				&addr2: InUse(2),
+			},
+			want: 1,
+		},
+		{
+			name: "all free",
+			addrSet: AddrSet{
+				&addr0: Unused(),
+				&addr1: Unused(),
+				&addr2: Unused(),
+			},
+			want: 3,
+		},
+		{
+			name: "all used",
+			addrSet: AddrSet{
+				&addr0: InUse(0),
+				&addr1: InUse(1),
+				&addr2: InUse(2),
+			},
+			want: 0,
+		},
+		{
+			name:    "empty",
+			addrSet: AddrSet{},
+			want:    0,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := tt.addrSet.AvailableAddrs(); got != tt.want {
+				t.Errorf("Region.AvailableAddrs() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestAddrSet_GetUnusedIP(t *testing.T) {
+	type args struct {
+		excluding *EdgeAddr
+	}
+	tests := []struct {
+		name    string
+		addrSet AddrSet
+		args    args
+		want    *EdgeAddr
+	}{
+		{
+			name: "happy test with excluding set",
+			addrSet: AddrSet{
+				&addr0: Unused(),
+				&addr1: Unused(),
+				&addr2: InUse(2),
+			},
+			args: args{excluding: &addr0},
+			want: &addr1,
+		},
+		{
+			name: "happy test with no excluding",
+			addrSet: AddrSet{
+				&addr0: InUse(0),
+				&addr1: Unused(),
+				&addr2: InUse(2),
+			},
+			args: args{excluding: nil},
+			want: &addr1,
+		},
+		{
+			name: "sad test with no excluding",
+			addrSet: AddrSet{
+				&addr0: InUse(0),
+				&addr1: InUse(1),
+				&addr2: InUse(2),
+			},
+			args: args{excluding: nil},
+			want: nil,
+		},
+		{
+			name: "sad test with excluding",
+			addrSet: AddrSet{
+				&addr0: Unused(),
+				&addr1: InUse(1),
+				&addr2: InUse(2),
+			},
+			args: args{excluding: &addr0},
+			want: nil,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := tt.addrSet.GetUnusedIP(tt.args.excluding); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("Region.GetUnusedIP() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestAddrSet_GiveBack(t *testing.T) {
+	type args struct {
+		addr *EdgeAddr
+	}
+	tests := []struct {
+		name           string
+		addrSet        AddrSet
+		args           args
+		wantOk         bool
+		availableAfter int
+	}{
+		{
+			name: "sad test with excluding",
+			addrSet: AddrSet{
+				&addr1: InUse(1),
+			},
+			args:           args{addr: &addr1},
+			wantOk:         true,
+			availableAfter: 1,
+		},
+		{
+			name: "sad test with excluding",
+			addrSet: AddrSet{
+				&addr1: InUse(1),
+			},
+			args:           args{addr: &addr2},
+			wantOk:         false,
+			availableAfter: 0,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if gotOk := tt.addrSet.GiveBack(tt.args.addr); gotOk != tt.wantOk {
+				t.Errorf("Region.GiveBack() = %v, want %v", gotOk, tt.wantOk)
+			}
+			if tt.availableAfter != tt.addrSet.AvailableAddrs() {
+				t.Errorf("Region.AvailableAddrs() = %v, want %v", tt.addrSet.AvailableAddrs(), tt.availableAfter)
+			}
+		})
+	}
+}
+
+func TestAddrSet_GetAnyAddress(t *testing.T) {
+	tests := []struct {
+		name    string
+		addrSet AddrSet
+		wantNil bool
+	}{
+		{
+			name:    "Sad test -- GetAnyAddress should only fail if the region is empty",
+			addrSet: AddrSet{},
+			wantNil: true,
+		},
+		{
+			name: "Happy test (all addresses unused)",
+			addrSet: AddrSet{
+				&addr0: Unused(),
+			},
+			wantNil: false,
+		},
+		{
+			name: "Happy test (GetAnyAddress can still return addresses used by proxy conns)",
+			addrSet: AddrSet{
+				&addr0: InUse(2),
+			},
+			wantNil: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := tt.addrSet.GetAnyAddress(); tt.wantNil != (got == nil) {
+				t.Errorf("Region.GetAnyAddress() = %v, but should it return nil? %v", got, tt.wantNil)
+			}
+		})
+	}
+}

edgediscovery/allregions/discovery.go

@@ -13,7 +13,7 @@ import (
 
 const (
 	// Used to discover HA origintunneld servers
-	srvService = "origintunneld"
+	srvService = "v2-origintunneld"
 	srvProto   = "tcp"
 	srvName    = "argotunnel.com"
 
@@ -115,6 +115,9 @@ func edgeDiscovery(log *zerolog.Logger, srvService string) ([][]*EdgeAddr, error
 		if err != nil {
 			return nil, err
 		}
+		for _, e := range edgeAddrs {
+			log.Debug().Msgf("Edge Address: %+v", *e)
+		}
 		resolvedAddrPerCNAME = append(resolvedAddrPerCNAME, edgeAddrs)
 	}
 
@@ -187,7 +190,6 @@ func ResolveAddrs(addrs []string, log *zerolog.Logger) (resolved []*EdgeAddr) {
 			UDP:       udpAddr,
 			IPVersion: version,
 		})
-
 	}
 	return
 }

edgediscovery/allregions/mocks_for_test.go

@@ -9,6 +9,115 @@ import (
 	"testing/quick"
 )
 
+var (
+	v4Addrs = []*EdgeAddr{&addr0, &addr1, &addr2, &addr3}
+	v6Addrs = []*EdgeAddr{&addr4, &addr5, &addr6, &addr7}
+	addr0   = EdgeAddr{
+		TCP: &net.TCPAddr{
+			IP:   net.ParseIP("123.4.5.0"),
+			Port: 8000,
+			Zone: "",
+		},
+		UDP: &net.UDPAddr{
+			IP:   net.ParseIP("123.4.5.0"),
+			Port: 8000,
+			Zone: "",
+		},
+		IPVersion: V4,
+	}
+	addr1 = EdgeAddr{
+		TCP: &net.TCPAddr{
+			IP:   net.ParseIP("123.4.5.1"),
+			Port: 8000,
+			Zone: "",
+		},
+		UDP: &net.UDPAddr{
+			IP:   net.ParseIP("123.4.5.1"),
+			Port: 8000,
+			Zone: "",
+		},
+		IPVersion: V4,
+	}
+	addr2 = EdgeAddr{
+		TCP: &net.TCPAddr{
+			IP:   net.ParseIP("123.4.5.2"),
+			Port: 8000,
+			Zone: "",
+		},
+		UDP: &net.UDPAddr{
+			IP:   net.ParseIP("123.4.5.2"),
+			Port: 8000,
+			Zone: "",
+		},
+		IPVersion: V4,
+	}
+	addr3 = EdgeAddr{
+		TCP: &net.TCPAddr{
+			IP:   net.ParseIP("123.4.5.3"),
+			Port: 8000,
+			Zone: "",
+		},
+		UDP: &net.UDPAddr{
+			IP:   net.ParseIP("123.4.5.3"),
+			Port: 8000,
+			Zone: "",
+		},
+		IPVersion: V4,
+	}
+	addr4 = EdgeAddr{
+		TCP: &net.TCPAddr{
+			IP:   net.ParseIP("2606:4700:a0::1"),
+			Port: 8000,
+			Zone: "",
+		},
+		UDP: &net.UDPAddr{
+			IP:   net.ParseIP("2606:4700:a0::1"),
+			Port: 8000,
+			Zone: "",
+		},
+		IPVersion: V6,
+	}
+	addr5 = EdgeAddr{
+		TCP: &net.TCPAddr{
+			IP:   net.ParseIP("2606:4700:a0::2"),
+			Port: 8000,
+			Zone: "",
+		},
+		UDP: &net.UDPAddr{
+			IP:   net.ParseIP("2606:4700:a0::2"),
+			Port: 8000,
+			Zone: "",
+		},
+		IPVersion: V6,
+	}
+	addr6 = EdgeAddr{
+		TCP: &net.TCPAddr{
+			IP:   net.ParseIP("2606:4700:a0::3"),
+			Port: 8000,
+			Zone: "",
+		},
+		UDP: &net.UDPAddr{
+			IP:   net.ParseIP("2606:4700:a0::3"),
+			Port: 8000,
+			Zone: "",
+		},
+		IPVersion: V6,
+	}
+	addr7 = EdgeAddr{
+		TCP: &net.TCPAddr{
+			IP:   net.ParseIP("2606:4700:a0::4"),
+			Port: 8000,
+			Zone: "",
+		},
+		UDP: &net.UDPAddr{
+			IP:   net.ParseIP("2606:4700:a0::4"),
+			Port: 8000,
+			Zone: "",
+		},
+		IPVersion: V6,
+	}
+)
+
 type mockAddrs struct {
 	// a set of synthetic SRV records
 	addrMap map[net.SRV][]*EdgeAddr

edgediscovery/allregions/region.go

@@ -1,79 +1,155 @@
 package allregions
 
+import "time"
+
+const (
+	timeoutDuration = 10 * time.Minute
+)
+
 // Region contains cloudflared edge addresses. The edge is partitioned into several regions for
 // redundancy purposes.
 type Region struct {
-	connFor map[*EdgeAddr]UsedBy
+	primaryIsActive bool
+	active          AddrSet
+	primary         AddrSet
+	secondary       AddrSet
+	primaryTimeout  time.Time
+	timeoutDuration time.Duration
 }
 
 // NewRegion creates a region with the given addresses, which are all unused.
-func NewRegion(addrs []*EdgeAddr) Region {
+func NewRegion(addrs []*EdgeAddr, overrideIPVersion ConfigIPVersion) Region {
 	// The zero value of UsedBy is Unused(), so we can just initialize the map's values with their
 	// zero values.
-	connFor := make(map[*EdgeAddr]UsedBy)
-	for _, addr := range addrs {
-		connFor[addr] = Unused()
+	connForv4 := make(AddrSet)
+	connForv6 := make(AddrSet)
+	systemPreference := V6
+	for i, addr := range addrs {
+		if i == 0 {
+			// First family of IPs returned is system preference of IP
+			systemPreference = addr.IPVersion
+		}
+		switch addr.IPVersion {
+		case V4:
+			connForv4[addr] = Unused()
+		case V6:
+			connForv6[addr] = Unused()
+		}
 	}
+
+	// Process as system preference
+	var primary AddrSet
+	var secondary AddrSet
+	switch systemPreference {
+	case V4:
+		primary = connForv4
+		secondary = connForv6
+	case V6:
+		primary = connForv6
+		secondary = connForv4
+	}
+
+	// Override with provided preference
+	switch overrideIPVersion {
+	case IPv4Only:
+		primary = connForv4
+		secondary = make(AddrSet) // empty
+	case IPv6Only:
+		primary = connForv6
+		secondary = make(AddrSet) // empty
+	case Auto:
+		// no change
+	default:
+		// no change
+	}
+
 	return Region{
-		connFor: connFor,
+		primaryIsActive: true,
+		active:          primary,
+		primary:         primary,
+		secondary:       secondary,
+		timeoutDuration: timeoutDuration,
 	}
 }
 
 // AddrUsedBy finds the address used by the given connection in this region.
 // Returns nil if the connection isn't using any IP.
 func (r *Region) AddrUsedBy(connID int) *EdgeAddr {
-	for addr, used := range r.connFor {
-		if used.Used && used.ConnID == connID {
-			return addr
-		}
+	edgeAddr := r.primary.AddrUsedBy(connID)
+	if edgeAddr == nil {
+		edgeAddr = r.secondary.AddrUsedBy(connID)
 	}
-	return nil
+	return edgeAddr
 }
 
 // AvailableAddrs counts how many unused addresses this region contains.
 func (r Region) AvailableAddrs() int {
-	n := 0
-	for _, usedby := range r.connFor {
-		if !usedby.Used {
-			n++
-		}
-	}
-	return n
+	return r.active.AvailableAddrs()
 }
 
-// GetUnusedIP returns a random unused address in this region.
-// Returns nil if all addresses are in use.
-func (r Region) GetUnusedIP(excluding *EdgeAddr) *EdgeAddr {
-	for addr, usedby := range r.connFor {
-		if !usedby.Used && addr != excluding {
-			return addr
-		}
-	}
-	return nil
-}
-
-// Use the address, assigning it to a proxy connection.
-func (r Region) Use(addr *EdgeAddr, connID int) {
-	if addr == nil {
-		return
-	}
-	r.connFor[addr] = InUse(connID)
-}
-
-// GetAnyAddress returns an arbitrary address from the region.
-func (r Region) GetAnyAddress() *EdgeAddr {
-	for addr := range r.connFor {
+// AssignAnyAddress returns a random unused address in this region now
+// assigned to the connID excluding the provided EdgeAddr.
+// Returns nil if all addresses are in use for the region.
+func (r Region) AssignAnyAddress(connID int, excluding *EdgeAddr) *EdgeAddr {
+	if addr := r.active.GetUnusedIP(excluding); addr != nil {
+		r.active.Use(addr, connID)
 		return addr
 	}
 	return nil
 }
 
+// GetAnyAddress returns an arbitrary address from the region.
+func (r Region) GetAnyAddress() *EdgeAddr {
+	return r.active.GetAnyAddress()
+}
+
 // GiveBack the address, ensuring it is no longer assigned to an IP.
 // Returns true if the address is in this region.
-func (r Region) GiveBack(addr *EdgeAddr) (ok bool) {
-	if _, ok := r.connFor[addr]; !ok {
-		return false
+func (r *Region) GiveBack(addr *EdgeAddr, hasConnectivityError bool) (ok bool) {
+	if ok = r.primary.GiveBack(addr); !ok {
+		// Attempt to give back the address in the secondary set
+		if ok = r.secondary.GiveBack(addr); !ok {
+			// Address is not in this region
+			return
+		}
 	}
-	r.connFor[addr] = Unused()
-	return true
+
+	// No connectivity error: no worry
+	if !hasConnectivityError {
+		return
+	}
+
+	// If using primary and returned address is IPv6 and secondary is available
+	if r.primaryIsActive && addr.IPVersion == V6 && len(r.secondary) > 0 {
+		r.active = r.secondary
+		r.primaryIsActive = false
+		r.primaryTimeout = time.Now().Add(r.timeoutDuration)
+		return
+	}
+
+	// Do nothing for IPv4 or if secondary is empty
+	if r.primaryIsActive {
+		return
+	}
+
+	// Immediately return to primary pool, regardless of current primary timeout
+	if addr.IPVersion == V4 {
+		activatePrimary(r)
+		return
+	}
+
+	// Timeout exceeded and can be reset to primary pool
+	if r.primaryTimeout.Before(time.Now()) {
+		activatePrimary(r)
+		return
+	}
+
+	return
+}
+
+// activatePrimary sets the primary set to the active set and resets the timeout.
+func activatePrimary(r *Region) {
+	r.active = r.primary
+	r.primaryIsActive = true
+	r.primaryTimeout = time.Now() // reset timeout
 }

edgediscovery/allregions/region_test.go

@@ -1,284 +1,357 @@
 package allregions
 
 import (
-	"reflect"
+	"net"
 	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
 )
 
+func makeAddrSet(addrs []*EdgeAddr) AddrSet {
+	addrSet := make(AddrSet, len(addrs))
+	for _, addr := range addrs {
+		addrSet[addr] = Unused()
+	}
+	return addrSet
+}
+
 func TestRegion_New(t *testing.T) {
-	r := NewRegion([]*EdgeAddr{&addr0, &addr1, &addr2})
-	if r.AvailableAddrs() != 3 {
-		t.Errorf("r.AvailableAddrs() == %v but want 3", r.AvailableAddrs())
-	}
-}
-
-func TestRegion_AddrUsedBy(t *testing.T) {
-	type fields struct {
-		connFor map[*EdgeAddr]UsedBy
-	}
-	type args struct {
-		connID int
-	}
 	tests := []struct {
-		name   string
-		fields fields
-		args   args
-		want   *EdgeAddr
+		name          string
+		addrs         []*EdgeAddr
+		mode          ConfigIPVersion
+		expectedAddrs int
+		primary       AddrSet
+		secondary     AddrSet
 	}{
 		{
-			name: "happy trivial test",
-			fields: fields{connFor: map[*EdgeAddr]UsedBy{
-				&addr0: InUse(0),
-			}},
-			args: args{connID: 0},
-			want: &addr0,
+			name:          "IPv4 addresses with IPv4Only",
+			addrs:         v4Addrs,
+			mode:          IPv4Only,
+			expectedAddrs: len(v4Addrs),
+			primary:       makeAddrSet(v4Addrs),
+			secondary:     AddrSet{},
 		},
 		{
-			name: "sad trivial test",
-			fields: fields{connFor: map[*EdgeAddr]UsedBy{
-				&addr0: InUse(0),
-			}},
-			args: args{connID: 1},
-			want: nil,
+			name:          "IPv6 addresses with IPv4Only",
+			addrs:         v6Addrs,
+			mode:          IPv4Only,
+			expectedAddrs: 0,
+			primary:       AddrSet{},
+			secondary:     AddrSet{},
 		},
 		{
-			name: "sad test",
-			fields: fields{connFor: map[*EdgeAddr]UsedBy{
-				&addr0: InUse(0),
-				&addr1: InUse(1),
-				&addr2: InUse(2),
-			}},
-			args: args{connID: 3},
-			want: nil,
+			name:          "IPv6 addresses with IPv6Only",
+			addrs:         v6Addrs,
+			mode:          IPv6Only,
+			expectedAddrs: len(v6Addrs),
+			primary:       makeAddrSet(v6Addrs),
+			secondary:     AddrSet{},
 		},
 		{
-			name: "happy test",
-			fields: fields{connFor: map[*EdgeAddr]UsedBy{
-				&addr0: InUse(0),
-				&addr1: InUse(1),
-				&addr2: InUse(2),
-			}},
-			args: args{connID: 1},
-			want: &addr1,
+			name:          "IPv6 addresses with IPv4Only",
+			addrs:         v6Addrs,
+			mode:          IPv4Only,
+			expectedAddrs: 0,
+			primary:       AddrSet{},
+			secondary:     AddrSet{},
+		},
+		{
+			name:          "IPv4 (first) and IPv6 addresses with Auto",
+			addrs:         append(v4Addrs, v6Addrs...),
+			mode:          Auto,
+			expectedAddrs: len(v4Addrs),
+			primary:       makeAddrSet(v4Addrs),
+			secondary:     makeAddrSet(v6Addrs),
+		},
+		{
+			name:          "IPv6 (first) and IPv4 addresses with Auto",
+			addrs:         append(v6Addrs, v4Addrs...),
+			mode:          Auto,
+			expectedAddrs: len(v6Addrs),
+			primary:       makeAddrSet(v6Addrs),
+			secondary:     makeAddrSet(v4Addrs),
+		},
+		{
+			name:          "IPv4 addresses with Auto",
+			addrs:         v4Addrs,
+			mode:          Auto,
+			expectedAddrs: len(v4Addrs),
+			primary:       makeAddrSet(v4Addrs),
+			secondary:     AddrSet{},
+		},
+		{
+			name:          "IPv6 addresses with Auto",
+			addrs:         v6Addrs,
+			mode:          Auto,
+			expectedAddrs: len(v6Addrs),
+			primary:       makeAddrSet(v6Addrs),
+			secondary:     AddrSet{},
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			r := &Region{
-				connFor: tt.fields.connFor,
+			r := NewRegion(tt.addrs, tt.mode)
+			assert.Equal(t, tt.expectedAddrs, r.AvailableAddrs())
+			assert.Equal(t, tt.primary, r.primary)
+			assert.Equal(t, tt.secondary, r.secondary)
+		})
+	}
+}
+
+func TestRegion_AnyAddress_EmptyActiveSet(t *testing.T) {
+	tests := []struct {
+		name  string
+		addrs []*EdgeAddr
+		mode  ConfigIPVersion
+	}{
+		{
+			name:  "IPv6 addresses with IPv4Only",
+			addrs: v6Addrs,
+			mode:  IPv4Only,
+		},
+		{
+			name:  "IPv4 addresses with IPv6Only",
+			addrs: v4Addrs,
+			mode:  IPv6Only,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			r := NewRegion(tt.addrs, tt.mode)
+			addr := r.GetAnyAddress()
+			assert.Nil(t, addr)
+			addr = r.AssignAnyAddress(0, nil)
+			assert.Nil(t, addr)
+		})
+	}
+}
+
+func TestRegion_AssignAnyAddress_FullyUsedActiveSet(t *testing.T) {
+	tests := []struct {
+		name  string
+		addrs []*EdgeAddr
+		mode  ConfigIPVersion
+	}{
+		{
+			name:  "IPv6 addresses with IPv6Only",
+			addrs: v6Addrs,
+			mode:  IPv6Only,
+		},
+		{
+			name:  "IPv4 addresses with IPv4Only",
+			addrs: v4Addrs,
+			mode:  IPv4Only,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			r := NewRegion(tt.addrs, tt.mode)
+			total := r.active.AvailableAddrs()
+			for i := 0; i < total; i++ {
+				addr := r.AssignAnyAddress(i, nil)
+				assert.NotNil(t, addr)
 			}
-			if got := r.AddrUsedBy(tt.args.connID); !reflect.DeepEqual(got, tt.want) {
-				t.Errorf("Region.AddrUsedBy() = %v, want %v", got, tt.want)
+			addr := r.AssignAnyAddress(9, nil)
+			assert.Nil(t, addr)
+		})
+	}
+}
+
+var giveBackTests = []struct {
+	name          string
+	addrs         []*EdgeAddr
+	mode          ConfigIPVersion
+	expectedAddrs int
+	primary       AddrSet
+	secondary     AddrSet
+	primarySwap   bool
+}{
+	{
+		name:          "IPv4 addresses with IPv4Only",
+		addrs:         v4Addrs,
+		mode:          IPv4Only,
+		expectedAddrs: len(v4Addrs),
+		primary:       makeAddrSet(v4Addrs),
+		secondary:     AddrSet{},
+		primarySwap:   false,
+	},
+	{
+		name:          "IPv6 addresses with IPv6Only",
+		addrs:         v6Addrs,
+		mode:          IPv6Only,
+		expectedAddrs: len(v6Addrs),
+		primary:       makeAddrSet(v6Addrs),
+		secondary:     AddrSet{},
+		primarySwap:   false,
+	},
+	{
+		name:          "IPv4 (first) and IPv6 addresses with Auto",
+		addrs:         append(v4Addrs, v6Addrs...),
+		mode:          Auto,
+		expectedAddrs: len(v4Addrs),
+		primary:       makeAddrSet(v4Addrs),
+		secondary:     makeAddrSet(v6Addrs),
+		primarySwap:   false,
+	},
+	{
+		name:          "IPv6 (first) and IPv4 addresses with Auto",
+		addrs:         append(v6Addrs, v4Addrs...),
+		mode:          Auto,
+		expectedAddrs: len(v6Addrs),
+		primary:       makeAddrSet(v6Addrs),
+		secondary:     makeAddrSet(v4Addrs),
+		primarySwap:   true,
+	},
+	{
+		name:          "IPv4 addresses with Auto",
+		addrs:         v4Addrs,
+		mode:          Auto,
+		expectedAddrs: len(v4Addrs),
+		primary:       makeAddrSet(v4Addrs),
+		secondary:     AddrSet{},
+		primarySwap:   false,
+	},
+	{
+		name:          "IPv6 addresses with Auto",
+		addrs:         v6Addrs,
+		mode:          Auto,
+		expectedAddrs: len(v6Addrs),
+		primary:       makeAddrSet(v6Addrs),
+		secondary:     AddrSet{},
+		primarySwap:   false,
+	},
+}
+
+func TestRegion_GiveBack_NoConnectivityError(t *testing.T) {
+	for _, tt := range giveBackTests {
+		t.Run(tt.name, func(t *testing.T) {
+			r := NewRegion(tt.addrs, tt.mode)
+			addr := r.AssignAnyAddress(0, nil)
+			assert.NotNil(t, addr)
+			assert.True(t, r.GiveBack(addr, false))
+		})
+	}
+}
+
+func TestRegion_GiveBack_ForeignAddr(t *testing.T) {
+	invalid := EdgeAddr{
+		TCP: &net.TCPAddr{
+			IP:   net.ParseIP("123.4.5.0"),
+			Port: 8000,
+			Zone: "",
+		},
+		UDP: &net.UDPAddr{
+			IP:   net.ParseIP("123.4.5.0"),
+			Port: 8000,
+			Zone: "",
+		},
+		IPVersion: V4,
+	}
+	for _, tt := range giveBackTests {
+		t.Run(tt.name, func(t *testing.T) {
+			r := NewRegion(tt.addrs, tt.mode)
+			assert.False(t, r.GiveBack(&invalid, false))
+			assert.False(t, r.GiveBack(&invalid, true))
+		})
+	}
+}
+
+func TestRegion_GiveBack_SwapPrimary(t *testing.T) {
+	for _, tt := range giveBackTests {
+		t.Run(tt.name, func(t *testing.T) {
+			r := NewRegion(tt.addrs, tt.mode)
+			addr := r.AssignAnyAddress(0, nil)
+			assert.NotNil(t, addr)
+			assert.True(t, r.GiveBack(addr, true))
+			assert.Equal(t, tt.primarySwap, !r.primaryIsActive)
+			if tt.primarySwap {
+				assert.Equal(t, r.secondary, r.active)
+				assert.False(t, r.primaryTimeout.IsZero())
+			} else {
+				assert.Equal(t, r.primary, r.active)
+				assert.True(t, r.primaryTimeout.IsZero())
 			}
 		})
 	}
 }
 
-func TestRegion_AvailableAddrs(t *testing.T) {
-	type fields struct {
-		connFor map[*EdgeAddr]UsedBy
-	}
-	tests := []struct {
-		name   string
-		fields fields
-		want   int
-	}{
-		{
-			name: "contains addresses",
-			fields: fields{connFor: map[*EdgeAddr]UsedBy{
-				&addr0: InUse(0),
-				&addr1: Unused(),
-				&addr2: InUse(2),
-			}},
-			want: 1,
-		},
-		{
-			name: "all free",
-			fields: fields{connFor: map[*EdgeAddr]UsedBy{
-				&addr0: Unused(),
-				&addr1: Unused(),
-				&addr2: Unused(),
-			}},
-			want: 3,
-		},
-		{
-			name: "all used",
-			fields: fields{connFor: map[*EdgeAddr]UsedBy{
-				&addr0: InUse(0),
-				&addr1: InUse(1),
-				&addr2: InUse(2),
-			}},
-			want: 0,
-		},
-		{
-			name:   "empty",
-			fields: fields{connFor: map[*EdgeAddr]UsedBy{}},
-			want:   0,
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			r := Region{
-				connFor: tt.fields.connFor,
-			}
-			if got := r.AvailableAddrs(); got != tt.want {
-				t.Errorf("Region.AvailableAddrs() = %v, want %v", got, tt.want)
-			}
-		})
-	}
+func TestRegion_GiveBack_IPv4_ResetPrimary(t *testing.T) {
+	r := NewRegion(append(v6Addrs, v4Addrs...), Auto)
+	// Exhaust all IPv6 addresses
+	a0 := r.AssignAnyAddress(0, nil)
+	a1 := r.AssignAnyAddress(1, nil)
+	a2 := r.AssignAnyAddress(2, nil)
+	a3 := r.AssignAnyAddress(3, nil)
+	assert.NotNil(t, a0)
+	assert.NotNil(t, a1)
+	assert.NotNil(t, a2)
+	assert.NotNil(t, a3)
+	// Give back the first IPv6 address to fallback to secondary IPv4 address set
+	assert.True(t, r.GiveBack(a0, true))
+	assert.False(t, r.primaryIsActive)
+	// Give back another IPv6 address
+	assert.True(t, r.GiveBack(a1, true))
+	// Primary shouldn't change
+	assert.False(t, r.primaryIsActive)
+	// Request an address (should be IPv4 from secondary)
+	a4_v4 := r.AssignAnyAddress(4, nil)
+	assert.NotNil(t, a4_v4)
+	assert.Equal(t, V4, a4_v4.IPVersion)
+	a5_v4 := r.AssignAnyAddress(5, nil)
+	assert.NotNil(t, a5_v4)
+	assert.Equal(t, V4, a5_v4.IPVersion)
+	a6_v4 := r.AssignAnyAddress(6, nil)
+	assert.NotNil(t, a6_v4)
+	assert.Equal(t, V4, a6_v4.IPVersion)
+	// Return IPv4 address (without failure)
+	// Primary shouldn't change because it is not a connectivity failure
+	assert.True(t, r.GiveBack(a4_v4, false))
+	assert.False(t, r.primaryIsActive)
+	// Return IPv4 address (with failure)
+	// Primary should change because it is a connectivity failure
+	assert.True(t, r.GiveBack(a5_v4, true))
+	assert.True(t, r.primaryIsActive)
+	// Return IPv4 address (with failure)
+	// Primary shouldn't change because the address is returned to the inactive
+	// secondary address set
+	assert.True(t, r.GiveBack(a6_v4, true))
+	assert.True(t, r.primaryIsActive)
+	// Return IPv6 address (without failure)
+	// Primary shoudn't change because it is not a connectivity failure
+	assert.True(t, r.GiveBack(a2, false))
+	assert.True(t, r.primaryIsActive)
 }
 
-func TestRegion_GetUnusedIP(t *testing.T) {
-	type fields struct {
-		connFor map[*EdgeAddr]UsedBy
-	}
-	type args struct {
-		excluding *EdgeAddr
-	}
-	tests := []struct {
-		name   string
-		fields fields
-		args   args
-		want   *EdgeAddr
-	}{
-		{
-			name: "happy test with excluding set",
-			fields: fields{connFor: map[*EdgeAddr]UsedBy{
-				&addr0: Unused(),
-				&addr1: Unused(),
-				&addr2: InUse(2),
-			}},
-			args: args{excluding: &addr0},
-			want: &addr1,
-		},
-		{
-			name: "happy test with no excluding",
-			fields: fields{connFor: map[*EdgeAddr]UsedBy{
-				&addr0: InUse(0),
-				&addr1: Unused(),
-				&addr2: InUse(2),
-			}},
-			args: args{excluding: nil},
-			want: &addr1,
-		},
-		{
-			name: "sad test with no excluding",
-			fields: fields{connFor: map[*EdgeAddr]UsedBy{
-				&addr0: InUse(0),
-				&addr1: InUse(1),
-				&addr2: InUse(2),
-			}},
-			args: args{excluding: nil},
-			want: nil,
-		},
-		{
-			name: "sad test with excluding",
-			fields: fields{connFor: map[*EdgeAddr]UsedBy{
-				&addr0: Unused(),
-				&addr1: InUse(1),
-				&addr2: InUse(2),
-			}},
-			args: args{excluding: &addr0},
-			want: nil,
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			r := Region{
-				connFor: tt.fields.connFor,
-			}
-			if got := r.GetUnusedIP(tt.args.excluding); !reflect.DeepEqual(got, tt.want) {
-				t.Errorf("Region.GetUnusedIP() = %v, want %v", got, tt.want)
-			}
-		})
-	}
-}
-
-func TestRegion_GiveBack(t *testing.T) {
-	type fields struct {
-		connFor map[*EdgeAddr]UsedBy
-	}
-	type args struct {
-		addr *EdgeAddr
-	}
-	tests := []struct {
-		name           string
-		fields         fields
-		args           args
-		wantOk         bool
-		availableAfter int
-	}{
-		{
-			name: "sad test with excluding",
-			fields: fields{connFor: map[*EdgeAddr]UsedBy{
-				&addr1: InUse(1),
-			}},
-			args:           args{addr: &addr1},
-			wantOk:         true,
-			availableAfter: 1,
-		},
-		{
-			name: "sad test with excluding",
-			fields: fields{connFor: map[*EdgeAddr]UsedBy{
-				&addr1: InUse(1),
-			}},
-			args:           args{addr: &addr2},
-			wantOk:         false,
-			availableAfter: 0,
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			r := Region{
-				connFor: tt.fields.connFor,
-			}
-			if gotOk := r.GiveBack(tt.args.addr); gotOk != tt.wantOk {
-				t.Errorf("Region.GiveBack() = %v, want %v", gotOk, tt.wantOk)
-			}
-			if tt.availableAfter != r.AvailableAddrs() {
-				t.Errorf("Region.AvailableAddrs() = %v, want %v", r.AvailableAddrs(), tt.availableAfter)
-			}
-		})
-	}
-}
-
-func TestRegion_GetAnyAddress(t *testing.T) {
-	type fields struct {
-		connFor map[*EdgeAddr]UsedBy
-	}
-	tests := []struct {
-		name    string
-		fields  fields
-		wantNil bool
-	}{
-		{
-			name:    "Sad test -- GetAnyAddress should only fail if the region is empty",
-			fields:  fields{connFor: map[*EdgeAddr]UsedBy{}},
-			wantNil: true,
-		},
-		{
-			name: "Happy test (all addresses unused)",
-			fields: fields{connFor: map[*EdgeAddr]UsedBy{
-				&addr0: Unused(),
-			}},
-			wantNil: false,
-		},
-		{
-			name: "Happy test (GetAnyAddress can still return addresses used by proxy conns)",
-			fields: fields{connFor: map[*EdgeAddr]UsedBy{
-				&addr0: InUse(2),
-			}},
-			wantNil: false,
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			r := Region{
-				connFor: tt.fields.connFor,
-			}
-			if got := r.GetAnyAddress(); tt.wantNil != (got == nil) {
-				t.Errorf("Region.GetAnyAddress() = %v, but should it return nil? %v", got, tt.wantNil)
-			}
-		})
-	}
+func TestRegion_GiveBack_Timeout(t *testing.T) {
+	r := NewRegion(append(v6Addrs, v4Addrs...), Auto)
+	a0 := r.AssignAnyAddress(0, nil)
+	a1 := r.AssignAnyAddress(1, nil)
+	a2 := r.AssignAnyAddress(2, nil)
+	assert.NotNil(t, a0)
+	assert.NotNil(t, a1)
+	assert.NotNil(t, a2)
+	// Give back IPv6 address to set timeout
+	assert.True(t, r.GiveBack(a0, true))
+	assert.False(t, r.primaryIsActive)
+	assert.False(t, r.primaryTimeout.IsZero())
+	// Request an address (should be IPv4 from secondary)
+	a3_v4 := r.AssignAnyAddress(3, nil)
+	assert.NotNil(t, a3_v4)
+	assert.Equal(t, V4, a3_v4.IPVersion)
+	assert.False(t, r.primaryIsActive)
+	// Give back IPv6 address inside timeout (no change)
+	assert.True(t, r.GiveBack(a2, true))
+	assert.False(t, r.primaryIsActive)
+	assert.False(t, r.primaryTimeout.IsZero())
+	// Accelerate timeout
+	r.primaryTimeout = time.Now().Add(-time.Minute)
+	// Return IPv6 address
+	assert.True(t, r.GiveBack(a1, true))
+	assert.True(t, r.primaryIsActive)
+	// Returning an IPv4 address after primary is active shouldn't change primary
+	// even with a connectivity error
+	assert.True(t, r.GiveBack(a3_v4, true))
+	assert.True(t, r.primaryIsActive)
 }

edgediscovery/allregions/regions.go

@@ -18,7 +18,7 @@ type Regions struct {
 // ------------------------------------
 
 // ResolveEdge resolves the Cloudflare edge, returning all regions discovered.
-func ResolveEdge(log *zerolog.Logger, region string) (*Regions, error) {
+func ResolveEdge(log *zerolog.Logger, region string, overrideIPVersion ConfigIPVersion) (*Regions, error) {
 	edgeAddrs, err := edgeDiscovery(log, getRegionalServiceName(region))
 	if err != nil {
 		return nil, err
@@ -27,8 +27,8 @@ func ResolveEdge(log *zerolog.Logger, region string) (*Regions, error) {
 		return nil, fmt.Errorf("expected at least 2 Cloudflare Regions regions, but SRV only returned %v", len(edgeAddrs))
 	}
 	return &Regions{
-		region1: NewRegion(edgeAddrs[0]),
-		region2: NewRegion(edgeAddrs[1]),
+		region1: NewRegion(edgeAddrs[0], overrideIPVersion),
+		region2: NewRegion(edgeAddrs[1], overrideIPVersion),
 	}, nil
 }
 
@@ -56,8 +56,8 @@ func NewNoResolve(addrs []*EdgeAddr) *Regions {
 	}
 
 	return &Regions{
-		region1: NewRegion(region1),
-		region2: NewRegion(region2),
+		region1: NewRegion(region1, Auto),
+		region2: NewRegion(region2, Auto),
 	}
 }
 
@@ -95,14 +95,12 @@ func (rs *Regions) GetUnusedAddr(excluding *EdgeAddr, connID int) *EdgeAddr {
 // getAddrs tries to grab address form `first` region, then `second` region
 // this is an unrolled loop over 2 element array
 func getAddrs(excluding *EdgeAddr, connID int, first *Region, second *Region) *EdgeAddr {
-	addr := first.GetUnusedIP(excluding)
+	addr := first.AssignAnyAddress(connID, excluding)
 	if addr != nil {
-		first.Use(addr, connID)
 		return addr
 	}
-	addr = second.GetUnusedIP(excluding)
+	addr = second.AssignAnyAddress(connID, excluding)
 	if addr != nil {
-		second.Use(addr, connID)
 		return addr
 	}
 
@@ -116,18 +114,18 @@ func (rs *Regions) AvailableAddrs() int {
 
 // GiveBack the address so that other connections can use it.
 // Returns true if the address is in this edge.
-func (rs *Regions) GiveBack(addr *EdgeAddr) bool {
-	if found := rs.region1.GiveBack(addr); found {
+func (rs *Regions) GiveBack(addr *EdgeAddr, hasConnectivityError bool) bool {
+	if found := rs.region1.GiveBack(addr, hasConnectivityError); found {
 		return found
 	}
-	return rs.region2.GiveBack(addr)
+	return rs.region2.GiveBack(addr, hasConnectivityError)
 }
 
 // Return regionalized service name if `region` isn't empty, otherwise return the global service name for origintunneld
 func getRegionalServiceName(region string) string {
 	if region != "" {
-		return region + "-" + srvService // Example: `us-origintunneld`
+		return region + "-" + srvService // Example: `us-v2-origintunneld`
 	}
 
-	return srvService // Global service is just `origintunneld`
+	return srvService // Global service is just `v2-origintunneld`
 }

edgediscovery/allregions/regions_test.go

@@ -1,134 +1,215 @@
 package allregions
 
 import (
-	"net"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 )
 
-var (
-	addr0 = EdgeAddr{
-		TCP: &net.TCPAddr{
-			IP:   net.ParseIP("123.4.5.0"),
-			Port: 8000,
-			Zone: "",
-		},
-		UDP: &net.UDPAddr{
-			IP:   net.ParseIP("123.4.5.0"),
-			Port: 8000,
-			Zone: "",
-		},
+func makeRegions(addrs []*EdgeAddr, mode ConfigIPVersion) Regions {
+	r1addrs := make([]*EdgeAddr, 0)
+	r2addrs := make([]*EdgeAddr, 0)
+	for i, addr := range addrs {
+		if i%2 == 0 {
+			r1addrs = append(r1addrs, addr)
+		} else {
+			r2addrs = append(r2addrs, addr)
+		}
 	}
-	addr1 = EdgeAddr{
-		TCP: &net.TCPAddr{
-			IP:   net.ParseIP("123.4.5.1"),
-			Port: 8000,
-			Zone: "",
-		},
-		UDP: &net.UDPAddr{
-			IP:   net.ParseIP("123.4.5.1"),
-			Port: 8000,
-			Zone: "",
-		},
-	}
-	addr2 = EdgeAddr{
-		TCP: &net.TCPAddr{
-			IP:   net.ParseIP("123.4.5.2"),
-			Port: 8000,
-			Zone: "",
-		},
-		UDP: &net.UDPAddr{
-			IP:   net.ParseIP("123.4.5.2"),
-			Port: 8000,
-			Zone: "",
-		},
-	}
-	addr3 = EdgeAddr{
-		TCP: &net.TCPAddr{
-			IP:   net.ParseIP("123.4.5.3"),
-			Port: 8000,
-			Zone: "",
-		},
-		UDP: &net.UDPAddr{
-			IP:   net.ParseIP("123.4.5.3"),
-			Port: 8000,
-			Zone: "",
-		},
-	}
-)
-
-func makeRegions() Regions {
-	r1 := NewRegion([]*EdgeAddr{&addr0, &addr1})
-	r2 := NewRegion([]*EdgeAddr{&addr2, &addr3})
+	r1 := NewRegion(r1addrs, mode)
+	r2 := NewRegion(r2addrs, mode)
 	return Regions{region1: r1, region2: r2}
 }
 
 func TestRegions_AddrUsedBy(t *testing.T) {
-	rs := makeRegions()
-	addr1 := rs.GetUnusedAddr(nil, 1)
-	assert.Equal(t, addr1, rs.AddrUsedBy(1))
-	addr2 := rs.GetUnusedAddr(nil, 2)
-	assert.Equal(t, addr2, rs.AddrUsedBy(2))
-	addr3 := rs.GetUnusedAddr(nil, 3)
-	assert.Equal(t, addr3, rs.AddrUsedBy(3))
+	tests := []struct {
+		name  string
+		addrs []*EdgeAddr
+		mode  ConfigIPVersion
+	}{
+		{
+			name:  "IPv4 addresses with IPv4Only",
+			addrs: v4Addrs,
+			mode:  IPv4Only,
+		},
+		{
+			name:  "IPv6 addresses with IPv6Only",
+			addrs: v6Addrs,
+			mode:  IPv6Only,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			rs := makeRegions(tt.addrs, tt.mode)
+			addr1 := rs.GetUnusedAddr(nil, 1)
+			assert.Equal(t, addr1, rs.AddrUsedBy(1))
+			addr2 := rs.GetUnusedAddr(nil, 2)
+			assert.Equal(t, addr2, rs.AddrUsedBy(2))
+			addr3 := rs.GetUnusedAddr(nil, 3)
+			assert.Equal(t, addr3, rs.AddrUsedBy(3))
+		})
+	}
 }
 
 func TestRegions_Giveback_Region1(t *testing.T) {
-	rs := makeRegions()
-	rs.region1.Use(&addr0, 0)
-	rs.region1.Use(&addr1, 1)
-	rs.region2.Use(&addr2, 2)
-	rs.region2.Use(&addr3, 3)
+	tests := []struct {
+		name  string
+		addrs []*EdgeAddr
+		mode  ConfigIPVersion
+	}{
+		{
+			name:  "IPv4 addresses with IPv4Only",
+			addrs: v4Addrs,
+			mode:  IPv4Only,
+		},
+		{
+			name:  "IPv6 addresses with IPv6Only",
+			addrs: v6Addrs,
+			mode:  IPv6Only,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			rs := makeRegions(tt.addrs, tt.mode)
+			addr := rs.region1.AssignAnyAddress(0, nil)
+			rs.region1.AssignAnyAddress(1, nil)
+			rs.region2.AssignAnyAddress(2, nil)
+			rs.region2.AssignAnyAddress(3, nil)
 
-	assert.Equal(t, 0, rs.AvailableAddrs())
+			assert.Equal(t, 0, rs.AvailableAddrs())
 
-	rs.GiveBack(&addr0)
-	assert.Equal(t, &addr0, rs.GetUnusedAddr(nil, 3))
+			rs.GiveBack(addr, false)
+			assert.Equal(t, addr, rs.GetUnusedAddr(nil, 0))
+		})
+	}
 }
 
 func TestRegions_Giveback_Region2(t *testing.T) {
-	rs := makeRegions()
-	rs.region1.Use(&addr0, 0)
-	rs.region1.Use(&addr1, 1)
-	rs.region2.Use(&addr2, 2)
-	rs.region2.Use(&addr3, 3)
+	tests := []struct {
+		name  string
+		addrs []*EdgeAddr
+		mode  ConfigIPVersion
+	}{
+		{
+			name:  "IPv4 addresses with IPv4Only",
+			addrs: v4Addrs,
+			mode:  IPv4Only,
+		},
+		{
+			name:  "IPv6 addresses with IPv6Only",
+			addrs: v6Addrs,
+			mode:  IPv6Only,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			rs := makeRegions(tt.addrs, tt.mode)
+			rs.region1.AssignAnyAddress(0, nil)
+			rs.region1.AssignAnyAddress(1, nil)
+			addr := rs.region2.AssignAnyAddress(2, nil)
+			rs.region2.AssignAnyAddress(3, nil)
 
-	assert.Equal(t, 0, rs.AvailableAddrs())
+			assert.Equal(t, 0, rs.AvailableAddrs())
 
-	rs.GiveBack(&addr2)
-	assert.Equal(t, &addr2, rs.GetUnusedAddr(nil, 2))
+			rs.GiveBack(addr, false)
+			assert.Equal(t, addr, rs.GetUnusedAddr(nil, 2))
+		})
+	}
 }
 
 func TestRegions_GetUnusedAddr_OneAddrLeft(t *testing.T) {
-	rs := makeRegions()
+	tests := []struct {
+		name  string
+		addrs []*EdgeAddr
+		mode  ConfigIPVersion
+	}{
+		{
+			name:  "IPv4 addresses with IPv4Only",
+			addrs: v4Addrs,
+			mode:  IPv4Only,
+		},
+		{
+			name:  "IPv6 addresses with IPv6Only",
+			addrs: v6Addrs,
+			mode:  IPv6Only,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			rs := makeRegions(tt.addrs, tt.mode)
+			rs.region1.AssignAnyAddress(0, nil)
+			rs.region1.AssignAnyAddress(1, nil)
+			rs.region2.AssignAnyAddress(2, nil)
+			addr := rs.region2.active.GetUnusedIP(nil)
 
-	rs.region1.Use(&addr0, 0)
-	rs.region1.Use(&addr1, 1)
-	rs.region2.Use(&addr2, 2)
-
-	assert.Equal(t, 1, rs.AvailableAddrs())
-	assert.Equal(t, &addr3, rs.GetUnusedAddr(nil, 3))
+			assert.Equal(t, 1, rs.AvailableAddrs())
+			assert.Equal(t, addr, rs.GetUnusedAddr(nil, 3))
+		})
+	}
 }
 
 func TestRegions_GetUnusedAddr_Excluding_Region1(t *testing.T) {
-	rs := makeRegions()
+	tests := []struct {
+		name  string
+		addrs []*EdgeAddr
+		mode  ConfigIPVersion
+	}{
+		{
+			name:  "IPv4 addresses with IPv4Only",
+			addrs: v4Addrs,
+			mode:  IPv4Only,
+		},
+		{
+			name:  "IPv6 addresses with IPv6Only",
+			addrs: v6Addrs,
+			mode:  IPv6Only,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			rs := makeRegions(tt.addrs, tt.mode)
 
-	rs.region1.Use(&addr0, 0)
-	rs.region1.Use(&addr1, 1)
+			rs.region1.AssignAnyAddress(0, nil)
+			rs.region1.AssignAnyAddress(1, nil)
+			addr := rs.region2.active.GetUnusedIP(nil)
+			a2 := rs.region2.active.GetUnusedIP(addr)
 
-	assert.Equal(t, 2, rs.AvailableAddrs())
-	assert.Equal(t, &addr3, rs.GetUnusedAddr(&addr2, 3))
+			assert.Equal(t, 2, rs.AvailableAddrs())
+			assert.Equal(t, addr, rs.GetUnusedAddr(a2, 3))
+		})
+	}
 }
 
 func TestRegions_GetUnusedAddr_Excluding_Region2(t *testing.T) {
-	rs := makeRegions()
+	tests := []struct {
+		name  string
+		addrs []*EdgeAddr
+		mode  ConfigIPVersion
+	}{
+		{
+			name:  "IPv4 addresses with IPv4Only",
+			addrs: v4Addrs,
+			mode:  IPv4Only,
+		},
+		{
+			name:  "IPv6 addresses with IPv6Only",
+			addrs: v6Addrs,
+			mode:  IPv6Only,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			rs := makeRegions(tt.addrs, tt.mode)
 
-	rs.region2.Use(&addr2, 0)
-	rs.region2.Use(&addr3, 1)
+			rs.region2.AssignAnyAddress(0, nil)
+			rs.region2.AssignAnyAddress(1, nil)
+			addr := rs.region1.active.GetUnusedIP(nil)
+			a2 := rs.region1.active.GetUnusedIP(addr)
 
-	assert.Equal(t, 2, rs.AvailableAddrs())
-	assert.Equal(t, &addr1, rs.GetUnusedAddr(&addr0, 1))
+			assert.Equal(t, 2, rs.AvailableAddrs())
+			assert.Equal(t, addr, rs.GetUnusedAddr(a2, 1))
+		})
+	}
 }
 
 func TestNewNoResolveBalancesRegions(t *testing.T) {