find optimal prefix to allocate to avoid fragmentation

ipv4/set.go

@@ -1,6 +1,7 @@
 package ipv4
 
 import (
+	"fmt"
 	"strings"
 )
 
@@ -339,3 +340,19 @@ func (me Set) Difference(other SetI) Set {
 func (me Set) isValid() bool {
 	return me.trie.isValid()
 }
+
+// FindAvailablePrefix returns a Prefix with a Mask of the given prefix length
+// that is contained by the current set but does not overlap the given reserved
+// set. The returned Prefix is optimally placed to avoid any further IP space
+// fragmentation. An error is returned if there is not enough space to allocate
+func (me Set) FindAvailablePrefix(reserved SetI, length uint32) (Prefix, error) {
+	prefix, err := me.trie.FindSmallestContainingPrefix(reserved.Set().trie, length)
+	if err != nil {
+		return Prefix{}, fmt.Errorf("no room for prefix of given length")
+	}
+
+	return Prefix{
+		addr:   prefix.Network().addr,
+		length: length,
+	}, nil
+}

ipv4/set_test.go

@@ -1,6 +1,7 @@
 package ipv4
 
 import (
+	"math"
 	"math/rand"
 	"strconv"
 	"sync"
@@ -953,3 +954,210 @@ func TestSetNumPrefixesStairs(t *testing.T) {
 		})
 	}
 }
+
+func TestFindAvailablePrefix(t *testing.T) {
+	tests := []struct {
+		description string
+		space       []SetI
+		reserved    []SetI
+		length      uint32
+		expected    Prefix
+		err         bool
+		change      int
+	}{
+		{
+			description: "empty",
+			space: []SetI{
+				_p("10.0.0.0/8"),
+			},
+			length: 24,
+			change: 1,
+		}, {
+			description: "find adjacent",
+			space: []SetI{
+				_p("10.0.0.0/8"),
+			},
+			reserved: []SetI{
+				_p("10.224.123.0/24"),
+			},
+			length:   24,
+			expected: _p("10.224.122.0/24"),
+		}, {
+			description: "many fewer prefixes",
+			space: []SetI{
+				_p("10.0.0.0/16"),
+			},
+			reserved: []SetI{
+				_p("10.0.1.0/24"),
+				_p("10.0.2.0/23"),
+				_p("10.0.4.0/22"),
+				_p("10.0.8.0/21"),
+				_p("10.0.16.0/20"),
+				_p("10.0.32.0/19"),
+				_p("10.0.64.0/18"),
+				_p("10.0.128.0/17"),
+			},
+			length: 24,
+			change: -7,
+		}, {
+			description: "toobig",
+			space: []SetI{
+				_p("10.0.0.0/8"),
+			},
+			reserved: []SetI{
+				_p("10.128.0.0/9"),
+				_p("10.64.0.0/10"),
+				_p("10.32.0.0/11"),
+				_p("10.16.0.0/12"),
+			},
+			length: 11,
+			err:    true,
+		}, {
+			description: "full",
+			space: []SetI{
+				_p("10.0.0.0/8"),
+			},
+			length: 7,
+			err:    true,
+		}, {
+			description: "random disjoint example",
+			space: []SetI{
+				_p("10.0.0.0/22"),
+				_p("192.168.0.0/21"),
+				_p("172.16.0.0/20"),
+			},
+			reserved: []SetI{
+				_p("192.168.0.0/21"),
+				_p("172.16.0.0/21"),
+				_p("172.16.8.0/22"),
+				_p("10.0.0.0/22"),
+				_p("172.16.12.0/24"),
+				_p("172.16.14.0/24"),
+				_p("172.16.15.0/24"),
+			},
+			length:   24,
+			expected: _p("172.16.13.0/24"),
+			change:   1,
+		}, {
+			description: "too fragmented",
+			space: []SetI{
+				_p("10.0.0.0/24"),
+			},
+			reserved: []SetI{
+				_p("10.0.0.0/27"),
+				_p("10.0.0.64/27"),
+				_p("10.0.0.128/27"),
+				_p("10.0.0.192/27"),
+			},
+			length: 25,
+			err:    true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.description, func(t *testing.T) {
+			// This is the full usable IP space
+			space := Set{}.Build(func(s_ Set_) bool {
+				for _, p := range tt.space {
+					s_.Insert(p)
+				}
+				return true
+			})
+			// This is the part of the usable space which has already been allocated
+			reserved := Set{}.Build(func(s_ Set_) bool {
+				for _, p := range tt.reserved {
+					s_.Insert(p)
+				}
+				return true
+			})
+
+			// Call the method under test to find the best allocation to avoid fragmentation.
+			prefix, err := space.FindAvailablePrefix(reserved, tt.length)
+
+			assert.Equal(t, tt.err, err != nil)
+			if err != nil {
+				return
+			}
+
+			assert.Equal(t, int64(0), reserved.Intersection(prefix).NumAddresses())
+
+			// Not all test cases care which prefix is returned but in some
+			// cases, there is only one right answer and so we might check it.
+			// This isn't strictly necessary but was handy with the first few.
+			if tt.expected.length != 0 {
+				assert.Equal(t, tt.expected.String(), prefix.String())
+			}
+
+			// What really matters is that fragmentation in the IP space is
+			// always avoided as much as possible. The `change` field in each
+			// test indicates what should happen to IP space fragmentation.
+			// This test framework measures fragmentation as the change in the
+			// minimal number of prefixes required to span the reserved set.
+			before := countPrefixes(reserved)
+			after := countPrefixes(reserved.Build(func(s_ Set_) bool {
+				s_.Insert(prefix)
+				return true
+			}))
+
+			diff := after - before
+			assert.LessOrEqual(t, diff, 1)
+			assert.LessOrEqual(t, diff, tt.change)
+		})
+	}
+
+	t.Run("randomized", func(t *testing.T) {
+		// Start with a space and an empty reserved set.
+		// This test will attempt to fragment the space by pulling out
+		space := _p("10.128.0.0/12").Set()
+		available := space.NumAddresses()
+
+		reserved := NewSet_()
+
+		rand.Seed(29)
+		for available > 0 {
+			// This is the most we can pull out. Assuming we avoid
+			// fragmentation, it should be the largest power of two that is
+			// less than or equal to the number of available addresses.
+			maxExponent := log2(available)
+
+			// Finding the maximum prefix here, proves we are avoiding fragmentation
+			maxPrefix, err := space.FindAvailablePrefix(reserved, 32-maxExponent)
+			assert.Nil(t, err)
+			assert.Equal(t, pow2(maxExponent), maxPrefix.NumAddresses())
+			assert.Equal(t, int64(0), reserved.Intersection(maxPrefix).NumAddresses())
+
+			// Pull out a random sized prefix up to the maximum size to attempt to further fragment the space.
+			randomSize := (rand.Uint32()%maxExponent + 1)
+			if randomSize > 12 {
+				randomSize = 12
+			}
+
+			randomSizePrefix, err := space.FindAvailablePrefix(reserved, 32-randomSize)
+			assert.Nil(t, err)
+			assert.Equal(t, pow2(randomSize), randomSizePrefix.NumAddresses())
+			assert.Equal(t, int64(0), reserved.Intersection(randomSizePrefix).NumAddresses())
+
+			// Reserve only the random sized one
+			reserved.Insert(randomSizePrefix)
+			available -= randomSizePrefix.NumAddresses()
+			assert.Equal(t, available, space.NumAddresses()-reserved.NumAddresses())
+		}
+	})
+}
+
+func pow2(x uint32) int64 {
+	return int64(math.Pow(2, float64(x)))
+}
+
+func log2(available_addresses int64) uint32 {
+	return uint32(math.Log2(float64(available_addresses)))
+}
+
+func countPrefixes(s Set) int {
+	var numPrefixes int
+	s.WalkPrefixes(func(_ Prefix) bool {
+		numPrefixes += 1
+		return true
+	})
+	return numPrefixes
+}

ipv4/setnode.go

@@ -1,6 +1,7 @@
 package ipv4
 
 import (
+	"fmt"
 	"math/bits"
 )
 
@@ -296,3 +297,110 @@ func (me *setNode) height() int {
 func (me *setNode) Walk(callback func(Prefix, interface{}) bool) bool {
 	return (*trieNode)(me).Walk(callback)
 }
+
+func best(left, right func() (Prefix, error), length uint32) (Prefix, error) {
+	lPrefix, lErr := left()
+	if lErr == nil {
+		if lPrefix.length == length {
+			return lPrefix, nil
+		}
+		rPrefix, rErr := right()
+		if rErr == nil {
+			if lPrefix.length < rPrefix.length {
+				return rPrefix, nil
+			} else {
+				return lPrefix, nil
+			}
+		}
+		return lPrefix, nil
+	}
+
+	rPrefix, rErr := right()
+	if rErr == nil {
+		return rPrefix, nil
+	}
+	return Prefix{}, fmt.Errorf("cannot find containing prefix")
+}
+
+func (me *setNode) findSmallestContainingPrefix(length uint32) (Prefix, error) {
+	if me == nil || length < me.Prefix.length {
+		return Prefix{}, fmt.Errorf("cannot find containing prefix")
+	}
+	if length == me.Prefix.length {
+		if me.isActive {
+			return me.Prefix, nil
+		}
+	}
+
+	l, r := (*setNode)(me.children[0]), (*setNode)(me.children[1])
+	bestPrefix, err := best(
+		func() (Prefix, error) { return l.findSmallestContainingPrefix(length) },
+		func() (Prefix, error) { return r.findSmallestContainingPrefix(length) },
+		length,
+	)
+	if err == nil {
+		return bestPrefix, nil
+	}
+	if !me.isActive {
+		return Prefix{}, fmt.Errorf("cannot find containing prefix")
+	}
+	return me.Prefix, nil
+}
+
+func (me *setNode) FindSmallestContainingPrefix(reserved *setNode, length uint32) (Prefix, error) {
+	if me == nil || length < me.Prefix.length {
+		return Prefix{}, fmt.Errorf("cannot find containing prefix")
+	}
+	if reserved == nil {
+		return me.findSmallestContainingPrefix(length)
+	}
+
+	result, _, _, child := compare(me.Prefix, reserved.Prefix)
+	switch result {
+	case compareIsContained:
+		if reserved.isActive {
+			return Prefix{}, fmt.Errorf("cannot find containing prefix")
+		}
+		return me.FindSmallestContainingPrefix((*setNode)(reserved.children[child]), length)
+	case compareDisjoint:
+		return me.findSmallestContainingPrefix(length)
+	}
+
+	if !me.isActive {
+		return best(
+			func() (Prefix, error) { return me.Left().FindSmallestContainingPrefix(reserved, length) },
+			func() (Prefix, error) { return me.Right().FindSmallestContainingPrefix(reserved, length) },
+			length,
+		)
+	}
+
+	// Assumes `me` is active as checked above
+	halves := func() (a, b *setNode) {
+		aPrefix, bPrefix := me.Prefix.Halves()
+		return setNodeFromPrefix(aPrefix), setNodeFromPrefix(bPrefix)
+	}
+
+	switch result {
+	case compareSame:
+		if reserved.isActive {
+			return Prefix{}, fmt.Errorf("cannot find containing prefix")
+		}
+		left, right := halves()
+		return best(
+			func() (Prefix, error) { return left.FindSmallestContainingPrefix(reserved.Left(), length) },
+			func() (Prefix, error) { return right.FindSmallestContainingPrefix(reserved.Right(), length) },
+			length,
+		)
+
+	case compareContains:
+		left, right := halves()
+		halves := [2]*setNode{left, right}
+		whole, partial := halves[(child+1)%2], halves[child]
+		return best(
+			func() (Prefix, error) { return whole.findSmallestContainingPrefix(length) },
+			func() (Prefix, error) { return partial.FindSmallestContainingPrefix(reserved, length) },
+			length,
+		)
+	}
+	panic("unreachable")
+}

ipv6/set.go

@@ -328,3 +328,19 @@ func (me Set) Difference(other SetI) Set {
 func (me Set) isValid() bool {
 	return me.trie.isValid()
 }
+
+// FindAvailablePrefix returns a Prefix with a Mask of the given prefix length
+// that is contained by the current set but does not overlap the given reserved
+// set. The returned Prefix is optimally placed to avoid any further IP space
+// fragmentation. An error is returned if there is not enough space to allocate
+func (me Set) FindAvailablePrefix(reserved SetI, length uint32) (Prefix, error) {
+	prefix, err := me.trie.FindSmallestContainingPrefix(reserved.Set().trie, length)
+	if err != nil {
+		return Prefix{}, fmt.Errorf("no room for prefix of given length")
+	}
+
+	return Prefix{
+		addr:   prefix.Network().addr,
+		length: length,
+	}, nil
+}