open-consul/sdk/freeport/freeport.go

460 lines
12 KiB
Go

// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0
// Package freeport provides a helper for reserving free TCP ports across multiple
// processes on the same machine. Each process reserves a block of ports outside
// the ephemeral port range. Tests can request one of these reserved ports
// and freeport will ensure that no other test uses that port until it is returned
// to freeport.
//
// Freeport is particularly useful when the code being tested does not accept
// a net.Listener. Any code that accepts a net.Listener (or uses net/http/httptest.Server)
// can use port 0 (ex: 127.0.0.1:0) to find an unused ephemeral port that will
// not conflict.
//
// Any code that does not accept a net.Listener or can not bind directly to port
// zero should use freeport to find an unused port.
package freeport
import (
"container/list"
"fmt"
"math/rand"
"net"
"os"
"runtime"
"sync"
"time"
)
const (
// maxBlocks is the number of available port blocks before exclusions.
maxBlocks = 30
// lowPort is the lowest port number that should be used.
lowPort = 10000
// attempts is how often we try to allocate a port block
// before giving up.
attempts = 10
)
var (
// blockSize is the size of the allocated port block. ports are given out
// consecutively from that block and after that point in a LRU fashion.
blockSize int
// effectiveMaxBlocks is the number of available port blocks.
// lowPort + effectiveMaxBlocks * blockSize must be less than 65535.
effectiveMaxBlocks int
// firstPort is the first port of the allocated block.
firstPort int
// lockLn is the system-wide mutex for the port block.
lockLn net.Listener
// mu guards:
// - pendingPorts
// - freePorts
// - total
mu sync.Mutex
// once is used to do the initialization on the first call to retrieve free
// ports
once sync.Once
// condNotEmpty is a condition variable to wait for freePorts to be not
// empty. Linked to 'mu'
condNotEmpty *sync.Cond
// freePorts is a FIFO of all currently free ports. Take from the front,
// and return to the back.
freePorts *list.List
// pendingPorts is a FIFO of recently freed ports that have not yet passed
// the not-in-use check.
pendingPorts *list.List
// total is the total number of available ports in the block for use.
total int
// seededRand is a random generator that is pre-seeded from the current time.
seededRand *rand.Rand
// stopCh is used to signal to background goroutines to terminate. Only
// really exists for the safety of reset() during unit tests.
stopCh chan struct{}
// stopWg is used to keep track of background goroutines that are still
// alive. Only really exists for the safety of reset() during unit tests.
stopWg sync.WaitGroup
// portLastUser associates ports with a test name in order to debug
// which test may be leaking unclosed TCP connections.
portLastUser map[int]string
)
// initialize is used to initialize freeport.
func initialize() {
var err error
blockSize = 1500
limit, err := systemLimit()
if err != nil {
panic("freeport: error getting system limit: " + err.Error())
}
if limit > 0 && limit < blockSize {
logf("INFO", "blockSize %d too big for system limit %d. Adjusting...", blockSize, limit)
blockSize = limit - 3
}
effectiveMaxBlocks, err = adjustMaxBlocks()
if err != nil {
panic("freeport: ephemeral port range detection failed: " + err.Error())
}
if effectiveMaxBlocks < 0 {
panic("freeport: no blocks of ports available outside of ephemeral range")
}
if lowPort+effectiveMaxBlocks*blockSize > 65535 {
panic("freeport: block size too big or too many blocks requested")
}
seededRand = rand.New(rand.NewSource(time.Now().UnixNano())) // This is compatible with go 1.19 but unnecessary in >= go1.20
firstPort, lockLn = alloc()
condNotEmpty = sync.NewCond(&mu)
freePorts = list.New()
pendingPorts = list.New()
// fill with all available free ports
for port := firstPort + 1; port < firstPort+blockSize; port++ {
if used := isPortInUse(port); !used {
freePorts.PushBack(port)
}
}
total = freePorts.Len()
stopWg.Add(1)
stopCh = make(chan struct{})
portLastUser = make(map[int]string)
// Note: we pass this param explicitly to the goroutine so that we can
// freely recreate the underlying stop channel during reset() after closing
// the original.
go checkFreedPorts(stopCh)
}
func shutdownGoroutine() {
mu.Lock()
if stopCh == nil {
mu.Unlock()
return
}
close(stopCh)
stopCh = nil
mu.Unlock()
stopWg.Wait()
}
// reset will reverse the setup from initialize() and then redo it (for tests)
func reset() {
logf("INFO", "resetting the freeport package state")
shutdownGoroutine()
mu.Lock()
defer mu.Unlock()
effectiveMaxBlocks = 0
firstPort = 0
if lockLn != nil {
lockLn.Close()
lockLn = nil
}
once = sync.Once{}
freePorts = nil
pendingPorts = nil
portLastUser = nil
total = 0
}
func checkFreedPorts(stopCh <-chan struct{}) {
defer stopWg.Done()
ticker := time.NewTicker(250 * time.Millisecond)
for {
select {
case <-stopCh:
logf("INFO", "Closing checkFreedPorts()")
return
case <-ticker.C:
checkFreedPortsOnce()
}
}
}
func checkFreedPortsOnce() {
mu.Lock()
defer mu.Unlock()
pending := pendingPorts.Len()
remove := make([]*list.Element, 0, pending)
for elem := pendingPorts.Front(); elem != nil; elem = elem.Next() {
port := elem.Value.(int)
if used := isPortInUse(port); !used {
freePorts.PushBack(port)
remove = append(remove, elem)
} else {
logf("WARN", "port %d still being used by %q", port, portLastUser[port])
}
}
retained := pending - len(remove)
if retained > 0 {
logf("WARN", "%d out of %d pending ports are still in use; something probably didn't wait around for the port to be closed!", retained, pending)
}
if len(remove) == 0 {
return
}
for _, elem := range remove {
pendingPorts.Remove(elem)
}
condNotEmpty.Broadcast()
}
// adjustMaxBlocks avoids having the allocation ranges overlap the ephemeral
// port range.
func adjustMaxBlocks() (int, error) {
ephemeralPortMin, ephemeralPortMax, err := getEphemeralPortRange()
if err != nil {
return 0, err
}
if ephemeralPortMin <= 0 || ephemeralPortMax <= 0 {
logf("INFO", "ephemeral port range detection not configured for GOOS=%q", runtime.GOOS)
return maxBlocks, nil
}
logf("INFO", "detected ephemeral port range of [%d, %d]", ephemeralPortMin, ephemeralPortMax)
for block := 0; block < maxBlocks; block++ {
min := lowPort + block*blockSize
max := min + blockSize
overlap := intervalOverlap(min, max-1, ephemeralPortMin, ephemeralPortMax)
if overlap {
logf("INFO", "reducing max blocks from %d to %d to avoid the ephemeral port range", maxBlocks, block)
return block, nil
}
}
return maxBlocks, nil
}
// alloc reserves a port block for exclusive use for the lifetime of the
// application. lockLn serves as a system-wide mutex for the port block and is
// implemented as a TCP listener which is bound to the firstPort and which will
// be automatically released when the application terminates.
func alloc() (int, net.Listener) {
for i := 0; i < attempts; i++ {
block := int(seededRand.Int31n(int32(effectiveMaxBlocks)))
firstPort := lowPort + block*blockSize
ln, err := net.ListenTCP("tcp", tcpAddr("127.0.0.1", firstPort))
if err != nil {
continue
}
// logf("DEBUG", "allocated port block %d (%d-%d)", block, firstPort, firstPort+blockSize-1)
return firstPort, ln
}
panic("freeport: cannot allocate port block")
}
// MustTake is the same as Take except it panics on error.
//
// Deprecated: Use GetN or GetOne instead.
func MustTake(n int) (ports []int) {
ports, err := Take(n)
if err != nil {
panic(err)
}
return ports
}
// Take returns a list of free ports from the reserved port block. It is safe
// to call this method concurrently. Ports have been tested to be available on
// 127.0.0.1 TCP but there is no guarantee that they will remain free in the
// future.
//
// Most callers should prefer GetN or GetOne.
func Take(n int) (ports []int, err error) {
if n <= 0 {
return nil, fmt.Errorf("freeport: cannot take %d ports", n)
}
mu.Lock()
defer mu.Unlock()
// Reserve a port block
once.Do(initialize)
if n > total {
return nil, fmt.Errorf("freeport: block size too small")
}
for len(ports) < n {
for freePorts.Len() == 0 {
if total == 0 {
return nil, fmt.Errorf("freeport: impossible to satisfy request; there are no actual free ports in the block anymore")
}
condNotEmpty.Wait()
}
elem := freePorts.Front()
freePorts.Remove(elem)
port := elem.Value.(int)
if used := isPortInUse(port); used {
// Something outside of the test suite has stolen this port, possibly
// due to assignment to an ephemeral port, remove it completely.
logf("WARN", "leaked port %d due to theft; removing from circulation", port)
total--
continue
}
ports = append(ports, port)
}
return ports, nil
}
// peekFree returns the next port that will be returned by Take to aid in testing.
func peekFree() int {
mu.Lock()
defer mu.Unlock()
return freePorts.Front().Value.(int)
}
// peekAllFree returns all free ports that could be returned by Take to aid in testing.
func peekAllFree() []int {
mu.Lock()
defer mu.Unlock()
var out []int
for elem := freePorts.Front(); elem != nil; elem = elem.Next() {
port := elem.Value.(int)
out = append(out, port)
}
return out
}
// stats returns diagnostic data to aid in testing
func stats() (numTotal, numPending, numFree int) {
mu.Lock()
defer mu.Unlock()
return total, pendingPorts.Len(), freePorts.Len()
}
// Return returns a block of ports back to the general pool. These ports should
// have been returned from a call to Take().
func Return(ports []int) {
if len(ports) == 0 {
return // convenience short circuit for test ergonomics
}
mu.Lock()
defer mu.Unlock()
for _, port := range ports {
if port > firstPort && port < firstPort+blockSize {
pendingPorts.PushBack(port)
}
}
}
func isPortInUse(port int) bool {
ln, err := net.ListenTCP("tcp", tcpAddr("127.0.0.1", port))
if err != nil {
return true
}
ln.Close()
return false
}
func tcpAddr(ip string, port int) *net.TCPAddr {
return &net.TCPAddr{IP: net.ParseIP(ip), Port: port}
}
// intervalOverlap returns true if the doubly-inclusive integer intervals
// represented by [min1, max1] and [min2, max2] overlap.
func intervalOverlap(min1, max1, min2, max2 int) bool {
if min1 > max1 {
logf("WARN", "interval1 is not ordered [%d, %d]", min1, max1)
return false
}
if min2 > max2 {
logf("WARN", "interval2 is not ordered [%d, %d]", min2, max2)
return false
}
return min1 <= max2 && min2 <= max1
}
func logf(severity string, format string, a ...interface{}) {
fmt.Fprintf(os.Stderr, "["+severity+"] freeport: "+format+"\n", a...)
}
// TestingT is the minimal set of methods implemented by *testing.T that are
// used by functions in freelist.
//
// In the future new methods may be added to this interface, but those methods
// should always be implemented by *testing.T
type TestingT interface {
Cleanup(func())
Helper()
Fatalf(format string, args ...interface{})
Name() string
}
// GetN returns n free ports from the reserved port block, and returns the
// ports to the pool when the test ends. See Take for more details.
func GetN(t TestingT, n int) []int {
t.Helper()
ports, err := Take(n)
if err != nil {
t.Fatalf("failed to take %v ports: %w", n, err)
}
logf("DEBUG", "Test %q took ports %v", t.Name(), ports)
mu.Lock()
for _, p := range ports {
portLastUser[p] = t.Name()
}
mu.Unlock()
t.Cleanup(func() {
Return(ports)
logf("DEBUG", "Test %q returned ports %v", t.Name(), ports)
})
return ports
}
// GetOne returns a single free port from the reserved port block, and returns the
// port to the pool when the test ends. See Take for more details.
// Use GetN if more than a single port is required.
func GetOne(t TestingT) int {
t.Helper()
return GetN(t, 1)[0]
}
// Deprecated: Please use Take/Return calls instead.
func Get(n int) (ports []int) { return MustTake(n) }
// Deprecated: Please use Take/Return calls instead.
func GetT(t TestingT, n int) (ports []int) { return MustTake(n) }
// Deprecated: Please use Take/Return calls instead.
func Free(n int) (ports []int, err error) { return MustTake(n), nil }