portallocator: rewrite to simplify, removes race condition

Docker-DCO-1.1-Signed-off-by: Erik Hollensbe <github@hollensbe.org> (github: erikh)

daemon/networkdriver/portallocator/portallocator.go

@@ -2,21 +2,21 @@ package portallocator
 
 import (
 	"errors"
-	"github.com/dotcloud/docker/pkg/collections"
 	"net"
 	"sync"
 )
 
+type (
+	portMap     map[int]bool
+	protocolMap map[string]portMap
+	ipMapping   map[string]protocolMap
+)
+
 const (
 	BeginPortRange = 49153
 	EndPortRange   = 65535
 )
 
-type (
-	portMappings map[string]*collections.OrderedIntSet
-	ipMapping    map[string]portMappings
-)
-
 var (
 	ErrAllPortsAllocated    = errors.New("all ports are allocated")
 	ErrPortAlreadyAllocated = errors.New("port has already been allocated")
@@ -24,165 +24,106 @@ var (
 )
 
 var (
-	currentDynamicPort = map[string]int{
-		"tcp": BeginPortRange - 1,
-		"udp": BeginPortRange - 1,
-	}
-	defaultIP             = net.ParseIP("0.0.0.0")
-	defaultAllocatedPorts = portMappings{}
-	otherAllocatedPorts   = ipMapping{}
-	lock                  = sync.Mutex{}
-)
+	mutex sync.Mutex
 
-func init() {
-	defaultAllocatedPorts["tcp"] = collections.NewOrderedIntSet()
-	defaultAllocatedPorts["udp"] = collections.NewOrderedIntSet()
-}
+	defaultIP = net.ParseIP("0.0.0.0")
+	globalMap = ipMapping{}
+)
 
-// 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(ip net.IP, proto string, port int) (int, error) {
-	lock.Lock()
-	defer lock.Unlock()
+	mutex.Lock()
+	defer mutex.Unlock()
 
-	if err := validateProtocol(proto); err != nil {
+	if err := validateProto(proto); err != nil {
 		return 0, err
 	}
 
-	// If the user requested a specific port to be allocated
+	ip = getDefault(ip)
+
+	mapping := getOrCreate(ip)
+
 	if port > 0 {
-		if err := registerSetPort(ip, proto, port); err != nil {
+		if !mapping[proto][port] {
+			mapping[proto][port] = true
+			return port, nil
+		} else {
+			return 0, ErrPortAlreadyAllocated
+		}
+	} else {
+		port, err := findPort(ip, proto)
+
+		if err != nil {
 			return 0, err
 		}
+
 		return port, nil
 	}
-	return registerDynamicPort(ip, proto)
 }
 
-// ReleasePort will return the provided port back into the
-// pool for reuse
 func ReleasePort(ip net.IP, proto string, port int) error {
-	lock.Lock()
-	defer lock.Unlock()
-
-	if err := validateProtocol(proto); err != nil {
-		return err
-	}
+	mutex.Lock()
+	defer mutex.Unlock()
 
-	allocated := defaultAllocatedPorts[proto]
-	allocated.Remove(port)
+	ip = getDefault(ip)
 
-	if !equalsDefault(ip) {
-		registerIP(ip)
+	mapping := getOrCreate(ip)
+	delete(mapping[proto], port)
 
-		// Remove the port for the specific ip address
-		allocated = otherAllocatedPorts[ip.String()][proto]
-		allocated.Remove(port)
-	}
 	return nil
 }
 
 func ReleaseAll() error {
-	lock.Lock()
-	defer lock.Unlock()
-
-	currentDynamicPort["tcp"] = BeginPortRange - 1
-	currentDynamicPort["udp"] = BeginPortRange - 1
-
-	defaultAllocatedPorts = portMappings{}
-	defaultAllocatedPorts["tcp"] = collections.NewOrderedIntSet()
-	defaultAllocatedPorts["udp"] = collections.NewOrderedIntSet()
+	mutex.Lock()
+	defer mutex.Unlock()
 
-	otherAllocatedPorts = ipMapping{}
+	globalMap = ipMapping{}
 
 	return nil
 }
 
-func registerDynamicPort(ip net.IP, proto string) (int, error) {
-
-	if !equalsDefault(ip) {
-		registerIP(ip)
-
-		ipAllocated := otherAllocatedPorts[ip.String()][proto]
-
-		port, err := findNextPort(proto, ipAllocated)
-		if err != nil {
-			return 0, err
-		}
-		ipAllocated.Push(port)
-		return port, nil
-
-	} else {
-
-		allocated := defaultAllocatedPorts[proto]
+func getOrCreate(ip net.IP) protocolMap {
+	ipstr := ip.String()
 
-		port, err := findNextPort(proto, allocated)
-		if err != nil {
-			return 0, err
+	if _, ok := globalMap[ipstr]; !ok {
+		globalMap[ipstr] = protocolMap{
+			"tcp": portMap{},
+			"udp": portMap{},
 		}
-		allocated.Push(port)
-		return port, nil
 	}
+
+	return globalMap[ipstr]
 }
 
-func registerSetPort(ip net.IP, proto string, port int) error {
-	allocated := defaultAllocatedPorts[proto]
-	if allocated.Exists(port) {
-		return ErrPortAlreadyAllocated
-	}
+func findPort(ip net.IP, proto string) (int, error) {
+	port := BeginPortRange
 
-	if !equalsDefault(ip) {
-		registerIP(ip)
+	mapping := getOrCreate(ip)
 
-		ipAllocated := otherAllocatedPorts[ip.String()][proto]
-		if ipAllocated.Exists(port) {
-			return ErrPortAlreadyAllocated
-		}
-		ipAllocated.Push(port)
-	} else {
-		allocated.Push(port)
-	}
-	return nil
-}
-
-func equalsDefault(ip net.IP) bool {
-	return ip == nil || ip.Equal(defaultIP)
-}
+	for mapping[proto][port] {
+		port++
 
-func findNextPort(proto string, allocated *collections.OrderedIntSet) (int, error) {
-	port := nextPort(proto)
-	startSearchPort := port
-	for allocated.Exists(port) {
-		port = nextPort(proto)
-		if startSearchPort == port {
+		if port > EndPortRange {
 			return 0, ErrAllPortsAllocated
 		}
 	}
+
+	mapping[proto][port] = true
+
 	return port, nil
 }
 
-func nextPort(proto string) int {
-	c := currentDynamicPort[proto] + 1
-	if c > EndPortRange {
-		c = BeginPortRange
+func getDefault(ip net.IP) net.IP {
+	if ip == nil {
+		return defaultIP
 	}
-	currentDynamicPort[proto] = c
-	return c
-}
 
-func registerIP(ip net.IP) {
-	if _, exists := otherAllocatedPorts[ip.String()]; !exists {
-		otherAllocatedPorts[ip.String()] = portMappings{
-			"tcp": collections.NewOrderedIntSet(),
-			"udp": collections.NewOrderedIntSet(),
-		}
-	}
+	return ip
 }
 
-func validateProtocol(proto string) error {
-	if _, exists := defaultAllocatedPorts[proto]; !exists {
+func validateProto(proto string) error {
+	if proto != "tcp" && proto != "udp" {
 		return ErrUnknownProtocol
 	}
+
 	return nil
 }