Rebalance auto-assigned egress IPs when necessary

pkg/network/common/egressip.go

@@ -490,10 +490,7 @@ func (eit *EgressIPTracker) findEgressIPAllocation(ip net.IP, allocation map[str
 		if node.offline {
 			continue
 		}
-		egressIPs, exists := allocation[node.nodeName]
-		if !exists {
-			continue
-		}
+		egressIPs := allocation[node.nodeName]
 		for _, parsed := range node.parsedCIDRs {
 			if parsed.Contains(ip) {
 				if bestNode != "" {
@@ -511,13 +508,13 @@ func (eit *EgressIPTracker) findEgressIPAllocation(ip net.IP, allocation map[str
 	return bestNode, otherNodes
 }
 
-// ReallocateEgressIPs returns a map from Node name to array-of-Egress-IP for all auto-allocated egress IPs
-func (eit *EgressIPTracker) ReallocateEgressIPs() map[string][]string {
-	eit.Lock()
-	defer eit.Unlock()
+func (eit *EgressIPTracker) makeEmptyAllocation() (map[string][]string, map[string]bool) {
+	return make(map[string][]string), make(map[string]bool)
+}
+
+func (eit *EgressIPTracker) allocateExistingEgressIPs(allocation map[string][]string, alreadyAllocated map[string]bool) bool {
+	removedEgressIPs := false
 
-	allocation := make(map[string][]string)
-	alreadyAllocated := make(map[string]bool)
 	for _, node := range eit.nodes {
 		if len(node.parsedCIDRs) > 0 {
 			allocation[node.nodeName] = make([]string, 0, node.requestedIPs.Len())
@@ -526,7 +523,7 @@ func (eit *EgressIPTracker) ReallocateEgressIPs() map[string][]string {
 	// For each active egress IP, if it still fits within some egress CIDR on its node,
 	// add it to that node's allocation.
 	for egressIP, eip := range eit.egressIPs {
-		if eip.assignedNodeIP == "" {
+		if eip.assignedNodeIP == "" || alreadyAllocated[egressIP] {
 			continue
 		}
 		node := eip.nodes[0]
@@ -539,13 +536,19 @@ func (eit *EgressIPTracker) ReallocateEgressIPs() map[string][]string {
 		}
 		if found && !node.offline {
 			allocation[node.nodeName] = append(allocation[node.nodeName], egressIP)
+		} else {
+			removedEgressIPs = true
 		}
 		// (We set alreadyAllocated even if the egressIP will be removed from
 		// its current node; we can't assign it to a new node until the next
 		// reallocation.)
 		alreadyAllocated[egressIP] = true
 	}
 
+	return removedEgressIPs
+}
+
+func (eit *EgressIPTracker) allocateNewEgressIPs(allocation map[string][]string, alreadyAllocated map[string]bool) {
 	// Allocate pending egress IPs that can only go to a single node
 	for egressIP, eip := range eit.egressIPs {
 		if alreadyAllocated[egressIP] {
@@ -567,6 +570,52 @@ func (eit *EgressIPTracker) ReallocateEgressIPs() map[string][]string {
 			allocation[nodeName] = append(allocation[nodeName], egressIP)
 		}
 	}
+}
+
+// ReallocateEgressIPs returns a map from Node name to array-of-Egress-IP for all auto-allocated egress IPs
+func (eit *EgressIPTracker) ReallocateEgressIPs() map[string][]string {
+	eit.Lock()
+	defer eit.Unlock()
+
+	allocation, alreadyAllocated := eit.makeEmptyAllocation()
+	removedEgressIPs := eit.allocateExistingEgressIPs(allocation, alreadyAllocated)
+	eit.allocateNewEgressIPs(allocation, alreadyAllocated)
+	if removedEgressIPs {
+		// Process the removals now; we'll get called again afterward and can
+		// check for balance then.
+		return allocation
+	}
+
+	// Compare the allocation to what we would have gotten if we started from scratch,
+	// to see if things have gotten too unbalanced. (In particular, if a node goes
+	// offline, gets emptied, and then comes back online, we want to move a bunch of
+	// egress IPs back onto that node.)
+	fullReallocation, alreadyAllocated := eit.makeEmptyAllocation()
+	eit.allocateNewEgressIPs(fullReallocation, alreadyAllocated)
+
+	emptyNodes := []string{}
+	for nodeName, fullEgressIPs := range fullReallocation {
+		incrementalEgressIPs := allocation[nodeName]
+		if len(incrementalEgressIPs) < len(fullEgressIPs) / 2 {
+			emptyNodes = append(emptyNodes, nodeName)
+		}
+	}
+
+	if len(emptyNodes) > 0 {
+		// Make a new incremental allocation, but skipping all of the egress IPs
+		// that got assigned to the "empty" nodes in the full reallocation; this
+		// will cause them to be dropped from their current nodes and then later
+		// reassigned (to one of the "empty" nodes, for balance).
+		allocation, alreadyAllocated = eit.makeEmptyAllocation()
+		for _, nodeName := range emptyNodes {
+			for _, egressIP := range fullReallocation[nodeName] {
+				alreadyAllocated[egressIP] = true
+			}
+		}
+		eit.allocateExistingEgressIPs(allocation, alreadyAllocated)
+		eit.allocateNewEgressIPs(allocation, alreadyAllocated)
+		eit.updateEgressCIDRs = true
+	}
 
 	return allocation
 }

pkg/network/common/egressip_test.go

@@ -1106,4 +1106,38 @@ func TestEgressCIDRAllocationOffline(t *testing.T) {
 		t.Fatalf("Bad IP allocation: %#v", allocation)
 	}
 	updateAllocations(eit, allocation)
+
+	// Bring node-3 back
+	eit.SetNodeOffline("172.17.0.3", false)
+	err = w.assertUpdateEgressCIDRsNotification()
+	if err != nil {
+		t.Fatalf("%v", err)
+	}
+
+	// First reallocation should remove some IPs from node-4 and node-5 but not add
+	// them to node-3. As above, the "balanced" allocation we're aiming for may not
+	// be perfect, but it has to be planning to assign at least 2 IPs to node-3.
+	allocation = eit.ReallocateEgressIPs()
+	node3ips = allocation["node-3"]
+	node4ips = allocation["node-4"]
+	node5ips = allocation["node-5"]
+	if len(node3ips) != 0 || len(node4ips) + len(node5ips) > 4 {
+		t.Fatalf("Bad IP allocation: %#v", allocation)
+	}
+	updateAllocations(eit, allocation)
+
+	err = w.assertUpdateEgressCIDRsNotification()
+	if err != nil {
+		t.Fatalf("%v", err)
+	}
+
+	// Next reallocation should reassign egress IPs to node-3
+	allocation = eit.ReallocateEgressIPs()
+	node3ips = allocation["node-3"]
+	node4ips = allocation["node-4"]
+	node5ips = allocation["node-5"]
+	if len(node3ips) < 2 || len(node4ips) == 0 || len(node5ips) == 0 {
+		t.Fatalf("Bad IP allocation: %#v", allocation)
+	}
+	updateAllocations(eit, allocation)
 }