Add support for IP exlusion to allocator

Mellanox · Jun 10, 2024 · 7314b4f · 7314b4f
1 parent b272d92
commit 7314b4f
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 15 deletions.
diff --git a/pkg/ipam-node/allocator/allocator.go b/pkg/ipam-node/allocator/allocator.go
@@ -39,17 +39,20 @@ type IPAllocator interface {
 }
 
 type allocator struct {
-	rangeSet *RangeSet
-	session  storePkg.Session
-	poolName string
+	rangeSet   *RangeSet
+	session    storePkg.Session
+	poolName   string
+	exclusions *RangeSet
 }
 
 // NewIPAllocator create and initialize a new instance of IP allocator
-func NewIPAllocator(s *RangeSet, poolName string, session storePkg.Session) IPAllocator {
+func NewIPAllocator(s *RangeSet, exclusions *RangeSet,
+	poolName string, session storePkg.Session) IPAllocator {
 	return &allocator{
-		rangeSet: s,
-		session:  session,
-		poolName: poolName,
+		rangeSet:   s,
+		session:    session,
+		poolName:   poolName,
+		exclusions: exclusions,
 	}
 }
 
@@ -64,6 +67,9 @@ func (a *allocator) Allocate(id string, ifName string, meta types.ReservationMet
 		if reservedIP == nil {
 			return nil, ErrNoFreeAddresses
 		}
+		if a.exclusions != nil && a.exclusions.Contains(reservedIP.IP) {
+			continue
+		}
 		err := a.session.Reserve(a.poolName, id, ifName, meta, reservedIP.IP)
 		if err == nil {
 			break

diff --git a/pkg/ipam-node/allocator/allocator_test.go b/pkg/ipam-node/allocator/allocator_test.go
@@ -50,7 +50,7 @@ func mkAlloc(session storePkg.Session) allocator.IPAllocator {
 		allocator.Range{Subnet: mustSubnet("192.168.1.0/29"), Gateway: net.ParseIP("192.168.1.1")},
 	}
 	Expect(p.Canonicalize()).NotTo(HaveOccurred())
-	return allocator.NewIPAllocator(&p, testPoolName, session)
+	return allocator.NewIPAllocator(&p, nil, testPoolName, session)
 }
 
 func newAllocatorWithMultiRanges(session storePkg.Session) allocator.IPAllocator {
@@ -59,7 +59,7 @@ func newAllocatorWithMultiRanges(session storePkg.Session) allocator.IPAllocator
 		allocator.Range{RangeStart: net.IP{192, 168, 2, 0}, RangeEnd: net.IP{192, 168, 2, 3}, Subnet: mustSubnet("192.168.2.0/30")},
 	}
 	Expect(p.Canonicalize()).NotTo(HaveOccurred())
-	return allocator.NewIPAllocator(&p, testPoolName, session)
+	return allocator.NewIPAllocator(&p, nil, testPoolName, session)
 }
 
 func (t AllocatorTestCase) run(idx int, session storePkg.Session) (*current.IPConfig, error) {
@@ -79,7 +79,7 @@ func (t AllocatorTestCase) run(idx int, session storePkg.Session) (*current.IPCo
 	session.SetLastReservedIP(testPoolName, net.ParseIP(t.lastIP))
 
 	Expect(p.Canonicalize()).To(Succeed())
-	alloc := allocator.NewIPAllocator(&p, testPoolName, session)
+	alloc := allocator.NewIPAllocator(&p, nil, testPoolName, session)
 	return alloc.Allocate(testContainerID, testIFName, types.ReservationMetadata{})
 }
 
@@ -391,7 +391,7 @@ var _ = Describe("allocator", func() {
 				allocator.Range{Subnet: mustSubnet("192.168.1.4/30")},
 			}
 			Expect(p.Canonicalize()).NotTo(HaveOccurred())
-			a := allocator.NewIPAllocator(&p, testPoolName, session)
+			a := allocator.NewIPAllocator(&p, nil, testPoolName, session)
 			// get range iterator and do the first Next
 			checkAlloc(a, "0", net.IP{192, 168, 1, 1})
 			checkAlloc(a, "1", net.IP{192, 168, 1, 2})
@@ -410,10 +410,36 @@ var _ = Describe("allocator", func() {
 				allocator.Range{Subnet: mustSubnet("192.168.1.0/31")},
 			}
 			Expect(p.Canonicalize()).NotTo(HaveOccurred())
-			a := allocator.NewIPAllocator(&p, testPoolName, session)
+			a := allocator.NewIPAllocator(&p, nil, testPoolName, session)
 			// get range iterator and do the first Next
 			checkAlloc(a, "0", net.IP{192, 168, 1, 0})
 			checkAlloc(a, "1", net.IP{192, 168, 1, 1})
 		})
 	})
+	Context("IP address exclusion", func() {
+		It("should exclude IPs", func() {
+			session, err := storePkg.New(
+				filepath.Join(GinkgoT().TempDir(), "test_store")).Open(context.Background())
+			Expect(err).NotTo(HaveOccurred())
+			defer func() {
+				_ = session.Commit()
+			}()
+			p := allocator.RangeSet{
+				allocator.Range{Subnet: mustSubnet("192.168.0.0/29")},
+			}
+			Expect(p.Canonicalize()).NotTo(HaveOccurred())
+			e := allocator.RangeSet{
+				allocator.Range{Subnet: mustSubnet("192.168.0.0/29"),
+					RangeStart: net.ParseIP("192.168.0.2"), RangeEnd: net.ParseIP("192.168.0.3")},
+				allocator.Range{Subnet: mustSubnet("192.168.0.0/29"),
+					RangeStart: net.ParseIP("192.168.0.5"), RangeEnd: net.ParseIP("192.168.0.5")},
+			}
+			Expect(e.Canonicalize()).NotTo(HaveOccurred())
+			a := allocator.NewIPAllocator(&p, &e, testPoolName, session)
+			// get range iterator and do the first Next
+			checkAlloc(a, "0", net.IP{192, 168, 0, 1})
+			checkAlloc(a, "1", net.IP{192, 168, 0, 4})
+			checkAlloc(a, "2", net.IP{192, 168, 0, 6})
+		})
+	})
 })
diff --git a/pkg/ipam-node/handlers/allocate.go b/pkg/ipam-node/handlers/allocate.go
@@ -121,7 +121,7 @@ func (h *Handlers) allocateInPool(pool string, reqLog logr.Logger,
 		return PoolAlloc{}, poolCfgError(poolLog, pool,
 			fmt.Sprintf("invalid range config: %s", err.Error()))
 	}
-	alloc := h.getAllocFunc(rangeSet, pool, session)
+	alloc := h.getAllocFunc(rangeSet, &allocator.RangeSet{}, pool, session)
 	allocMeta := types.ReservationMetadata{
 		CreateTime:         time.Now().Format(time.RFC3339Nano),
 		PoolConfigSnapshot: poolCfg.String(),

diff --git a/pkg/ipam-node/handlers/handlers.go b/pkg/ipam-node/handlers/handlers.go
@@ -27,7 +27,8 @@ import (
 	poolPkg "github.com/Mellanox/nvidia-k8s-ipam/pkg/pool"
 )
 
-type GetAllocatorFunc = func(s *allocator.RangeSet, poolName string, session storePkg.Session) allocator.IPAllocator
+type GetAllocatorFunc = func(s *allocator.RangeSet, exclusions *allocator.RangeSet,
+	poolName string, session storePkg.Session) allocator.IPAllocator
 
 // New create and initialize new instance of grpc Handlers
 func New(poolConfReader poolPkg.ConfigReader, store storePkg.Store, getAllocFunc GetAllocatorFunc) *Handlers {

diff --git a/pkg/ipam-node/handlers/handlers_test.go b/pkg/ipam-node/handlers/handlers_test.go
@@ -95,7 +95,8 @@ var _ = Describe("Handlers", func() {
 		allocators = map[string]*allocatorMockPkg.IPAllocator{
 			testPoolName1: allocatorMockPkg.NewIPAllocator(GinkgoT()),
 			testPoolName2: allocatorMockPkg.NewIPAllocator(GinkgoT())}
-		getAllocFunc = func(s *allocatorPkg.RangeSet, poolName string, store storePkg.Session) allocatorPkg.IPAllocator {
+		getAllocFunc = func(s *allocatorPkg.RangeSet, exclusions *allocatorPkg.RangeSet,
+			poolName string, store storePkg.Session) allocatorPkg.IPAllocator {
 			return allocators[poolName]
 		}
 		handlers = handlersPkg.New(poolManager, store, getAllocFunc)