110 lines
2.7 KiB
Go
110 lines
2.7 KiB
Go
package connect
|
||
|
||
import (
|
||
"errors"
|
||
"fmt"
|
||
"net"
|
||
"sync"
|
||
|
||
"FireLeave_tool/logger"
|
||
)
|
||
|
||
// PortManager 管理端口分配
|
||
type PortManager struct {
|
||
mu sync.Mutex
|
||
allocated map[string]int // DeviceUUID 到推理端口的映射
|
||
portPool []int // 推理端口池（HTTP）
|
||
available map[int]bool // 可用端口标记
|
||
basePort int
|
||
maxPort int
|
||
}
|
||
|
||
func NewPortManager() *PortManager {
|
||
// 推理端口范围：58881-58890（HTTP），支持最多10个设备
|
||
ports := []int{58881, 58882, 58883, 58884, 58885, 58886, 58887, 58888, 58889, 58890}
|
||
|
||
// 初始化可用端口映射
|
||
available := make(map[int]bool)
|
||
for _, port := range ports {
|
||
available[port] = true
|
||
}
|
||
|
||
return &PortManager{
|
||
allocated: make(map[string]int),
|
||
portPool: ports,
|
||
available: available,
|
||
}
|
||
}
|
||
|
||
func (pm *PortManager) AllocatePort(deviceUUID string) (int, error) {
|
||
pm.mu.Lock()
|
||
defer pm.mu.Unlock()
|
||
|
||
// 检查设备是否已分配端口
|
||
if port, exists := pm.allocated[deviceUUID]; exists {
|
||
logger.Logger.Printf("Device %s already allocated port: %d", deviceUUID, port)
|
||
return port, nil
|
||
}
|
||
|
||
// 查找可用端口
|
||
for _, port := range pm.portPool {
|
||
if !pm.available[port] {
|
||
continue // 端口已被分配
|
||
}
|
||
|
||
tempPort := port + 1000 // 温度端口（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] = port
|
||
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")
|
||
}
|
||
|
||
func (pm *PortManager) ReleasePort(deviceUUID string) {
|
||
pm.mu.Lock()
|
||
defer pm.mu.Unlock()
|
||
|
||
// 删除设备端口映射
|
||
if port, exists := pm.allocated[deviceUUID]; exists {
|
||
delete(pm.allocated, deviceUUID)
|
||
pm.available[port] = true // 标记端口为可用
|
||
logger.Logger.Printf("Port released for device %s: %d", deviceUUID, port)
|
||
}
|
||
}
|
||
|
||
// GlobalPortManager 全局端口管理器
|
||
var GlobalPortManager = NewPortManager()
|
||
|
||
func (pm *PortManager) GetPort(deviceUUID string) (int, error) {
|
||
pm.mu.Lock()
|
||
defer pm.mu.Unlock()
|
||
|
||
if port, exists := pm.allocated[deviceUUID]; exists {
|
||
return port, nil
|
||
}
|
||
|
||
return 0, fmt.Errorf("device %s has no allocated port", deviceUUID)
|
||
}
|