Add port allocator and move ipset into orderedintset

Docker-DCO-1.1-Signed-off-by: Michael Crosby <michael@crosbymichael.com> (github: crosbymichael)

networkdriver/ipallocator/allocator.go

@@ -4,11 +4,13 @@ import (
 	"encoding/binary"
 	"errors"
 	"github.com/dotcloud/docker/networkdriver"
+	"github.com/dotcloud/docker/pkg/collections"
+	"github.com/dotcloud/docker/pkg/netlink"
 	"net"
 	"sync"
 )
 
-type networkSet map[string]*iPSet
+type networkSet map[iPNet]*collections.OrderedIntSet
 
 var (
 	ErrNoAvailableIPs     = errors.New("no available ip addresses on network")

networkdriver/portallocator/portallocator.go

@@ -0,0 +1,103 @@
+package portallocator
+
+import (
+	"errors"
+	"github.com/dotcloud/docker/pkg/collections"
+	"sync"
+)
+
+type portMappings map[string]*collections.OrderedIntSet
+
+const (
+	BeginPortRange = 49153
+	EndPortRange   = 65535
+)
+
+var (
+	ErrPortAlreadyAllocated = errors.New("port has already been allocated")
+	ErrPortExceedsRange     = errors.New("port exceeds upper range")
+	ErrUnknownProtocol      = errors.New("unknown protocol")
+)
+
+var (
+	lock           = sync.Mutex{}
+	allocatedPorts = portMappings{}
+	availablePorts = portMappings{}
+)
+
+func init() {
+	allocatedPorts["udp"] = collections.NewOrderedIntSet()
+	availablePorts["udp"] = collections.NewOrderedIntSet()
+	allocatedPorts["tcp"] = collections.NewOrderedIntSet()
+	availablePorts["tcp"] = collections.NewOrderedIntSet()
+}
+
+// RequestPort returns an available port if the port is 0
+// If the provided port is not 0 then it will be checked if
+// it is available for allocation
+func RequestPort(proto string, port int) (int, error) {
+	lock.Lock()
+	defer lock.Unlock()
+
+	if err := validateProtocol(proto); err != nil {
+		return 0, err
+	}
+
+	var (
+		allocated = allocatedPorts[proto]
+		available = availablePorts[proto]
+	)
+
+	if port != 0 {
+		if allocated.Exists(port) {
+			return 0, ErrPortAlreadyAllocated
+		}
+		available.Remove(port)
+		allocated.Push(port)
+		return port, nil
+	}
+
+	next := available.Pop()
+	if next == 0 {
+		next = allocated.PullBack()
+		if next == 0 {
+			next = BeginPortRange
+		} else {
+			next++
+		}
+		if next > EndPortRange {
+			return 0, ErrPortExceedsRange
+		}
+	}
+
+	allocated.Push(next)
+	return next, nil
+}
+
+// ReleasePort will return the provided port back into the
+// pool for reuse
+func ReleasePort(proto string, port int) error {
+	lock.Lock()
+	defer lock.Unlock()
+
+	if err := validateProtocol(proto); err != nil {
+		return err
+	}
+
+	var (
+		allocated = allocatedPorts[proto]
+		available = availablePorts[proto]
+	)
+
+	allocated.Remove(port)
+	available.Push(port)
+
+	return nil
+}
+
+func validateProtocol(proto string) error {
+	if _, exists := allocatedPorts[proto]; !exists {
+		return ErrUnknownProtocol
+	}
+	return nil
+}

networkdriver/portallocator/portallocator_test.go

@@ -0,0 +1,150 @@
+package portallocator
+
+import (
+	"github.com/dotcloud/docker/pkg/collections"
+	"testing"
+)
+
+func reset() {
+	lock.Lock()
+	defer lock.Unlock()
+
+	allocatedPorts = portMappings{}
+	availablePorts = portMappings{}
+
+	allocatedPorts["udp"] = collections.NewOrderedIntSet()
+	availablePorts["udp"] = collections.NewOrderedIntSet()
+	allocatedPorts["tcp"] = collections.NewOrderedIntSet()
+	availablePorts["tcp"] = collections.NewOrderedIntSet()
+}
+
+func TestRequestNewPort(t *testing.T) {
+	defer reset()
+
+	port, err := RequestPort("tcp", 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if expected := BeginPortRange; port != expected {
+		t.Fatalf("Expected port %d got %d", expected, port)
+	}
+}
+
+func TestRequestSpecificPort(t *testing.T) {
+	defer reset()
+
+	port, err := RequestPort("tcp", 5000)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if port != 5000 {
+		t.Fatalf("Expected port 5000 got %d", port)
+	}
+}
+
+func TestReleasePort(t *testing.T) {
+	defer reset()
+
+	port, err := RequestPort("tcp", 5000)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if port != 5000 {
+		t.Fatalf("Expected port 5000 got %d", port)
+	}
+
+	if err := ReleasePort("tcp", 5000); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestReuseReleasedPort(t *testing.T) {
+	defer reset()
+
+	port, err := RequestPort("tcp", 5000)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if port != 5000 {
+		t.Fatalf("Expected port 5000 got %d", port)
+	}
+
+	if err := ReleasePort("tcp", 5000); err != nil {
+		t.Fatal(err)
+	}
+
+	port, err = RequestPort("tcp", 5000)
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestReleaseUnreadledPort(t *testing.T) {
+	defer reset()
+
+	port, err := RequestPort("tcp", 5000)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if port != 5000 {
+		t.Fatalf("Expected port 5000 got %d", port)
+	}
+
+	port, err = RequestPort("tcp", 5000)
+	if err != ErrPortAlreadyAllocated {
+		t.Fatalf("Expected error %s got %s", ErrPortAlreadyAllocated, err)
+	}
+}
+
+func TestUnknowProtocol(t *testing.T) {
+	defer reset()
+
+	if _, err := RequestPort("tcpp", 0); err != ErrUnknownProtocol {
+		t.Fatalf("Expected error %s got %s", ErrUnknownProtocol, err)
+	}
+}
+
+func TestAllocateAllPorts(t *testing.T) {
+	defer reset()
+
+	for i := 0; i <= EndPortRange-BeginPortRange; i++ {
+		port, err := RequestPort("tcp", 0)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if expected := BeginPortRange + i; port != expected {
+			t.Fatalf("Expected port %d got %d", expected, port)
+		}
+	}
+
+	if _, err := RequestPort("tcp", 0); err != ErrPortExceedsRange {
+		t.Fatalf("Expected error %s got %s", ErrPortExceedsRange, err)
+	}
+
+	_, err := RequestPort("udp", 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func BenchmarkAllocatePorts(b *testing.B) {
+	defer reset()
+
+	b.StartTimer()
+	for i := 0; i < b.N; i++ {
+		for i := 0; i <= EndPortRange-BeginPortRange; i++ {
+			port, err := RequestPort("tcp", 0)
+			if err != nil {
+				b.Fatal(err)
+			}
+
+			if expected := BeginPortRange + i; port != expected {
+				b.Fatalf("Expected port %d got %d", expected, port)
+			}
+		}
+		reset()
+	}
+	b.StopTimer()
+}

pkg/collections/orderedintset.go

@@ -1,18 +1,23 @@
-package ipallocator
+package collections
 
 import (
 	"sort"
 	"sync"
 )
 
-// iPSet is a thread-safe sorted set and a stack.
-type iPSet struct {
+// OrderedIntSet is a thread-safe sorted set and a stack.
+type OrderedIntSet struct {
 	sync.RWMutex
 	set []int
 }
 
+// NewOrderedSet returns an initialized OrderedSet
+func NewOrderedIntSet() *OrderedIntSet {
+	return &OrderedIntSet{}
+}
+
 // Push takes a string and adds it to the set. If the elem aready exists, it has no effect.
-func (s *iPSet) Push(elem int) {
+func (s *OrderedIntSet) Push(elem int) {
 	s.RLock()
 	for _, e := range s.set {
 		if e == elem {
@@ -30,13 +35,13 @@ func (s *iPSet) Push(elem int) {
 }
 
 // Pop is an alias to PopFront()
-func (s *iPSet) Pop() int {
+func (s *OrderedIntSet) Pop() int {
 	return s.PopFront()
 }
 
 // Pop returns the first elemen from the list and removes it.
 // If the list is empty, it returns 0
-func (s *iPSet) PopFront() int {
+func (s *OrderedIntSet) PopFront() int {
 	s.RLock()
 
 	for i, e := range s.set {
@@ -55,7 +60,7 @@ func (s *iPSet) PopFront() int {
 // PullBack retrieve the last element of the list.
 // The element is not removed.
 // If the list is empty, an empty element is returned.
-func (s *iPSet) PullBack() int {
+func (s *OrderedIntSet) PullBack() int {
 	if len(s.set) == 0 {
 		return 0
 	}
@@ -63,7 +68,7 @@ func (s *iPSet) PullBack() int {
 }
 
 // Exists checks if the given element present in the list.
-func (s *iPSet) Exists(elem int) bool {
+func (s *OrderedIntSet) Exists(elem int) bool {
 	for _, e := range s.set {
 		if e == elem {
 			return true
@@ -74,7 +79,7 @@ func (s *iPSet) Exists(elem int) bool {
 
 // Remove removes an element from the list.
 // If the element is not found, it has no effect.
-func (s *iPSet) Remove(elem int) {
+func (s *OrderedIntSet) Remove(elem int) {
 	for i, e := range s.set {
 		if e == elem {
 			s.set = append(s.set[:i], s.set[i+1:]...)