first commit

main.go

@@ -23,10 +23,6 @@ import (
 	"time"
 )
 
-var (
-	inferenceLastHeartbeat sync.Map // key: deviceUUID, value: time.Time
-)
-
 type deviceAlarmState struct {
 	InCondition     bool
 	StartTime       int64
@@ -250,12 +246,12 @@ func processDevice(ctx context.Context, deviceData connect.DeviceData) {
 	}
 
 	// 关键：保存进程对象到全局 map
-	inferenceProcesses.Store(deviceData.DeviceUUID, inferenceCmd)
+	connect.InferenceProcesses.Store(deviceData.DeviceUUID, inferenceCmd)
 
 	logger.Logger.Printf("推理程序已启动，PID=%d (DeviceUUID=%s)", inferenceCmd.Process.Pid, deviceData.DeviceUUID)
 
 	// 初始化心跳
-	UpdateInferenceHeartbeat(deviceData.DeviceUUID)
+	connect.UpdateInferenceHeartbeat(deviceData.DeviceUUID)
 
 	// 检查推理程序是否启动成功
 	time.Sleep(2 * time.Second)
@@ -492,10 +488,7 @@ func waitForShutdown(ctx context.Context, cancel context.CancelFunc, wg *sync.Wa
 	case <-stop:
 		logger.Logger.Printf("Upon receiving the stop signal, start cleaning...")
 
-		// First cancel all contexts
 		cancel()
-
-		// Wait for all goroutines to complete, but with timeout
 		done := make(chan struct{})
 		go func() {
 			wg.Wait()
@@ -523,17 +516,14 @@ func waitForShutdown(ctx context.Context, cancel context.CancelFunc, wg *sync.Wa
 	os.Exit(0)
 }
 
-// initWebSocketAsync Asynchronously initializes WebSocket
 func initWebSocketAsync(ctx context.Context, wg *sync.WaitGroup) {
 	logger.Logger.Printf("Start the initialization of asynchronous WebSocket...")
 
-	// 1. Initialize WebSocket channel
 	if err := connect.InitWSChannel(); err != nil {
 		logger.Logger.Printf("Failed to initialize WebSocket channel: %v", err)
 		return
 	}
 
-	// 2. Wait for connection to be ready (with context cancellation)
 	logger.Logger.Printf("Waiting for WebSocket connection to be established...")
 	containerName := "fireleave-container"
 
@@ -545,7 +535,7 @@ func initWebSocketAsync(ctx context.Context, wg *sync.WaitGroup) {
 		case <-time.After(1 * time.Second):
 			client := connect.GetWSClient()
 			if client != nil && client.IsConnected() {
-				// Connection successful → proceed to register
+
 				goto register
 			}
 		}
@@ -555,7 +545,7 @@ func initWebSocketAsync(ctx context.Context, wg *sync.WaitGroup) {
 	return
 
 register:
-	// 3. Register service
+
 	serviceID, err := connect.RegisterService(containerName)
 	if err != nil {
 		logger.Logger.Printf("Service registration failed: %v", err)
@@ -571,12 +561,10 @@ register:
 	}
 }
 
-// GetAllDevices Get all devices (new method)
 func (dm *DeviceManager) GetAllDevices() map[string]*connect.DeviceData {
 	dm.mu.Lock()
 	defer dm.mu.Unlock()
 
-	// Return copy to avoid concurrent modification
 	result := make(map[string]*connect.DeviceData)
 	for k, v := range dm.devices {
 		result[k] = v
@@ -701,99 +689,29 @@ func triggerAlarmWithDelay(deviceData connect.DeviceData, alarmTime int64, coold
 	video_server.CleanupOldRawVideos(deviceData.DeviceUUID)
 	video_server.CleanupOldVideos(deviceData.DeviceUUID)
 }
-func UpdateInferenceHeartbeat(deviceUUID string) {
-	inferenceLastHeartbeat.Store(deviceUUID, time.Now())
-}
-func ShouldRestartInference(deviceUUID string) bool {
-	if val, ok := inferenceLastHeartbeat.Load(deviceUUID); ok {
-		last := val.(time.Time)
-		return time.Since(last) > 90*time.Second
-	}
-	return true // 从未收到过，也要启动
-}
 func (dm *DeviceManager) RangeDevices(fn func(uuid string, device *connect.DeviceData) bool) {
 	dm.mu.Lock()
 	defer dm.mu.Unlock()
-
 	for uuid, device := range dm.devices {
-		// 拷贝一份指针，防止外部修改影响遍历
-		devCopy := device
-		if !fn(uuid, devCopy) {
+
+		if !fn(uuid, device) {
 			return
 		}
 	}
 }
-
-// StopInferenceProcess 停止指定设备的推理进程（你现在 defer 中就是这么干的，抽出来统一管理）
-func (dm *DeviceManager) StopInferenceProcess(deviceUUID string) {
-	// 这里你之前没有保存 cmd，我们需要一个地方存起来！
-	// 所以我们要加一个 map 保存每个设备的推理进程
-	if cmd, ok := inferenceProcesses.Load(deviceUUID); ok {
-		if process, ok2 := cmd.(*exec.Cmd); ok2 && process.Process != nil {
-			logger.Logger.Printf("正在停止推理进程 (DeviceUUID=%s, PID=%d)", deviceUUID, process.Process.Pid)
-			process.Process.Kill()
-			process.Wait() // 等待退出，防止僵尸进程
-		}
-		inferenceProcesses.Delete(deviceUUID)
-	}
-}
-
-// StartInferenceProcess 启动推理进程（把你 processDevice 里启动推理的部分抽出来）
-func (dm *DeviceManager) StartInferenceProcess(deviceData *connect.DeviceData) {
-	dm.StopInferenceProcess(deviceData.DeviceUUID)
-
-	// 正确：只查询，不分配
-	port, err := connect.GlobalPortManager.GetPort(deviceData.DeviceUUID)
-	if err != nil {
-		logger.Logger.Printf("获取端口失败，无法重启推理程序 (DeviceUUID=%s): %v", deviceData.DeviceUUID, err)
-		return
-	}
-
-	detectAreaStr := connect.ProcessDetectAreaForInference(deviceData.DetectArea)
-
-	inferenceCmd := exec.Command(
-		"./yolov5",
-		"-s", deviceData.CameraRTSP,
-		"-m", "model",
-		"-c", deviceData.DeviceUUID,
-		"-p", fmt.Sprintf("%d", port), // 用老端口！
-		"-w", detectAreaStr,
-		"-r", fmt.Sprintf("%d", deviceData.Confidence),
-	)
-	inferenceCmd.Stdout = os.Stdout
-	inferenceCmd.Stderr = os.Stderr
-
-	logger.Logger.Printf("正在重启推理程序 (使用原端口 %d): %v", port, inferenceCmd.Args)
-
-	if err := inferenceCmd.Start(); err != nil {
-		logger.Logger.Printf("推理程序重启失败 (DeviceUUID=%s): %v", deviceData.DeviceUUID, err)
-		return
-	}
-
-	inferenceProcesses.Store(deviceData.DeviceUUID, inferenceCmd)
-	logger.Logger.Printf("推理程序已成功重启，PID=%d (DeviceUUID=%s, Port=%d)",
-		inferenceCmd.Process.Pid, deviceData.DeviceUUID, port)
-
-	UpdateInferenceHeartbeat(deviceData.DeviceUUID)
-}
-
-var inferenceProcesses sync.Map
-
 func main() {
-	// Define command line parameters (remove ws parameter)
+
 	var rtspURL string
 	var deviceUUID string
 	flag.StringVar(&rtspURL, "rtsp", "", "RTSP stream URL for testing")
 	flag.StringVar(&deviceUUID, "uuid", "", "Device UUID for testing")
 	flag.Parse()
 
-	// Initialize logger
 	if err := logger.InitLogger(); err != nil {
 		log.Fatalf("Logger initialization failed: %v", err)
 	}
 	defer logger.CloseLogger()
 
-	// Log cleanup
 	go func() {
 		for {
 			logger.CleanupOldLogs()
@@ -801,20 +719,18 @@ func main() {
 		}
 	}()
 
-	// Directly enter normal mode (remove wsMode judgment)
 	logger.Logger.Printf("=== Starting FireLeave Tool ===")
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 	var wg sync.WaitGroup
 
-	// 1. Initialize device data directly from service config
 	logger.Logger.Printf("Step 1: Initialize device data from service config")
 	deviceDataList, err := connect.InitializeDevicesFromServiceConfig()
 	if err != nil {
 		logger.Logger.Printf("Device data initialization failed: %v", err)
 		logger.Logger.Printf("Will try to use existing device files...")
-		// If initialization fails, try to load from existing files
+
 		deviceDataList = connect.GetDeviceDataList()
 	}
 
@@ -833,7 +749,6 @@ func main() {
 
 	logger.Logger.Printf("Step 3: All devices started")
 
-	// 3. Asynchronously initialize WebSocket
 	logger.Logger.Printf("Step 4: Asynchronously initialize WebSocket connection")
 	go initWebSocketAsync(ctx, &wg)
 
@@ -848,12 +763,12 @@ func main() {
 		defer ticker.Stop()
 		for range ticker.C {
 			GlobalDeviceManager.RangeDevices(func(uuid string, device *connect.DeviceData) bool {
-				if ShouldRestartInference(uuid) {
+				if connect.ShouldRestartInference(uuid) {
 					logger.Logger.Printf("推理程序心跳超时（>90s），准备重启: DeviceUUID=%s", uuid)
-					GlobalDeviceManager.StopInferenceProcess(uuid)
-					time.Sleep(2 * time.Second)
-					GlobalDeviceManager.StartInferenceProcess(device)
-					UpdateInferenceHeartbeat(uuid) // 防止重复重启
+					connect.StopInferenceProcess(uuid)
+					time.Sleep(6 * time.Second)
+					connect.StartInferenceProcess(device)
+					connect.UpdateInferenceHeartbeat(uuid)
 				}
 				return true
 			})
@@ -867,7 +782,7 @@ func main() {
 			}
 		}
 	}()
-	// Wait for termination signal
+
 	waitForShutdown(ctx, cancel, &wg)
 }
 func printGlobalDeviceStatus() {
@@ -883,7 +798,7 @@ func printGlobalDeviceStatus() {
 
 	i := 1
 	for uuid, data := range activeDevices {
-		// 从 deviceStates 获取冷却时间
+
 		stateInf, _ := deviceStates.Load(uuid)
 		state := stateInf.(*deviceAlarmState)
 		cooldownLeft := int64(0)
@@ -930,7 +845,6 @@ func monitorDeviceStatus(ctx context.Context, initialDevices []connect.DeviceDat
 	}
 }
 
-// debugDeviceData Debug device data
 func debugDeviceData(deviceData connect.DeviceData) {
 	logger.Logger.Printf("Debug device data: UUID=%s", deviceData.DeviceUUID)
 	logger.Logger.Printf("  TaskID: %q", deviceData.TaskID)