github.com/outbrain/consul@v1.4.5/lib/freeport/freeport.go (about)

     1  // Package freeport provides a helper for allocating free ports across multiple
     2  // processes on the same machine.
     3  package freeport
     4  
     5  import (
     6  	"fmt"
     7  	"math/rand"
     8  	"net"
     9  	"sync"
    10  	"time"
    11  
    12  	"github.com/mitchellh/go-testing-interface"
    13  )
    14  
    15  const (
    16  	// blockSize is the size of the allocated port block. ports are given out
    17  	// consecutively from that block with roll-over for the lifetime of the
    18  	// application/test run.
    19  	blockSize = 1500
    20  
    21  	// maxBlocks is the number of available port blocks.
    22  	// lowPort + maxBlocks * blockSize must be less than 65535.
    23  	maxBlocks = 30
    24  
    25  	// lowPort is the lowest port number that should be used.
    26  	lowPort = 10000
    27  
    28  	// attempts is how often we try to allocate a port block
    29  	// before giving up.
    30  	attempts = 10
    31  )
    32  
    33  var (
    34  	// firstPort is the first port of the allocated block.
    35  	firstPort int
    36  
    37  	// lockLn is the system-wide mutex for the port block.
    38  	lockLn net.Listener
    39  
    40  	// mu guards nextPort
    41  	mu sync.Mutex
    42  
    43  	// once is used to do the initialization on the first call to retrieve free
    44  	// ports
    45  	once sync.Once
    46  
    47  	// port is the last allocated port.
    48  	port int
    49  )
    50  
    51  // initialize is used to initialize freeport.
    52  func initialize() {
    53  	if lowPort+maxBlocks*blockSize > 65535 {
    54  		panic("freeport: block size too big or too many blocks requested")
    55  	}
    56  
    57  	rand.Seed(time.Now().UnixNano())
    58  	firstPort, lockLn = alloc()
    59  }
    60  
    61  // alloc reserves a port block for exclusive use for the lifetime of the
    62  // application. lockLn serves as a system-wide mutex for the port block and is
    63  // implemented as a TCP listener which is bound to the firstPort and which will
    64  // be automatically released when the application terminates.
    65  func alloc() (int, net.Listener) {
    66  	for i := 0; i < attempts; i++ {
    67  		block := int(rand.Int31n(int32(maxBlocks)))
    68  		firstPort := lowPort + block*blockSize
    69  		ln, err := net.ListenTCP("tcp", tcpAddr("127.0.0.1", firstPort))
    70  		if err != nil {
    71  			continue
    72  		}
    73  		// log.Printf("[DEBUG] freeport: allocated port block %d (%d-%d)", block, firstPort, firstPort+blockSize-1)
    74  		return firstPort, ln
    75  	}
    76  	panic("freeport: cannot allocate port block")
    77  }
    78  
    79  func tcpAddr(ip string, port int) *net.TCPAddr {
    80  	return &net.TCPAddr{IP: net.ParseIP(ip), Port: port}
    81  }
    82  
    83  // Get wraps the Free function and panics on any failure retrieving ports.
    84  func Get(n int) (ports []int) {
    85  	ports, err := Free(n)
    86  	if err != nil {
    87  		panic(err)
    88  	}
    89  
    90  	return ports
    91  }
    92  
    93  // GetT is suitable for use when retrieving unused ports in tests. If there is
    94  // an error retrieving free ports, the test will be failed.
    95  func GetT(t testing.T, n int) (ports []int) {
    96  	ports, err := Free(n)
    97  	if err != nil {
    98  		t.Fatalf("Failed retrieving free port: %v", err)
    99  	}
   100  
   101  	return ports
   102  }
   103  
   104  // Free returns a list of free ports from the allocated port block. It is safe
   105  // to call this method concurrently. Ports have been tested to be available on
   106  // 127.0.0.1 TCP but there is no guarantee that they will remain free in the
   107  // future.
   108  func Free(n int) (ports []int, err error) {
   109  	mu.Lock()
   110  	defer mu.Unlock()
   111  
   112  	if n > blockSize-1 {
   113  		return nil, fmt.Errorf("freeport: block size too small")
   114  	}
   115  
   116  	// Reserve a port block
   117  	once.Do(initialize)
   118  
   119  	for len(ports) < n {
   120  		port++
   121  
   122  		// roll-over the port
   123  		if port < firstPort+1 || port >= firstPort+blockSize {
   124  			port = firstPort + 1
   125  		}
   126  
   127  		// if the port is in use then skip it
   128  		ln, err := net.ListenTCP("tcp", tcpAddr("127.0.0.1", port))
   129  		if err != nil {
   130  			// log.Println("[DEBUG] freeport: port already in use: ", port)
   131  			continue
   132  		}
   133  		ln.Close()
   134  
   135  		ports = append(ports, port)
   136  	}
   137  	// log.Println("[DEBUG] freeport: free ports:", ports)
   138  	return ports, nil
   139  }