PortPool: add check port before alloc

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
 }

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)