github.com/mysteriumnetwork/node@v0.0.0-20240516044423-365054f76801/core/port/pool.go (about)

     1  /*
     2   * Copyright (C) 2019 The "MysteriumNetwork/node" Authors.
     3   *
     4   * This program is free software: you can redistribute it and/or modify
     5   * it under the terms of the GNU General Public License as published by
     6   * the Free Software Foundation, either version 3 of the License, or
     7   * (at your option) any later version.
     8   *
     9   * This program is distributed in the hope that it will be useful,
    10   * but WITHOUT ANY WARRANTY; without even the implied warranty of
    11   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    12   * GNU General Public License for more details.
    13   *
    14   * You should have received a copy of the GNU General Public License
    15   * along with this program.  If not, see <http://www.gnu.org/licenses/>.
    16   */
    17  
    18  package port
    19  
    20  import (
    21  	"fmt"
    22  	"math/rand"
    23  
    24  	"github.com/pkg/errors"
    25  	"github.com/rs/zerolog/log"
    26  
    27  	"github.com/mysteriumnetwork/node/utils/random"
    28  )
    29  
    30  // Pool hands out ports for service use
    31  type Pool struct {
    32  	start, capacity int
    33  	rand            *rand.Rand
    34  }
    35  
    36  // ServicePortSupplier provides port needed to run a service on
    37  type ServicePortSupplier interface {
    38  	Acquire() (Port, error)
    39  	AcquireMultiple(n int) (ports []Port, err error)
    40  }
    41  
    42  // NewFixedRangePool creates a fixed size pool from port.Range
    43  func NewFixedRangePool(r Range) *Pool {
    44  	return &Pool{
    45  		start:    r.Start,
    46  		capacity: r.Capacity(),
    47  		rand:     random.NewTimeSeededRand(),
    48  	}
    49  }
    50  
    51  // Acquire returns an unused port in pool's range
    52  func (pool *Pool) Acquire() (Port, error) {
    53  	p, err := pool.seekAvailablePort()
    54  	log.Info().Err(err).Msgf("Supplying port %d", p)
    55  	return Port(p), err
    56  }
    57  
    58  func (pool *Pool) seekAvailablePort() (int, error) {
    59  	randomOffset := pool.rand.Intn(pool.capacity)
    60  	for i := 0; i < pool.capacity; i++ {
    61  		p := pool.start + (randomOffset+i)%pool.capacity
    62  		available, err := available(p)
    63  		if available || err != nil {
    64  			return p, err
    65  		}
    66  	}
    67  	return 0, errors.New("port pool is exhausted")
    68  }
    69  
    70  // AcquireMultiple returns n unused ports from pool's range.
    71  func (pool *Pool) AcquireMultiple(n int) (ports []Port, err error) {
    72  	if n > pool.capacity {
    73  		return nil, fmt.Errorf("requested more ports (%d) than pool capacity (%d)", n, pool.capacity)
    74  	}
    75  
    76  	portSet := make(map[Port]struct{})
    77  	for i := 0; i < 10*n && len(portSet) < n; i++ {
    78  		p, err := pool.Acquire()
    79  		if err != nil {
    80  			continue
    81  		}
    82  
    83  		portSet[p] = struct{}{}
    84  
    85  		if len(portSet) == n {
    86  			for port := range portSet {
    87  				ports = append(ports, port)
    88  			}
    89  			return ports, nil
    90  		}
    91  	}
    92  
    93  	return nil, errors.New("too many collisions")
    94  }