github.com/cloudwego/kitex@v0.9.0/internal/test/port.go (about)

     1  // Copyright 2023 CloudWeGo Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package test
    16  
    17  import (
    18  	"fmt"
    19  	"hash/fnv"
    20  	"net"
    21  	"os"
    22  	"runtime/debug"
    23  	"strconv"
    24  	"strings"
    25  	"sync/atomic"
    26  	"time"
    27  )
    28  
    29  const (
    30  	UnixUserPortStart = 1023
    31  	UnixUserPortEnd   = 49151
    32  )
    33  
    34  func hashInt(n int) uint32 {
    35  	h := fnv.New32a()
    36  	_, _ = h.Write([]byte(strconv.Itoa(n)))
    37  	return h.Sum32()
    38  }
    39  
    40  // Go test may start multiple process for testing, so it's easy to get same port with same starting port
    41  var curPort uint32 = UnixUserPortStart + hashInt(os.Getpid())%200*200
    42  
    43  // GetLocalAddress return a local address starting from 1024
    44  // This API ensures no repeated addr returned in one UNIX OS
    45  func GetLocalAddress() string {
    46  	for {
    47  		port := atomic.AddUint32(&curPort, 1)
    48  		addr := "127.0.0.1:" + strconv.Itoa(int(port))
    49  		if !IsAddressInUse(addr) {
    50  			trace := strings.Split(string(debug.Stack()), "\n")
    51  			if len(trace) > 6 {
    52  				println(fmt.Sprintf("%s: GetLocalAddress = %v", trace[6], addr))
    53  			}
    54  			return addr
    55  		}
    56  	}
    57  }
    58  
    59  // tells if a net address is already in use.
    60  func IsAddressInUse(address string) bool {
    61  	// Attempt to establish a TCP connection to the address
    62  	conn, err := net.DialTimeout("tcp", address, 1*time.Second)
    63  	if err != nil {
    64  		// If there's an error, the address is likely not in use or not reachable
    65  		return false
    66  	}
    67  	_ = conn.Close()
    68  	// If the connection is successful, the address is in use
    69  	return true
    70  }
    71  
    72  // WaitServerStart waits for server to start for at most 1 second
    73  func WaitServerStart(addr string) {
    74  	for begin := time.Now(); time.Since(begin) < time.Second; {
    75  		if _, err := net.Dial("tcp", addr); err == nil {
    76  			println("server is up at", addr)
    77  			return
    78  		}
    79  		time.Sleep(time.Millisecond * 10)
    80  	}
    81  }