github.com/matrixorigin/matrixone@v1.2.0/pkg/tests/network.go (about)

     1  // Copyright 2021 - 2022 Matrix Origin
     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 tests
    16  
    17  import (
    18  	"net"
    19  	"strconv"
    20  	"sync"
    21  
    22  	"github.com/matrixorigin/matrixone/pkg/common/moerr"
    23  )
    24  
    25  var (
    26  	maxPort   = 65535
    27  	curPort   = 10000 // curPort indicates allocated port.
    28  	curPortMu sync.Mutex
    29  )
    30  
    31  // GetAddressBatch generates service addresses by batch.
    32  func GetAddressBatch(host string, batch int) ([]string, error) {
    33  	addrs := make([]string, batch)
    34  	for i := 0; i < batch; i++ {
    35  		port, err := GetAvailablePort(host)
    36  		if err != nil {
    37  			return nil, err
    38  		}
    39  		addrs[i] = net.JoinHostPort(host, port)
    40  	}
    41  	return addrs, nil
    42  }
    43  
    44  // GetAvailablePort gets available port on host address.
    45  func GetAvailablePort(host string) (string, error) {
    46  	curPortMu.Lock()
    47  	defer curPortMu.Unlock()
    48  
    49  	port := 0
    50  	for curPort < maxPort {
    51  		curPort++
    52  
    53  		ln, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP(host), Port: curPort})
    54  		if err != nil {
    55  			continue
    56  		}
    57  		ln.Close()
    58  
    59  		port = curPort
    60  		break
    61  	}
    62  
    63  	if port == 0 {
    64  		return "", moerr.NewInternalErrorNoCtx("failed to allocate")
    65  	}
    66  	return strconv.Itoa(port), nil
    67  }