naza/pkg/nazanet/avail_udp_conn_pool.go

// Copyright 2020, Chef.  All rights reserved.
// https://github.com/q191201771/naza
//
// Use of this source code is governed by a MIT-style license
// that can be found in the License file.
//
// Author: Chef (191201771@qq.com)

package nazanet

import (
	"net"
	"sync"
)

// 从指定的UDP端口范围内，寻找可绑定监听的端口，绑定监听并返回
// Pool只提供Acquire获取接口，不提供释放接口，连接资源是标准*net.UDPConn对象，需要释放时，外部直接Close即可
//
type AvailUDPConnPool struct {
	minPort uint16
	maxPort uint16

	m        sync.Mutex
	lastPort uint16
}

func NewAvailUDPConnPool(minPort uint16, maxPort uint16) *AvailUDPConnPool {
	return &AvailUDPConnPool{
		minPort:  minPort,
		maxPort:  maxPort,
		lastPort: minPort,
	}
}

func (a *AvailUDPConnPool) Acquire() (*net.UDPConn, uint16, error) {
	a.m.Lock()
	defer a.m.Unlock()

	loopFirstFlag := true
	p := a.lastPort
	for {
		// 找了一轮也没有可用的，返回错误
		if !loopFirstFlag && p == a.lastPort {
			return nil, 0, ErrNazaNet
		}
		loopFirstFlag = false

		conn, err := listenUDPWithPort(p)

		// 绑定失败，尝试下一个端口
		if err != nil {
			p = a.nextPort(p)
			continue
		}

		// 绑定成功，更新last，返回结果
		a.lastPort = a.nextPort(p)
		return conn, p, nil
	}
}

// 有的业务场景，需要返回两个可用的端口，并且必须是连续的
// @return 前面的是端口小的，后面的是端口+1的
//
func (a *AvailUDPConnPool) Acquire2() (*net.UDPConn, uint16, *net.UDPConn, uint16, error) {
	a.m.Lock()
	defer a.m.Unlock()

	loopFirstFlag := true
	p := a.lastPort
	for {
		// 找了一轮也没有可用的，返回错误
		if !loopFirstFlag && p == a.lastPort {
			return nil, 0, nil, 0, ErrNazaNet
		}
		loopFirstFlag = false

		// 因为第一个端口如果为最大值，那么和第二个端口肯定不是线性连续了
		if p == a.maxPort {
			p = a.minPort
			continue
		}

		conn, err := listenUDPWithPort(p)

		// 第一个就绑定失败，尝试下一个端口
		if err != nil {
			p = a.nextPort(p)
			continue
		}

		// 绑定成功，因为我们需要两个，所以我们还要找第二个

		// 因为前面已经有判断最大值了，所以直接+1
		conn2, err := listenUDPWithPort(p + 1)

		// 第二个失败了，关闭第一个，然后从第二个的下一个重新尝试
		if err != nil {
			_ = conn.Close()
			p = a.nextPort(p + 1)
			continue
		}

		// 绑定成功，更新last，返回结果
		a.lastPort = a.nextPort(p + 1)
		return conn, p, conn2, p + 1, nil
	}
}

// 通过Acquire获取到可用net.UDPConn对象后，将对象关闭，只返回可用的端口
func (a *AvailUDPConnPool) Peek() (uint16, error) {
	conn, port, err := a.Acquire()
	if err == nil {
		err = conn.Close()
	}
	return port, err
}

func (a *AvailUDPConnPool) nextPort(p uint16) uint16 {
	if p == a.maxPort {
		return a.minPort
	}

	return p + 1
}