From 08a75bcb4bc71e720f866134340d918735bffc78 Mon Sep 17 00:00:00 2001 From: Akagi201 Date: Fri, 30 Sep 2016 10:45:59 +0800 Subject: [PATCH] PortPool: add check port support --- shell/ports.go | 104 ++++++++++++++++++++++++++++++++++---------- shell/shell_test.go | 38 ++++++++-------- 2 files changed, 101 insertions(+), 41 deletions(-) diff --git a/shell/ports.go b/shell/ports.go index 8da6bc0..fff24ba 100644 --- a/shell/ports.go +++ b/shell/ports.go @@ -27,49 +27,107 @@ SOFTWARE. */ package main -import "fmt" +import ( + "fmt" + "net" + "strconv" +) -// The port pool manage available ports. +// PortPool The port pool manage available ports. type PortPool struct { - shell *ShellBoss - // alloc new port from ports, fill ports from left, - // release to ports when free port. + // Alloc new port from ports. Alloc new port from the head of the list. + // Release to ports when free port. + // If check failed, put the port to the tail of the list ports []int - left []int + // The ports in use. fill used when alloc + used []int } -// alloc port in [start,stop] +// NewPortPool alloc port in [start,stop] func NewPortPool(start, stop int) *PortPool { v := &PortPool{} + v.ports = make([]int, stop-start+1) for i := start; i <= stop; i++ { - if len(v.ports) < 64 { - v.ports = append(v.ports, i) - } else { - v.left = append(v.left, i) - } + v.ports[i-start] = i } return v } -func (v *PortPool) Alloc(nbPort int) (ports []int, err error) { +// allocOnePort alloc a port from the port pool +func (v *PortPool) allocOnePort() (int, error) { + if len(v.ports) <= 0 { + return 0, fmt.Errorf("empty port pool") + } + + for i, p := range v.ports { + if checkPort(p) { + // delete i-index element + v.ports = append(v.ports[:i], v.ports[i+1:]...) + v.used = append(v.used, p) + return p, nil + } + } + + return 0, fmt.Errorf("No available port") +} + +// Alloc alloc nbPort ports from the port pool +func (v *PortPool) Alloc(nbPort int) ([]int, error) { if nbPort <= 0 { return nil, fmt.Errorf("invalid ports %v", nbPort) } - if len(v.ports)+len(v.left) < nbPort { - return nil, fmt.Errorf("no %v port available, left %v", nbPort, len(v.ports)+len(v.left)) - } if len(v.ports) < nbPort { - cp := nbPort - len(v.ports) - v.ports = append(v.ports, v.left[0:cp]...) - v.left = v.left[cp:] + return nil, fmt.Errorf("no %v port available, left %v", nbPort, len(v.ports)) } - ports = v.ports[0:nbPort] - v.ports = v.ports[nbPort:] - return + ports := []int{} + for i := 0; i < nbPort; i++ { + p, err := v.allocOnePort() + if err != nil { + return nil, err + } else { + ports = append(ports, p) + } + } + return ports, nil } +// Free free port from the port pool func (v *PortPool) Free(port int) { - v.ports = append(v.ports, port) + // Free used ports to the head of the ports list for better port reused. + v.ports = append([]int{port}, v.ports...) + for i, p := range v.used { + if port == p { + // delete i-index element + v.used = append(v.used[:i], v.used[i+1:]...) + return + } + } +} + +// GetPortsInUse get a list of port in use +func (v *PortPool) GetPortsInUse() []int { + return v.used +} + +// checkPort check if a port is available +func checkPort(port int) bool { + // Concatenate a colon and the port + host := ":" + strconv.Itoa(port) + + // Try to create a server with the port + listener, err := net.Listen("tcp", host) + + // if it fails then the port is likely taken + if err != nil { + return false + } + + // close the server + listener.Close() + + // we successfully used and closed the port + // so it's now available to be used again + return true } diff --git a/shell/shell_test.go b/shell/shell_test.go index dd8ee01..859091c 100644 --- a/shell/shell_test.go +++ b/shell/shell_test.go @@ -24,55 +24,57 @@ SOFTWARE. package main -import "testing" +import ( + "testing" +) func TestPortPool_Alloc(t *testing.T) { - p := NewPortPool(1, 10) + p := NewPortPool(6001, 6010) if _, err := p.Alloc(0); err == nil { t.Error("should error") } if ps, err := p.Alloc(1); err != nil { - t.Errorf("alloc failed, err is %v", err) - } else if len(ps) != 1 || ps[0] != 1 { - t.Errorf("alloc failed, ports=%v", ps) + t.Errorf("alloc 1 failed, err is %v", err) + } else if len(ps) != 1 || ps[0] != 6001 { + t.Errorf("alloc 1 failed, ports=%v", ps) } else if ps, err = p.Alloc(9); err != nil { - t.Errorf("alloc failed, err is %v", err) - } else if len(ps) != 9 || ps[0] != 2 || ps[8] != 10 { - t.Errorf("alloc failed, ports=%v", ps) + t.Errorf("alloc 9 failed, err is %v", err) + } else if len(ps) != 9 || ps[0] != 6002 || ps[8] != 6010 { + t.Errorf("alloc 9 failed, ports=%v", ps) } else if ps, err = p.Alloc(1); err == nil { t.Errorf("should error, ports=%v", ps) } } func TestPortPool_Free(t *testing.T) { - p := NewPortPool(1, 10) + p := NewPortPool(6001, 6010) - if ps, err := p.Alloc(10); err != nil || len(ps) != 10 || ps[0] != 1 || ps[9] != 10 { + if ps, err := p.Alloc(10); err != nil || len(ps) != 10 || ps[0] != 6001 || ps[9] != 6010 { t.Errorf("alloc failed, ports=%v, err is %v", ps, err) } - p.Free(11) - if ps, err := p.Alloc(1); err != nil || len(ps) != 1 || ps[0] != 11 { + p.Free(6010) + if ps, err := p.Alloc(1); err != nil || len(ps) != 1 { t.Error("free failed, ports=%v, err is %v", ps, err) - } else if ps[0] != 11 { + } else if ps[0] != 6010 { t.Errorf("invalid port=%v", ps) } } func TestPortPool_Alloc2(t *testing.T) { - p := NewPortPool(1, 100) + p := NewPortPool(9001, 9100) - if ps, err := p.Alloc(64); err != nil || len(ps) != 64 || ps[0] != 1 || ps[63] != 64 { + if ps, err := p.Alloc(64); err != nil || len(ps) != 64 || ps[0] != 9001 || ps[63] != 9064 { t.Errorf("alloc failed, ports=%v, err is %v", ps, err) - } else if ps, err := p.Alloc(36); err != nil || len(ps) != 36 || ps[0] != 65 || ps[35] != 100 { + } else if ps, err := p.Alloc(36); err != nil || len(ps) != 36 || ps[0] != 9065 || ps[35] != 9100 { t.Errorf("alloc failed, ports=%v, err is %v", ps, err) } else if ps, err = p.Alloc(1); err == nil { t.Errorf("should error, ports=%v", ps) } - p.Free(11) - if ps, err := p.Alloc(1); err != nil || len(ps) != 1 || ps[0] != 11 { + p.Free(9011) + if ps, err := p.Alloc(1); err != nil || len(ps) != 1 || ps[0] != 9011 { t.Errorf("alloc failed, ports=%v, err is %v", ps, err) } else if ps, err = p.Alloc(1); err == nil { t.Errorf("should error, ports=%v", ps)