From 303ed3c8300183bab09a36b58e0c1db89d12424a Mon Sep 17 00:00:00 2001 From: Michael Crosby Date: Thu, 23 Jan 2014 06:43:50 -0800 Subject: [PATCH] Add port allocator and move ipset into orderedintset Docker-DCO-1.1-Signed-off-by: Michael Crosby (github: crosbymichael) --- networkdriver/ipallocator/allocator.go | 4 +- networkdriver/portallocator/portallocator.go | 103 ++++++++++++ .../portallocator/portallocator_test.go | 150 ++++++++++++++++++ .../collections/orderedintset.go | 23 +-- 4 files changed, 270 insertions(+), 10 deletions(-) create mode 100644 networkdriver/portallocator/portallocator.go create mode 100644 networkdriver/portallocator/portallocator_test.go rename networkdriver/ipallocator/ipset.go => pkg/collections/orderedintset.go (72%) diff --git a/networkdriver/ipallocator/allocator.go b/networkdriver/ipallocator/allocator.go index 09319d1332..602c9d220b 100644 --- a/networkdriver/ipallocator/allocator.go +++ b/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") diff --git a/networkdriver/portallocator/portallocator.go b/networkdriver/portallocator/portallocator.go new file mode 100644 index 0000000000..a0a3ebd8d2 --- /dev/null +++ b/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 +} diff --git a/networkdriver/portallocator/portallocator_test.go b/networkdriver/portallocator/portallocator_test.go new file mode 100644 index 0000000000..603bf5a15a --- /dev/null +++ b/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() +} diff --git a/networkdriver/ipallocator/ipset.go b/pkg/collections/orderedintset.go similarity index 72% rename from networkdriver/ipallocator/ipset.go rename to pkg/collections/orderedintset.go index 43d54691d1..456975cbb0 100644 --- a/networkdriver/ipallocator/ipset.go +++ b/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:]...)