Refactoring portallocator
Faster, more documented, less code. Docker-DCO-1.1-Signed-off-by: Alexandr Morozov <lk4d4math@gmail.com> (github: LK4D4)
This commit is contained in:
parent
bc33ce0d8f
commit
f387cc1205
|
@ -12,10 +12,22 @@ type portMap struct {
|
||||||
last int
|
last int
|
||||||
}
|
}
|
||||||
|
|
||||||
type (
|
func newPortMap() *portMap {
|
||||||
protocolMap map[string]*portMap
|
return &portMap{
|
||||||
ipMapping map[string]protocolMap
|
p: map[int]struct{}{},
|
||||||
)
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type protoMap map[string]*portMap
|
||||||
|
|
||||||
|
func newProtoMap() protoMap {
|
||||||
|
return protoMap{
|
||||||
|
"tcp": newPortMap(),
|
||||||
|
"udp": newPortMap(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type ipMapping map[string]protoMap
|
||||||
|
|
||||||
const (
|
const (
|
||||||
BeginPortRange = 49153
|
BeginPortRange = 49153
|
||||||
|
@ -62,107 +74,83 @@ func (e ErrPortAlreadyAllocated) Error() string {
|
||||||
return fmt.Sprintf("Bind for %s:%d failed: port is already allocated", e.ip, e.port)
|
return fmt.Sprintf("Bind for %s:%d failed: port is already allocated", e.ip, e.port)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// RequestPort requests new port from global ports pool for specified ip and proto.
|
||||||
|
// If port is 0 it returns first free port. Otherwise it cheks port availability
|
||||||
|
// in pool and return that port or error if port is already busy.
|
||||||
func RequestPort(ip net.IP, proto string, port int) (int, error) {
|
func RequestPort(ip net.IP, proto string, port int) (int, error) {
|
||||||
mutex.Lock()
|
mutex.Lock()
|
||||||
defer mutex.Unlock()
|
defer mutex.Unlock()
|
||||||
|
|
||||||
if err := validateProto(proto); err != nil {
|
if proto != "tcp" && proto != "udp" {
|
||||||
return 0, err
|
return 0, ErrUnknownProtocol
|
||||||
}
|
}
|
||||||
|
|
||||||
ip = getDefault(ip)
|
if ip == nil {
|
||||||
|
ip = defaultIP
|
||||||
mapping := getOrCreate(ip)
|
}
|
||||||
|
ipstr := ip.String()
|
||||||
|
protomap, ok := globalMap[ipstr]
|
||||||
|
if !ok {
|
||||||
|
protomap = newProtoMap()
|
||||||
|
globalMap[ipstr] = protomap
|
||||||
|
}
|
||||||
|
mapping := protomap[proto]
|
||||||
if port > 0 {
|
if port > 0 {
|
||||||
if _, ok := mapping[proto].p[port]; !ok {
|
if _, ok := mapping.p[port]; !ok {
|
||||||
mapping[proto].p[port] = struct{}{}
|
mapping.p[port] = struct{}{}
|
||||||
return port, nil
|
return port, nil
|
||||||
} else {
|
|
||||||
return 0, NewErrPortAlreadyAllocated(ip.String(), port)
|
|
||||||
}
|
}
|
||||||
} else {
|
return 0, NewErrPortAlreadyAllocated(ipstr, port)
|
||||||
port, err := findPort(ip, proto)
|
}
|
||||||
|
|
||||||
|
port, err := mapping.findPort()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return 0, err
|
return 0, err
|
||||||
}
|
}
|
||||||
|
|
||||||
return port, nil
|
return port, nil
|
||||||
}
|
}
|
||||||
}
|
|
||||||
|
|
||||||
|
// ReleasePort releases port from global ports pool for specified ip and proto.
|
||||||
func ReleasePort(ip net.IP, proto string, port int) error {
|
func ReleasePort(ip net.IP, proto string, port int) error {
|
||||||
mutex.Lock()
|
mutex.Lock()
|
||||||
defer mutex.Unlock()
|
defer mutex.Unlock()
|
||||||
|
|
||||||
ip = getDefault(ip)
|
if ip == nil {
|
||||||
|
ip = defaultIP
|
||||||
mapping := getOrCreate(ip)[proto]
|
}
|
||||||
delete(mapping.p, port)
|
protomap, ok := globalMap[ip.String()]
|
||||||
|
if !ok {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
delete(protomap[proto].p, port)
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ReleaseAll releases all ports for all ips.
|
||||||
func ReleaseAll() error {
|
func ReleaseAll() error {
|
||||||
mutex.Lock()
|
mutex.Lock()
|
||||||
defer mutex.Unlock()
|
|
||||||
|
|
||||||
globalMap = ipMapping{}
|
globalMap = ipMapping{}
|
||||||
|
mutex.Unlock()
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func getOrCreate(ip net.IP) protocolMap {
|
func (pm *portMap) findPort() (int, error) {
|
||||||
ipstr := ip.String()
|
if pm.last == 0 {
|
||||||
|
pm.p[BeginPortRange] = struct{}{}
|
||||||
if _, ok := globalMap[ipstr]; !ok {
|
pm.last = BeginPortRange
|
||||||
globalMap[ipstr] = protocolMap{
|
|
||||||
"tcp": &portMap{p: map[int]struct{}{}, last: 0},
|
|
||||||
"udp": &portMap{p: map[int]struct{}{}, last: 0},
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return globalMap[ipstr]
|
|
||||||
}
|
|
||||||
|
|
||||||
func findPort(ip net.IP, proto string) (int, error) {
|
|
||||||
mapping := getOrCreate(ip)[proto]
|
|
||||||
|
|
||||||
if mapping.last == 0 {
|
|
||||||
mapping.p[BeginPortRange] = struct{}{}
|
|
||||||
mapping.last = BeginPortRange
|
|
||||||
return BeginPortRange, nil
|
return BeginPortRange, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
for port := mapping.last + 1; port != mapping.last; port++ {
|
for port := pm.last + 1; port != pm.last; port++ {
|
||||||
if port > EndPortRange {
|
if port > EndPortRange {
|
||||||
port = BeginPortRange
|
port = BeginPortRange
|
||||||
}
|
}
|
||||||
|
|
||||||
if _, ok := mapping.p[port]; !ok {
|
if _, ok := pm.p[port]; !ok {
|
||||||
mapping.p[port] = struct{}{}
|
pm.p[port] = struct{}{}
|
||||||
mapping.last = port
|
pm.last = port
|
||||||
return port, nil
|
return port, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|
||||||
return 0, ErrAllPortsAllocated
|
return 0, ErrAllPortsAllocated
|
||||||
}
|
}
|
||||||
|
|
||||||
func getDefault(ip net.IP) net.IP {
|
|
||||||
if ip == nil {
|
|
||||||
return defaultIP
|
|
||||||
}
|
|
||||||
|
|
||||||
return ip
|
|
||||||
}
|
|
||||||
|
|
||||||
func validateProto(proto string) error {
|
|
||||||
if proto != "tcp" && proto != "udp" {
|
|
||||||
return ErrUnknownProtocol
|
|
||||||
}
|
|
||||||
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
|
|
Loading…
Reference in New Issue