Add a ForEach() to IP allocator

This commit is contained in:
Tim Hockin 2016-12-23 00:46:14 -08:00
parent f75bed8682
commit 2b5a10f8ff
2 changed files with 46 additions and 0 deletions

View File

@ -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 {

View File

@ -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 {