Create ip_allocator that copes with IPv6

It uses a set (via a map) of allocated IPs
This commit is contained in:
Justin SB 2014-11-06 21:26:26 -08:00
parent ddaa716de1
commit 04313ff173
3 changed files with 145 additions and 124 deletions

View File

@ -18,17 +18,67 @@ package service
import (
"fmt"
math_rand "math/rand"
"net"
"sync"
"time"
"github.com/golang/glog"
)
type ipAllocator struct {
subnet *net.IPNet
// TODO: This could be smarter, but for now a bitmap will suffice.
lock sync.Mutex // protects 'used'
used []byte // a bitmap of allocated IPs
subnet net.IPNet
ipSpaceSize int64 // Size of subnet, or -1 if it does not fit in an int64
used ipAddrSet
randomAttempts int
random *math_rand.Rand
}
type ipAddrSet struct {
// We are pretty severely restricted in the types of things we can use as a key
ips map[string]bool
}
func (s *ipAddrSet) Init() {
s.ips = map[string]bool{}
}
// Adds to the ipAddrSet; returns true iff it was added (was not already in set)
func (s *ipAddrSet) Size() int {
return len(s.ips)
}
func (s *ipAddrSet) Contains(ip net.IP) bool {
key := ip.String()
exists := s.ips[key]
return exists
}
// Adds to the ipAddrSet; returns true iff it was added (was not already in set)
func (s *ipAddrSet) Add(ip net.IP) bool {
key := ip.String()
exists := s.ips[key]
if exists {
return false
}
s.ips[key] = true
return true
}
// Removes from the ipAddrSet; returns true iff it was removed (was already in set)
func (s *ipAddrSet) Remove(ip net.IP) bool {
key := ip.String()
exists := s.ips[key]
if !exists {
return false
}
delete(s.ips, key)
// TODO: We probably should add this IP to an 'embargo' list for a limited amount of time
return true
}
// The smallest number of IPs we accept.
@ -40,20 +90,42 @@ func newIPAllocator(subnet *net.IPNet) *ipAllocator {
return nil
}
seed := time.Now().UTC().UnixNano()
r := math_rand.New(math_rand.NewSource(seed))
ipSpaceSize := int64(-1)
ones, bits := subnet.Mask.Size()
// TODO: some settings with IPv6 address could cause this to take
// an excessive amount of memory.
numIps := 1 << uint(bits-ones)
if numIps < minIPSpace {
glog.Errorf("IPAllocator requires at least %d IPs", minIPSpace)
return nil
if (bits-ones) < 63 {
ipSpaceSize = int64(1)<<uint(bits - ones)
if ipSpaceSize < minIPSpace {
glog.Errorf("IPAllocator requires at least %d IPs", minIPSpace)
return nil
}
}
ipa := &ipAllocator{
subnet: subnet,
used: make([]byte, numIps/8),
subnet: *subnet,
ipSpaceSize: ipSpaceSize,
random: r,
randomAttempts: 1000,
}
ipa.used[0] = 0x01 // block the network addr
ipa.used[(numIps/8)-1] = 0x80 // block the broadcast addr
ipa.used.Init()
zero := make(net.IP, len(subnet.IP), len(subnet.IP))
for i := 0; i < len(subnet.IP); i++ {
zero[i] = subnet.IP[i]&subnet.Mask[i]
}
ipa.used.Add(zero) // block the zero addr
ipa.used.Add(subnet.IP) // block the network addr
broadcast := make(net.IP, len(subnet.IP), len(subnet.IP))
for i := 0; i < len(subnet.IP); i++ {
broadcast[i] = subnet.IP[i]|^subnet.Mask[i]
}
ipa.used.Add(broadcast) // block the broadcast addr
return ipa
}
@ -65,13 +137,11 @@ func (ipa *ipAllocator) Allocate(ip net.IP) error {
if !ipa.subnet.Contains(ip) {
return fmt.Errorf("IP %s does not fall within subnet %s", ip, ipa.subnet)
}
offset := ipSub(ip, ipa.subnet.IP)
i := offset / 8
m := byte(1 << byte(offset%8))
if ipa.used[i]&m != 0 {
if !ipa.used.Add(ip) {
return fmt.Errorf("IP %s is already allocated", ip)
}
ipa.used[i] |= m
return nil
}
@ -80,33 +150,48 @@ func (ipa *ipAllocator) AllocateNext() (net.IP, error) {
ipa.lock.Lock()
defer ipa.lock.Unlock()
for i := range ipa.used {
if ipa.used[i] != 0xff {
freeMask := ^ipa.used[i]
nextBit, err := ffs(freeMask)
if err != nil {
// If this happens, something really weird is going on.
glog.Errorf("ffs(%#x) had an unexpected error: %s", freeMask, err)
return nil, err
}
ipa.used[i] |= 1 << nextBit
offset := (i * 8) + int(nextBit)
ip := ipAdd(ipa.subnet.IP, offset)
if int64(ipa.used.Size()) == ipa.ipSpaceSize {
return nil, fmt.Errorf("can't find a free IP in %s", ipa.subnet)
}
// Try randomly first
for i := 0; i < ipa.randomAttempts; i++ {
ip := ipa.createRandomIp()
if ipa.used.Add(ip) {
return ip, nil
}
}
// If that doesn't work, try a linear search
ip := copyIP(ipa.subnet.IP)
for ipa.subnet.Contains(ip) {
ip = ipAdd(ip, 1)
if ipa.used.Add(ip) {
return ip, nil
}
}
return nil, fmt.Errorf("can't find a free IP in %s", ipa.subnet)
}
// This is a really dumb implementation of find-first-set-bit.
func ffs(val byte) (uint, error) {
if val == 0 {
return 0, fmt.Errorf("Can't find-first-set on 0")
func (ipa *ipAllocator) createRandomIp() net.IP {
ip := ipa.subnet.IP
mask := ipa.subnet.Mask
n := len(ip)
randomIp := make(net.IP, n, n)
for i := 0; i < n; i++ {
if mask[i] == 0xff {
randomIp[i] = ipa.subnet.IP[i]
} else {
b := byte(ipa.random.Intn(256))
randomIp[i] = (ipa.subnet.IP[i]&mask[i])|(b&^mask[i])
}
}
i := uint(0)
for ; i < 8 && (val&(1<<i) == 0); i++ {
}
return i, nil
return randomIp
}
// Add an offset to an IP address - used for joining network addr and host addr parts.
@ -118,38 +203,11 @@ func ipAdd(ip net.IP, offset int) net.IP {
result := int(out[i]) + add
out[i] = byte(result % 256)
offset >>= 8
offset += result / 256 // carry
offset += result/256 // carry
}
return out
}
// Subtract two IPs, returning the difference as an offset - used or splitting an IP into
// network addr and host addr parts.
func ipSub(lhs, rhs net.IP) int {
// If they are not the same length, normalize them. Make copies because net.IP is
// a slice underneath. Sneaky sneaky.
lhs = simplifyIP(lhs)
rhs = simplifyIP(rhs)
if len(lhs) != len(rhs) {
lhs = copyIP(lhs).To16()
rhs = copyIP(rhs).To16()
}
offset := 0
borrow := 0
// Loop from most-significant to least.
for i := range lhs {
offset <<= 8
result := (int(lhs[i]) - borrow) - int(rhs[i])
if result < 0 {
borrow = 1
} else {
borrow = 0
}
offset += result
}
return offset
}
// Get the optimal slice for an IP. IPv4 addresses will come back in a 4 byte slice. IPv6
// addresses will come back in a 16 byte slice. Non-IP arguments will produce nil.
func simplifyIP(in net.IP) net.IP {
@ -176,9 +234,6 @@ func (ipa *ipAllocator) Release(ip net.IP) error {
if !ipa.subnet.Contains(ip) {
return fmt.Errorf("IP %s does not fall within subnet %s", ip, ipa.subnet)
}
offset := ipSub(ip, ipa.subnet.IP)
i := offset / 8
m := byte(1 << byte(offset%8))
ipa.used[i] &^= m
ipa.used.Remove(ip)
return nil
}

View File

@ -36,13 +36,16 @@ func TestNew(t *testing.T) {
if ipa == nil {
t.Errorf("expected non-nil")
}
if len(ipa.used) != 128 { // a /22 has 1024 IPs, 8 per byte = 128
t.Errorf("wrong size for ipa.used")
if ipa.ipSpaceSize != 1024 {
t.Errorf("wrong size for ipa.ipSpaceSize")
}
if ipa.used[0] != 0x01 {
if ipa.used.Size() != 2 {
t.Errorf("wrong size() for ipa.used")
}
if !ipa.used.Contains(net.ParseIP("93.76.0.0")) {
t.Errorf("network address was not reserved")
}
if ipa.used[127] != 0x80 {
if !ipa.used.Contains(net.ParseIP("93.76.3.255")) {
t.Errorf("broadcast address was not reserved")
}
}
@ -68,6 +71,9 @@ func TestAllocateNext(t *testing.T) {
_, ipnet, _ := net.ParseCIDR("93.76.0.0/22")
ipa := newIPAllocator(ipnet)
// Turn off random allocation attempts, so we just allocate in sequence
ipa.randomAttempts = 0
ip1, err := ipa.AllocateNext()
if err != nil {
t.Error(err)
@ -128,6 +134,8 @@ func TestRelease(t *testing.T) {
_, ipnet, _ := net.ParseCIDR("93.76.0.0/24")
ipa := newIPAllocator(ipnet)
ipa.randomAttempts = 0
err := ipa.Release(net.ParseIP("1.2.3.4"))
if err == nil {
t.Errorf("Expected an error")
@ -172,31 +180,6 @@ func TestRelease(t *testing.T) {
}
}
func TestFFS(t *testing.T) {
_, err := ffs(0)
if err == nil {
t.Errorf("Expected error")
}
testCases := []struct {
value byte
expected uint
}{
{0x01, 0}, {0x02, 1}, {0x04, 2}, {0x08, 3},
{0x10, 4}, {0x20, 5}, {0x40, 6}, {0x80, 7},
{0x22, 1}, {0xa0, 5}, {0xfe, 1}, {0xff, 0},
}
for _, tc := range testCases {
r, err := ffs(tc.value)
if err != nil {
t.Error(err)
}
if r != tc.expected {
t.Errorf("Expected %d, got %d", tc.expected, r)
}
}
}
func TestIPAdd(t *testing.T) {
testCases := []struct {
ip string
@ -223,31 +206,6 @@ func TestIPAdd(t *testing.T) {
}
}
func TestIPSub(t *testing.T) {
testCases := []struct {
lhs string
rhs string
expected int
}{
{"1.2.3.0", "1.2.3.0", 0},
{"1.2.3.1", "1.2.3.0", 1},
{"1.2.3.255", "1.2.3.0", 255},
{"1.2.4.0", "1.2.3.0", 256},
{"1.2.4.0", "1.2.3.255", 1},
{"1.2.4.1", "1.2.3.0", 257},
{"1.3.3.0", "1.2.3.0", 65536},
{"1.2.3.5", "1.2.3.4", 1},
{"0.0.0.0", "0.0.0.1", -1},
{"0.0.1.0", "0.0.0.1", 255},
}
for _, tc := range testCases {
r := ipSub(net.ParseIP(tc.lhs), net.ParseIP(tc.rhs))
if r != tc.expected {
t.Errorf("Expected %v, got %v", tc.expected, r)
}
}
}
func TestCopyIP(t *testing.T) {
ip1 := net.ParseIP("1.2.3.4")
ip2 := copyIP(ip1)

View File

@ -44,6 +44,8 @@ func TestServiceRegistryCreate(t *testing.T) {
fakeCloud := &cloud.FakeCloud{}
machines := []string{"foo", "bar", "baz"}
storage := NewREST(registry, fakeCloud, registrytest.NewMinionRegistry(machines, api.NodeResources{}), makeIPNet(t))
storage.portalMgr.randomAttempts = 0
svc := &api.Service{
Port: 6502,
ObjectMeta: api.ObjectMeta{Name: "foo"},
@ -414,6 +416,7 @@ func TestServiceRegistryIPAllocation(t *testing.T) {
fakeCloud := &cloud.FakeCloud{}
machines := []string{"foo", "bar", "baz"}
rest := NewREST(registry, fakeCloud, registrytest.NewMinionRegistry(machines, api.NodeResources{}), makeIPNet(t))
rest.portalMgr.randomAttempts = 0
svc1 := &api.Service{
Port: 6502,
@ -467,6 +470,7 @@ func TestServiceRegistryIPReallocation(t *testing.T) {
fakeCloud := &cloud.FakeCloud{}
machines := []string{"foo", "bar", "baz"}
rest := NewREST(registry, fakeCloud, registrytest.NewMinionRegistry(machines, api.NodeResources{}), makeIPNet(t))
rest.portalMgr.randomAttempts = 0
svc1 := &api.Service{
Port: 6502,
@ -509,6 +513,7 @@ func TestServiceRegistryIPUpdate(t *testing.T) {
fakeCloud := &cloud.FakeCloud{}
machines := []string{"foo", "bar", "baz"}
rest := NewREST(registry, fakeCloud, registrytest.NewMinionRegistry(machines, api.NodeResources{}), makeIPNet(t))
rest.portalMgr.randomAttempts = 0
svc := &api.Service{
Port: 6502,
@ -561,6 +566,7 @@ func TestServiceRegistryIPExternalLoadBalancer(t *testing.T) {
fakeCloud := &cloud.FakeCloud{}
machines := []string{"foo", "bar", "baz"}
rest := NewREST(registry, fakeCloud, registrytest.NewMinionRegistry(machines, api.NodeResources{}), makeIPNet(t))
rest.portalMgr.randomAttempts = 0
svc := &api.Service{
Port: 6502,
@ -588,6 +594,7 @@ func TestServiceRegistryIPReloadFromStorage(t *testing.T) {
fakeCloud := &cloud.FakeCloud{}
machines := []string{"foo", "bar", "baz"}
rest1 := NewREST(registry, fakeCloud, registrytest.NewMinionRegistry(machines, api.NodeResources{}), makeIPNet(t))
rest1.portalMgr.randomAttempts = 0
svc := &api.Service{
Port: 6502,
@ -607,6 +614,7 @@ func TestServiceRegistryIPReloadFromStorage(t *testing.T) {
// This will reload from storage, finding the previous 2
rest2 := NewREST(registry, fakeCloud, registrytest.NewMinionRegistry(machines, api.NodeResources{}), makeIPNet(t))
rest2.portalMgr.randomAttempts = 0
svc = &api.Service{
Port: 6502,