Set difference when picking random ports

nomad/structs/bitmap.go

@@ -6,8 +6,8 @@ import "fmt"
 type Bitmap []byte
 
 // NewBitmap returns a bitmap with up to size indexes
-func NewBitmap(size int) (Bitmap, error) {
-	if size <= 0 {
+func NewBitmap(size uint) (Bitmap, error) {
+	if size == 0 {
 		return nil, fmt.Errorf("bitmap must be positive size")
 	}
 	if size&7 != 0 {
@@ -17,6 +17,22 @@ func NewBitmap(size int) (Bitmap, error) {
 	return Bitmap(b), nil
 }
 
+// Copy returns a copy of the Bitmap
+func (b Bitmap) Copy() (Bitmap, error) {
+	if b == nil {
+		return nil, fmt.Errorf("can't copy nil Bitmap")
+	}
+
+	raw := make([]byte, len(b))
+	copy(raw, b)
+	return Bitmap(raw), nil
+}
+
+// Size returns the size of the bitmap
+func (b Bitmap) Size() uint {
+	return uint(len(b) << 3)
+}
+
 // Set is used to set the given index of the bitmap
 func (b Bitmap) Set(idx uint) {
 	bucket := idx >> 3
@@ -37,3 +53,17 @@ func (b Bitmap) Clear() {
 		b[i] = 0
 	}
 }
+
+// IndexesInRange returns the indexes in which the values are either set or unset based
+// on the passed parameter in the passed range
+func (b Bitmap) IndexesInRange(set bool, from, to uint) []int {
+	var indexes []int
+	for i := from; i < to; i++ {
+		c := b.Check(i)
+		if c && set || !c && !set {
+			indexes = append(indexes, int(i))
+		}
+	}
+
+	return indexes
+}

nomad/structs/bitmap_test.go

@@ -1,6 +1,9 @@
 package structs
 
-import "testing"
+import (
+	"reflect"
+	"testing"
+)
 
 func TestBitmap(t *testing.T) {
 	// Check invalid sizes
@@ -14,11 +17,16 @@ func TestBitmap(t *testing.T) {
 	}
 
 	// Create a normal bitmap
-	b, err := NewBitmap(256)
+	var s uint = 256
+	b, err := NewBitmap(s)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
 
+	if b.Size() != s {
+		t.Fatalf("bad size")
+	}
+
 	// Set a few bits
 	b.Set(0)
 	b.Set(255)
@@ -46,6 +54,39 @@ func TestBitmap(t *testing.T) {
 		}
 	}
 
+	// Check the indexes
+	idxs := b.IndexesInRange(true, 0, 256)
+	expected := []int{0, 255}
+	if !reflect.DeepEqual(idxs, expected) {
+		t.Fatalf("bad: got %#v; want %#v", idxs, expected)
+	}
+
+	idxs = b.IndexesInRange(true, 1, 256)
+	expected = []int{255}
+	if !reflect.DeepEqual(idxs, expected) {
+		t.Fatalf("bad: got %#v; want %#v", idxs, expected)
+	}
+
+	idxs = b.IndexesInRange(false, 0, 256)
+	if len(idxs) != 254 {
+		t.Fatalf("bad")
+	}
+
+	idxs = b.IndexesInRange(false, 100, 200)
+	if len(idxs) != 100 {
+		t.Fatalf("bad")
+	}
+
+	// Check the copy is correct
+	b2, err := b.Copy()
+	if err != nil {
+		t.Fatalf("bad: %v", err)
+	}
+
+	if !reflect.DeepEqual(b, b2) {
+		t.Fatalf("bad")
+	}
+
 	// Clear
 	b.Clear()
 

nomad/structs/network.go

@@ -183,6 +183,8 @@ func (idx *NetworkIndex) AssignNetwork(ask *NetworkResource) (out *NetworkResour
 			return
 		}
 
+		used := idx.UsedPorts[ipStr]
+
 		// Check if any of the reserved ports are in use
 		for _, port := range ask.ReservedPorts {
 			// Guard against invalid port
@@ -192,7 +194,6 @@ func (idx *NetworkIndex) AssignNetwork(ask *NetworkResource) (out *NetworkResour
 			}
 
 			// Check if in use
-			used := idx.UsedPorts[ipStr]
 			if used != nil && used.Check(uint(port.Value)) {
 				err = fmt.Errorf("reserved port collision")
 				return
@@ -208,28 +209,25 @@ func (idx *NetworkIndex) AssignNetwork(ask *NetworkResource) (out *NetworkResour
 			DynamicPorts:  ask.DynamicPorts,
 		}
 
-		// Check if we need to generate any ports
-		for i := 0; i < len(ask.DynamicPorts); i++ {
-			attempts := 0
-		PICK:
-			attempts++
-			if attempts > maxRandPortAttempts {
-				err = fmt.Errorf("dynamic port selection failed")
-				return
-			}
+		// Try to stochastically pick the dynamic ports as it is faster and
+		// lower memory usage.
+		var dynPorts []int
+		var dynErr error
+		dynPorts, dynErr = getDynamicPortsStochastic(used, ask)
+		if err == nil {
+			goto BUILD_OFFER
+		}
 
-			randPort := MinDynamicPort + rand.Intn(MaxDynamicPort-MinDynamicPort)
-			used := idx.UsedPorts[ipStr]
-			if used != nil && used.Check(uint(randPort)) {
-				goto PICK
-			}
+		// Fall back to the precise method if the random sampling failed.
+		dynPorts, dynErr = getDynamicPortsPrecise(used, ask)
+		if dynErr != nil {
+			err = dynErr
+			return
+		}
 
-			for _, ports := range [][]Port{offer.ReservedPorts, offer.DynamicPorts} {
-				if isPortReserved(ports, randPort) {
-					goto PICK
-				}
-			}
-			offer.DynamicPorts[i].Value = randPort
+	BUILD_OFFER:
+		for i, port := range dynPorts {
+			offer.DynamicPorts[i].Value = port
 		}
 
 		// Stop, we have an offer!
@@ -240,10 +238,86 @@ func (idx *NetworkIndex) AssignNetwork(ask *NetworkResource) (out *NetworkResour
 	return
 }
 
+// getDynamicPortsPrecise takes the nodes used port bitmap which may be nil if
+// no ports have been allocated yet, the network ask and returns a set of unused
+// ports to fullfil the ask's DynamicPorts or an error if it failed. An error
+// means the ask can not be satisfied as the method does a precise search.
+func getDynamicPortsPrecise(nodeUsed Bitmap, ask *NetworkResource) ([]int, error) {
+	// Create a copy of the used ports and apply the new reserves
+	var usedSet Bitmap
+	var err error
+	if nodeUsed != nil {
+		usedSet, err = nodeUsed.Copy()
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		usedSet, err = NewBitmap(maxValidPort)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	for _, port := range ask.ReservedPorts {
+		usedSet.Set(uint(port.Value))
+	}
+
+	// Get the indexes of the unset
+	availablePorts := usedSet.IndexesInRange(false, MinDynamicPort, MaxDynamicPort)
+
+	// Randomize the amount we need
+	numDyn := len(ask.DynamicPorts)
+	if len(availablePorts) < numDyn {
+		return nil, fmt.Errorf("dynamic port selection failed")
+	}
+
+	for i := 0; i < numDyn; i++ {
+		j := rand.Intn(numDyn)
+		availablePorts[i], availablePorts[j] = availablePorts[j], availablePorts[i]
+	}
+
+	return availablePorts[:numDyn], nil
+}
+
+// getDynamicPortsStochastic takes the nodes used port bitmap which may be nil if
+// no ports have been allocated yet, the network ask and returns a set of unused
+// ports to fullfil the ask's DynamicPorts or an error if it failed. An error
+// does not mean the ask can not be satisfied as the method has a fixed amount
+// of random probes and if these fail, the search is aborted.
+func getDynamicPortsStochastic(nodeUsed Bitmap, ask *NetworkResource) ([]int, error) {
+	var reserved, dynamic []int
+	for _, port := range ask.ReservedPorts {
+		reserved = append(reserved, port.Value)
+	}
+
+	for i := 0; i < len(ask.DynamicPorts); i++ {
+		attempts := 0
+	PICK:
+		attempts++
+		if attempts > maxRandPortAttempts {
+			return nil, fmt.Errorf("stochastic dynamic port selection failed")
+		}
+
+		randPort := MinDynamicPort + rand.Intn(MaxDynamicPort-MinDynamicPort)
+		if nodeUsed != nil && nodeUsed.Check(uint(randPort)) {
+			goto PICK
+		}
+
+		for _, ports := range [][]int{reserved, dynamic} {
+			if isPortReserved(ports, randPort) {
+				goto PICK
+			}
+		}
+		dynamic = append(dynamic, randPort)
+	}
+
+	return dynamic, nil
+}
+
 // IntContains scans an integer slice for a value
-func isPortReserved(haystack []Port, needle int) bool {
+func isPortReserved(haystack []int, needle int) bool {
 	for _, item := range haystack {
-		if item.Value == needle {
+		if item == needle {
 			return true
 		}
 	}

nomad/structs/network_test.go

@@ -343,8 +343,62 @@ func TestNetworkIndex_AssignNetwork(t *testing.T) {
 	}
 }
 
+// This test ensures that even with a small domain of available ports we are
+// able to make a dynamic port allocation.
+func TestNetworkIndex_AssignNetwork_Dynamic_Contention(t *testing.T) {
+
+	// Create a node that only has one free port
+	idx := NewNetworkIndex()
+	n := &Node{
+		Resources: &Resources{
+			Networks: []*NetworkResource{
+				&NetworkResource{
+					Device: "eth0",
+					CIDR:   "192.168.0.100/32",
+					MBits:  1000,
+				},
+			},
+		},
+		Reserved: &Resources{
+			Networks: []*NetworkResource{
+				&NetworkResource{
+					Device: "eth0",
+					IP:     "192.168.0.100",
+					MBits:  1,
+				},
+			},
+		},
+	}
+	for i := MinDynamicPort; i < MaxDynamicPort-1; i++ {
+		n.Reserved.Networks[0].ReservedPorts = append(n.Reserved.Networks[0].ReservedPorts, Port{Value: i})
+	}
+
+	idx.SetNode(n)
+
+	// Ask for dynamic ports
+	ask := &NetworkResource{
+		DynamicPorts: []Port{{"http", 0}},
+	}
+	offer, err := idx.AssignNetwork(ask)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if offer == nil {
+		t.Fatalf("bad")
+	}
+	if offer.IP != "192.168.0.100" {
+		t.Fatalf("bad: %#v", offer)
+	}
+	if len(offer.DynamicPorts) != 1 {
+		t.Fatalf("There should be three dynamic ports")
+	}
+	if p := offer.DynamicPorts[0].Value; p == MaxDynamicPort {
+		t.Fatalf("Dynamic Port: should have been assigned %d; got %d", p, MaxDynamicPort)
+	}
+}
+
 func TestIntContains(t *testing.T) {
-	l := []Port{{"one", 1}, {"two", 2}, {"ten", 10}, {"twenty", 20}}
+	l := []int{1, 2, 10, 20}
 	if isPortReserved(l, 50) {
 		t.Fatalf("bad")
 	}