Allow subnet allocator to mark assigned subnet dynamically

- Currently, we need to pass all allocated subnets during
subnet allocator creation time (inUse arg to NewSubnetAllocator()).
This means we need to know all existing subnets beforehand.

- This change exposes additional method so that we can mark a
specific subnet as already allocated dynamically (after the subnet
allocated is created).

Precursor for #18911

Author: Ravi Sankar Penta

pkg/util/netutils/subnet_allocator.go

@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"net"
 	"sync"
+
+	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
 )
 
 var ErrSubnetAllocatorFull = fmt.Errorf("No subnets available.")
@@ -23,14 +25,14 @@ type SubnetAllocator struct {
 func NewSubnetAllocator(network string, hostBits uint32, inUse []string) (*SubnetAllocator, error) {
 	_, netIP, err := net.ParseCIDR(network)
 	if err != nil {
-		return nil, fmt.Errorf("Failed to parse network address: %q", network)
+		return nil, fmt.Errorf("failed to parse network address: %q", network)
 	}
 
 	netMaskSize, _ := netIP.Mask.Size()
 	if hostBits == 0 {
-		return nil, fmt.Errorf("Host capacity cannot be zero.")
+		return nil, fmt.Errorf("host capacity cannot be zero.")
 	} else if hostBits > (32 - uint32(netMaskSize)) {
-		return nil, fmt.Errorf("Subnet capacity cannot be larger than number of networks available.")
+		return nil, fmt.Errorf("subnet capacity cannot be larger than number of networks available.")
 	}
 	subnetBits := 32 - uint32(netMaskSize) - hostBits
 
@@ -61,29 +63,41 @@ func NewSubnetAllocator(network string, hostBits uint32, inUse []string) (*Subne
 		rightMask = 0
 	}
 
-	amap := make(map[string]bool)
-	for _, netStr := range inUse {
-		_, nIp, err := net.ParseCIDR(netStr)
-		if err != nil {
-			fmt.Println("Failed to parse network address: ", netStr)
-			continue
-		}
-		if !netIP.Contains(nIp.IP) {
-			fmt.Println("Provided subnet doesn't belong to network: ", nIp)
-			continue
-		}
-		amap[nIp.String()] = true
-	}
-	return &SubnetAllocator{
+	sa := &SubnetAllocator{
 		network:    netIP,
 		hostBits:   hostBits,
 		leftShift:  leftShift,
 		leftMask:   leftMask,
 		rightShift: rightShift,
 		rightMask:  rightMask,
 		next:       0,
-		allocMap:   amap,
-	}, nil
+		allocMap:   make(map[string]bool),
+	}
+	for _, netStr := range inUse {
+		_, ipNet, err := net.ParseCIDR(netStr)
+		if err != nil {
+			utilruntime.HandleError(fmt.Errorf("failed to parse network address: %s", netStr))
+			continue
+		}
+		if err = sa.AllocateNetwork(ipNet); err != nil {
+			utilruntime.HandleError(err)
+			continue
+		}
+	}
+	return sa, nil
+}
+
+func (sna *SubnetAllocator) AllocateNetwork(ipNet *net.IPNet) error {
+	sna.mutex.Lock()
+	defer sna.mutex.Unlock()
+
+	if !sna.network.Contains(ipNet.IP) {
+		return fmt.Errorf("provided subnet doesn't belong to network: %v", ipNet)
+	}
+	if !sna.allocMap[ipNet.String()] {
+		sna.allocMap[ipNet.String()] = true
+	}
+	return nil
 }
 
 func (sna *SubnetAllocator) GetNetwork() (*net.IPNet, error) {
@@ -120,16 +134,16 @@ func (sna *SubnetAllocator) GetNetwork() (*net.IPNet, error) {
 func (sna *SubnetAllocator) ReleaseNetwork(ipnet *net.IPNet) error {
 	sna.mutex.Lock()
 	defer sna.mutex.Unlock()
+
 	if !sna.network.Contains(ipnet.IP) {
-		return fmt.Errorf("Provided subnet %v doesn't belong to the network %v.", ipnet, sna.network)
+		return fmt.Errorf("provided subnet %v doesn't belong to the network %v.", ipnet, sna.network)
 	}
 
 	ipnetStr := ipnet.String()
 	if !sna.allocMap[ipnetStr] {
-		return fmt.Errorf("Provided subnet %v is already available.", ipnet)
+		return fmt.Errorf("provided subnet %v is already available.", ipnet)
+	} else {
+		sna.allocMap[ipnetStr] = false
 	}
-
-	sna.allocMap[ipnetStr] = false
-
 	return nil
 }

pkg/util/netutils/subnet_allocator_test.go

@@ -205,6 +205,59 @@ func TestAllocateSubnetInUse(t *testing.T) {
 	}
 }
 
+func TestAllocateNetwork(t *testing.T) {
+	sna, err := NewSubnetAllocator("10.1.0.0/16", 14, nil)
+	if err != nil {
+		t.Fatal("Failed to initialize IP allocator: ", err)
+	}
+
+	allocSubnets := make([]*net.IPNet, 4)
+	for i := 0; i < 4; i++ {
+		if allocSubnets[i], err = sna.GetNetwork(); err != nil {
+			t.Fatal("Failed to get network: ", err)
+		}
+	}
+
+	if sn, err := sna.GetNetwork(); err == nil {
+		t.Fatalf("Unexpectedly succeeded in getting network (sn=%s)", sn.String())
+	}
+	if err := sna.ReleaseNetwork(allocSubnets[2]); err != nil {
+		t.Fatalf("Failed to release the subnet (allocSubnets[2]=%s): %v", allocSubnets[2].String(), err)
+	}
+	for i := 0; i < 2; i++ {
+		if err := sna.AllocateNetwork(allocSubnets[2]); err != nil {
+			t.Fatalf("Failed to allocate the subnet (allocSubnets[2]=%s): %v", allocSubnets[2].String(), err)
+		}
+	}
+	if sn, err := sna.GetNetwork(); err == nil {
+		t.Fatalf("Unexpectedly succeeded in getting network (sn=%s)", sn.String())
+	}
+
+	// Test subnet does not belong to network
+	var sn *net.IPNet
+	_, sn, err = net.ParseCIDR("10.2.3.4/24")
+	if err != nil {
+		t.Fatal("Failed to parse given network: ", err)
+	}
+	if err := sna.AllocateNetwork(sn); err == nil {
+		t.Fatalf("Unexpectedly succeeded in allocating subnet that doesn't belong to network (sn=%s)", sn.String())
+	}
+
+	// Test AllocateNetwork usage in NewSubnetAllocator
+	var subnetStrs []string
+	for _, sn := range allocSubnets {
+		subnetStrs = append(subnetStrs, sn.String())
+	}
+	var sa *SubnetAllocator
+	sa, err = NewSubnetAllocator("10.1.0.0/16", 14, subnetStrs)
+	if err != nil {
+		t.Fatal("Failed to initialize IP allocator: ", err)
+	}
+	if sn, err = sa.GetNetwork(); err == nil {
+		t.Fatalf("Unexpectedly succeeded in getting network (sn=%s)", sn.String())
+	}
+}
+
 func TestAllocateReleaseSubnet(t *testing.T) {
 	sna, err := NewSubnetAllocator("10.1.0.0/16", 14, nil)
 	if err != nil {