AI-tianyan/internal/stream/ingestor.go

package stream

import (
	"context"
	"fmt"
	"io"
	"log"
	"os/exec"
	"sync"
	"time"

	"bytes"

	"tianyan-edge/internal/config"
)

type Frame struct {
	StreamID int
	DeviceID string
	URL      string
	JPEG     []byte
	TS       float64
}

type StreamManager struct {
	cfg       *config.Config
	frames    chan<- Frame
	processes map[string]*exec.Cmd
	mu        sync.Mutex
	updates   <-chan *config.Config
}

func NewStreamManager(cfg *config.Config, frames chan<- Frame, updates <-chan *config.Config) *StreamManager {
	return &StreamManager{
		cfg:       cfg,
		frames:    frames,
		processes: make(map[string]*exec.Cmd),
		updates:   updates,
	}
}

func (sm *StreamManager) Run(ctx context.Context) {
	// Initial start
	sm.applyStreams(ctx, sm.cfg.RTSPURLs)

	for {
		select {
		case <-ctx.Done():
			sm.stopAll()
			return
		case newCfg := <-sm.updates:
			log.Printf("stream: applying new config, urls=%d", len(newCfg.RTSPURLs))
			sm.applyStreams(ctx, newCfg.RTSPURLs)
		}
	}
}

func (sm *StreamManager) applyStreams(ctx context.Context, urls []string) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	// Identify URLs to remove (present in sm.processes but not in new urls)
	toRemove := map[string]*exec.Cmd{}
	for url, cmd := range sm.processes {
		found := false
		for _, u := range urls {
			if u == url {
				found = true
				break
			}
		}
		if !found {
			toRemove[url] = cmd
		}
	}

	// Stop removed streams
	for url, cmd := range toRemove {
		log.Printf("stream: stopping %s", url)
		cmd.Process.Kill()
		cmd.Wait()
		delete(sm.processes, url)
	}

	// Start new streams
	for i, url := range urls {
		if _, exists := sm.processes[url]; !exists {
			go sm.streamLoop(ctx, i, url)
		}
	}
}

func (sm *StreamManager) streamLoop(ctx context.Context, idx int, url string) {
	deviceID := fmt.Sprintf("cam-%03d", idx)
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}

		cmd := exec.CommandContext(ctx, "ffmpeg",
			"-i", url,
			"-vf", fmt.Sprintf("fps=%d", sm.cfg.InferFPS),
			"-f", "image2pipe", "-vcodec", "mjpeg", "-q:v", "5", "-",
		)
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			log.Printf("stream[%d] stdout pipe failed: %v", idx, err)
			time.Sleep(5 * time.Second)
			continue
		}

		sm.mu.Lock()
		sm.processes[url] = cmd
		sm.mu.Unlock()

		if err := cmd.Start(); err != nil {
			log.Printf("stream[%d] start failed: %v", idx, err)
			time.Sleep(5 * time.Second)
			continue
		}

		log.Printf("stream[%d] connected  url=%s", idx, url)
		sm.readFrames(ctx, stdout, idx, deviceID, url)
		cmd.Wait()

		select {
		case <-ctx.Done():
			return
		default:
			log.Printf("stream[%d] lost, retry 5s", idx)
			time.Sleep(5 * time.Second)
		}
	}
}

func (sm *StreamManager) readFrames(ctx context.Context, r io.Reader, idx int, deviceID, url string) {
	soi, eoi := []byte{0xFF, 0xD8}, []byte{0xFF, 0xD9}
	buf, tmp := make([]byte, 0, 1<<20), make([]byte, 32768)
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}
		n, err := r.Read(tmp)
		if n > 0 {
			buf = append(buf, tmp[:n]...)
			for {
				s := bytes.Index(buf, soi)
				if s < 0 {
					buf = buf[:0]
					break
				}
				buf = buf[s:]
				e := bytes.Index(buf[2:], eoi)
				if e < 0 {
					break
				}
				end := e + 4
				frame := make([]byte, end)
				copy(frame, buf[:end])
				buf = buf[end:]
				select {
				case sm.frames <- Frame{StreamID: idx, DeviceID: deviceID, URL: url,
					JPEG: frame, TS: float64(time.Now().UnixMilli()) / 1000.0}:
				default:
				}
			}
		}
		if err != nil {
			return
		}
	}
}

func (sm *StreamManager) stopAll() {
	sm.mu.Lock()
	defer sm.mu.Unlock()
	for url, cmd := range sm.processes {
		log.Printf("stream: stopping all %s", url)
		cmd.Process.Kill()
		delete(sm.processes, url)
	}
}