m7s/m7s.live/plugin/gb28181/v4@v4.4.4/portmanager.go

package gb28181

import (
	"io"
	"fmt"
	"net"
	"sync"
)

type PortManager struct {
	recycle chan uint16
	max     uint16
	pos     uint16
	Valid   bool
}

var portMutex sync.Mutex

func (pm *PortManager) Init(start, end uint16) {
	pm.pos = start - 1
	pm.max = end
	if pm.pos > 0 && pm.max > pm.pos {
		pm.Valid = true
		pm.recycle = make(chan uint16, pm.Range())
	}
}

func (pm *PortManager) Range() uint16 {
	return pm.max - pm.pos
}

func (pm *PortManager) Recycle(p uint16) (err error) {

	select {
	case pm.recycle <- p:
		//fmt.Println("recycle port 00000000000000000000000000",p)
		return nil
	default:
		return io.EOF //TODO: 换一个Error
	}
}

func tryTcpPortUnUse(port uint16) bool {

	ln, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		// 如果监听失败，则说明端口已被占用
		fmt.Printf("Port %d is NOT available,err=%s\n", port,err.Error())
		return false
	}
	defer ln.Close()
	//fmt.Println("端口可用   ------------------------",port)
	return true
}

func (pm *PortManager) GetUnUseTcpPort() (p uint16, err error){
	for i:=0;i<3;i++{
		p, err = pm.GetPort()
		if err != nil {
			return
		}
		if tryTcpPortUnUse(p){
			return
		}
	}

	return 0 ,fmt.Errorf("尝试三次端口被占用")
}

func (pm *PortManager) GetPort() (p uint16, err error) {
	portMutex.Lock()
	defer portMutex.Unlock()
	select {
	case p = <-pm.recycle:
		return
	default:
		if pm.Range() > 0 {
			pm.pos++
			p = pm.pos
			return
		} else {
			return 0 , fmt.Errorf("无可用端口")
			//return 0, io.EOF //TODO: 换一个Error
		}
	}
}