From 2b5a10f8ff90979acd345ca4b03a4582c8c7f444 Mon Sep 17 00:00:00 2001
From: Tim Hockin <thockin@google.com>
Date: Fri, 23 Dec 2016 00:46:14 -0800
Subject: [PATCH] Add a ForEach() to IP allocator

---
 .../core/service/ipallocator/allocator.go     |  9 +++++
 .../service/ipallocator/allocator_test.go     | 37 +++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/pkg/registry/core/service/ipallocator/allocator.go b/pkg/registry/core/service/ipallocator/allocator.go
index 85e0b5f6797..2ef4c6edda2 100644
--- a/pkg/registry/core/service/ipallocator/allocator.go
+++ b/pkg/registry/core/service/ipallocator/allocator.go
@@ -32,6 +32,7 @@ type Interface interface {
 	Allocate(net.IP) error
 	AllocateNext() (net.IP, error)
 	Release(net.IP) error
+	ForEach(func(net.IP))
 }
 
 var (
@@ -146,6 +147,14 @@ func (r *Range) Release(ip net.IP) error {
 	return r.alloc.Release(offset)
 }
 
+// ForEach calls the provided function for each allocated IP.
+func (r *Range) ForEach(fn func(net.IP)) {
+	r.alloc.ForEach(func(offset int) {
+		ip, _ := GetIndexedIP(r.net, offset+1) // +1 because Range doesn't store IP 0
+		fn(ip)
+	})
+}
+
 // Has returns true if the provided IP is already allocated and a call
 // to Allocate(ip) would fail with ErrAllocated.
 func (r *Range) Has(ip net.IP) bool {
diff --git a/pkg/registry/core/service/ipallocator/allocator_test.go b/pkg/registry/core/service/ipallocator/allocator_test.go
index d8e5bfb73d2..f87707cd771 100644
--- a/pkg/registry/core/service/ipallocator/allocator_test.go
+++ b/pkg/registry/core/service/ipallocator/allocator_test.go
@@ -167,6 +167,43 @@ func TestRangeSize(t *testing.T) {
 	}
 }
 
+func TestForEach(t *testing.T) {
+	_, cidr, err := net.ParseCIDR("192.168.1.0/24")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	testCases := []sets.String{
+		sets.NewString(),
+		sets.NewString("192.168.1.1"),
+		sets.NewString("192.168.1.1", "192.168.1.254"),
+		sets.NewString("192.168.1.1", "192.168.1.128", "192.168.1.254"),
+	}
+
+	for i, tc := range testCases {
+		r := NewCIDRRange(cidr)
+		for ips := range tc {
+			ip := net.ParseIP(ips)
+			if err := r.Allocate(ip); err != nil {
+				t.Errorf("[%d] error allocating IP %v: %v", i, ip, err)
+			}
+			if !r.Has(ip) {
+				t.Errorf("[%d] expected IP %v allocated", i, ip)
+			}
+		}
+		calls := sets.NewString()
+		r.ForEach(func(ip net.IP) {
+			calls.Insert(ip.String())
+		})
+		if len(calls) != len(tc) {
+			t.Errorf("[%d] expected %d calls, got %d", i, len(tc), len(calls))
+		}
+		if !calls.Equal(tc) {
+			t.Errorf("[%d] expected calls to equal testcase: %v vs %v", i, calls.List(), tc.List())
+		}
+	}
+}
+
 func TestSnapshot(t *testing.T) {
 	_, cidr, err := net.ParseCIDR("192.168.1.0/24")
 	if err != nil {