gqc
279 lines
7.0 KiB
Go
279 lines
7.0 KiB
Go
package portmanager
|
||
|
||
import (
|
||
"errors"
|
||
"fmt"
|
||
"net"
|
||
"sync"
|
||
"time"
|
||
|
||
"FireLeave_tool/logger"
|
||
)
|
||
|
||
// PortManager 管理端口分配
|
||
type PortManager struct {
|
||
mu sync.Mutex
|
||
allocated map[string]DevicePorts // DeviceUUID 到端口映射的映射
|
||
portPool []int // 推理端口池（HTTP）
|
||
available map[int]bool // 可用端口标记
|
||
}
|
||
|
||
// NewPortManager 创建端口管理器
|
||
func NewPortManager() *PortManager {
|
||
// 初始化端口池
|
||
var ports []int
|
||
for port := BasePortStart; port <= BasePortEnd; port++ {
|
||
ports = append(ports, port)
|
||
}
|
||
|
||
// 初始化可用端口映射
|
||
available := make(map[int]bool)
|
||
for _, port := range ports {
|
||
available[port] = true
|
||
}
|
||
|
||
return &PortManager{
|
||
allocated: make(map[string]DevicePorts),
|
||
portPool: ports,
|
||
available: available,
|
||
}
|
||
}
|
||
|
||
// AllocatePort 为设备分配基础端口
|
||
func (pm *PortManager) AllocatePort(deviceUUID string) (int, error) {
|
||
pm.mu.Lock()
|
||
defer pm.mu.Unlock()
|
||
|
||
// 检查设备是否已分配端口
|
||
if ports, exists := pm.allocated[deviceUUID]; exists {
|
||
logger.Logger.Printf("Device %s already allocated port: %d", deviceUUID, ports.Inference)
|
||
return ports.Inference, nil
|
||
}
|
||
|
||
// 查找可用端口
|
||
for _, port := range pm.portPool {
|
||
if !pm.available[port] {
|
||
continue // 端口已被分配
|
||
}
|
||
|
||
tempPort := port + TemperaturePortOffset // 温度端口（TCP）
|
||
|
||
// 检查推理端口（HTTP）是否可用
|
||
inferenceLn, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
|
||
if err != nil {
|
||
logger.Logger.Printf("Port %d (inference) is not available: %v", port, err)
|
||
continue
|
||
}
|
||
inferenceLn.Close()
|
||
|
||
// 检查温度端口（TCP）是否可用
|
||
tempLn, err := net.Listen("tcp", fmt.Sprintf(":%d", tempPort))
|
||
if err != nil {
|
||
logger.Logger.Printf("Port %d (temperature) is not available: %v", tempPort, err)
|
||
continue
|
||
}
|
||
tempLn.Close()
|
||
|
||
// 端口可用，记录分配
|
||
pm.allocated[deviceUUID] = DevicePorts{
|
||
Inference: port,
|
||
Temperature: tempPort,
|
||
}
|
||
pm.available[port] = false
|
||
|
||
logger.Logger.Printf("Port allocated for device %s: inference=%d, temperature=%d",
|
||
deviceUUID, port, tempPort)
|
||
return port, nil
|
||
}
|
||
|
||
return 0, errors.New("no available ports for inference and temperature")
|
||
}
|
||
|
||
// AllocateFireCheckPort 为火焰检测程序分配端口
|
||
func (pm *PortManager) AllocateFireCheckPort(deviceUUID string) (int, error) {
|
||
pm.mu.Lock()
|
||
defer pm.mu.Unlock()
|
||
|
||
// 检查设备是否已分配端口
|
||
ports, exists := pm.allocated[deviceUUID]
|
||
if !exists {
|
||
// 先分配基础端口
|
||
_, err := pm.AllocatePort(deviceUUID)
|
||
if err != nil {
|
||
return 0, fmt.Errorf("failed to allocate base port: %v", err)
|
||
}
|
||
// 重新获取更新后的端口映射
|
||
ports = pm.allocated[deviceUUID]
|
||
}
|
||
|
||
// 计算火焰检测端口
|
||
fireCheckPort := ports.Inference + FireCheckPortOffset
|
||
|
||
// 检查火焰检测端口是否可用
|
||
fireCheckLn, err := net.Listen("tcp", fmt.Sprintf(":%d", fireCheckPort))
|
||
if err != nil {
|
||
logger.Logger.Printf("Port %d (fire_check) is not available: %v", fireCheckPort, err)
|
||
return 0, fmt.Errorf("fire check port %d not available: %v", fireCheckPort, err)
|
||
}
|
||
fireCheckLn.Close()
|
||
|
||
// 更新端口映射
|
||
ports.FireCheck = fireCheckPort
|
||
pm.allocated[deviceUUID] = ports
|
||
|
||
logger.Logger.Printf("Fire check port allocated for device %s: %d", deviceUUID, fireCheckPort)
|
||
return fireCheckPort, nil
|
||
}
|
||
|
||
// ReleasePort 释放设备端口
|
||
func (pm *PortManager) ReleasePort(deviceUUID string) {
|
||
pm.mu.Lock()
|
||
defer pm.mu.Unlock()
|
||
|
||
// 删除设备端口映射
|
||
if ports, exists := pm.allocated[deviceUUID]; exists {
|
||
delete(pm.allocated, deviceUUID)
|
||
pm.available[ports.Inference] = true // 标记基础端口为可用
|
||
logger.Logger.Printf("Port released for device %s: inference=%d, temperature=%d, fire_check=%d",
|
||
deviceUUID, ports.Inference, ports.Temperature, ports.FireCheck)
|
||
}
|
||
}
|
||
|
||
// GetPort 获取设备基础端口
|
||
func (pm *PortManager) GetPort(deviceUUID string) (int, error) {
|
||
pm.mu.Lock()
|
||
defer pm.mu.Unlock()
|
||
|
||
if ports, exists := pm.allocated[deviceUUID]; exists {
|
||
return ports.Inference, nil
|
||
}
|
||
|
||
return 0, fmt.Errorf("device %s has no allocated port", deviceUUID)
|
||
}
|
||
|
||
// GetFireCheckPort 获取火焰检测程序端口
|
||
func (pm *PortManager) GetFireCheckPort(deviceUUID string) (int, error) {
|
||
pm.mu.Lock()
|
||
defer pm.mu.Unlock()
|
||
|
||
if ports, exists := pm.allocated[deviceUUID]; exists {
|
||
if ports.FireCheck > 0 {
|
||
return ports.FireCheck, nil
|
||
}
|
||
return 0, fmt.Errorf("device %s has no allocated fire check port", deviceUUID)
|
||
}
|
||
|
||
return 0, fmt.Errorf("device %s has no allocated ports", deviceUUID)
|
||
}
|
||
|
||
// GetTemperaturePort 获取温度程序端口
|
||
func (pm *PortManager) GetTemperaturePort(deviceUUID string) (int, error) {
|
||
pm.mu.Lock()
|
||
defer pm.mu.Unlock()
|
||
|
||
if ports, exists := pm.allocated[deviceUUID]; exists {
|
||
return ports.Temperature, nil
|
||
}
|
||
|
||
return 0, fmt.Errorf("device %s has no allocated ports", deviceUUID)
|
||
}
|
||
|
||
// GetAllocatedPorts 获取所有已分配的端口信息
|
||
func (pm *PortManager) GetAllocatedPorts() []PortInfo {
|
||
pm.mu.Lock()
|
||
defer pm.mu.Unlock()
|
||
|
||
var portInfos []PortInfo
|
||
now := time.Now().Unix()
|
||
|
||
for deviceUUID, ports := range pm.allocated {
|
||
// 基础推理端口
|
||
portInfos = append(portInfos, PortInfo{
|
||
DeviceUUID: deviceUUID,
|
||
PortType: PortTypeInference,
|
||
Port: ports.Inference,
|
||
AllocatedAt: now,
|
||
})
|
||
|
||
// 温度端口
|
||
portInfos = append(portInfos, PortInfo{
|
||
DeviceUUID: deviceUUID,
|
||
PortType: PortTypeTemperature,
|
||
Port: ports.Temperature,
|
||
AllocatedAt: now,
|
||
})
|
||
|
||
// 火焰检测端口
|
||
if ports.FireCheck > 0 {
|
||
portInfos = append(portInfos, PortInfo{
|
||
DeviceUUID: deviceUUID,
|
||
PortType: PortTypeFireCheck,
|
||
Port: ports.FireCheck,
|
||
AllocatedAt: now,
|
||
})
|
||
}
|
||
|
||
// 新推理端口
|
||
if ports.NewInference > 0 {
|
||
portInfos = append(portInfos, PortInfo{
|
||
DeviceUUID: deviceUUID,
|
||
PortType: PortTypeNewInference,
|
||
Port: ports.NewInference,
|
||
AllocatedAt: now,
|
||
})
|
||
}
|
||
}
|
||
|
||
return portInfos
|
||
}
|
||
|
||
// GetDevicePorts 获取设备的所有端口
|
||
func (pm *PortManager) GetDevicePorts(deviceUUID string) (DevicePorts, error) {
|
||
pm.mu.Lock()
|
||
defer pm.mu.Unlock()
|
||
|
||
if ports, exists := pm.allocated[deviceUUID]; exists {
|
||
return ports, nil
|
||
}
|
||
|
||
return DevicePorts{}, fmt.Errorf("device %s has no allocated ports", deviceUUID)
|
||
}
|
||
|
||
// IsPortAvailable 检查端口是否可用
|
||
func (pm *PortManager) IsPortAvailable(port int) bool {
|
||
pm.mu.Lock()
|
||
defer pm.mu.Unlock()
|
||
|
||
// 检查是否在端口池中
|
||
for _, p := range pm.portPool {
|
||
if p == port {
|
||
return pm.available[port]
|
||
}
|
||
}
|
||
|
||
// 对于非端口池中的端口，直接检查可用性
|
||
ln, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
|
||
if err != nil {
|
||
return false
|
||
}
|
||
ln.Close()
|
||
return true
|
||
}
|
||
|
||
// GetAvailablePortCount 获取可用端口数量
|
||
func (pm *PortManager) GetAvailablePortCount() int {
|
||
pm.mu.Lock()
|
||
defer pm.mu.Unlock()
|
||
|
||
count := 0
|
||
for _, available := range pm.available {
|
||
if available {
|
||
count++
|
||
}
|
||
}
|
||
return count
|
||
}
|
||
|
||
// GlobalPortManager 全局端口管理器
|
||
var GlobalPortManager = NewPortManager()
|